Skip to content

Add a hook so that downstream can use to wrap more fixtures - #191

Open
seberg wants to merge 10 commits into
Quansight-Labs:mainfrom
seberg:feature/prepare-args-hook
Open

Add a hook so that downstream can use to wrap more fixtures#191
seberg wants to merge 10 commits into
Quansight-Labs:mainfrom
seberg:feature/prepare-args-hook

Conversation

@seberg

@seberg seberg commented Jul 18, 2026

Copy link
Copy Markdown

This uses a hook to setup wrapping of fixtures and also uses that for tmp_dir and tmp_path.

Discussed and worked in with @agriyakhetarpal. We decided on this approach of returning a dict, discarding the idea of calling the original fixture function again, because that requires reaching into private pytest API (probably works in practice, but...).

(Large written with agent, at least the tests probably need one iteration, the code should be OK enough.).

Closes gh-189

seberg and others added 2 commits July 18, 2026 14:28
This uses a hook to setup wrapping of fixtures and also uses that
for tmp_dir and tmp_path.

Co-authored-by: Agriya Khetarpal <74401230+agriyakhetarpal@users.noreply.github.com>
@agriyakhetarpal
agriyakhetarpal self-requested a review July 18, 2026 15:14
@agriyakhetarpal

Copy link
Copy Markdown
Member

Thanks a lot, @seberg, it was pretty nice to work together on this! I got my side of the Fable implementation running on top of this, and it found a few bugs here for us, along with a few general improvements. Listing the AI output below; note that I distilled it down quite a bit to make it much less verbose than it usually is (😅):

  1. A deadlock (confirmed by a 30-second hang). If a transform was raised in one worker thread, that thread died before reaching barrier.wait(), so the other threads waited forever, and the error was never reported. The closure now records the error, aborts the barrier, and the other threads catch BrokenBarrierError and exit. See 6dfaaf8.
  2. A regression vs main: the hook was called once per session with the configured --parallel-threads value, so a test marked force_parallel_threads(2) under the default CLI got no tmp_path/tmpdir isolation - threads shared one directory, which main handled correctly. The hook is now called with each test's actual thread count, which is cached per distinct count, so it still runs once per count rather than once per test. This also makes n_workers more meaningful for downstream hookimpls.
  3. Smaller fixes:
    • The ordering comment in _get_wrap_fixtures said pluggy finds the most specific hookimpl last - it's the opposite, and our code was right but for the wrong reason, so I corrected the comment and two similar ones in the tests.
    • A teardown error during a failed setup no longer shadows the original transform error.
    • The hookspec registration has been moved from pytest_configure to pytest_addhooks, which is the canonical place.
  4. The per-test n_workers change alters the hook's calling convention slightly (it can now fire more than once per session, once per distinct count). I think it's the right call since it fixes the force_parallel_threads regression without touching the agreed dict-of-transforms shape.

I also got it to port the tests from my implementation to yours and added a few more in aa2e60d. Then I finished by adding the change in f854f9e that we both liked, which is to use *thread_setups* in the API naming. I split the commits so that it's easier for you to take a look!

agriyakhetarpal added a commit to HIPS/autograd that referenced this pull request Jul 19, 2026
Replaces the per-test `rng = npr.RandomState(42)` locals (401 of them)
with an `rng` fixture defined once in tests/conftest.py. Under
pytest-run-parallel, the new `pytest_run_parallel_get_thread_setups`
hook swaps the fixture for a fresh, identically seeded instance in
every thread, so each thread reproduces the single-threaded stream.
The hookimpl uses optionalhook, so plain pytest runs without the
plugin are unaffected.

