slide2vec 5.9.2__tar.gz → 6.0.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 5.9.2
+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
@@ -80,7 +80,7 @@ Requires-Dist: wandb; extra == "fm"
 Requires-Dist: torch<2.8,>=2.3; extra == "fm"
 Requires-Dist: torchvision>=0.18.0; extra == "fm"
 Requires-Dist: einops>=0.8.0; extra == "fm"
-Requires-Dist: timm>=1.0.3; extra == "fm"
+Requires-Dist: timm>=1.0.20; extra == "fm"
 Requires-Dist: huggingface_hub<1.0,>=0.30.0; extra == "fm"
 Requires-Dist: environs; extra == "fm"
 Requires-Dist: einops-exts>=0.0.4; extra == "fm"
@@ -119,162 +119,88 @@ Dynamic: license-file
 [![PyPI version](https://img.shields.io/pypi/v/slide2vec?label=pypi&logo=pypi&color=3776AB)](https://pypi.org/project/slide2vec/)
 [![Docs](https://img.shields.io/badge/docs-website-blue)](https://clemsgrs.github.io/slide2vec/)
 
-`slide2vec` is a Python package for efficient encoding of whole-slide images using publicly available foundation models. It builds on [`hs2p`](https://pypi.org/project/hs2p/) for fast preprocessing and exposes a focused surface around `Model`, `Pipeline`, and `ExecutionOptions`.
+`slide2vec` encodes whole-slide images with publicly available pathology foundation models. It uses [`hs2p`](https://pypi.org/project/hs2p/) for tissue detection and tiling, and handles batching, multi-GPU execution, and embedding storage.
 
-Documentation site: [https://clemsgrs.github.io/slide2vec/](https://clemsgrs.github.io/slide2vec/)
+## Install
 
-## Installation
+Python 3.10 or newer is required:
 
 ```shell
 pip install slide2vec
-pip install "slide2vec[fm]"
 ```
 
-`slide2vec` keeps the base install focused on the core package surface. Use `slide2vec[fm]` when you want the PyPI-hosted FM dependencies.
+Many models need additional dependencies available through `pip install "slide2vec[fm]"`. See the [model installation guide](https://clemsgrs.github.io/slide2vec/models.html#model-installation) for model-specific extras, separate environments, and upstream packages.
 
-Some model backends still rely on upstream Git repositories that PyPI will not accept as package metadata. Install those separately when needed:
+For gated models such as Virchow2, request access on the model's Hugging Face page and authenticate with `hf auth login` or an `HF_TOKEN` environment variable.
 
-```shell
-pip install git+https://github.com/lilab-stanford/MUSK.git
-pip install git+https://github.com/Mahmoodlab/CONCH.git
-pip install git+https://github.com/prov-gigapath/prov-gigapath.git
-```
-
-AtlasPatch-backed tissue segmentation is available through hs2p's `sam2` path in the bundled install.
-
-Waiv encoders use a separately tested Transformers 5 runtime. Install them in
-their own environment with `pip install "slide2vec[waiv]"`; the `waiv` extra is
-incompatible with the existing `fm`, `prism`, and `titan` dependency pins.
-
-## Python API
+## Embed a slide
 
 ```python
-from slide2vec import Model
-from slide2vec.utils.config import hf_login
-
-hf_login()
+from slide2vec import Model, PreprocessingConfig
 
 model = Model.from_preset("virchow2")
-embedded = model.embed_slide("/path/to/slide.svs")
+preprocessing = PreprocessingConfig(requested_spacing_um=0.5)
+embedded = model.embed_slide("/path/to/slide.svs", preprocessing=preprocessing)
 
-tile_embeddings = embedded.tile_embeddings
-x = embedded.x
-y = embedded.y
+tile_embeddings = embedded.tile_embeddings  # (N, 2560)
+x, y = embedded.x, embedded.y               # level-0 tile coordinates
 ```
 
-Use `list_models()` when you want to inspect the shipped presets programmatically:
+The preset supplies tile size and precision defaults. Declare spacing explicitly for models such as Virchow2 that support several scales. Use `list_models()` to list presets, or filter with `list_models("tile")`, `list_models("slide")`, or `list_models("patient")`.
 
-```python
-from slide2vec import list_models
+See [getting started](https://clemsgrs.github.io/slide2vec/getting-started.html) for preprocessing and execution settings, and the [API guide](https://clemsgrs.github.io/slide2vec/api.html) for patient embeddings, image inputs, and dense grids.
+
+## Save a batch
 
-all_models = list_models()
-tile_models = list_models("tile")
-slide_models = list_models("slide")
-patient_models = list_models("patient")
+Create a CSV manifest:
+
+```csv
+sample_id,image_path
+slide-1,/data/slide-1.svs
+slide-2,/data/slide-2.svs
 ```
 
-Use `Pipeline(...)` for manifest-driven batch processing when you want artifacts written to disk instead of only in-memory outputs:
+Optional `mask_path` and `spacing_at_level_0` columns supply a mask or correct missing or incorrect level-0 spacing. Patient-level models also require `patient_id`; see the [manifest schema](https://clemsgrs.github.io/slide2vec/manifest.html).
 
 ```python
-from slide2vec import ExecutionOptions, Pipeline, PreprocessingConfig
+from slide2vec import ExecutionOptions, Model, Pipeline, PreprocessingConfig
 
 pipeline = Pipeline(
-    model=model,
-    preprocessing=PreprocessingConfig(
-        requested_spacing_um=0.5,
-        requested_tile_size_px=224,
-        masks={"min_coverage": {"tissue": 0.1}},
-    ),
-    execution=ExecutionOptions(output_dir="outputs/demo"),
+    model=Model.from_preset("virchow2"),
+    preprocessing=PreprocessingConfig(requested_spacing_um=0.5),
+    execution=ExecutionOptions(output_dir="outputs/run"),
 )
 result = pipeline.run(manifest_path="/path/to/slides.csv")
 ```
 
-By default, `ExecutionOptions()` uses all available GPUs. Set `ExecutionOptions(num_gpus=4)` when you want to cap the sharding explicitly.
-
-### Hierarchical Feature Extraction
-
-Tile embeddings can be spatially grouped into regions for downstream models that consume region-level structure. Enable it by setting `region_tile_multiple` on `PreprocessingConfig`:
-
-```python
-preprocessing = PreprocessingConfig(
-    requested_spacing_um=0.5,
-    requested_tile_size_px=224,
-    region_tile_multiple=6,  # 6x6 tiles per region
-)
-embedded = model.embed_slide("/path/to/slide.svs", preprocessing=preprocessing)
-```
-
-Hierarchical outputs have shape `(num_regions, tiles_per_region, feature_dim)` and are written to `hierarchical_embeddings/` when persisted.
-
-See the [hierarchical features guide](https://clemsgrs.github.io/slide2vec/hierarchical.html) for details.
-
-### Input Manifest
-
-Manifest-driven runs use the schema below. `mask_path` and `spacing_at_level_0` are optional.
-
-```csv
-sample_id,image_path,mask_path,spacing_at_level_0
-slide-1,/path/to/slide-1.svs,/path/to/mask-1.png,0.25
-slide-2,/path/to/slide-2.svs,,
-...
-```
-
-Use `spacing_at_level_0` when the slide file reports a missing or incorrect level-0 spacing and you want to override it.
-
+Runs use all available GPUs by default; set `ExecutionOptions(num_gpus=2)` to limit them. Embeddings are saved as `.pt` tensors with metadata sidecars. Use `ExecutionOptions(output_format="npz")` for NumPy archives. The [output guide](https://clemsgrs.github.io/slide2vec/output-layout.html) describes directories, shapes, coordinates, and progress records.
 
-### Outputs
+Add `region_tile_multiple=6` to the preprocessing config to group tiles into 6×6 regions. These produce `(num_regions, 36, feature_dim)` tensors in `hierarchical_embeddings/`; see [hierarchical features](https://clemsgrs.github.io/slide2vec/hierarchical.html).
 
-The package writes explicit artifact directories:
-
-- `tile_embeddings/<sample_id>.pt` or `.npz`
-- `tile_embeddings/<sample_id>.meta.json`
-- `hierarchical_embeddings/<sample_id>.pt` or `.npz` (when `region_tile_multiple` is set)
-- `hierarchical_embeddings/<sample_id>.meta.json`
-- `slide_embeddings/<sample_id>.pt` or `.npz`
-- `slide_embeddings/<sample_id>.meta.json`
-- optional `slide_latents/<sample_id>.pt` or `.npz`
-
-`.pt` remains the default format. `.npz` is available through `ExecutionOptions(output_format="npz")`.
-
-### Supported Models
-
-`slide2vec` currently ships presets for 28 tile-level models, 4 slide-level models,
-and 1 patient-level model.
-For the full catalog and preset names, see the [model zoo](https://clemsgrs.github.io/slide2vec/models.html).
-
-## CLI
-
-The CLI is a thin wrapper over the package API.  
-Bundled configs live under `slide2vec/configs/preprocessing/` and `slide2vec/configs/models/`.
+The same batch workflow is available from the terminal:
 
 ```shell
 slide2vec /path/to/config.yaml
 ```
 
-By default, manifest-driven CLI runs use all available GPUs. Set `speed.num_gpus=4` when you want to cap the sharding explicitly.
-
-New to the CLI or doing batch runs to disk? Start with the [CLI guide](https://clemsgrs.github.io/slide2vec/cli.html) for the config-driven workflow and common run patterns.
+The [CLI guide](https://clemsgrs.github.io/slide2vec/cli.html) provides a complete config example, overrides, and resume instructions.
 
 ## Docker
 
 [![Docker Version](https://img.shields.io/docker/v/waticlems/slide2vec?sort=semver&label=docker&logo=docker&color=2496ED)](https://hub.docker.com/r/waticlems/slide2vec)
 
-Docker remains available when you prefer a containerized runtime:
-
 ```shell
 docker pull waticlems/slide2vec:latest
 docker run --rm -it \
     -v /path/to/your/data:/data \
-    -e HF_TOKEN=<your-huggingface-api-token> \
+    -e HF_TOKEN \
     waticlems/slide2vec:latest
 ```
 
-## Documentation
+Set `HF_TOKEN` in your shell before starting the container.
+
+## More documentation
 
-- [Documentation website](https://clemsgrs.github.io/slide2vec/)
-- [API guide](https://clemsgrs.github.io/slide2vec/api.html)
-- [CLI guide](https://clemsgrs.github.io/slide2vec/cli.html)
 - [Model zoo](https://clemsgrs.github.io/slide2vec/models.html)
-- [Performance benchmarks and QA](docs/performance.md)
-- [`tutorials/api_walkthrough.ipynb`](tutorials/api_walkthrough.ipynb) for a notebook walkthrough of the API
+- [Performance benchmarks and QA](https://clemsgrs.github.io/slide2vec/performance.html)
+- [API walkthrough notebook](tutorials/api_walkthrough.ipynb)

slide2vec-6.0.1/README.md

@@ -0,0 +1,90 @@
+# slide2vec
+
+[![PyPI version](https://img.shields.io/pypi/v/slide2vec?label=pypi&logo=pypi&color=3776AB)](https://pypi.org/project/slide2vec/)
+[![Docs](https://img.shields.io/badge/docs-website-blue)](https://clemsgrs.github.io/slide2vec/)
+
+`slide2vec` encodes whole-slide images with publicly available pathology foundation models. It uses [`hs2p`](https://pypi.org/project/hs2p/) for tissue detection and tiling, and handles batching, multi-GPU execution, and embedding storage.
+
+## Install
+
+Python 3.10 or newer is required:
+
+```shell
+pip install slide2vec
+```
+
+Many models need additional dependencies available through `pip install "slide2vec[fm]"`. See the [model installation guide](https://clemsgrs.github.io/slide2vec/models.html#model-installation) for model-specific extras, separate environments, and upstream packages.
+
+For gated models such as Virchow2, request access on the model's Hugging Face page and authenticate with `hf auth login` or an `HF_TOKEN` environment variable.
+
+## Embed a slide
+
+```python
+from slide2vec import Model, PreprocessingConfig
+
+model = Model.from_preset("virchow2")
+preprocessing = PreprocessingConfig(requested_spacing_um=0.5)
+embedded = model.embed_slide("/path/to/slide.svs", preprocessing=preprocessing)
+
+tile_embeddings = embedded.tile_embeddings  # (N, 2560)
+x, y = embedded.x, embedded.y               # level-0 tile coordinates
+```
+
+The preset supplies tile size and precision defaults. Declare spacing explicitly for models such as Virchow2 that support several scales. Use `list_models()` to list presets, or filter with `list_models("tile")`, `list_models("slide")`, or `list_models("patient")`.
+
+See [getting started](https://clemsgrs.github.io/slide2vec/getting-started.html) for preprocessing and execution settings, and the [API guide](https://clemsgrs.github.io/slide2vec/api.html) for patient embeddings, image inputs, and dense grids.
+
+## Save a batch
+
+Create a CSV manifest:
+
+```csv
+sample_id,image_path
+slide-1,/data/slide-1.svs
+slide-2,/data/slide-2.svs
+```
+
+Optional `mask_path` and `spacing_at_level_0` columns supply a mask or correct missing or incorrect level-0 spacing. Patient-level models also require `patient_id`; see the [manifest schema](https://clemsgrs.github.io/slide2vec/manifest.html).
+
+```python
+from slide2vec import ExecutionOptions, Model, Pipeline, PreprocessingConfig
+
+pipeline = Pipeline(
+    model=Model.from_preset("virchow2"),
+    preprocessing=PreprocessingConfig(requested_spacing_um=0.5),
+    execution=ExecutionOptions(output_dir="outputs/run"),
+)
+result = pipeline.run(manifest_path="/path/to/slides.csv")
+```
+
+Runs use all available GPUs by default; set `ExecutionOptions(num_gpus=2)` to limit them. Embeddings are saved as `.pt` tensors with metadata sidecars. Use `ExecutionOptions(output_format="npz")` for NumPy archives. The [output guide](https://clemsgrs.github.io/slide2vec/output-layout.html) describes directories, shapes, coordinates, and progress records.
+
+Add `region_tile_multiple=6` to the preprocessing config to group tiles into 6×6 regions. These produce `(num_regions, 36, feature_dim)` tensors in `hierarchical_embeddings/`; see [hierarchical features](https://clemsgrs.github.io/slide2vec/hierarchical.html).
+
+The same batch workflow is available from the terminal:
+
+```shell
+slide2vec /path/to/config.yaml
+```
+
+The [CLI guide](https://clemsgrs.github.io/slide2vec/cli.html) provides a complete config example, overrides, and resume instructions.
+
+## Docker
+
+[![Docker Version](https://img.shields.io/docker/v/waticlems/slide2vec?sort=semver&label=docker&logo=docker&color=2496ED)](https://hub.docker.com/r/waticlems/slide2vec)
+
+```shell
+docker pull waticlems/slide2vec:latest
+docker run --rm -it \
+    -v /path/to/your/data:/data \
+    -e HF_TOKEN \
+    waticlems/slide2vec:latest
+```
+
+Set `HF_TOKEN` in your shell before starting the container.
+
+## More documentation
+
+- [Model zoo](https://clemsgrs.github.io/slide2vec/models.html)
+- [Performance benchmarks and QA](https://clemsgrs.github.io/slide2vec/performance.html)
+- [API walkthrough notebook](tutorials/api_walkthrough.ipynb)

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "slide2vec"
-version = "5.9.2"
+version = "6.0.1"
 description = "Embedding of whole slide images with Foundation Models"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -105,7 +105,7 @@ fm = [
     "torch>=2.3,<2.8",
     "torchvision>=0.18.0",
     "einops>=0.8.0",
-    "timm>=1.0.3",
+    "timm>=1.0.20",
     "huggingface_hub>=0.30.0,<1.0",
     "environs",
     "einops-exts>=0.0.4",
@@ -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 = "5.9.2"
+current_version = "6.0.1"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit = false
 tag = false

{slide2vec-5.9.2 → 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__ = "5.9.2"
+__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-5.9.2 → slide2vec-6.0.1}/slide2vec/api.py

@@ -1010,14 +1010,14 @@ class Model:
         )
         self._encoder_input = contract
         plan = contract.plan
-        if emit_run_info and plan.requires_variable_model_input:
+        if emit_run_info:
             logging.getLogger("slide2vec").info(
-                "Pooled encoder input for '%s': preset %dpx, requested %dpx, "
-                "exact encoder input %dpx; using normalization-only preprocessing.",
+                "Pooled encoder input for '%s': preset %dpx, requested %dpx; encoding "
+                "exactly %dpx with geometry-preserving preprocessing.",
                 self.name,
                 plan.preset_input_size_px,
                 plan.requested_tile_size_px,
-                plan.expected_encoder_input_size_px,
+                plan.requested_tile_size_px,
             )
         return contract
 

slide2vec-6.0.1/slide2vec/cufile.json

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

{slide2vec-5.9.2 → 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-5.9.2 → 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-5.9.2 → 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-5.9.2 → 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}"
         )
     )