Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,5 @@ cython_debug/
output/
outputs/
archive/
tasks/
docs/20*-*.md
data/
29 changes: 14 additions & 15 deletions slide2vec/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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),
Expand All @@ -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:
Expand Down
5 changes: 2 additions & 3 deletions slide2vec/configs/default.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down
20 changes: 12 additions & 8 deletions slide2vec/distributed/direct_embed_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down
20 changes: 10 additions & 10 deletions slide2vec/distributed/pipeline_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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():
Expand Down
22 changes: 13 additions & 9 deletions slide2vec/encoders/validation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
)
Loading
Loading