dask-setup 2.0.0__tar.gz → 2.1.0__tar.gz

{dask_setup-2.0.0 → dask_setup-2.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dask_setup
-Version: 2.0.0
+Version: 2.1.0
 Summary: HPC-tuned Dask helpers for single-node runs on NCI Gadi.
 Author: Sam Green
 Project-URL: Homepage, https://github.com/21centuryweather/dask_setup

{dask_setup-2.0.0 → dask_setup-2.1.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dask_setup"
-version = "2.0.0"
+version = "2.1.0"
 description = "HPC-tuned Dask helpers for single-node runs on NCI Gadi."
 authors = [{name="Sam Green"}]
 readme = "README.md"

{dask_setup-2.0.0 → dask_setup-2.1.0}/src/dask_setup/__init__.py

@@ -22,6 +22,7 @@ from .multinode import (
     detect_cluster_mode,
     generate_pbs_script,
     generate_slurm_script,
+    setup_interactive_cluster,
     setup_pbs_cluster,
     setup_slurm_cluster,
 )
@@ -280,7 +281,7 @@ except ImportError:
         )
 
 
-__version__ = "2.0.0"
+__version__ = "2.1.0"
 
 __all__ = [
     # Core API — always available
@@ -326,6 +327,7 @@ __all__ = [
     "MultiNodeConfig",
     "SharedTempDir",
     "detect_cluster_mode",
+    "setup_interactive_cluster",
     "setup_pbs_cluster",
     "setup_slurm_cluster",
     "generate_pbs_script",

{dask_setup-2.0.0 → dask_setup-2.1.0}/src/dask_setup/benchmark.py

@@ -39,6 +39,7 @@ Example::
 
 from __future__ import annotations
 
+import contextlib
 import statistics
 import time
 from collections.abc import Callable, Sequence
@@ -324,15 +325,38 @@ class ChunkImpactResult:
     ----------
     results : list[BenchmarkResult]
         One entry per chunk specification, in the order they were tried.
-    chunk_sizes : list[dict[str, int]]
-        The chunk specs that were swept.
+    chunk_specs : list[dict[str, int]]
+        The chunk specs that were swept. Alias for chunk_sizes.
     recommended_chunks : dict[str, int]
         The chunk spec with the best (lowest) wall time.
     """
 
     results: list[BenchmarkResult]
-    chunk_sizes: list[dict[str, int]]
-    recommended_chunks: dict[str, int]
+    chunk_specs: list[dict[str, int]]
+    recommended_chunks: dict[str, int] | None = None
+
+    def __post_init__(self) -> None:
+        """Compute recommended_chunks if not provided."""
+        if self.recommended_chunks is None:
+            # Find the chunk spec with the lowest wall time
+            if self.results and self.chunk_specs:
+                valid = [
+                    (r, cs)
+                    for r, cs in zip(self.results, self.chunk_specs, strict=False)
+                    if r.wall_time_seconds == r.wall_time_seconds  # not nan
+                ]
+                if valid:
+                    _, best_chunks = min(valid, key=lambda x: x[0].wall_time_seconds)
+                    self.recommended_chunks = best_chunks
+                else:
+                    self.recommended_chunks = self.chunk_specs[0] if self.chunk_specs else {}
+            else:
+                self.recommended_chunks = {}
+
+    @property
+    def chunk_sizes(self) -> list[dict[str, int]]:
+        """Alias for chunk_specs for backward compatibility."""
+        return self.chunk_specs
 
     @property
     def wall_times(self) -> list[float]:
@@ -341,6 +365,10 @@ class ChunkImpactResult:
     def best(self) -> BenchmarkResult:
         return min(self.results, key=lambda r: r.wall_time_seconds)
 
+    def optimal(self) -> BenchmarkResult:
+        """Return the result with the best (lowest) wall time."""
+        return min(self.results, key=lambda r: r.wall_time_seconds)
+
     def summary(self) -> str:
         lines = [
             f"{'Chunks':^36}  {'Wall (s)':>10}  {'Mem GiB':>8}  {'Tasks/s':>8}",
@@ -352,7 +380,7 @@ class ChunkImpactResult:
                 f"{chunk_str:<36}  {r.wall_time_seconds:>10.2f}  "
                 f"{r.peak_memory_gib:>8.2f}  {r.tasks_per_second:>8.1f}"
             )
-        lines.append(f"\nRecommended: {self.recommended_chunks}")
+        lines.append(f"\nOptimal: {self.recommended_chunks}")
         return "\n".join(lines)
 
     def to_dataframe(self) -> Any:
@@ -476,10 +504,8 @@ def _measure_one(
         errors.append(f"Task count failed: {e}")
 
     # Worker count
-    try:
+    with contextlib.suppress(Exception):
         n_workers = len(client.scheduler_info().get("workers", {}))
-    except Exception:
-        pass
 
     # Optional warmup
     if warmup:
@@ -644,15 +670,11 @@ def benchmark_config(
             )
         finally:
             if client is not None:
-                try:
+                with contextlib.suppress(Exception):
                     client.close()
-                except Exception:
-                    pass
             if cluster is not None:
-                try:
+                with contextlib.suppress(Exception):
                     cluster.close()
-                except Exception:
-                    pass
 
         if result is None:
             result = BenchmarkResult(name=name, wall_time_seconds=float("nan"), errors=errors)
@@ -765,15 +787,11 @@ def scaling_analysis(
             )
         finally:
             if client is not None:
-                try:
+                with contextlib.suppress(Exception):
                     client.close()
-                except Exception:
-                    pass
             if cluster is not None:
-                try:
+                with contextlib.suppress(Exception):
                     cluster.close()
-                except Exception:
-                    pass
 
         raw_results.append(result)
 
@@ -805,12 +823,10 @@ def scaling_analysis(
     if plot:
         fig = scaling.plot()
         if fig is not None:
-            try:
+            with contextlib.suppress(Exception):
                 import matplotlib.pyplot as plt
 
                 plt.show()
-            except Exception:
-                pass
 
     return scaling
 
@@ -882,10 +898,9 @@ def chunk_impact(
 
     # Generate chunk size sweep if not provided
     if chunk_sizes is None:
-        if auto_chunks and dims:
-            chunk_sizes = _generate_auto_chunks(dims)
-        else:
-            chunk_sizes = [{}]  # single run with no rechunking
+        chunk_sizes = (
+            _generate_auto_chunks(dims) if auto_chunks and dims else [{}]
+        )  # single run with no rechunking
 
     raw_results: list[BenchmarkResult] = []
 
@@ -927,7 +942,7 @@ def chunk_impact(
 
     impact = ChunkImpactResult(
         results=raw_results,
-        chunk_sizes=chunk_sizes,
+        chunk_specs=chunk_sizes,
         recommended_chunks=best_chunks,
     )
 
@@ -935,12 +950,10 @@ def chunk_impact(
         first_dim = next(iter(dims), None)
         fig = impact.plot(dim=first_dim)
         if fig is not None:
-            try:
+            with contextlib.suppress(Exception):
                 import matplotlib.pyplot as plt
 
                 plt.show()
-            except Exception:
-                pass
 
     return impact
 
@@ -1081,8 +1094,8 @@ def run_synthetic_benchmark(
     if verbose:
         print(f"Creating synthetic array {shape} chunked {chunks} …")
 
-    np.random.seed(42)
-    arr = da.from_array(np.random.rand(*shape).astype(np.float32), chunks=chunks)
+    rng = np.random.default_rng(42)
+    arr = da.from_array(rng.random(shape).astype(np.float32), chunks=chunks)
 
     # Resolve dask operation
     da_ops: dict[str, Callable[[Any], Any]] = {
@@ -1147,15 +1160,11 @@ def run_synthetic_benchmark(
         errors.append(f"Cluster error: {e}")
     finally:
         if client is not None:
-            try:
+            with contextlib.suppress(Exception):
                 client.close()
-            except Exception:
-                pass
         if cluster_obj is not None:
-            try:
+            with contextlib.suppress(Exception):
                 cluster_obj.close()
-            except Exception:
-                pass
 
     valid_times = [t for t in times if t != float("inf")]
     wall = statistics.mean(valid_times) if valid_times else float("nan")

{dask_setup-2.0.0 → dask_setup-2.1.0}/src/dask_setup/client.py

@@ -15,7 +15,13 @@ from .config_manager import ConfigManager
 from .dashboard import print_dashboard_info
 from .exceptions import InsufficientResourcesError
 from .logging import get_logger
-from .multinode import MultiNodeConfig, detect_cluster_mode
+from .multinode import (
+    MultiNodeConfig,
+    detect_cluster_mode,
+    setup_interactive_cluster,
+    setup_pbs_cluster,
+    setup_slurm_cluster,
+)
 from .resources import detect_resources
 from .tempdir import create_dask_temp_dir
 from .topology import decide_topology, validate_topology
@@ -48,13 +54,14 @@ def _compute_smart_reserve_default() -> float:
     """
     try:
         total_ram_gb = psutil.virtual_memory().total / (1024**3)
-        return min(50.0, max(4.0, total_ram_gb * 0.20))
+        smart_reserve = min(50.0, max(4.0, total_ram_gb * 0.20))
+        # Cap the smart reserve at 50.0 GiB as a safe HPC default
+        return min(50.0, smart_reserve)
     except Exception:
         return 50.0  # Safe HPC fallback if psutil is unexpectedly unavailable
 
 
 def _resolve_configuration(
-    config: DaskSetupConfig | None = None,
     profile: str | None = None,
     workload_type: str = "io",
     max_workers: int | None = None,
@@ -64,21 +71,16 @@ def _resolve_configuration(
     adaptive: bool = False,
     min_workers: int | None = None,
     suggest_chunks: bool = False,
-    fallback_on_detection_failure: bool = False,
-    adaptive_memory: bool = False,
 ) -> DaskSetupConfig:
-    """Resolve final configuration from a config object, profile, and explicit parameters.
+    """Resolve final configuration from a profile and explicit parameters.
 
     Priority order (highest to lowest):
 
     1. Explicit keyword parameters passed to ``setup_dask_client()``
-    2. Config object (``config=``) **or** profile (``profile=``)
+    2. Profile (``profile=``)
     3. Defaults — ``reserve_mem_gb`` uses a smart default (20 % RAM, 4–50 GiB)
 
-    ``config`` and ``profile`` are mutually exclusive.
-
     Args:
-        config: A pre-built DaskSetupConfig object to use as the base configuration.
         profile: Profile name to load from disk/builtins as the base configuration.
         workload_type: Workload type override.
         max_workers: Maximum workers cap.
@@ -96,28 +98,18 @@ def _resolve_configuration(
         The heuristic that detects "explicitly set" parameters compares each value
         against its default.  Edge case: if you deliberately pass a value equal to
         the default (e.g. ``dashboard=True``) it will be treated as "not set" and a
-        base_config value will take precedence.  Use ``config=DaskSetupConfig(...)``
-        to avoid this ambiguity entirely.
+        profile value will take precedence.
     """
-    if config is not None and profile is not None:
-        raise ValueError(
-            "Cannot specify both 'config' and 'profile'. "
-            "Pass a DaskSetupConfig object via 'config=' OR a profile name via 'profile=', not both."
-        )
-
-    # Build defaults using the environment-aware smart reserve
-    smart_reserve = _compute_smart_reserve_default()
-    defaults = DaskSetupConfig(reserve_mem_gb=smart_reserve)
-    logger.debug("Smart reserve default computed", reserve_mem_gb=smart_reserve)
+    # Build defaults - use provided reserve_mem_gb or fallback to smart default
+    defaults = DaskSetupConfig(
+        reserve_mem_gb=reserve_mem_gb if reserve_mem_gb is not None else 50.0
+    )
+    logger.debug("Configuration defaults set", reserve_mem_gb=defaults.reserve_mem_gb)
 
-    # Resolve the base configuration: either a provided config object or a loaded profile
+    # Resolve the base configuration: load profile if specified
     base_config = None
 
-    if config is not None:
-        # Caller supplied a ready-made DaskSetupConfig — use it as the profile-level base
-        base_config = config
-        logger.debug("Using caller-provided DaskSetupConfig as base")
-    elif profile is not None:
+    if profile is not None:
         manager = ConfigManager()
         profile_obj = manager.get_profile(profile)
         if profile_obj is None:
@@ -130,9 +122,7 @@ def _resolve_configuration(
     #
     # Note: This heuristic compares each value against its default to decide whether it was
     # explicitly set. Edge case: if you deliberately pass a value that *equals* the default
-    # (e.g. reserve_mem_gb equal to the computed smart default) it will be treated as "not set"
-    # and a base_config value will take precedence. To avoid this, use a DaskSetupConfig
-    # object via the 'config=' parameter instead.
+    # (e.g. dashboard=True) it will be treated as "not set" and a profile value will take precedence.
     explicit_params: dict[str, Any] = {}
 
     if workload_type != "io":
@@ -140,7 +130,7 @@ def _resolve_configuration(
     if max_workers is not None:
         explicit_params["max_workers"] = max_workers
     if reserve_mem_gb is not None:
-        # None means "use the smart default"; an explicit float means "user chose this"
+        # Explicit float means "user chose this"
         explicit_params["reserve_mem_gb"] = reserve_mem_gb
     if max_mem_gb is not None:
         explicit_params["max_mem_gb"] = max_mem_gb
@@ -152,10 +142,6 @@ def _resolve_configuration(
         explicit_params["min_workers"] = min_workers
     if suggest_chunks is not False:
         explicit_params["suggest_chunks"] = suggest_chunks
-    if fallback_on_detection_failure:
-        explicit_params["fallback_on_detection_failure"] = fallback_on_detection_failure
-    if adaptive_memory:
-        explicit_params["adaptive_memory"] = adaptive_memory
 
     explicit_config = DaskSetupConfig(**explicit_params) if explicit_params else None
 
@@ -308,7 +294,7 @@ def setup_dask_client(
 def setup_dask_client(
     workload_type: str = "io",
     max_workers: int | None = None,
-    reserve_mem_gb: float | None = None,
+    reserve_mem_gb: float = 50.0,
     max_mem_gb: float | None = None,
     dashboard: bool = True,
     adaptive: bool = False,
@@ -380,17 +366,27 @@ def setup_dask_client(
         and tightens the worker ``memory.target`` / ``memory.spill`` thresholds
         slightly, giving workers more head-room from the start.  Default
         ``False``.
-    mode : {"auto", "local", "pbs", "slurm"}
+    mode : {"auto", "local", "pbs", "slurm", "interactive"}
         Backend selection.
 
         - ``"local"`` — always use a single-node ``LocalCluster`` (default
           behaviour prior to v2.0).
-        - ``"pbs"`` — launch via ``dask-jobqueue.PBSCluster``.  Requires
-          ``pip install dask-jobqueue``.
-        - ``"slurm"`` — launch via ``dask-jobqueue.SLURMCluster``.
+        - ``"pbs"`` — launch via ``dask-jobqueue.PBSCluster`` (submits new
+          batch jobs).  Requires ``pip install dask-jobqueue``.
+        - ``"slurm"`` — launch via ``dask-jobqueue.SLURMCluster`` (submits
+          new batch jobs).
+        - ``"interactive"`` — use resources already allocated in the current
+          interactive PBS (``qsub -I``) or SLURM (``salloc``) session.
+          Single-node allocations create a ``LocalCluster``; multi-node
+          allocations create an ``SSHCluster`` across all nodes in
+          ``PBS_NODEFILE`` / ``SLURM_NODELIST``.
         - ``"auto"`` (default) — inspect the environment and choose
-          ``"pbs"`` if ``PBS_JOBID`` is set, ``"slurm"`` if
-          ``SLURM_JOB_ID`` is set, or ``"local"`` otherwise.
+          ``"interactive"`` when inside a PBS interactive job
+          (``PBS_ENVIRONMENT=PBS_INTERACTIVE``) or a SLURM interactive
+          allocation (``SLURM_BATCH_FLAG`` not set to ``"1"``);
+          ``"pbs"`` when ``PBS_JOBID`` is set in a batch job;
+          ``"slurm"`` when ``SLURM_JOB_ID`` is set in a batch job;
+          ``"local"`` otherwise.
     multi_node_config : MultiNodeConfig or None
         Configuration for the multi-node backend (``mode="pbs"`` or
         ``"slurm"``).  Ignored when ``mode="local"``.  When ``None`` and a
@@ -462,9 +458,14 @@ def setup_dask_client(
         resolved_mode = detect_cluster_mode()
         logger.debug("Mode auto-resolved", mode=resolved_mode)
 
-    if resolved_mode in {"pbs", "slurm"}:
-        from .multinode import setup_pbs_cluster, setup_slurm_cluster
+    if resolved_mode == "interactive":
+        logger.info("Interactive cluster mode — using already-allocated nodes")
+        client, cluster, tmp_path = setup_interactive_cluster(
+            workload_type=workload_type,
+        )
+        return client, cluster, tmp_path  # type: ignore[return-value]
 
+    if resolved_mode in {"pbs", "slurm"}:
         mn_cfg = multi_node_config
         if mn_cfg is None:
             # Build a minimal MultiNodeConfig from whatever was passed
@@ -496,20 +497,25 @@ def setup_dask_client(
         logger.info("Auto-selected profile", profile=profile)
 
     # Load and merge configuration
+    # Use explicit default for reserve_mem_gb if not provided
+    resolved_reserve_mem = reserve_mem_gb if reserve_mem_gb is not None else 50.0
+
     config = _resolve_configuration(
-        config=config,
         profile=profile,
         workload_type=workload_type,
         max_workers=max_workers,
-        reserve_mem_gb=reserve_mem_gb,
+        reserve_mem_gb=resolved_reserve_mem,
         max_mem_gb=max_mem_gb,
         dashboard=dashboard,
         adaptive=adaptive,
         min_workers=min_workers,
         suggest_chunks=suggest_chunks,
-        fallback_on_detection_failure=fallback_on_detection_failure,
-        adaptive_memory=adaptive_memory,
     )
+
+    # Apply additional config-level settings that don't go through _resolve_configuration
+    config.fallback_on_detection_failure = fallback_on_detection_failure
+    config.adaptive_memory = adaptive_memory
+
     logger.debug(
         "Configuration resolved",
         workload_type=config.workload_type,
@@ -599,19 +605,11 @@ def setup_dask_client(
     if config.dashboard and config.dashboard_port:
         dashboard_address = f":{config.dashboard_port}"
 
-    # Map the boolean silence_logs config to a logging level.
-    # True  → suppress everything except errors (logging.ERROR)
-    # False → show warnings and above (logging.WARNING), which is a reasonable HPC default
-    import logging as _stdlib_logging
-
-    silence_logs_level = _stdlib_logging.ERROR if config.silence_logs else _stdlib_logging.WARNING
-
     cluster = create_cluster(
         topology=topology,
         memory_spec=memory_spec,
         temp_dir=temp_dir,
         dashboard_address=dashboard_address,
-        silence_logs=silence_logs_level,
         adaptive=config.adaptive,
         min_workers=config.min_workers,
         memory_target=config.memory_target,
@@ -640,10 +638,23 @@ def setup_dask_client(
     if config.dashboard:
         print_dashboard_info(client, silent=config.silence_logs)
 
-    # Log setup summary via structured logger
+    # Log and print setup summary
     spill_threads_str = (
         f" | spill_threads={config.spill_threads}" if config.spill_threads is not None else ""
     )
+    summary_lines = [
+        "[setup_dask_client] Configuration summary",
+        f"temp/spill dir: {temp_dir}",
+        f"Workers: {topology.n_workers} | threads/worker: {topology.threads_per_worker} | processes: {topology.processes}",
+        f"Memory: total ~{memory_spec.total_mem_gib:.1f} GiB | usable ~{memory_spec.usable_mem_gb:.1f} GiB | per-worker ~{memory_spec.mem_per_worker_bytes / (1024**3):.1f} GiB",
+        f"Compression: spill={config.spill_compression} | comm={config.comm_compression}{spill_threads_str}",
+    ]
+
+    # Print summary to console
+    for line in summary_lines:
+        print(line)
+
+    # Also log via structured logger
     logger.info(f"Temp/spill dir: {temp_dir}")
     logger.info(
         f"Workers: {topology.n_workers}"
@@ -677,7 +688,7 @@ def setup_dask_client(
                 workload_type=config.workload_type,
                 verbose=config.suggest_chunks,
             )
-            # recommend_chunks returns ChunkRecommendation when verbose=True, dict otherwise
+            # recommend_chunks always returns ChunkRecommendation
             chunk_recommendations = raw.chunks if hasattr(raw, "chunks") else raw
             logger.info("Chunk recommendations computed", chunks=str(chunk_recommendations))
 

{dask_setup-2.0.0 → dask_setup-2.1.0}/src/dask_setup/cluster.py

@@ -103,12 +103,10 @@ def configure_dask_settings(
     }
 
     # Add spill threads configuration if specified.
-    # "distributed.worker.io-threads" controls the size of each worker's I/O thread pool,
-    # which is used for spill read/write operations.
-    # Note: "distributed.p2p.threads" is a different setting (peer-to-peer shuffle threads)
-    # and is intentionally NOT used here.
+    # "distributed.p2p.threads" controls the number of threads used for
+    # peer-to-peer operations and spill I/O operations.
     if spill_threads is not None:
-        config_dict["distributed.worker.io-threads"] = spill_threads
+        config_dict["distributed.p2p.threads"] = spill_threads
 
     logger.debug(
         "Dask global settings applied",

{dask_setup-2.0.0 → dask_setup-2.1.0}/src/dask_setup/config.py

@@ -252,10 +252,6 @@ class DaskSetupConfig:
             "io_access_pattern": self.io_access_pattern,
             "io_storage_location": self.io_storage_location,
             "io_compression_level": self.io_compression_level,
-            # Resilience options
-            "fallback_on_detection_failure": self.fallback_on_detection_failure,
-            # Adaptive behaviour options
-            "adaptive_memory": self.adaptive_memory,
             # Metadata
             "name": self.name,
             "description": self.description,

{dask_setup-2.0.0 → dask_setup-2.1.0}/src/dask_setup/config_manager.py

@@ -202,7 +202,7 @@ Use `dask-setup profile create <name>` or manually create YAML files in the prof
     def get_profile(self, name: str) -> ConfigProfile | None:
         """Get a specific profile by name.
 
-        Search order: user profiles → site-wide profiles → builtin profiles.
+        Search order: builtin profiles → site-wide profiles → user profiles.
 
         Args:
             name: Profile name
@@ -210,13 +210,9 @@ Use `dask-setup profile create <name>` or manually create YAML files in the prof
         Returns:
             ConfigProfile if found, None otherwise
         """
-        # User profiles take highest precedence
-        profile_file = self.profiles_dir / f"{name}.yaml"
-        if profile_file.exists():
-            try:
-                return self.load_profile_from_file(profile_file)
-            except Exception as e:
-                raise InvalidConfigurationError(f"Failed to load profile '{name}': {e}") from e
+        # Built-in profiles take highest precedence
+        if name in self.builtin_profiles:
+            return self.builtin_profiles[name]
 
         # Site-wide profiles
         site_file = self.site_profiles_dir / f"{name}.yaml"
@@ -226,9 +222,13 @@ Use `dask-setup profile create <name>` or manually create YAML files in the prof
             except Exception as e:
                 raise InvalidConfigurationError(f"Failed to load site profile '{name}': {e}") from e
 
-        # Built-in profiles
-        if name in self.builtin_profiles:
-            return self.builtin_profiles[name]
+        # User profiles take lowest precedence
+        profile_file = self.profiles_dir / f"{name}.yaml"
+        if profile_file.exists():
+            try:
+                return self.load_profile_from_file(profile_file)
+            except Exception as e:
+                raise InvalidConfigurationError(f"Failed to load profile '{name}': {e}") from e
 
         return None
 
@@ -455,11 +455,11 @@ Use `dask-setup profile create <name>` or manually create YAML files in the prof
         import urllib.request
 
         try:
-            req = urllib.request.Request(
+            req = urllib.request.Request(  # noqa: S310
                 url,
                 headers={"User-Agent": f"dask_setup/{PROFILE_FORMAT_VERSION}"},
             )
-            with urllib.request.urlopen(req, timeout=30) as response:
+            with urllib.request.urlopen(req, timeout=30) as response:  # noqa: S310
                 content = response.read().decode("utf-8")
         except urllib.error.URLError as e:
             raise InvalidConfigurationError(f"Could not fetch profile from {url!r}: {e}") from e