fix floating point precision and port (#123)

Author: clemsgrs

.gitignore

@@ -165,6 +165,5 @@ cython_debug/
 output/
 outputs/
 archive/
-tasks/
 docs/20*-*.md
 data/

slide2vec/api.py

@@ -159,8 +159,9 @@ class ExecutionOptions:
     output_format: str = "pt"
     #: Number of tiles per forward pass.
     batch_size: int = 32
-    #: DataLoader worker count. ``None`` means auto (capped by CPU / SLURM limit).
-    num_workers: int | None = None
+    #: DataLoader worker count per GPU rank. ``None`` means auto
+    #: (capped by CPU / SLURM limit, then split across the resolved GPU count).
+    num_workers_per_gpu: int | None = None
     #: Tiling worker count. ``None`` means auto (capped by CPU / SLURM limit).
     num_preprocessing_workers: int | None = None
     #: Number of GPUs to use. ``None`` defaults to all available GPUs.
@@ -170,8 +171,6 @@ class ExecutionOptions:
     precision: str | None = None
     #: DataLoader prefetch queue depth per worker (default ``4``).
     prefetch_factor: int = 4
-    #: Keep DataLoader workers alive between batches (default ``True``).
-    persistent_workers: bool = True
     #: Persist tile embeddings to disk when running a slide-level model.
     save_tile_embeddings: bool = False
     #: Persist slide embeddings to disk when running a patient-level model.
@@ -183,14 +182,13 @@ class ExecutionOptions:
     def from_config(cls, cfg: Any, *, run_on_cpu: bool = False) -> "ExecutionOptions":
         configured_num_gpus = cfg.speed.num_gpus
         requested_precision = normalize_precision_name(cfg.speed.precision)
-        num_workers = cfg.speed.num_dataloader_workers
+        num_workers_per_gpu = cfg.speed.num_dataloader_workers
         prefetch_factor = int(cfg.speed.prefetch_factor_embedding)
-        persistent_workers = bool(cfg.speed.persistent_workers_embedding)
         return cls(
             output_dir=Path(cfg.output_dir),
             output_format="pt",
             batch_size=int(cfg.model.batch_size),
-            num_workers=int(num_workers) if num_workers is not None else None,
+            num_workers_per_gpu=int(num_workers_per_gpu) if num_workers_per_gpu is not None else None,
             num_preprocessing_workers=(
                 int(cfg.speed.num_preprocessing_workers)
                 if cfg.speed.num_preprocessing_workers is not None
@@ -199,7 +197,6 @@ def from_config(cls, cfg: Any, *, run_on_cpu: bool = False) -> "ExecutionOptions
             num_gpus=1 if run_on_cpu else (int(configured_num_gpus) if configured_num_gpus is not None else None),
             precision="fp32" if run_on_cpu else requested_precision,
             prefetch_factor=prefetch_factor,
-            persistent_workers=persistent_workers,
             save_tile_embeddings=bool(cfg.model.save_tile_embeddings),
             save_slide_embeddings=bool(cfg.model.save_slide_embeddings),
             save_latents=bool(cfg.model.save_latents),
@@ -222,23 +219,25 @@ def __post_init__(self) -> None:
         object.__setattr__(self, "num_preprocessing_workers", capped_num_preprocessing_workers)
         logger = logging.getLogger(__name__)
         cap_source = f"slurm_cpu_limit={slurm_limit}" if slurm_limit is not None else f"cpu_count={cpu_count}"
-        resolved_num_workers = self.resolved_num_workers()
-        num_workers_label = (
+        resolved_num_workers = self.resolved_num_workers_per_gpu()
+        num_workers_per_gpu_label = (
             f"{resolved_num_workers} (requested=auto)"
-            if self.num_workers is None
+            if self.num_workers_per_gpu is None
             else str(resolved_num_workers)
         )
         logger.info(
-            "ExecutionOptions: num_workers=%s, num_preprocessing_workers=%d "
+            "ExecutionOptions: num_workers_per_gpu=%s, num_preprocessing_workers=%d "
             "(preprocessing cap=%d via %s)",
-            num_workers_label,
+            num_workers_per_gpu_label,
             capped_num_preprocessing_workers,
             cap,
             cap_source,
         )
 
-    def resolved_num_workers(self) -> int:
-        return cpu_worker_limit() if self.num_workers is None else int(self.num_workers)
+    def resolved_num_workers_per_gpu(self) -> int:
+        if self.num_workers_per_gpu is not None:
+            return self.num_workers_per_gpu
+        return max(1, cpu_worker_limit() // self.num_gpus)
 
     def with_output_dir(self, output_dir: PathLike | None) -> "ExecutionOptions":
         if output_dir is None:

slide2vec/configs/default.yaml

@@ -42,7 +42,7 @@ tiling:
     sthresh_up: 255 # upper threshold value for scaling the binary mask
     mthresh: 7 # median filter size (positive, odd integer)
     close: 4 # additional morphological closing to apply following initial thresholding (positive integer)
-    method: "hsv" # tissue segmentation method: "hsv", "otsu", "threshold", or "sam2"
+    method: # tissue segmentation method: "hsv", "otsu", "threshold", or "sam2"; ignored when precomputed tissue masks are provided
     sam2_checkpoint_path: # optional when method="sam2"; if empty, hs2p downloads the default AtlasPatch checkpoint from Hugging Face
     sam2_config_path: # optional local override for the SAM2 model config; if empty, hs2p downloads the default AtlasPatch config from Hugging Face
     sam2_device: "cpu" # device for SAM2 inference, e.g. "cpu", "cuda", or "cuda:0"
@@ -71,12 +71,11 @@ tiling:
 
 speed:
   precision: # model inference precision ["fp32", "fp16", "bf16"]; if not set, determined automatically based on model recommendations
-  num_dataloader_workers: # number of DataLoader worker processes for reading tiles during embedding; defaults to auto (job CPU budget, except cuCIM on-the-fly uses cpu_budget // speed.num_cucim_workers)
+  num_dataloader_workers: # number of DataLoader worker processes per GPU rank for reading tiles during embedding; defaults to auto (job CPU budget split across GPUs, except cuCIM on-the-fly uses per-GPU budget // speed.num_cucim_workers)
   num_gpus: # number of GPUs to use for feature extraction; defaults to all available GPUs
   num_preprocessing_workers: # number of workers for hs2p tiling (WSI reading, JPEG encoding, tar writing); defaults to the runtime CPU budget capped at 64
   num_cucim_workers: 4 # number of internal cucim threads per read_region call (embedding path, on-the-fly only); DataLoader workers are auto-set to cpu_count // num_cucim_workers
   prefetch_factor_embedding: 4 # prefetch factor for tile embedding dataloaders
-  persistent_workers_embedding: true # keep DataLoader workers alive across epochs/batches
 
 wandb:
   enable: false

slide2vec/distributed/direct_embed_worker.py

@@ -119,20 +119,24 @@ def main(argv=None) -> int:
                 return 0
             assigned_slides = [paired_by_sample[sample_id][0] for sample_id in assigned_ids]
             assigned_tiling_results = [paired_by_sample[sample_id][1] for sample_id in assigned_ids]
-            embedded_slides = _compute_embedded_slides(
-                model,
-                assigned_slides,
-                assigned_tiling_results,
-                preprocessing=preprocessing,
-                execution=execution,
-            )
-            for embedded_slide in embedded_slides:
+
+            def _persist_embedded_slide(slide, tiling_result, embedded_slide) -> None:
                 payload = {
                     "tile_embeddings": _to_cpu_payload(embedded_slide.tile_embeddings),
                     "slide_embedding": _to_cpu_payload(embedded_slide.slide_embedding),
                     "latents": _to_cpu_payload(embedded_slide.latents),
                 }
                 torch.save(payload, coordination_dir / f"{embedded_slide.sample_id}.embedded.pt")
+
+            _compute_embedded_slides(
+                model,
+                assigned_slides,
+                assigned_tiling_results,
+                preprocessing=preprocessing,
+                execution=execution,
+                on_embedded_slide=_persist_embedded_slide,
+                collect_results=False,
+            )
             return 0
     finally:
         if dist.is_available() and dist.is_initialized():

slide2vec/distributed/pipeline_worker.py

@@ -19,8 +19,8 @@ def main(argv=None) -> int:
     import slide2vec.distributed as distributed
     from slide2vec.api import Model
     from slide2vec.inference import (
+        _build_incremental_persist_callback,
         _compute_embedded_slides,
-        _persist_embedded_slide,
         load_successful_tiled_slides,
     )
     from slide2vec.progress import JsonlProgressReporter, activate_progress_reporter
@@ -70,21 +70,21 @@ def main(argv=None) -> int:
         )
         context = activate_progress_reporter(reporter) if reporter is not None else nullcontext()
         with context:
-            embedded_slides = _compute_embedded_slides(
+            persist_callback, _, _ = _build_incremental_persist_callback(
+                model=model,
+                preprocessing=preprocessing,
+                execution=execution,
+                process_list_path=None,
+            )
+            _compute_embedded_slides(
                 model,
                 assigned_slides,
                 assigned_tiling_results,
                 preprocessing=preprocessing,
                 execution=execution,
+                on_embedded_slide=persist_callback,
+                collect_results=False,
             )
-            for embedded_slide, tiling_result in zip(embedded_slides, assigned_tiling_results):
-                _persist_embedded_slide(
-                    model,
-                    embedded_slide,
-                    tiling_result,
-                    preprocessing=preprocessing,
-                    execution=execution,
-                )
         return 0
     finally:
         if dist.is_available() and dist.is_initialized():

slide2vec/encoders/validation.py

@@ -63,14 +63,18 @@ def validate_encoder_config(
     if not mismatches:
         return
 
-    message = (
-        f"Model '{encoder_name}' is configured with "
-        f"{'; '.join(mismatches)}. "
-        "Set `model.allow_non_recommended_settings=true` in YAML/CLI or "
-        "`allow_non_recommended_settings=True` in `Model.from_preset(...)` "
-        "to continue with a warning."
-    )
     if allow_non_recommended:
-        logger.warning(message)
+        logger.warning(
+            f"Model '{encoder_name}' is configured with "
+            f"{'; '.join(mismatches)}. "
+            "Warning-only mode is enabled because "
+            "`allow_non_recommended_settings=True`."
+        )
     else:
-        raise ValueError(message)
+        raise ValueError(
+            f"Model '{encoder_name}' is configured with "
+            f"{'; '.join(mismatches)}. "
+            "Set `model.allow_non_recommended_settings=true` in YAML/CLI or "
+            "`allow_non_recommended_settings=True` in `Model.from_preset(...)` "
+            "to continue."
+        )