slide2vec 6.0.0__tar.gz → 6.0.1__tar.gz

{slide2vec-6.0.0 → slide2vec-6.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 6.0.0
+Version: 6.0.1
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0

{slide2vec-6.0.0 → slide2vec-6.0.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "slide2vec"
-version = "6.0.0"
+version = "6.0.1"
 description = "Embedding of whole slide images with Foundation Models"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -149,7 +149,7 @@ zip-safe = false
 include = ["slide2vec*"]
 
 [tool.setuptools.package-data]
-slide2vec = ["py.typed"]
+slide2vec = ["py.typed", "cufile.json"]
 "slide2vec.configs" = ["*.yaml", "models/*.yaml", "preprocessing/*.yaml"]
 
 [tool.pytest.ini_options]
@@ -180,7 +180,7 @@ no_implicit_reexport = true
 max-line-length = 160
 
 [tool.bumpver]
-current_version = "6.0.0"
+current_version = "6.0.1"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit = false
 tag = false

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/__init__.py

@@ -1,3 +1,7 @@
+from slide2vec._cufile import disable_cufile_compat_mode
+
+disable_cufile_compat_mode()  # Before anything can import cucim.
+
 from slide2vec.api import (
     DenseImageOptions,
     DenseOptions,
@@ -27,7 +31,7 @@ from slide2vec.encoders import (
 )
 
 
-__version__ = "6.0.0"
+__version__ = "6.0.1"
 
 __all__ = [
     "Model",

slide2vec-6.0.1/slide2vec/_cufile.py

@@ -0,0 +1,21 @@
+"""Keep ``import cucim`` from opening a CUDA context on every visible GPU (issue #326).
+
+cuCIM opens the cuFile (GPUDirect Storage) driver when it is imported. Without the
+``nvidia-fs`` kernel module the driver falls back to "compat mode", which still creates a
+~520 MiB CUDA context on each visible device although every read then goes through plain
+POSIX I/O. Pointing cuFile at a config that disallows compat mode makes the driver open fail
+instead: cuCIM reads through POSIX exactly as before, and no context is created. Where
+``nvidia-fs`` is loaded the driver opens as usual and GDS keeps working.
+
+cuFile reads ``CUFILE_ENV_PATH_JSON`` once, when cuCIM is imported, so this must run before
+that import: ``slide2vec/__init__.py`` calls it first. A value the user already set wins.
+"""
+
+import os
+from pathlib import Path
+
+CUFILE_ENV = "CUFILE_ENV_PATH_JSON"
+
+
+def disable_cufile_compat_mode(environ=os.environ) -> None:
+    environ.setdefault(CUFILE_ENV, str(Path(__file__).with_name("cufile.json")))

slide2vec-6.0.1/slide2vec/cufile.json

@@ -0,0 +1,5 @@
+{
+    "properties": {
+        "allow_compat_mode": false
+    }
+}

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/distributed/__init__.py

@@ -2,6 +2,8 @@ import os
 
 import torch
 
+from slide2vec.distributed.worker_entry import PINNED_ENV
+
 _RANK = -1
 _WORLD_SIZE = -1
 _LOCAL_RANK = -1
@@ -43,6 +45,19 @@ def get_local_rank() -> int:
     return _LOCAL_RANK
 
 
+def get_device_ordinal() -> int:
+    """
+    Returns:
+        The CUDA ordinal of this process's GPU: 0 once :mod:`slide2vec.distributed.pin_gpu`
+        has left it the only visible device, else the local rank.
+    """
+    return _device_ordinal(get_local_rank())
+
+
+def _device_ordinal(local_rank: int) -> int:
+    return 0 if os.environ.get(PINNED_ENV) == "1" else local_rank
+
+
 def get_local_size() -> int:
     """
     Returns:
@@ -181,7 +196,7 @@ def enable(
 
     Args:
         set_cuda_current_device: If True, call torch.cuda.set_device() to set the
-            current PyTorch CUDA device to the one matching the local rank.
+            current PyTorch CUDA device to this rank's GPU (see get_device_ordinal()).
         overwrite: If True, overwrites already set variables. Else fails.
     """
 
@@ -192,7 +207,7 @@ def enable(
     torch_env.export(overwrite=overwrite)
 
     if set_cuda_current_device:
-        torch.cuda.set_device(torch_env.local_rank)
+        torch.cuda.set_device(_device_ordinal(torch_env.local_rank))
 
     _RANK = torch_env.rank
     _WORLD_SIZE = torch_env.world_size

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/distributed/direct_embed_worker.py

@@ -42,7 +42,7 @@ def main(argv=None) -> int:
     model_spec = dict(request["model"])
     model = Model.from_preset(
         model_spec["name"],
-        device=f"cuda:{local_rank}",
+        device=f"cuda:{distributed.get_device_ordinal()}",
         output_variant=model_spec.get("output_variant"),
         allow_non_recommended_settings=bool(model_spec["allow_non_recommended_settings"]),
     )

slide2vec-6.0.1/slide2vec/distributed/pin_gpu.py

@@ -0,0 +1,47 @@
+"""Bootstrap that shows a torchrun worker only its own GPU, then runs the worker (issue #326).
+
+``import cucim`` opens a CUDA context on every visible device, so a rank that can see its
+siblings' GPUs holds an idle context on each of them. CUDA reads ``CUDA_VISIBLE_DEVICES``
+once, when it initialises, and ``import slide2vec`` already initialises it (``transformers``
+calls ``torch.cuda.is_available()`` at import). A worker module therefore cannot narrow its
+own visibility: ``python -m slide2vec.distributed.<worker>`` has imported the package before
+the worker's first line runs. torchrun runs this file **by path** instead, so the narrowing
+happens before anything is imported. Keep this file free of slide2vec and torch imports.
+
+Only for workers that run no collectives: a pinned rank cannot see its siblings' GPUs.
+
+Usage: ``torchrun ... pin_gpu.py <worker module> [worker args...]``
+"""
+
+import os
+import runpy
+import sys
+
+PINNED_ENV = "SLIDE2VEC_WORKER_GPU_PINNED"  # Mirrored in worker_entry.PINNED_ENV.
+
+
+def pin_to_own_gpu(environ) -> None:
+    """Narrow ``CUDA_VISIBLE_DEVICES`` to the ``LOCAL_RANK``-th GPU this process inherited."""
+    local_rank = int(environ["LOCAL_RANK"])
+    visible = [gpu.strip() for gpu in environ.get("CUDA_VISIBLE_DEVICES", "").split(",") if gpu.strip()]
+    if visible and local_rank >= len(visible):
+        raise RuntimeError(
+            f"LOCAL_RANK={local_rank} but CUDA_VISIBLE_DEVICES lists only {len(visible)} GPU(s)"
+        )
+    environ["CUDA_VISIBLE_DEVICES"] = visible[local_rank] if visible else str(local_rank)
+    environ[PINNED_ENV] = "1"
+
+
+def main() -> None:
+    pin_to_own_gpu(os.environ)
+    module, *worker_args = sys.argv[1:]
+    # Prefer the package that supplied this bootstrap over a competing checkout in cwd.
+    # Keep cwd importable without changing how relative data and output paths resolve.
+    package_root = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+    sys.path[:1] = [package_root, os.getcwd()]
+    sys.argv = [module, *worker_args]
+    runpy.run_module(module, run_name="__main__", alter_sys=True)
+
+
+if __name__ == "__main__":
+    main()

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/distributed/pipeline_worker.py

@@ -36,7 +36,7 @@ def main(argv=None) -> int:
     model_spec = dict(request["model"])
     model = Model.from_preset(
         model_spec["name"],
-        device=f"cuda:{local_rank}",
+        device=f"cuda:{distributed.get_device_ordinal()}",
         output_variant=model_spec.get("output_variant"),
         allow_non_recommended_settings=bool(model_spec["allow_non_recommended_settings"]),
     )

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/distributed/worker_entry.py

@@ -18,6 +18,10 @@ from contextlib import nullcontext
 from dataclasses import dataclass
 
 
+#: Set by :mod:`slide2vec.distributed.pin_gpu` once this process sees only its own GPU.
+PINNED_ENV = "SLIDE2VEC_WORKER_GPU_PINNED"
+
+
 @dataclass(frozen=True, kw_only=True)
 class WorkerRank:
     """This process's place in the torchrun world: which shard, and which GPU."""
@@ -26,12 +30,14 @@ class WorkerRank:
     global_rank: int
     #: Number of shards the work list is cut into.
     world_size: int
-    #: Local device ordinal, i.e. ``cuda:<local_rank>``.
+    #: This process's index among the ranks on its node; names it in progress labels.
     local_rank: int
+    #: CUDA ordinal of this rank's GPU: 0 once the process is pinned to it, else ``local_rank``.
+    device_ordinal: int
 
     @property
     def device(self) -> str:
-        return f"cuda:{self.local_rank}"
+        return f"cuda:{self.device_ordinal}"
 
 
 def worker_args_parser(module: str, *, add_help: bool = True) -> argparse.ArgumentParser:
@@ -46,12 +52,15 @@ def resolve_worker_rank() -> WorkerRank:
     """Read this process's rank from the env torchrun exports.
 
     Defaults to a single in-process rank when the module is launched bare, so a worker can
-    be run by hand (or in a test) without a torchrun agent.
+    be run by hand (or in a test) without a torchrun agent. A worker launched through
+    :mod:`slide2vec.distributed.pin_gpu` sees only its own GPU, which is then ``cuda:0``.
     """
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
     return WorkerRank(
         global_rank=int(os.environ.get("RANK", "0")),
         world_size=int(os.environ.get("WORLD_SIZE", "1")),
-        local_rank=int(os.environ.get("LOCAL_RANK", "0")),
+        local_rank=local_rank,
+        device_ordinal=0 if os.environ.get(PINNED_ENV) == "1" else local_rank,
     )
 
 
@@ -82,6 +91,6 @@ def worker_progress_context(request: dict, *, rank: WorkerRank):
         return nullcontext()
     return activate_progress_reporter(
         JsonlProgressReporter(
-            progress_events_path, rank=rank.global_rank, progress_label=rank.device
+            progress_events_path, rank=rank.global_rank, progress_label=f"cuda:{rank.local_rank}"
         )
     )

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/runtime/dense_image_stage.py

@@ -275,6 +275,7 @@ def _run_dense_images_distributed(
         request_path.write_text(json.dumps(request, indent=2, sort_keys=True), encoding="utf-8")
         run_torchrun_worker(
             module="slide2vec.distributed.dense_image_worker",
+            pin_gpus=True,  # No collectives: each rank needs only its own GPU.
             num_gpus=execution.num_gpus,
             output_dir=out_dir,
             request_path=request_path,

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/runtime/dense_stage.py

@@ -297,6 +297,7 @@ def _run_dense_distributed(model, specs, *, dense, execution, out_dir) -> None:
         request_path.write_text(json.dumps(request, indent=2, sort_keys=True), encoding="utf-8")
         run_torchrun_worker(
             module="slide2vec.distributed.dense_worker",
+            pin_gpus=True,  # No collectives: each rank needs only its own GPU.
             num_gpus=execution.num_gpus,
             output_dir=out_dir,
             request_path=request_path,

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/runtime/distributed.py

@@ -158,18 +158,28 @@ def run_torchrun_worker(
     output_dir: Path,
     request_path: Path,
     failure_title: str,
+    pin_gpus: bool = False,
     progress_events_path: Path | None = None,
     progress_event_callback: Callable[[Any], None] | None = None,
     popen_factory=subprocess.Popen,
 ) -> None:
+    """Run ``module`` once per GPU under torchrun.
+
+    ``pin_gpus`` shows each rank only its own GPU (see :mod:`slide2vec.distributed.pin_gpu`).
+    Use it for workers that run no collectives.
+    """
+    if pin_gpus:
+        # By path, not ``-m``: the bootstrap must run before the slide2vec package is imported.
+        worker = [str(Path(__file__).resolve().parents[1] / "distributed" / "pin_gpu.py"), module]
+    else:
+        worker = ["-m", module]
     command = [
         sys.executable,
         "-m",
         "torch.distributed.run",
         "--standalone",
         f"--nproc_per_node={num_gpus}",
-        "-m",
-        module,
+        *worker,
         "--output-dir",
         str(output_dir),
         "--request-path",
@@ -188,9 +198,15 @@ def run_torchrun_worker(
             previous_sigterm = signal.signal(signal.SIGTERM, _raise_on_sigterm)
         except (ValueError, OSError):
             previous_sigterm = None
+    # Workers keep the parent's working directory, so a relative output_dir (and every path
+    # derived from it: the request, the tiling artifacts, the progress log) means the same
+    # place in both. PYTHONPATH, not cwd, is what lets an uninstalled checkout be imported.
+    package_root = str(Path(__file__).resolve().parents[2])
+    env = dict(os.environ)
+    env["PYTHONPATH"] = os.pathsep.join(filter(None, [package_root, env.get("PYTHONPATH")]))
     process = popen_factory(
         command,
-        cwd=str(Path(__file__).resolve().parents[2]),
+        env=env,
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
         text=True,

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/runtime/distributed_stage.py

@@ -153,6 +153,7 @@ def run_distributed_embedding_stage(
         output_dir=output_dir,
         request_path=request_path,
         failure_title="Distributed feature extraction failed",
+        pin_gpus=True,
         progress_events_path=progress_events_path,
         progress_event_callback=on_progress_event,
         popen_factory=Popen,
@@ -190,6 +191,7 @@ def run_distributed_direct_embedding_stage(
         output_dir=output_dir,
         request_path=request_path,
         failure_title="Distributed direct embedding failed",
+        pin_gpus=True,
         progress_events_path=progress_events_path,
         popen_factory=Popen,
     )

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec/runtime/image_stage.py

@@ -147,6 +147,7 @@ def _run_images_distributed(model, specs, *, execution, out_dir) -> None:
         request_path.write_text(json.dumps(request, indent=2, sort_keys=True), encoding="utf-8")
         run_torchrun_worker(
             module="slide2vec.distributed.image_worker",
+            pin_gpus=True,  # No collectives: each rank needs only its own GPU.
             num_gpus=execution.num_gpus,
             output_dir=out_dir,
             request_path=request_path,

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 6.0.0
+Version: 6.0.1
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0

{slide2vec-6.0.0 → slide2vec-6.0.1}/slide2vec.egg-info/SOURCES.txt

@@ -3,9 +3,11 @@ README.md
 pyproject.toml
 slide2vec/__init__.py
 slide2vec/__main__.py
+slide2vec/_cufile.py
 slide2vec/api.py
 slide2vec/artifacts.py
 slide2vec/cli.py
+slide2vec/cufile.json
 slide2vec/inference.py
 slide2vec/progress.py
 slide2vec.egg-info/PKG-INFO
@@ -27,6 +29,7 @@ slide2vec/distributed/dense_image_worker.py
 slide2vec/distributed/dense_worker.py
 slide2vec/distributed/direct_embed_worker.py
 slide2vec/distributed/image_worker.py
+slide2vec/distributed/pin_gpu.py
 slide2vec/distributed/pipeline_worker.py
 slide2vec/distributed/worker_entry.py
 slide2vec/encoders/__init__.py
@@ -113,6 +116,7 @@ slide2vec/utils/utils.py
 tests/test_architecture_runtime_split.py
 tests/test_attention_extraction.py
 tests/test_benchmark_tooling.py
+tests/test_cufile_compat_mode.py
 tests/test_dense_encode_kit.py
 tests/test_dense_encoder_input.py
 tests/test_dense_extraction.py
@@ -163,4 +167,5 @@ tests/test_slide_coordinate_preparation.py
 tests/test_soma_migration.py
 tests/test_tile_store.py
 tests/test_tiling_pipeline.py
-tests/test_titan.py
+tests/test_titan.py
+tests/test_worker_gpu_pinning.py

slide2vec-6.0.1/tests/test_cufile_compat_mode.py

@@ -0,0 +1,50 @@
+"""``import cucim`` must not open a CUDA context on every visible GPU (issue #326).
+
+Without ``nvidia-fs``, cuFile's compat mode creates an idle context per device at import.
+slide2vec points cuFile at a config that disallows compat mode, before cucim can load.
+"""
+
+from __future__ import annotations
+
+import ast
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import slide2vec
+from slide2vec import _cufile
+
+
+def test_packaged_config_disallows_compat_mode():
+    config = json.loads(Path(slide2vec.__file__).with_name("cufile.json").read_text())
+    assert config == {"properties": {"allow_compat_mode": False}}
+
+
+def test_cufile_is_pointed_at_the_packaged_config():
+    environ: dict[str, str] = {}
+    _cufile.disable_cufile_compat_mode(environ)
+    assert Path(environ[_cufile.CUFILE_ENV]) == Path(slide2vec.__file__).with_name("cufile.json")
+
+
+def test_a_config_the_user_chose_is_left_alone():
+    environ = {_cufile.CUFILE_ENV: "/etc/my-cufile.json"}
+    _cufile.disable_cufile_compat_mode(environ)
+    assert environ[_cufile.CUFILE_ENV] == "/etc/my-cufile.json"
+
+
+def test_package_import_sets_it_before_any_other_import():
+    # cuFile reads the variable once, when cucim is imported; nothing may get there first.
+    first, second = ast.parse(Path(slide2vec.__file__).read_text()).body[:2]
+    assert isinstance(first, ast.ImportFrom) and first.module == "slide2vec._cufile"
+    assert isinstance(second, ast.Expr) and second.value.func.id == "disable_cufile_compat_mode"
+
+
+def test_importing_slide2vec_exports_the_variable():
+    env = {k: v for k, v in os.environ.items() if k != _cufile.CUFILE_ENV}
+    out = subprocess.run(
+        [sys.executable, "-c", f"import os, slide2vec; print(os.environ['{_cufile.CUFILE_ENV}'])"],
+        env=env, capture_output=True, text=True, check=True, cwd=Path(slide2vec.__file__).parents[1],
+    ).stdout.strip().splitlines()[-1]
+    assert Path(out).name == "cufile.json" and Path(out).is_file()

slide2vec-6.0.1/tests/test_worker_gpu_pinning.py

@@ -0,0 +1,234 @@
+"""Each sharded torchrun worker must see only its own GPU (issue #326).
+
+``import cucim`` opens a CUDA context on every visible device, so a rank that can see its
+siblings' GPUs holds an idle context on each of them.
+"""
+
+from __future__ import annotations
+
+import ast
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+import pytest
+
+from slide2vec.distributed import pin_gpu, worker_entry
+from slide2vec.distributed.worker_entry import resolve_worker_rank
+
+
+def test_rank_sees_only_its_own_gpu():
+    environ = {"LOCAL_RANK": "2"}
+    pin_gpu.pin_to_own_gpu(environ)
+    assert environ["CUDA_VISIBLE_DEVICES"] == "2"
+    assert environ[pin_gpu.PINNED_ENV] == "1"
+
+
+def test_rank_picks_its_entry_of_an_inherited_gpu_list():
+    environ = {"LOCAL_RANK": "1", "CUDA_VISIBLE_DEVICES": "GPU-aaa, GPU-bbb,GPU-ccc"}
+    pin_gpu.pin_to_own_gpu(environ)
+    assert environ["CUDA_VISIBLE_DEVICES"] == "GPU-bbb"
+
+
+def test_more_ranks_than_visible_gpus_is_an_error():
+    with pytest.raises(RuntimeError, match="LOCAL_RANK=2"):
+        pin_gpu.pin_to_own_gpu({"LOCAL_RANK": "2", "CUDA_VISIBLE_DEVICES": "0,1"})
+
+
+def test_bootstrap_imports_nothing_that_could_initialise_cuda():
+    # It runs by path, before the package import that freezes CUDA's visible-device list.
+    tree = ast.parse(Path(pin_gpu.__file__).read_text())
+    imported = {
+        alias.name.split(".")[0]
+        for node in ast.walk(tree) if isinstance(node, ast.Import)
+        for alias in node.names
+    } | {
+        (node.module or "").split(".")[0]
+        for node in ast.walk(tree) if isinstance(node, ast.ImportFrom)
+    }
+    assert imported == {"os", "runpy", "sys"}
+
+
+def test_pinned_worker_uses_cuda_0_but_keeps_its_rank_label(monkeypatch):
+    assert worker_entry.PINNED_ENV == pin_gpu.PINNED_ENV
+    monkeypatch.setenv("RANK", "3")
+    monkeypatch.setenv("WORLD_SIZE", "4")
+    monkeypatch.setenv("LOCAL_RANK", "3")
+    monkeypatch.setenv(worker_entry.PINNED_ENV, "1")
+    rank = resolve_worker_rank()
+    assert rank.device == "cuda:0"
+    assert (rank.global_rank, rank.local_rank) == (3, 3)
+
+
+def test_unpinned_worker_uses_its_local_rank_device(monkeypatch):
+    monkeypatch.setenv("LOCAL_RANK", "3")
+    monkeypatch.delenv(worker_entry.PINNED_ENV, raising=False)
+    assert resolve_worker_rank().device == "cuda:3"
+
+
+@pytest.mark.parametrize("pin_gpus", [True, False])
+def test_launcher_runs_the_bootstrap_by_path_only_when_pinning(tmp_path, pin_gpus):
+    from slide2vec.runtime.distributed import run_torchrun_worker
+
+    commands = []
+
+    class _Done:
+        pid = None
+        stdout = stderr = None
+
+        def __init__(self, command, **kwargs):
+            commands.append(command)
+
+        def poll(self):
+            return 0
+
+        def wait(self, timeout=None):
+            return 0
+
+    run_torchrun_worker(
+        module="slide2vec.distributed.dense_worker", num_gpus=4, output_dir=tmp_path,
+        request_path=tmp_path / "request.json", failure_title="failed",
+        pin_gpus=pin_gpus, popen_factory=_Done,
+    )
+    (command,) = commands
+    tail = command[command.index("--nproc_per_node=4") + 1:]
+    if pin_gpus:
+        assert Path(tail[0]) == Path(pin_gpu.__file__).resolve()
+        assert tail[1] == "slide2vec.distributed.dense_worker"
+    else:
+        assert tail[:2] == ["-m", "slide2vec.distributed.dense_worker"]
+
+
+@pytest.mark.parametrize("pinned, ordinal", [(True, 0), (False, 3)])
+def test_enable_binds_a_pinned_rank_to_cuda_0(monkeypatch, pinned, ordinal):
+    # pipeline_worker and direct_embed_worker go through distributed.enable(); once pinned,
+    # cuda:<local_rank> no longer exists for any rank but 0.
+    import torch
+
+    import slide2vec.distributed as distributed
+
+    for name in ("_RANK", "_WORLD_SIZE", "_LOCAL_RANK", "_LOCAL_WORLD_SIZE"):
+        monkeypatch.setattr(distributed, name, -1)
+    for key in ("RANK", "WORLD_SIZE", "LOCAL_RANK", "LOCAL_WORLD_SIZE"):
+        monkeypatch.setenv(key, "4" if "SIZE" in key else "3")
+    if pinned:
+        monkeypatch.setenv(worker_entry.PINNED_ENV, "1")
+    else:
+        monkeypatch.delenv(worker_entry.PINNED_ENV, raising=False)
+    monkeypatch.setattr(distributed, "_restrict_print_to_main_process", lambda: None)
+    bound = []
+    monkeypatch.setattr(torch.cuda, "set_device", bound.append)
+
+    distributed.enable(overwrite=True)
+
+    assert bound == [ordinal]
+    assert distributed.get_device_ordinal() == ordinal
+    assert distributed.get_local_rank() == 3
+
+
+@pytest.mark.parametrize("worker", ["pipeline_worker", "direct_embed_worker"])
+def test_enable_based_workers_build_the_model_on_the_device_ordinal(worker):
+    source = (Path(pin_gpu.__file__).parent / f"{worker}.py").read_text()
+    assert 'device=f"cuda:{distributed.get_device_ordinal()}"' in source
+    assert 'progress_label=f"cuda:{local_rank}"' in source
+
+
+def test_every_torchrun_stage_pins_its_workers():
+    runtime = Path(pin_gpu.__file__).resolve().parents[1] / "runtime"
+    launches = [
+        node
+        for path in runtime.glob("*.py")
+        for node in ast.walk(ast.parse(path.read_text()))
+        if isinstance(node, ast.Call) and getattr(node.func, "id", None) == "run_torchrun_worker"
+    ]
+    assert len(launches) == 5
+    for call in launches:
+        (pin,) = [kw.value for kw in call.keywords if kw.arg == "pin_gpus"]
+        assert pin.value is True
+
+
+def test_workers_keep_the_parents_working_directory(tmp_path, monkeypatch):
+    # A relative output_dir must mean the same place in the parent and in every rank, so the
+    # checkout is made importable through PYTHONPATH rather than by moving the workers' cwd.
+    import os
+
+    import slide2vec
+    from slide2vec.runtime.distributed import run_torchrun_worker
+
+    monkeypatch.setenv("PYTHONPATH", "/somewhere/else")
+    launches = []
+
+    class _Done:
+        pid = None
+        stdout = stderr = None
+
+        def __init__(self, command, **kwargs):
+            launches.append((command, kwargs))
+
+        def poll(self):
+            return 0
+
+        def wait(self, timeout=None):
+            return 0
+
+    monkeypatch.chdir(tmp_path)
+    run_torchrun_worker(
+        module="slide2vec.distributed.dense_worker", num_gpus=2, output_dir=Path("out"),
+        request_path=Path("out") / "request.json", failure_title="failed", popen_factory=_Done,
+    )
+    ((command, kwargs),) = launches
+    assert kwargs.get("cwd") is None
+    assert command[command.index("--output-dir") + 1] == "out"
+    package_root = str(Path(slide2vec.__file__).resolve().parents[1])
+    assert kwargs["env"]["PYTHONPATH"].split(os.pathsep) == [package_root, "/somewhere/else"]
+
+
+def test_bootstrap_prefers_its_checkout_and_keeps_relative_paths(tmp_path):
+    # Minimal competing packages let a fresh interpreter exercise import precedence
+    # without loading CUDA or model dependencies.
+    package = tmp_path / "parent-checkout" / "slide2vec"
+    bootstrap = package / "distributed" / "pin_gpu.py"
+    bootstrap.parent.mkdir(parents=True)
+    bootstrap.write_text(Path(pin_gpu.__file__).read_text())
+    (package / "__init__.py").write_text("")
+
+    work_dir = tmp_path / "other-checkout"
+    shadow_package = work_dir / "slide2vec"
+    shadow_package.mkdir(parents=True)
+    (shadow_package / "__init__.py").write_text(
+        "raise AssertionError('worker imported the competing checkout')\n"
+    )
+    (work_dir / "input.txt").write_text("relative input")
+    (work_dir / "out").mkdir()
+    (work_dir / "out" / "request.json").write_text(
+        json.dumps({"input": "input.txt", "output": "out/report.json"})
+    )
+    # This module is only available in the working directory, which must remain
+    # importable after the parent's package root.
+    (work_dir / "probe_worker.py").write_text(
+        "import json, os, sys\n"
+        "from pathlib import Path\n"
+        "import slide2vec\n"
+        "request = json.loads(Path(sys.argv[1]).read_text())\n"
+        "Path(request['output']).write_text(json.dumps({\n"
+        "    'package': slide2vec.__file__,\n"
+        "    'cwd': os.getcwd(),\n"
+        "    'input': Path(request['input']).read_text(),\n"
+        "    'visible': os.environ['CUDA_VISIBLE_DEVICES'],\n"
+        "}))\n"
+    )
+    env = dict(os.environ)
+    env.update(LOCAL_RANK="1", CUDA_VISIBLE_DEVICES="2,3", PYTHONPATH=str(package.parent))
+    result = subprocess.run(
+        [sys.executable, str(bootstrap), "probe_worker", "out/request.json"],
+        cwd=work_dir, env=env, capture_output=True, text=True, timeout=15,
+    )
+    assert result.returncode == 0, result.stderr
+    assert json.loads((work_dir / "out" / "report.json").read_text()) == {
+        "package": str(package / "__init__.py"),
+        "cwd": str(work_dir),
+        "input": "relative input",
+        "visible": "3",
+    }