slurm cpu

Author: michielkallenberg

cybench/runs/run_experiments.py

@@ -21,7 +21,13 @@
 from cybench.datasets.dataset import PandasDataset
 from cybench.evaluation.eval import evaluate_predictions
 from cybench.evaluation.aggregated_metrics import compute_report_metrics, format_report_metrics
-from cybench.util.config_utils import adjust_model_cfg_to_dataset, set_seed, remove_search_keys
+from cybench.util.config_utils import (
+    adjust_model_cfg_to_dataset,
+    apply_force_cpu_to_frozen_model_cfg,
+    remove_search_keys,
+    set_seed,
+    walk_forward_force_cpu,
+)
 from cybench.util.optuna_hyper_opt import OptunaOptimizer
 from cybench.util.store_and_cache import make_folder, save_preds, save_meta_dict
 from cybench.util.feature_selection import (
@@ -218,6 +224,14 @@ def main(cfg):
         fs_cfg: DictConfig | None = None
         if is_walk_forward:
             model_cfg = cast(DictConfig, OmegaConf.create(OmegaConf.to_container(frozen_model_cfg)))
+            if walk_forward_force_cpu(cfg):
+                prev_device = OmegaConf.select(model_cfg, "device")
+                model_cfg = apply_force_cpu_to_frozen_model_cfg(model_cfg)
+                if prev_device and prev_device != "cpu":
+                    log.info(
+                        "Walk-forward force CPU: frozen model device %s → cpu",
+                        prev_device,
+                    )
             fs_cfg = (
                 cast(DictConfig, OmegaConf.create(OmegaConf.to_container(frozen_fs_cfg)))
                 if frozen_fs_cfg is not None

cybench/util/config_utils.py

@@ -52,10 +52,46 @@ def _clean(obj):
 
     return cast(DictConfig, OmegaConf.create(_clean(cfg_dict)))
 
+
 def remove_search_keys(model_cfg: DictConfig) -> DictConfig:
     return remove_keys(model_cfg, key="_search_")
 
 
+def is_cybench_force_cpu() -> bool:
+    """True when Slurm submit_array.sh --cpu exported CYBENCH_FORCE_CPU=1."""
+    val = os.environ.get("CYBENCH_FORCE_CPU", "").strip().lower()
+    return val in ("1", "true", "yes", "on")
+
+
+def walk_forward_force_cpu(cfg: DictConfig) -> bool:
+    """Whether walk-forward should override device in frozen optimal_model.yaml."""
+    if is_cybench_force_cpu():
+        return True
+    if OmegaConf.select(cfg, "model.device") == "cpu":
+        return True
+    return OmegaConf.select(cfg, "experiment.device") == "cpu"
+
+
+def apply_force_cpu_to_frozen_model_cfg(model_cfg: DictConfig) -> DictConfig:
+    """
+    Override device fields baked into screening optimal_model.yaml for CPU runs.
+
+    Walk-forward loads the frozen config wholesale; without this, TabDPT and other
+    GPU-screened models keep device=cuda even on the main partition.
+    """
+    cfg_out = cast(DictConfig, OmegaConf.create(OmegaConf.to_container(model_cfg)))
+    with open_dict(cfg_out):
+        if "device" in cfg_out:
+            cfg_out.device = "cpu"
+        if "allow_cpu_fallback" in cfg_out:
+            cfg_out.allow_cpu_fallback = True
+        torch_model = cfg_out.get("torch_model")
+        if torch_model is not None and "device" in torch_model:
+            with open_dict(torch_model):
+                torch_model.device = "cpu"
+    return cfg_out
+
+
 def reload_config_with_overrides(
         config_dir: Path,
         config_name: str,
@@ -74,7 +110,6 @@ def reload_config_with_overrides(
     return cfg
 
 
-
 def get_run_description(overrides_path: pathlib.Path) -> str:
     """Loads and formats the list of Hydra overrides into a unique string."""
     try:

tests/util/test_config_utils.py

@@ -8,7 +8,12 @@
 from hydra import compose, initialize_config_dir
 from omegaconf import OmegaConf
 
-from cybench.util.config_utils import adjust_model_cfg_to_dataset
+from cybench.util.config_utils import (
+    adjust_model_cfg_to_dataset,
+    apply_force_cpu_to_frozen_model_cfg,
+    is_cybench_force_cpu,
+    walk_forward_force_cpu,
+)
 
 
 class _FakeTorchDataset:
@@ -45,3 +50,49 @@ def test_adjust_model_cfg_sets_missing_placeholder_dims():
     assert adjusted.torch_model.temporal_in_dim == 12
     assert adjusted.torch_model.context_in_dim == 5
     assert not OmegaConf.is_missing(adjusted.torch_model, "temporal_in_dim")
+
+
+def test_is_cybench_force_cpu(monkeypatch):
+    monkeypatch.delenv("CYBENCH_FORCE_CPU", raising=False)
+    assert is_cybench_force_cpu() is False
+    monkeypatch.setenv("CYBENCH_FORCE_CPU", "1")
+    assert is_cybench_force_cpu() is True
+    monkeypatch.setenv("CYBENCH_FORCE_CPU", "yes")
+    assert is_cybench_force_cpu() is True
+
+
+def test_walk_forward_force_cpu_from_hydra_overrides(monkeypatch):
+    monkeypatch.delenv("CYBENCH_FORCE_CPU", raising=False)
+    cfg = OmegaConf.create({"model": {"device": "cpu"}, "experiment": {"device": "cuda"}})
+    assert walk_forward_force_cpu(cfg) is True
+    cfg = OmegaConf.create({"model": {"device": "auto"}, "experiment": {"device": "cpu"}})
+    assert walk_forward_force_cpu(cfg) is True
+    cfg = OmegaConf.create({"model": {"device": "auto"}, "experiment": {"device": "cuda"}})
+    assert walk_forward_force_cpu(cfg) is False
+
+
+def test_apply_force_cpu_to_frozen_model_cfg():
+    frozen = OmegaConf.create(
+        {
+            "_target_": "cybench.models.tabular_foundation_model.TabDPTModel",
+            "name": "tabdpt",
+            "device": "cuda",
+            "allow_cpu_fallback": False,
+        }
+    )
+    out = apply_force_cpu_to_frozen_model_cfg(frozen)
+    assert out.device == "cpu"
+    assert out.allow_cpu_fallback is True
+
+
+def test_apply_force_cpu_to_frozen_torch_model_cfg():
+    frozen = OmegaConf.create(
+        {
+            "framework": "torch",
+            "device": "cuda",
+            "torch_model": {"input_size": 6, "device": "cuda"},
+        }
+    )
+    out = apply_force_cpu_to_frozen_model_cfg(frozen)
+    assert out.device == "cpu"
+    assert out.torch_model.device == "cpu"