second round of PR feedback

Author: robin-janssen

.github/workflows/ci.yaml

@@ -13,7 +13,6 @@ on:
       - opened
       - synchronize
       - reopened
-      - closed
 
 permissions:
   contents: write
@@ -172,4 +171,4 @@ jobs:
           publish_branch: gh-pages
           user_name: "GitHub Actions"
           user_email: "actions@github.com"
-          force_orphan: false
+          force_orphan: false

codes/benchmark/bench_fcts.py

@@ -103,7 +103,7 @@ def run_benchmark(surr_name: str, surrogate_class, conf: dict) -> dict[str, Any]
         log=dataset_conf.get("log10_transform", True),
         log_params=dataset_conf.get("log10_transform_params", True),
         normalisation_mode=dataset_conf.get("normalise", "minmax"),
-        tolerance=dataset_conf.get("tolerance", 1e-25),
+        tolerance=dataset_conf.get("tolerance", None),
         per_species=dataset_conf.get("normalise_per_species", False),
     )
 

codes/tune/optuna_fcts.py

@@ -258,7 +258,7 @@ def training_run(
         log=dataset_cfg.get("log10_transform", True),
         log_params=dataset_cfg.get("log10_transform_params", True),
         normalisation_mode=dataset_cfg.get("normalise", "minmax"),
-        tolerance=dataset_cfg.get("tolerance", 1e-25),
+        tolerance=dataset_cfg.get("tolerance", None),
         per_species=dataset_cfg.get("normalise_per_species", False),
     )
 

codes/tune/tune_utils.py

@@ -96,7 +96,7 @@ def apply_tuning_defaults(config: dict) -> dict:
     dataset.setdefault("log10_transform_params", True)
     dataset.setdefault("normalise", "minmax")
     dataset.setdefault("normalise_per_species", False)
-    dataset.setdefault("tolerance", 1e-25)
+    dataset.setdefault("tolerance", None)
     dataset.setdefault("subset_factor", 1)
     dataset.setdefault("log_timesteps", False)
     cfg["dataset"] = dataset

docs/source/reference/configuration.md

@@ -91,7 +91,7 @@ dataset:
   log10_transform_params: true
   normalise: "minmax"
   normalise_per_species: false
-  tolerance: 1e-25
+  tolerance: null
   subset_factor: 1
   log_timesteps: false
   use_optimal_params: true

docs/source/tutorials/index.md

@@ -11,7 +11,7 @@ Interactive tutorials live in this folder so they can be rendered directly on th
 To execute a notebook locally:
 
 ```bash
-pip install -r requirements.txt  # ensures docs dependencies
+uv pip install --group dev  # installs docs + notebook deps
 jupyter lab docs/source/tutorials/benchmark_quickstart.ipynb
 ```
 

run_training.py

@@ -35,8 +35,10 @@ def main(args):
     torch.use_deterministic_algorithms(True)
     config = load_and_save_config(config_path=args.config, save=False)
 
-    if torch.cuda.is_available():
-        print(torch.cuda.device_count(), torch.cuda.current_device())
+    if torch.cuda.is_available() and config.get("verbose", False):
+        nice_print(
+            f"CUDA devices detected: count={torch.cuda.device_count()}, current={torch.cuda.current_device()}"
+        )
 
     download_data(config["dataset"]["name"], verbose=config.get("verbose", False))
     task_list_filepath, copy_config = check_training_status(config)

run_tuning.py

@@ -56,7 +56,9 @@ def run_single_study(config: dict, study_name: str, db_url: str, sqlite_backend:
     if not config.get("optuna_logging", False):
         optuna.logging.set_verbosity(optuna.logging.WARNING)
 
-    if config.get("multi_objective", False):
+    multi_objective = config.get("multi_objective", False)
+
+    if multi_objective:
         sampler = optuna.samplers.NSGAIISampler(
             seed=config["seed"], population_size=config["population_size"]
         )
@@ -151,7 +153,7 @@ def trial_complete_callback(study_: optuna.Study, trial_: optuna.trial.FrozenTri
         try:
             study.optimize(
                 objective_fn,
-                n_trials=n_trials * 2 if config["multi_objective"] else n_trials,
+                n_trials=n_trials,
                 n_jobs=n_jobs,
                 callbacks=[
                     MaxValidTrialsCallback(n_trials),

test/test_bench_main.py

@@ -147,18 +147,47 @@ def test_time_inference(monkeypatch):
     assert result["mean_inference_time_per_prediction"] == pytest.approx(3.0 / 2)
 
 
-def test_evaluate_compute(monkeypatch):
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA unavailable, skipping measure_memory_footprint test.")
+def test_evaluate_compute_cpu_path(monkeypatch):
+    # Force non-CUDA execution
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
 
-    # patch memory footprint and parameter count
+    # measure_memory_footprint should not be called
+    called = {"mem": False}
+
+    def fake_measure(model, inputs, device):
+        called["mem"] = True
+        return {}, model
+
+    monkeypatch.setattr(bf, "measure_memory_footprint", fake_measure)
+    monkeypatch.setattr(bf, "count_trainable_parameters", lambda m: 12345)
+
+    class DummyLoader:
+        def __iter__(self):
+            yield ("inp",)
+
+    loader = DummyLoader()
+    model = FakeModel()
+    surr = "SurrB"
+    conf = {"training_id": "TID"}
+    out = bf.evaluate_compute(model, surr, test_loader=loader, conf=conf)
+    assert model.load_calls == [("TID", surr, f"{surr.lower()}_main")]
+    assert out["num_trainable_parameters"] == 12345
+    assert out["memory_footprint"] == {}
+    assert not called["mem"]
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA unavailable for memory profiling test."
+)
+def test_evaluate_compute_cuda_path(monkeypatch):
     fake_mem = {"model_memory": 100, "forward_memory_nograd": 50}
-    monkeypatch.setattr(
-        bf, "measure_memory_footprint", lambda m, inp, device: (fake_mem, m)
-    )
+
+    def fake_measure(model, inputs, device):
+        return fake_mem, model
+
+    monkeypatch.setattr(bf, "measure_memory_footprint", fake_measure)
     monkeypatch.setattr(bf, "count_trainable_parameters", lambda m: 12345)
 
-    # test_loader yields one tuple of inputs
     class DummyLoader:
         def __iter__(self):
             yield ("inp",)
@@ -168,7 +197,6 @@ def __iter__(self):
     surr = "SurrB"
     conf = {"training_id": "TID"}
     out = bf.evaluate_compute(model, surr, test_loader=loader, conf=conf)
-    # load main was invoked
     assert model.load_calls == [("TID", surr, f"{surr.lower()}_main")]
     assert out["num_trainable_parameters"] == 12345
     assert out["memory_footprint"] is fake_mem