diff --git a/slide2vec/distributed/__init__.py b/slide2vec/distributed/__init__.py index 4d0c299..8b55b94 100644 --- a/slide2vec/distributed/__init__.py +++ b/slide2vec/distributed/__init__.py @@ -1,11 +1,9 @@ -import datetime import os -import random -import socket import torch -import torch.distributed as dist +_RANK = -1 +_WORLD_SIZE = -1 _LOCAL_RANK = -1 _LOCAL_WORLD_SIZE = -1 @@ -13,41 +11,31 @@ def is_enabled() -> bool: """ Returns: - True if distributed training is enabled + True if distributed mode has been enabled (the process topology is known). """ - return dist.is_available() and dist.is_initialized() - - -def is_enabled_and_multiple_gpus() -> bool: - """ - Returns: - True if distributed training is enabled and there are multiple GPUs available. - """ - return ( - dist.is_available() and dist.is_initialized() and torch.cuda.device_count() > 1 - ) + return _RANK >= 0 def get_global_size() -> int: """ Returns: - The number of processes in the process group + The number of processes in the job. """ - return dist.get_world_size() if is_enabled() else 1 + return _WORLD_SIZE if is_enabled() else 1 def get_global_rank() -> int: """ Returns: - The rank of the current process within the global process group. + The rank of the current process in the job. """ - return dist.get_rank() if is_enabled() else 0 + return _RANK if is_enabled() else 0 def get_local_rank() -> int: """ Returns: - The rank of the current process within the local (per-machine) process group. + The rank of the current process on its machine. """ if not is_enabled(): return 0 @@ -58,8 +46,7 @@ def get_local_rank() -> int: def get_local_size() -> int: """ Returns: - The size of the per-machine process group, - i.e. the number of processes per machine. + The number of processes on the current machine. """ if not is_enabled(): return 1 @@ -91,18 +78,10 @@ def print(*args, **kwargs): __builtin__.print = print -def _get_available_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - # A "" host address means INADDR_ANY i.e. binding to all interfaces. - # Note this is not compatible with IPv6. - s.bind(("", 0)) - port = s.getsockname()[1] - return port - - -_TORCH_DISTRIBUTED_ENV_VARS = ( - "MASTER_ADDR", - "MASTER_PORT", +# The process-topology variables torchrun exports. MASTER_ADDR / MASTER_PORT are +# deliberately ignored: they only locate a process-group rendezvous, and no process +# group is ever created (issue #219). +_TORCHRUN_ENV_VARS = ( "RANK", "WORLD_SIZE", "LOCAL_RANK", @@ -113,7 +92,7 @@ def _get_available_port() -> int: def _collect_env_vars() -> dict[str, str]: return { env_var: os.environ[env_var] - for env_var in _TORCH_DISTRIBUTED_ENV_VARS + for env_var in _TORCHRUN_ENV_VARS if env_var in os.environ } @@ -128,8 +107,6 @@ def _check_env_variable(key: str, new_value: str): class _TorchDistributedEnvironment: def __init__(self): - self.master_addr = "127.0.0.1" - self.master_port = 0 self.rank = -1 self.world_size = -1 self.local_rank = -1 @@ -139,7 +116,7 @@ def __init__(self): if not env_vars: # Environment is not set pass - elif len(env_vars) == len(_TORCH_DISTRIBUTED_ENV_VARS): + elif len(env_vars) == len(_TORCHRUN_ENV_VARS): # Environment is fully set return self._set_from_preset_env() else: @@ -150,13 +127,13 @@ def __init__(self): if torch.cuda.device_count() > 0: return self._set_from_local() - raise RuntimeError("Can't initialize PyTorch distributed environment") + raise RuntimeError( + "Can't determine the process topology: no torchrun environment and no CUDA device" + ) # Single node job with preset environment (i.e. torchrun) def _set_from_preset_env(self): # logger.info("Initialization from preset environment") - self.master_addr = os.environ["MASTER_ADDR"] - self.master_port = os.environ["MASTER_PORT"] self.rank = int(os.environ["RANK"]) self.world_size = int(os.environ["WORLD_SIZE"]) assert self.rank < self.world_size @@ -167,20 +144,16 @@ def _set_from_preset_env(self): # Single node and GPU job (i.e. local script run) def _set_from_local(self): # logger.info("Initialization from local") - self.master_addr = "127.0.0.1" - self.master_port = _get_available_port() self.rank = 0 self.world_size = 1 self.local_rank = 0 self.local_world_size = 1 def export(self, *, overwrite: bool) -> "_TorchDistributedEnvironment": - # See the "Environment variable initialization" section from - # https://pytorch.org/docs/stable/distributed.html for the complete list of - # environment variables required for the env:// initialization method. + # Export the topology so child processes and later lookups see one consistent + # view. There is no env:// rendezvous to feed (no process group is created), so + # MASTER_ADDR / MASTER_PORT are not exported. env_vars = { - "MASTER_ADDR": self.master_addr, - "MASTER_PORT": str(self.master_port), "RANK": str(self.rank), "WORLD_SIZE": str(self.world_size), "LOCAL_RANK": str(self.local_rank), @@ -198,9 +171,15 @@ def enable( *, set_cuda_current_device: bool = True, overwrite: bool = False, - allow_nccl_timeout: bool = False, ): - """Enable distributed mode + """Enable distributed mode. + + Reads the process topology (RANK / WORLD_SIZE / LOCAL_RANK / LOCAL_WORLD_SIZE) that + torchrun exports — or single-process defaults when launched bare — and records it for + the ``get_*`` helpers. No ``torch.distributed`` process group is created: the workers + run no collectives (each rank writes its own artifacts and nobody gathers), so an NCCL + group would only add init cost and a 14-day-timeout hang surface for no benefit. See + issue #219. Args: set_cuda_current_device: If True, call torch.cuda.set_device() to set the @@ -208,8 +187,8 @@ def enable( overwrite: If True, overwrites already set variables. Else fails. """ - global _LOCAL_RANK, _LOCAL_WORLD_SIZE - if _LOCAL_RANK >= 0 or _LOCAL_WORLD_SIZE >= 0: + global _RANK, _WORLD_SIZE, _LOCAL_RANK, _LOCAL_WORLD_SIZE + if is_enabled(): raise RuntimeError("Distributed mode has already been enabled") torch_env = _TorchDistributedEnvironment() torch_env.export(overwrite=overwrite) @@ -217,20 +196,9 @@ def enable( if set_cuda_current_device: torch.cuda.set_device(torch_env.local_rank) - if allow_nccl_timeout: - # This allows to use torch distributed timeout in a NCCL backend - key, value = "NCCL_ASYNC_ERROR_HANDLING", "1" - if not overwrite: - _check_env_variable(key, value) - os.environ[key] = value - - dist.init_process_group( - backend="nccl", - timeout=datetime.timedelta(days=14), - device_id=torch.device(f"cuda:{torch_env.local_rank}") - ) - # Finalize setup + _RANK = torch_env.rank + _WORLD_SIZE = torch_env.world_size _LOCAL_RANK = torch_env.local_rank _LOCAL_WORLD_SIZE = torch_env.local_world_size _restrict_print_to_main_process() diff --git a/slide2vec/distributed/direct_embed_worker.py b/slide2vec/distributed/direct_embed_worker.py index 67dbe34..08f14ae 100644 --- a/slide2vec/distributed/direct_embed_worker.py +++ b/slide2vec/distributed/direct_embed_worker.py @@ -15,7 +15,6 @@ def get_args_parser(add_help: bool = True) -> argparse.ArgumentParser: def main(argv=None) -> int: import torch - import torch.distributed as dist import slide2vec.distributed as distributed import slide2vec.inference as inference @@ -34,149 +33,146 @@ def main(argv=None) -> int: output_dir = Path(args.output_dir) coordination_dir = Path(request["coordination_dir"]) + # Reads the torchrun-exported rank/world-size; no NCCL process group (issue #219). distributed.enable(overwrite=True) - try: - global_rank = distributed.get_global_rank() - world_size = distributed.get_global_size() - local_rank = distributed.get_local_rank() - - model_spec = dict(request["model"]) - model = Model.from_preset( - model_spec["name"], - device=f"cuda:{local_rank}", - output_variant=model_spec.get("output_variant"), - allow_non_recommended_settings=bool(model_spec["allow_non_recommended_settings"]), - ) - preprocessing = deserialize_preprocessing(request["preprocessing"]) - execution = deserialize_execution(request["execution"]) - from slide2vec.runtime.distributed import ( - decode_work_unit, - encode_work_unit, - work_unit_shard_stem, - ) - from slide2vec.runtime.embedding import tiling_result_annotation - - load_successful_tiled_slides_fn = getattr(inference, "load_successful_tiled_slides", None) - if not callable(load_successful_tiled_slides_fn): - from slide2vec.runtime.manifest import load_successful_tiled_slides as load_successful_tiled_slides_fn - slide_records, tiling_results = load_successful_tiled_slides_fn(output_dir) - # Key by the composite (sample_id, annotation) work unit so a multi-class slide's sibling - # classes never overwrite each other; flat units collapse to the bare sample_id key. - paired_by_unit = { - encode_work_unit(slide.sample_id, tiling_result_annotation(tiling_result)): (slide, tiling_result) - for slide, tiling_result in zip(slide_records, tiling_results) - } - progress_events_path = request.get("progress_events_path") - reporter = ( - JsonlProgressReporter( - progress_events_path, - rank=global_rank, - progress_label=f"cuda:{local_rank}", - ) - if progress_events_path - else None + global_rank = distributed.get_global_rank() + world_size = distributed.get_global_size() + local_rank = distributed.get_local_rank() + + model_spec = dict(request["model"]) + model = Model.from_preset( + model_spec["name"], + device=f"cuda:{local_rank}", + output_variant=model_spec.get("output_variant"), + allow_non_recommended_settings=bool(model_spec["allow_non_recommended_settings"]), + ) + preprocessing = deserialize_preprocessing(request["preprocessing"]) + execution = deserialize_execution(request["execution"]) + from slide2vec.runtime.distributed import ( + decode_work_unit, + encode_work_unit, + work_unit_shard_stem, + ) + from slide2vec.runtime.embedding import tiling_result_annotation + + load_successful_tiled_slides_fn = getattr(inference, "load_successful_tiled_slides", None) + if not callable(load_successful_tiled_slides_fn): + from slide2vec.runtime.manifest import load_successful_tiled_slides as load_successful_tiled_slides_fn + slide_records, tiling_results = load_successful_tiled_slides_fn(output_dir) + # Key by the composite (sample_id, annotation) work unit so a multi-class slide's sibling + # classes never overwrite each other; flat units collapse to the bare sample_id key. + paired_by_unit = { + encode_work_unit(slide.sample_id, tiling_result_annotation(tiling_result)): (slide, tiling_result) + for slide, tiling_result in zip(slide_records, tiling_results) + } + progress_events_path = request.get("progress_events_path") + reporter = ( + JsonlProgressReporter( + progress_events_path, + rank=global_rank, + progress_label=f"cuda:{local_rank}", ) - context = activate_progress_reporter(reporter) if reporter is not None else nullcontext() - - with context: - if request["strategy"] == "tile_shard": - work_unit = request["work_unit"] - shard_stem = work_unit_shard_stem(*decode_work_unit(work_unit)) - slide, tiling_result = paired_by_unit[work_unit] - loaded = model._load_backend() - if is_hierarchical_preprocessing(preprocessing): - geometry = resolve_hierarchical_geometry(preprocessing, tiling_result) - index = build_hierarchical_index( - tiling_result, - region_tile_multiple=int(preprocessing.region_tile_multiple), - tile_size_lv0=int(geometry["tile_size_lv0"]), - ) - flat_indices = np.array_split(index.flat_index, world_size)[global_rank] - compute_hierarchical_embedding_shard_for_slide_fn = getattr( - inference, - "_compute_hierarchical_embedding_shard_for_slide", - None, - ) - if not callable(compute_hierarchical_embedding_shard_for_slide_fn): - from slide2vec.runtime.embedding_pipeline import ( - compute_hierarchical_embedding_shard_for_slide as compute_hierarchical_embedding_shard_for_slide_fn, - ) - shard_indices, tile_embeddings = compute_hierarchical_embedding_shard_for_slide_fn( - loaded, - slide, - tiling_result, - preprocessing=preprocessing, - execution=execution, - flat_indices=flat_indices, - ) - payload = { - "flat_index": torch.as_tensor(shard_indices, dtype=torch.long), - "tile_embeddings": tile_embeddings.detach().cpu() if torch.is_tensor(tile_embeddings) else torch.as_tensor(tile_embeddings), - } - torch.save(payload, coordination_dir / f"{shard_stem}.hier.rank{global_rank}.pt") - else: - num_tiles = len(tiling_result.x) - tile_indices = np.array_split(np.arange(num_tiles, dtype=np.int64), world_size)[global_rank] - compute_tile_embeddings_for_slide_fn = getattr( - inference, - "_compute_tile_embeddings_for_slide", - None, - ) - if not callable(compute_tile_embeddings_for_slide_fn): - from slide2vec.runtime.embedding_pipeline import ( - compute_tile_embeddings_for_slide as compute_tile_embeddings_for_slide_fn, - ) - tile_embeddings = compute_tile_embeddings_for_slide_fn( - loaded, - model, - slide, - tiling_result, - preprocessing=preprocessing, - execution=execution, - tile_indices=tile_indices, + if progress_events_path + else None + ) + context = activate_progress_reporter(reporter) if reporter is not None else nullcontext() + + with context: + if request["strategy"] == "tile_shard": + work_unit = request["work_unit"] + shard_stem = work_unit_shard_stem(*decode_work_unit(work_unit)) + slide, tiling_result = paired_by_unit[work_unit] + loaded = model._load_backend() + if is_hierarchical_preprocessing(preprocessing): + geometry = resolve_hierarchical_geometry(preprocessing, tiling_result) + index = build_hierarchical_index( + tiling_result, + region_tile_multiple=int(preprocessing.region_tile_multiple), + tile_size_lv0=int(geometry["tile_size_lv0"]), + ) + flat_indices = np.array_split(index.flat_index, world_size)[global_rank] + compute_hierarchical_embedding_shard_for_slide_fn = getattr( + inference, + "_compute_hierarchical_embedding_shard_for_slide", + None, + ) + if not callable(compute_hierarchical_embedding_shard_for_slide_fn): + from slide2vec.runtime.embedding_pipeline import ( + compute_hierarchical_embedding_shard_for_slide as compute_hierarchical_embedding_shard_for_slide_fn, ) - payload = { - "tile_index": torch.as_tensor(tile_indices, dtype=torch.long), - "tile_embeddings": tile_embeddings.detach().cpu() if torch.is_tensor(tile_embeddings) else torch.as_tensor(tile_embeddings), - } - torch.save(payload, coordination_dir / f"{shard_stem}.tiles.rank{global_rank}.pt") - return 0 - - assigned_ids = list(request.get("assignments", {}).get(str(global_rank), [])) - if not assigned_ids: - return 0 - assigned_slides = [paired_by_unit[unit_key][0] for unit_key in assigned_ids] - assigned_tiling_results = [paired_by_unit[unit_key][1] for unit_key in assigned_ids] - - def _persist_embedded_slide(slide, tiling_result, embedded_slide) -> None: + shard_indices, tile_embeddings = compute_hierarchical_embedding_shard_for_slide_fn( + loaded, + slide, + tiling_result, + preprocessing=preprocessing, + execution=execution, + flat_indices=flat_indices, + ) 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), + "flat_index": torch.as_tensor(shard_indices, dtype=torch.long), + "tile_embeddings": tile_embeddings.detach().cpu() if torch.is_tensor(tile_embeddings) else torch.as_tensor(tile_embeddings), } - # Stem by (sample_id, annotation) so two classes of one slide never overwrite each - # other; flat units keep the bare-sample_id filename for backward compatibility. - stem = work_unit_shard_stem( - embedded_slide.sample_id, tiling_result_annotation(tiling_result) + torch.save(payload, coordination_dir / f"{shard_stem}.hier.rank{global_rank}.pt") + else: + num_tiles = len(tiling_result.x) + tile_indices = np.array_split(np.arange(num_tiles, dtype=np.int64), world_size)[global_rank] + compute_tile_embeddings_for_slide_fn = getattr( + inference, + "_compute_tile_embeddings_for_slide", + None, ) - torch.save(payload, coordination_dir / f"{stem}.embedded.pt") - - compute_embedded_slides_fn = getattr(inference, "_compute_embedded_slides", None) - if not callable(compute_embedded_slides_fn): - from slide2vec.runtime.embedding_pipeline import compute_embedded_slides as compute_embedded_slides_fn - compute_embedded_slides_fn( - model, - assigned_slides, - assigned_tiling_results, - preprocessing=preprocessing, - execution=execution, - on_embedded_slide=_persist_embedded_slide, - collect_results=False, - ) + if not callable(compute_tile_embeddings_for_slide_fn): + from slide2vec.runtime.embedding_pipeline import ( + compute_tile_embeddings_for_slide as compute_tile_embeddings_for_slide_fn, + ) + tile_embeddings = compute_tile_embeddings_for_slide_fn( + loaded, + model, + slide, + tiling_result, + preprocessing=preprocessing, + execution=execution, + tile_indices=tile_indices, + ) + payload = { + "tile_index": torch.as_tensor(tile_indices, dtype=torch.long), + "tile_embeddings": tile_embeddings.detach().cpu() if torch.is_tensor(tile_embeddings) else torch.as_tensor(tile_embeddings), + } + torch.save(payload, coordination_dir / f"{shard_stem}.tiles.rank{global_rank}.pt") + return 0 + + assigned_ids = list(request.get("assignments", {}).get(str(global_rank), [])) + if not assigned_ids: return 0 - finally: - if dist.is_available() and dist.is_initialized(): - dist.destroy_process_group() + assigned_slides = [paired_by_unit[unit_key][0] for unit_key in assigned_ids] + assigned_tiling_results = [paired_by_unit[unit_key][1] for unit_key in assigned_ids] + + 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), + } + # Stem by (sample_id, annotation) so two classes of one slide never overwrite each + # other; flat units keep the bare-sample_id filename for backward compatibility. + stem = work_unit_shard_stem( + embedded_slide.sample_id, tiling_result_annotation(tiling_result) + ) + torch.save(payload, coordination_dir / f"{stem}.embedded.pt") + + compute_embedded_slides_fn = getattr(inference, "_compute_embedded_slides", None) + if not callable(compute_embedded_slides_fn): + from slide2vec.runtime.embedding_pipeline import compute_embedded_slides as compute_embedded_slides_fn + compute_embedded_slides_fn( + model, + assigned_slides, + assigned_tiling_results, + preprocessing=preprocessing, + execution=execution, + on_embedded_slide=_persist_embedded_slide, + collect_results=False, + ) + return 0 def _to_cpu_payload(value): diff --git a/slide2vec/distributed/pipeline_worker.py b/slide2vec/distributed/pipeline_worker.py index c1f43c8..a4201e3 100644 --- a/slide2vec/distributed/pipeline_worker.py +++ b/slide2vec/distributed/pipeline_worker.py @@ -15,8 +15,6 @@ def get_args_parser(add_help: bool = True) -> argparse.ArgumentParser: def main(argv=None) -> int: - import torch.distributed as dist - import slide2vec.distributed as distributed import slide2vec.inference as inference from slide2vec.api import Model @@ -29,84 +27,81 @@ def main(argv=None) -> int: request = json.loads(Path(args.request_path).read_text(encoding="utf-8")) output_dir = Path(args.output_dir) + # Reads the torchrun-exported rank/world-size; no NCCL process group (issue #219). distributed.enable(overwrite=True) - try: - local_rank = distributed.get_local_rank() - global_rank = distributed.get_global_rank() - world_size = distributed.get_global_size() + local_rank = distributed.get_local_rank() + global_rank = distributed.get_global_rank() + world_size = distributed.get_global_size() - model_spec = dict(request["model"]) - model = Model.from_preset( - model_spec["name"], - device=f"cuda:{local_rank}", - output_variant=model_spec.get("output_variant"), - allow_non_recommended_settings=bool(model_spec["allow_non_recommended_settings"]), - ) - preprocessing = deserialize_preprocessing(request["preprocessing"]) - execution = deserialize_execution(request["execution"]) - tiling_input_dir = Path(request.get("tiling_input_dir", str(output_dir))) - load_successful_tiled_slides_fn = getattr(inference, "load_successful_tiled_slides", None) - if not callable(load_successful_tiled_slides_fn): - from slide2vec.runtime.manifest import load_successful_tiled_slides as load_successful_tiled_slides_fn - slide_records, tiling_results = load_successful_tiled_slides_fn(tiling_input_dir) - # Each (sample_id, annotation) row is an independent work unit; key by the composite so a - # multi-class slide's sibling classes never overwrite each other. Flat units (None / tissue / - # merged) encode to the bare sample_id, byte-identical to pre-#168 single-class runs. + model_spec = dict(request["model"]) + model = Model.from_preset( + model_spec["name"], + device=f"cuda:{local_rank}", + output_variant=model_spec.get("output_variant"), + allow_non_recommended_settings=bool(model_spec["allow_non_recommended_settings"]), + ) + preprocessing = deserialize_preprocessing(request["preprocessing"]) + execution = deserialize_execution(request["execution"]) + tiling_input_dir = Path(request.get("tiling_input_dir", str(output_dir))) + load_successful_tiled_slides_fn = getattr(inference, "load_successful_tiled_slides", None) + if not callable(load_successful_tiled_slides_fn): + from slide2vec.runtime.manifest import load_successful_tiled_slides as load_successful_tiled_slides_fn + slide_records, tiling_results = load_successful_tiled_slides_fn(tiling_input_dir) + # Each (sample_id, annotation) row is an independent work unit; key by the composite so a + # multi-class slide's sibling classes never overwrite each other. Flat units (None / tissue / + # merged) encode to the bare sample_id, byte-identical to pre-#168 single-class runs. + paired_by_unit = { + encode_work_unit(slide.sample_id, tiling_result_annotation(tiling_result)): (slide, tiling_result) + for slide, tiling_result in zip(slide_records, tiling_results) + } + requested_work_units = request.get("work_units") + if requested_work_units is not None: + requested_unit_set = {str(unit) for unit in requested_work_units} paired_by_unit = { - encode_work_unit(slide.sample_id, tiling_result_annotation(tiling_result)): (slide, tiling_result) - for slide, tiling_result in zip(slide_records, tiling_results) + unit_key: pair + for unit_key, pair in paired_by_unit.items() + if unit_key in requested_unit_set } - requested_work_units = request.get("work_units") - if requested_work_units is not None: - requested_unit_set = {str(unit) for unit in requested_work_units} - paired_by_unit = { - unit_key: pair - for unit_key, pair in paired_by_unit.items() - if unit_key in requested_unit_set - } - slide_records = [slide for slide, _ in paired_by_unit.values()] - tiling_results = [tiling_result for _, tiling_result in paired_by_unit.values()] - assignments = assign_slides_to_ranks(slide_records, tiling_results, num_gpus=world_size) - assigned_ids = assignments.get(global_rank, []) - if not assigned_ids: - return 0 - assigned_slides = [paired_by_unit[unit_key][0] for unit_key in assigned_ids] - assigned_tiling_results = [paired_by_unit[unit_key][1] for unit_key in assigned_ids] - progress_events_path = request.get("progress_events_path") - reporter = ( - JsonlProgressReporter( - progress_events_path, - rank=global_rank, - progress_label=f"cuda:{local_rank}", - ) - if progress_events_path - else None - ) - context = activate_progress_reporter(reporter) if reporter is not None else nullcontext() - with context: - build_incremental_persist_callback_fn = getattr(inference, "_build_incremental_persist_callback", build_incremental_persist_callback) - persist_callback, _, _ = build_incremental_persist_callback_fn( - model=model, - preprocessing=preprocessing, - execution=execution, - process_list_path=None, - ) - compute_embedded_slides_fn = getattr(inference, "_compute_embedded_slides", None) - if not callable(compute_embedded_slides_fn): - from slide2vec.runtime.embedding_pipeline import compute_embedded_slides as compute_embedded_slides_fn - compute_embedded_slides_fn( - model, - assigned_slides, - assigned_tiling_results, - preprocessing=preprocessing, - execution=execution, - on_embedded_slide=persist_callback, - collect_results=False, - ) + slide_records = [slide for slide, _ in paired_by_unit.values()] + tiling_results = [tiling_result for _, tiling_result in paired_by_unit.values()] + assignments = assign_slides_to_ranks(slide_records, tiling_results, num_gpus=world_size) + assigned_ids = assignments.get(global_rank, []) + if not assigned_ids: return 0 - finally: - if dist.is_available() and dist.is_initialized(): - dist.destroy_process_group() + assigned_slides = [paired_by_unit[unit_key][0] for unit_key in assigned_ids] + assigned_tiling_results = [paired_by_unit[unit_key][1] for unit_key in assigned_ids] + progress_events_path = request.get("progress_events_path") + reporter = ( + JsonlProgressReporter( + progress_events_path, + rank=global_rank, + progress_label=f"cuda:{local_rank}", + ) + if progress_events_path + else None + ) + context = activate_progress_reporter(reporter) if reporter is not None else nullcontext() + with context: + build_incremental_persist_callback_fn = getattr(inference, "_build_incremental_persist_callback", build_incremental_persist_callback) + persist_callback, _, _ = build_incremental_persist_callback_fn( + model=model, + preprocessing=preprocessing, + execution=execution, + process_list_path=None, + ) + compute_embedded_slides_fn = getattr(inference, "_compute_embedded_slides", None) + if not callable(compute_embedded_slides_fn): + from slide2vec.runtime.embedding_pipeline import compute_embedded_slides as compute_embedded_slides_fn + compute_embedded_slides_fn( + model, + assigned_slides, + assigned_tiling_results, + preprocessing=preprocessing, + execution=execution, + on_embedded_slide=persist_callback, + collect_results=False, + ) + return 0 if __name__ == "__main__": diff --git a/slide2vec/utils/config.py b/slide2vec/utils/config.py index 1553854..f6afc06 100644 --- a/slide2vec/utils/config.py +++ b/slide2vec/utils/config.py @@ -6,10 +6,7 @@ from pathlib import Path from omegaconf import OmegaConf -from slide2vec.distributed import ( - is_enabled_and_multiple_gpus, - is_main_process, -) +from slide2vec.distributed import is_main_process from slide2vec.runtime.model_settings import canonicalize_model_name from slide2vec.utils import initialize_wandb, fix_random_seeds, get_sha, setup_logging from slide2vec.configs import default_config @@ -170,9 +167,5 @@ def hf_login(): prompted = True if token is None: return - if is_enabled_and_multiple_gpus(): - import torch.distributed as dist - - dist.barrier() if is_main_process() and prompted: login(token) diff --git a/tests/test_regression_core.py b/tests/test_regression_core.py index 68cea68..7e36770 100644 --- a/tests/test_regression_core.py +++ b/tests/test_regression_core.py @@ -670,7 +670,6 @@ def _fake_login(*args, **kwargs): monkeypatch.setenv("HF_TOKEN", "token-from-env") monkeypatch.setattr(config, "is_main_process", lambda: True) - monkeypatch.setattr(config, "is_enabled_and_multiple_gpus", lambda: False) monkeypatch.setattr("huggingface_hub.login", _fake_login) config.hf_login() diff --git a/tests/test_regression_inference.py b/tests/test_regression_inference.py index 78a28ff..8131018 100644 --- a/tests/test_regression_inference.py +++ b/tests/test_regression_inference.py @@ -1680,8 +1680,6 @@ def test_compute_embedded_slides_skips_retaining_results_when_collect_results_is def test_pipeline_worker_disables_result_collection_when_streaming(monkeypatch, tmp_path: Path): - import torch.distributed as dist - import slide2vec.distributed as distributed import slide2vec.inference as inference import slide2vec.runtime.serialization as serialization @@ -1713,8 +1711,6 @@ def test_pipeline_worker_disables_result_collection_when_streaming(monkeypatch, monkeypatch.setattr(distributed, "get_local_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_size", lambda: 1) - monkeypatch.setattr(dist, "is_available", lambda: False) - monkeypatch.setattr(dist, "is_initialized", lambda: False) monkeypatch.setattr(Model, "from_preset", lambda *args, **kwargs: SimpleNamespace()) monkeypatch.setattr(serialization, "deserialize_preprocessing", lambda payload: DEFAULT_PREPROCESSING) monkeypatch.setattr( @@ -1747,8 +1743,6 @@ def fake_compute_embedded_slides(*args, **kwargs): def test_pipeline_worker_filters_to_requested_sample_ids(monkeypatch, tmp_path: Path): - import torch.distributed as dist - import slide2vec.distributed as distributed import slide2vec.runtime.serialization as serialization from slide2vec.api import Model @@ -1781,8 +1775,6 @@ def test_pipeline_worker_filters_to_requested_sample_ids(monkeypatch, tmp_path: monkeypatch.setattr(distributed, "get_local_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_size", lambda: 1) - monkeypatch.setattr(dist, "is_available", lambda: False) - monkeypatch.setattr(dist, "is_initialized", lambda: False) monkeypatch.setattr(Model, "from_preset", lambda *args, **kwargs: SimpleNamespace()) monkeypatch.setattr(serialization, "deserialize_preprocessing", lambda payload: DEFAULT_PREPROCESSING) monkeypatch.setattr( @@ -1811,7 +1803,6 @@ def fake_compute_embedded_slides(_model, slides, _tiling_results, **kwargs): def test_direct_embed_worker_streams_payloads_without_retaining_results(monkeypatch, tmp_path: Path): import torch - import torch.distributed as dist import slide2vec.distributed as distributed import slide2vec.runtime.serialization as serialization @@ -1846,8 +1837,6 @@ def test_direct_embed_worker_streams_payloads_without_retaining_results(monkeypa monkeypatch.setattr(distributed, "get_local_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_size", lambda: 1) - monkeypatch.setattr(dist, "is_available", lambda: False) - monkeypatch.setattr(dist, "is_initialized", lambda: False) monkeypatch.setattr(Model, "from_preset", lambda *args, **kwargs: SimpleNamespace()) monkeypatch.setattr(serialization, "deserialize_preprocessing", lambda payload: DEFAULT_PREPROCESSING) monkeypatch.setattr( @@ -3106,46 +3095,89 @@ def test_make_embedded_slide_validates_coordinates_and_supports_tile_and_slide_o tile_embeddings=np.zeros((1, 4), dtype=np.float32), ) -def test_distributed_enable_keeps_explicit_device_binding_without_unconditional_barrier(): +def test_distributed_enable_drops_nccl_process_group_but_keeps_device_binding(): + # Issue #219: the dense/pooled workers run no collectives, so enable() must NOT + # build an NCCL process group. It reads rank/world-size from the torchrun-exported + # env and only binds the CUDA device. source = (ROOT / "slide2vec" / "distributed" / "__init__.py").read_text(encoding="utf-8") tree = ast.parse(source) - init_process_group_calls = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and isinstance(node.func.value, ast.Name) - and node.func.value.id == "dist" - and node.func.attr == "init_process_group" - ] - assert init_process_group_calls + def _attr_calls(attr: str) -> list[ast.Call]: + return [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == attr + ] - device_keyword = next( - ( - keyword - for keyword in init_process_group_calls[0].keywords - if keyword.arg == "device_id" - ), - None, - ) - assert device_keyword is not None - assert isinstance(device_keyword.value, ast.Call) - assert isinstance(device_keyword.value.func, ast.Attribute) - assert isinstance(device_keyword.value.func.value, ast.Name) - assert device_keyword.value.func.value.id == "torch" - assert device_keyword.value.func.attr == "device" + # No process group is created or torn down, and no collective is issued. + assert not _attr_calls("init_process_group") + assert not _attr_calls("destroy_process_group") + assert not _attr_calls("barrier") + assert not _attr_calls("all_gather") + assert not _attr_calls("all_reduce") - barrier_calls = [ + # The explicit device binding survives. + set_device_calls = [ node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and isinstance(node.func.value, ast.Name) - and node.func.value.id == "dist" - and node.func.attr == "barrier" + for node in _attr_calls("set_device") + if isinstance(node.func.value, ast.Attribute) + and node.func.value.attr == "cuda" ] - assert not barrier_calls + assert set_device_calls + + +def test_global_rank_and_size_read_module_state_not_a_process_group(monkeypatch): + # Issue #219: get_global_rank/get_global_size derive from the module state that + # enable() records from the env, not from a live torch.distributed process group. + import slide2vec.distributed as distributed + + monkeypatch.setattr(distributed, "_RANK", 3, raising=False) + monkeypatch.setattr(distributed, "_WORLD_SIZE", 4, raising=False) + monkeypatch.setattr(distributed, "_LOCAL_RANK", 1, raising=False) + monkeypatch.setattr(distributed, "_LOCAL_WORLD_SIZE", 2, raising=False) + + assert distributed.is_enabled() is True + assert distributed.get_global_rank() == 3 + assert distributed.get_global_size() == 4 + assert distributed.get_local_rank() == 1 + assert distributed.get_local_size() == 2 + assert distributed.is_main_process() is False + + +def test_global_helpers_return_single_process_defaults_when_not_enabled(monkeypatch): + import slide2vec.distributed as distributed + + monkeypatch.setattr(distributed, "_RANK", -1, raising=False) + monkeypatch.setattr(distributed, "_WORLD_SIZE", -1, raising=False) + monkeypatch.setattr(distributed, "_LOCAL_RANK", -1, raising=False) + monkeypatch.setattr(distributed, "_LOCAL_WORLD_SIZE", -1, raising=False) + + assert distributed.is_enabled() is False + assert distributed.get_global_rank() == 0 + assert distributed.get_global_size() == 1 + assert distributed.get_local_rank() == 0 + assert distributed.get_local_size() == 1 + assert distributed.is_main_process() is True + + +def test_distributed_module_has_no_torch_distributed_dependency(): + # The whole point of #219: slide2vec.distributed no longer imports or calls into + # torch.distributed — nor datetime (the NCCL timeout) or socket (the MASTER_PORT + # pick), which only the process-group setup needed. + source = (ROOT / "slide2vec" / "distributed" / "__init__.py").read_text(encoding="utf-8") + tree = ast.parse(source) + imported = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + imported.add(node.module) + assert "torch.distributed" not in imported + assert "datetime" not in imported + assert "socket" not in imported + def test_setup_distributed_helper_has_been_removed(): source = (ROOT / "slide2vec" / "utils" / "config.py").read_text(encoding="utf-8") @@ -6028,8 +6060,6 @@ def test_work_unit_shard_stem_real_class_is_unique_and_filesystem_safe(): def test_direct_embed_worker_slide_shard_writes_per_class_payloads_without_collision(monkeypatch, tmp_path: Path): """Two classes of one slide assigned to the same rank must persist to distinct .embedded.pt files (per (sample_id, annotation) stem), so a sibling class is never overwritten.""" - import torch.distributed as dist - import slide2vec.distributed as distributed import slide2vec.runtime.serialization as serialization from slide2vec.api import Model @@ -6063,8 +6093,6 @@ def test_direct_embed_worker_slide_shard_writes_per_class_payloads_without_colli monkeypatch.setattr(distributed, "get_local_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_size", lambda: 1) - monkeypatch.setattr(dist, "is_available", lambda: False) - monkeypatch.setattr(dist, "is_initialized", lambda: False) monkeypatch.setattr(Model, "from_preset", lambda *args, **kwargs: SimpleNamespace()) monkeypatch.setattr(serialization, "deserialize_preprocessing", lambda payload: DEFAULT_PREPROCESSING) monkeypatch.setattr(serialization, "deserialize_execution", lambda payload: ExecutionOptions(output_dir=tmp_path)) @@ -6146,8 +6174,6 @@ def test_compute_embedded_slides_finished_event_carries_annotation(monkeypatch): def test_pipeline_worker_embeds_every_class_of_a_multi_class_slide(monkeypatch, tmp_path: Path): """Pipeline-worker path: a multi-label slide's classes must NOT collapse by sample_id — every (sample_id, annotation) unit is computed (the bug #168 fixes).""" - import torch.distributed as dist - import slide2vec.distributed as distributed import slide2vec.runtime.serialization as serialization from slide2vec.api import Model @@ -6179,8 +6205,6 @@ def test_pipeline_worker_embeds_every_class_of_a_multi_class_slide(monkeypatch, monkeypatch.setattr(distributed, "get_local_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_rank", lambda: 0) monkeypatch.setattr(distributed, "get_global_size", lambda: 1) - monkeypatch.setattr(dist, "is_available", lambda: False) - monkeypatch.setattr(dist, "is_initialized", lambda: False) monkeypatch.setattr(Model, "from_preset", lambda *args, **kwargs: SimpleNamespace()) monkeypatch.setattr(serialization, "deserialize_preprocessing", lambda payload: DEFAULT_PREPROCESSING) monkeypatch.setattr(serialization, "deserialize_execution", lambda payload: ExecutionOptions(output_dir=tmp_path))