The nox session and the TSAN workflow temporarily install the plugin
from the Quansight-Labs/pytest-run-parallel#191 branch. This commit
will be reverted once that PR is settled and released.
agriyakhetarpal added a commit to HIPS/autograd that referenced this pull request Jul 19, 2026
Replaces the per-test `rng = npr.RandomState(42)` locals (401 of them)
with an `rng` fixture defined once in tests/conftest.py. Under
pytest-run-parallel, the new `pytest_run_parallel_get_thread_setups`
hook swaps the fixture for a fresh, identically seeded instance in
every thread, so each thread reproduces the single-threaded stream.
The hookimpl uses optionalhook, so plain pytest runs without the
plugin are unaffected.

The nox session and the TSAN workflow temporarily install the plugin
from the Quansight-Labs/pytest-run-parallel#191 branch. This commit
will be reverted once that PR is settled and released.

Co-Authored-By: Sebastian Berg <sebastian@sipsolutions.net>

@agriyakhetarpal agriyakhetarpal left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tested this downstream in HIPS/autograd#659 (see the linked commit above and the conftest.py file), and it works pretty well! @seberg, could you please try this with cuda-python now? Once you confirm the new changes, of course.

Also, @ngoldbaum, I guess this is also worth testing downstream with scikit-image as a follow-up to scikit-image/scikit-image#8065, which kind of motivated my Autograd changes after your suggestion about it, and motivated this PR too.

Thank you!

@ngoldbaum

Copy link
Copy Markdown
Collaborator

@bwhitt7 is there any chance I can interest you in taking a look at this? It generalizes the idea you implemented last year to fix thread safety issues in fixtures like tmp_path by allowing library authors to register new fixtures to handle.

@ngoldbaum ngoldbaum left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I left a few comments, the deadlock issue is the biggest one.

I think the feature itself is worth having but maybe we can do this in a more simple way?

Alternative 1: a per-thread setup hook instead of fixture transforms.

The original ask in gh-189 was "run some setup, like cuInit, in every thread". This PR answers that indirectly: you can only attach setup to a fixture, and only when a test asks for that fixture by name in its signature. Tests that don't use such a fixture get nothing, and autouse / usefixtures fixtures don't count. A more direct answer is a hook that just runs inside every worker thread, before the barrier, written like a fixture:

# conftest.py
import threading, pytest

_cuda = threading.local()

@pytest.hookimpl(optionalhook=True)
def pytest_run_parallel_thread_setup(thread_index, n_workers):
    _cuda.ctx = make_context(thread_index)   # runs in the worker thread
    yield                                    # test body runs here
    _cuda.ctx.destroy()                      # per-thread teardown

Because it runs inside the worker at test time, almost all of the bookkeeping in this PR disappears: no collection-time hook call, no caching per thread count, no rules for merging dicts from several conftests, no per-fixture filtering. The plugin calls the hook in each thread and drives the generators. If we want, the transform hook from this PR can be added later on top of the same plumbing — the hard parts (running user code in each thread, the teardown stack, breaking the barrier when setup fails) are shared.

Alternative 2: keep fixture transforms, but call the hook once and simplify the rules.

If per-thread fixture values are the important part, the current shape can still shed weight. Call the hook once per session with no n_workers argument, and give that information to the transform when it runs, bundled in one context object so fields can be added later without breaking existing transforms:

def transform(value, ctx):   # ctx.thread_index, ctx.n_workers, ...
    ...

That deletes the per-thread-count cache and the "hook may run several times, once per distinct count" rule — the built-in tmp_path transform just returns value unchanged when ctx.n_workers == 1. It also fixes a trap in the current calling convention: transforms are invoked as transform(value, thread_index=...), so nothing new (say, iteration_index) can ever be passed without breaking every released transform. Separately, I'd consider "most specific wins" instead of chaining when two implementations register the same fixture name — that's the rule fixture overrides already use, it's easier to explain, and it makes it possible to replace the built-in tmp_path handling rather than only wrap it.

