Fix tune trial timeout handling

Author: Donglai Wei

connectomics/decoding/tuning/optuna_tuner.py

@@ -27,6 +27,7 @@
 
 import h5py
 import numpy as np
+import torch
 from omegaconf import DictConfig, OmegaConf
 
 try:
@@ -335,8 +336,14 @@ def _trial_evaluation_worker(send_conn, evaluation_kind: str, payload: Dict[str,
 
 
 def _get_trial_process_context() -> mp.context.BaseContext:
-    """Prefer fork to avoid copying large prediction arrays when available."""
-    for method in ("fork", "spawn"):
+    """Choose a multiprocessing start method for timeout-enforced trials."""
+    methods = ("fork", "spawn")
+    if torch.cuda.is_available() and torch.cuda.is_initialized():
+        # Tune mode runs inference before Optuna; once CUDA is initialized,
+        # forking the parent process is unsafe and can hang.
+        methods = ("spawn", "fork")
+
+    for method in methods:
         try:
             return mp.get_context(method)
         except ValueError:
@@ -795,6 +802,12 @@ def optimize(self) -> optuna.Study:
 
         metric = self.tune_cfg.optimization["single_objective"]["metric"]
         trial_timeout = self._get_trial_timeout_seconds()
+        if timeout is not None and trial_timeout is None:
+            logger.warning(
+                "tune.timeout=%s limits the whole Optuna study, not one trial. "
+                "Long WaterZ runs will still block unless tune.trial_timeout is set.",
+                timeout,
+            )
         logger.info(
             "Starting Optuna optimization: %s | Trials: %s | Metric: %s | "
             "Direction: %s | Trial timeout: %s",

connectomics/training/lightning/utils.py

@@ -157,6 +157,18 @@ def parse_args():
         default=None,
         help="Number of Optuna trials (overrides config, use with --mode tune or tune-test)",
     )
+    tune_group.add_argument(
+        "--tune-timeout",
+        type=int,
+        default=None,
+        help="Whole-study Optuna timeout in seconds (overrides tune.timeout)",
+    )
+    tune_group.add_argument(
+        "--tune-trial-timeout",
+        type=int,
+        default=None,
+        help="Per-trial tuning timeout in seconds (overrides tune.trial_timeout)",
+    )
     parser.add_argument(
         "overrides",
         nargs="*",
@@ -217,6 +229,28 @@ def setup_config(args) -> Config:
     # Resolve data paths on merged runtime data section.
     cfg = resolve_data_paths(cfg)
 
+    if cfg.tune is not None:
+        if args.tune_trials is not None:
+            logger.info("Overriding tune.n_trials: %s -> %s", cfg.tune.n_trials, args.tune_trials)
+            cfg.tune.n_trials = args.tune_trials
+
+        if args.tune_timeout is not None:
+            logger.info("Overriding tune.timeout: %s -> %s", cfg.tune.timeout, args.tune_timeout)
+            cfg.tune.timeout = args.tune_timeout
+
+        if args.tune_trial_timeout is not None:
+            logger.info(
+                "Overriding tune.trial_timeout: %s -> %s",
+                cfg.tune.trial_timeout,
+                args.tune_trial_timeout,
+            )
+            cfg.tune.trial_timeout = args.tune_trial_timeout
+    elif any(
+        value is not None
+        for value in (args.tune_trials, args.tune_timeout, args.tune_trial_timeout)
+    ):
+        logger.warning("Ignoring --tune-* CLI overrides because the config has no tune section")
+
     # Override max_epochs if --reset-max-epochs is specified
     if args.reset_max_epochs is not None:
         logger.info(

tests/unit/test_lit_utils.py

@@ -5,6 +5,7 @@
 import pytest
 
 from connectomics.config import Config, save_config
+from connectomics.config.schema.stages import TuneConfig
 from connectomics.training.lightning.data_factory import _calculate_validation_steps_per_epoch
 from connectomics.training.lightning.path_utils import (
     expand_file_paths as canonical_expand_file_paths,
@@ -27,6 +28,9 @@ def _make_args(
     fast_dev_run: int = 0,
     mode: str = "train",
     nnunet_preprocess: bool = False,
+    tune_timeout: int | None = None,
+    tune_trial_timeout: int | None = None,
+    tune_trials: int | None = None,
 ):
     return argparse.Namespace(
         config=str(config_path),
@@ -42,7 +46,9 @@ def _make_args(
         external_prefix=None,
         params=None,
         param_source=None,
-        tune_trials=None,
+        tune_trials=tune_trials,
+        tune_timeout=tune_timeout,
+        tune_trial_timeout=tune_trial_timeout,
         nnunet_preprocess=nnunet_preprocess,
         overrides=overrides or [],
     )
@@ -84,6 +90,27 @@ def test_setup_config_enables_nnunet_preprocess_from_cli_switch(tmp_path):
     assert updated.data.nnunet_preprocessing.enabled is True
 
 
+def test_setup_config_applies_tune_timeout_cli_overrides(tmp_path):
+    cfg = Config()
+    cfg.tune = TuneConfig()
+
+    cfg_path = tmp_path / "config.yaml"
+    save_config(cfg, cfg_path)
+
+    args = _make_args(
+        cfg_path,
+        mode="tune",
+        tune_trials=17,
+        tune_timeout=3600,
+        tune_trial_timeout=300,
+    )
+    updated = setup_config(args)
+
+    assert updated.tune.n_trials == 17
+    assert updated.tune.timeout == 3600
+    assert updated.tune.trial_timeout == 300
+
+
 def test_expand_file_paths_handles_globs_and_lists(tmp_path):
     data_dir = tmp_path / "data"
     data_dir.mkdir()

tests/unit/test_main_cli_contract.py

@@ -59,3 +59,13 @@ def test_parse_args_demo_mode_requires_no_config(monkeypatch):
     args = _parse_with_argv(monkeypatch, ["--demo"])
     assert args.demo is True
     assert args.config is None
+
+
+def test_parse_args_accepts_tune_timeout_flags(monkeypatch):
+    args = _parse_with_argv(
+        monkeypatch,
+        ["--tune-timeout", "3600", "--tune-trial-timeout", "300"],
+    )
+
+    assert args.tune_timeout == 3600
+    assert args.tune_trial_timeout == 300

tests/unit/test_main_runtime_stage_switch.py

@@ -9,8 +9,8 @@
 from scripts.main import (
     _is_test_evaluation_enabled,
     has_assigned_test_shard,
-    maybe_limit_test_devices,
     maybe_enable_independent_test_sharding,
+    maybe_limit_test_devices,
     resolve_test_stage_runtime,
 )
 
@@ -32,6 +32,8 @@ def _make_args(config_path: Path, mode: str = "test"):
         params=None,
         param_source=None,
         tune_trials=None,
+        tune_timeout=None,
+        tune_trial_timeout=None,
         nnunet_preprocess=False,
         overrides=[],
         shard_id=None,
@@ -163,9 +165,7 @@ def test_maybe_enable_independent_test_sharding_uses_explicit_shard_args(tmp_pat
     assert cfg.system.num_gpus == (1 if torch.cuda.is_available() else 0)
 
 
-def test_maybe_enable_independent_test_sharding_skips_single_volume_tests(
-    tmp_path, monkeypatch
-):
+def test_maybe_enable_independent_test_sharding_skips_single_volume_tests(tmp_path, monkeypatch):
     cfg = Config()
     cfg.system.num_gpus = 4
     args = _make_args(tmp_path / "config.yaml")
@@ -182,9 +182,7 @@ def test_maybe_enable_independent_test_sharding_skips_single_volume_tests(
     assert cfg.system.num_gpus == 4
 
 
-def test_has_assigned_test_shard_returns_false_for_empty_slice(
-    tmp_path, monkeypatch
-):
+def test_has_assigned_test_shard_returns_false_for_empty_slice(tmp_path, monkeypatch):
     args = _make_args(tmp_path / "config.yaml")
     cfg = Config()
     args.shard_id = 3

tests/unit/test_optuna_tuner.py

@@ -10,6 +10,7 @@
 from connectomics.decoding.tuning.optuna_tuner import (
     OptunaDecodingTuner,
     TrialEvaluationTimeoutError,
+    _get_trial_process_context,
     load_and_apply_best_params,
     run_tuning,
 )
@@ -361,3 +362,63 @@ def _raise_timeout(_evaluation_kind, _payload):
     assert trial.user_attrs["timed_out"] is True
     assert trial.user_attrs["timeout_stage"] == "waterz_batch"
     assert trial.user_attrs["trial_timeout"] == 30.0
+
+
+def test_get_trial_process_context_prefers_spawn_after_cuda_init(monkeypatch):
+    observed = []
+
+    class _DummyContext:
+        pass
+
+    def _fake_get_context(method=None):
+        observed.append(method)
+        if method == "spawn":
+            return _DummyContext()
+        raise ValueError(f"unsupported: {method}")
+
+    monkeypatch.setattr(
+        "connectomics.decoding.tuning.optuna_tuner.torch.cuda.is_available",
+        lambda: True,
+    )
+    monkeypatch.setattr(
+        "connectomics.decoding.tuning.optuna_tuner.torch.cuda.is_initialized",
+        lambda: True,
+    )
+    monkeypatch.setattr(
+        "connectomics.decoding.tuning.optuna_tuner.mp.get_context", _fake_get_context
+    )
+
+    ctx = _get_trial_process_context()
+
+    assert isinstance(ctx, _DummyContext)
+    assert observed == ["spawn"]
+
+
+def test_get_trial_process_context_prefers_fork_without_cuda_init(monkeypatch):
+    observed = []
+
+    class _DummyContext:
+        pass
+
+    def _fake_get_context(method=None):
+        observed.append(method)
+        if method == "fork":
+            return _DummyContext()
+        raise ValueError(f"unsupported: {method}")
+
+    monkeypatch.setattr(
+        "connectomics.decoding.tuning.optuna_tuner.torch.cuda.is_available",
+        lambda: False,
+    )
+    monkeypatch.setattr(
+        "connectomics.decoding.tuning.optuna_tuner.torch.cuda.is_initialized",
+        lambda: False,
+    )
+    monkeypatch.setattr(
+        "connectomics.decoding.tuning.optuna_tuner.mp.get_context", _fake_get_context
+    )
+
+    ctx = _get_trial_process_context()
+
+    assert isinstance(ctx, _DummyContext)
+    assert observed == ["fork"]