Comment thread README.md
```python
import pytest

@pytest.hookimpl

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should use optionalhook=True and say why, otherwise people will have issues if pytest-run-parallel isn't installed

Comment thread README.md

@pytest.hookimpl
def pytest_run_parallel_get_thread_setups(n_workers):
def transform_db(value, *, thread_index):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not take **kwargs to allow extensibility?

Comment thread README.md
Comment on lines +164 to +166
chained. Transforms from more specific hookimpls (which pluggy calls first,
such as one in a nested `conftest.py`) are applied last, so they can wrap or
override the others. The built-in implementation wraps `tmp_path` and `tmpdir`

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My agent reviewer says this claim isn't true and to do that you'd need to use the item.ihook pytest hook for this and a different cache key to implement this scoping rule.

Comment thread README.md
When using the fixtures `thread_index` and `iteration_index`, they should be
requested directly by tests, and will return 0 when requested by other fixtures.

### Per-thread fixture setups

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The docs should target a Python developer working with a test suite who doesn't necessarily know pytest or pytest-run-parallel all that well. Currently these docs target a pytest/pluggy expert. The very specific docs are good too but they shouldn't be the only docs and we shouldn't lead with them IMO.

iteration_index=iteration_index,
n_workers=n_workers,
)
except Exception as e:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs to be BaseException, otherwise there's a possible deadlock. This script deadlocks for me on this PR:

import subprocess
import sys
import tempfile
from pathlib import Path

tmp = Path(tempfile.mkdtemp())

(tmp / "conftest.py").write_text("""
import pytest

@pytest.hookimpl
def pytest_run_parallel_get_thread_setups(n_workers):
    def transform(value, *, thread_index):
        if thread_index == 1:
            pytest.skip("no device in this thread")
        return value
    return {"resource": transform}
""")

(tmp / "test_repro.py").write_text("""
import pytest

@pytest.fixture
def resource():
    return "base"

def test_resource(resource):
    pass
""")

try:
    subprocess.run(
        [sys.executable, "-m", "pytest", "--parallel-threads=2"],
        cwd=tmp, timeout=5,
    )
except subprocess.TimeoutExpired:
    print("\nDEADLOCK: pytest still hanging after 5 seconds")

@seberg

seberg commented Jul 24, 2026

Copy link
Copy Markdown
Author

So with the alternatives. A single function was my first thought. It is nice, but you have to pass in which fixtures and mutate values in-place (or return a dict there). It may also make things harder if you set up multiple fixtures (but don't know which ones you need always).
Now, one thing that I do have and that I missed, is that sometimes you have fixtures that take other fixtures. In that case, in theory you'd want to call this setup in the right order and pass the previous values into the transformer.
I'll note that I suspect I don't need this in practice and of course one can always work around, although a single setup function might make that more convenient (it wouldn't solve it fully but we don't need to aim for it).

Maybe, you can pass in everything into the transform, one example that would be harder with this is if you want to have a barrier fixture, as that needs to know the workers before setting up the threads. No real opinion, it removes the need to call it more than once possibly, but means the args you probably ignore are passed to every transformer.

@ngoldbaum

Copy link
Copy Markdown
Collaborator

Thanks for the extra context! I think given those constraints this is alright (after fixing the BaseException catching to fix the deadlock).

Can you also do one more thing to make this extensible in the future? Instead of locking in the signature as def transform(value, *, thread_index), can you use inspect to get the signature at runtime? Like so:

params = inspect.signature(transform).parameters
accepts_any = any(p.kind is p.VAR_KEYWORD for p in params.values())
kwargs = offered if accepts_any else {k: v for k, v in offered.items() if k in params}

This is apparently how pluggy does it and it allows future extensibility. For example, we could add support for a new iteration_index argument without breaking anything.

agriyakhetarpal added a commit to HIPS/autograd that referenced this pull request Aug 3, 2026
@ngoldbaum

Copy link
Copy Markdown
Collaborator

@agriyakhetarpal @seberg did either of you ever want to finish this off?

@seberg

seberg commented Sep 8, 2026

Copy link
Copy Markdown
Author

Yeah, I'll do an iteration in the next days (and confirm that it covers enough of my use-case).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add per thread/test setup hook?

3 participants