extending tests

Author: TheEimer

hydra_plugins/hyper_smac/hyper_smac.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from copy import deepcopy
 from pathlib import Path
 
 import pandas as pd
@@ -10,6 +11,7 @@
 from omegaconf import OmegaConf
 from smac import Scenario
 from smac.facade import HyperparameterOptimizationFacade
+from smac.intensifier.hyperband import Hyperband
 from smac.runhistory.dataclasses import TrialInfo, TrialValue
 
 
@@ -43,12 +45,27 @@ class HyperSMACAdapter:
     def __init__(self, smac):
         """Initialize the adapter."""
         self.smac = smac
+        if isinstance(self.smac._intensifier, Hyperband):
+            self.hyperband = True
+            self.total_configs = 0
+            self.smac._intensifier.__post_init__()
+            self.n_configs_per_stage = deepcopy(self.smac._intensifier._n_configs_in_stage)
+            self.current_bracket = 0
+            self.total_stages = len(self.n_configs_per_stage.keys())
 
     def ask(self):
         """Ask for the next configuration."""
         smac_info = self.smac.ask()
         info = Info(config=smac_info.config, budget=smac_info.budget, load_path=None, seed=smac_info.seed)
-        return info, False, False
+        terminate = False
+        if self.hyperband:
+            self.total_configs += 1
+            if self.total_configs >= sum(self.n_configs_per_stage[self.current_bracket]):
+                self.current_bracket += 1
+                self.total_configs = 0
+                terminate = True
+        print(info, terminate)
+        return info, terminate, False
 
     def tell(self, info, value):
         """Tell the result of the configuration."""

hydra_plugins/hypersweeper/hypersweeper_backend.py

@@ -167,7 +167,7 @@ def sweep(self, arguments: list[str]) -> list | None:
             if isinstance(v, np.integer | np.floating):
                 v = v.item()  # noqa: PLW2901
             if isinstance(v, np.str_):
-                v = str(v)
+                v = str(v)  # noqa: PLW2901
             OmegaConf.update(final_config, k, v, force_add=True)
 
         with open(Path(optimizer.output_dir) / "final_config.yaml", "w+") as fp:

hydra_plugins/hypersweeper/hypersweeper_sweeper.py

@@ -462,6 +462,7 @@ def run(self, verbose=False):
             infos = []
             t = 0
             terminate = False
+            print(t < self.max_parallel, terminate, trial_termination, optimizer_termination)
             while t < self.max_parallel and not terminate and not trial_termination and not optimizer_termination:
                 if not any(b is None for b in self.history["budget"]) and self.budget is not None:
                     budget_termination = sum(self.history["budget"]) >= self.budget

tests/test_initialization.py

@@ -192,12 +192,16 @@ def test_init_without_n_trials(config):
     config["n_trials"] = None
     sweeper = HypersweeperSweeper(**config)
     assert sweeper.n_trials is None, "Number of trials should be set to None"
-    assert sweeper.max_parallel == config.get("job_array_size_limit", 100), (
-        "Max parallelization should match job array size limit without n_trials"
+    job_array_limit = config.get("job_array_size_limit", 100)
+    if "seeds" in config and config["seeds"] is not None:
+        job_array_limit = max(job_array_limit // len(config["seeds"]), 1)
+    assert sweeper.max_parallel == job_array_limit, (
+        f"Max parallelization should match job array size limit without n_trials (is {sweeper.max_parallel})"
     )
     shutil.rmtree("some_dir", ignore_errors=True)  # Clean up the base directory if it was created
     shutil.rmtree("checkpoints", ignore_errors=True)
 
+
 @mark.skipif(IN_GITHUB_ACTIONS, reason="Test doesn't work in Github Actions due to wandb logging.")
 @mark.parametrize("config", [TEST_CONFIG1, TEST_CONFIG2])
 def test_init_with_wandb(config):

tests/test_job_spawning.py

@@ -6,11 +6,12 @@
 import re
 import shutil
 import subprocess
-import pytest
 from pathlib import Path
 
+import pytest
+
 
-@pytest.mark.parametrize("max_parallel, job_limit", [(0.05, 10), (0.5, 9)])
+@pytest.mark.parametrize(("max_parallel", "job_limit"), [(0.05, 10), (0.5, 9)])
 def test_max_parallel(max_parallel, job_limit):
     if Path("branin_non_max").exists():
         shutil.rmtree(Path("branin_max_parallel"))
@@ -30,16 +31,21 @@ def test_max_parallel(max_parallel, job_limit):
     keyword = "Launching "
     all_keyword_indices = [m.start() for m in re.finditer(keyword, process_logs)]
     n_jobs_spawned = [int(process_logs[occ + len(keyword)]) for occ in all_keyword_indices]
-    n_times_spawned = process_logs.count(keyword) 
+    n_times_spawned = process_logs.count(keyword)
     total_jobs_spawned = sum(n_jobs_spawned)
 
-    assert total_jobs_spawned == job_limit, f"Total number of spawned jobs doesn't match the overall budget. Used: {total_jobs_spawned}, expected: {job_limit}."
+    assert total_jobs_spawned == job_limit, (
+        f"Total number of spawned jobs doesn't match the budget. Used: {total_jobs_spawned}, expected: {job_limit}."
+    )
     for i, n_jobs in enumerate(n_jobs_spawned):
-        assert n_jobs > 0, f"No jobs spawned in execution {i+1}/{n_times_spawned}."
-        assert n_jobs <= max_parallel*job_limit or n_jobs == 1, f"Number of spawned jobs exceeds the maximum parallelization limit in execution {i+1}/{n_times_spawned}: {n_jobs}."
+        assert n_jobs > 0, f"No jobs spawned in execution {i + 1}/{n_times_spawned}."
+        assert n_jobs <= max_parallel * job_limit or n_jobs == 1, (
+            f"Number of spawned jobs exceeds the maximum parallelization limit {i + 1}/{n_times_spawned}: {n_jobs}."
+        )
     shutil.rmtree(Path("branin_max_parallel"))
 
-@pytest.mark.parametrize("max_parallel, job_limit, seeds", [(0.05, 10, []), (0.5, 9, [1,2])])
+
+@pytest.mark.parametrize(("max_parallel", "job_limit", "seeds"), [(0.05, 10, []), (0.5, 9, [1, 2])])
 def test_max_parallel_with_seeds(max_parallel, job_limit, seeds):
     if Path("branin_non_max").exists():
         shutil.rmtree(Path("mlp_max_parallel_with_seeds"))
@@ -60,11 +66,15 @@ def test_max_parallel_with_seeds(max_parallel, job_limit, seeds):
     keyword = "Launching "
     all_keyword_indices = [m.start() for m in re.finditer(keyword, process_logs)]
     n_jobs_spawned = [int(process_logs[occ + len(keyword)]) for occ in all_keyword_indices]
-    n_times_spawned = process_logs.count(keyword) 
+    n_times_spawned = process_logs.count(keyword)
     total_jobs_spawned = sum(n_jobs_spawned)
 
-    assert total_jobs_spawned == job_limit*max(1, len(seeds)), f"Total number of spawned jobs doesn't match the overall budget. Used: {total_jobs_spawned}, expected: {job_limit*max(1, len(seeds))}."
+    assert total_jobs_spawned == job_limit * max(1, len(seeds)), (
+        f"Number of jobs doesn't match budget. Used: {total_jobs_spawned}, expected: {job_limit * max(1, len(seeds))}."
+    )
     for i, n_jobs in enumerate(n_jobs_spawned):
-        assert n_jobs > 0, f"No jobs spawned in execution {i+1}/{n_times_spawned}."
-        assert n_jobs <= max_parallel*job_limit or n_jobs == 1, f"Number of spawned jobs exceeds the maximum parallelization limit in execution {i+1}/{n_times_spawned}: {n_jobs}."
-    shutil.rmtree(Path("mlp_max_parallel_with_seeds"))
+        assert n_jobs > 0, f"No jobs spawned in execution {i + 1}/{n_times_spawned}."
+        assert n_jobs <= max_parallel * job_limit or n_jobs == 1, (
+            f"Number of spawned jobs exceeds the maximum parallelization limit {i + 1}/{n_times_spawned}: {n_jobs}."
+        )
+    shutil.rmtree(Path("mlp_max_parallel_with_seeds"))

tests/test_logging.py

@@ -5,6 +5,7 @@
 import shutil
 import subprocess
 from pathlib import Path
+
 import pandas as pd
 
 
@@ -45,7 +46,8 @@ def test_logfiles():
     assert "budget" in incumbent.columns, "Incumbent file missing expected column 'budget'."
     assert "config_id" in incumbent.columns, "Incumbent file missing expected column 'config_id'."
     assert "total_wallclock_time" in incumbent.columns, "Incumbent file missing expected column 'wall_clock_time'."
-    assert "total_optimization_time" in incumbent.columns, "Incumbent file missing expected column 'total_optimization_time'."
-
+    assert "total_optimization_time" in incumbent.columns, (
+        "Incumbent file missing expected column 'total_optimization_time'."
+    )
 
-    shutil.rmtree(Path("branin_logging"))
+    shutil.rmtree(Path("branin_logging"))

tests/test_runs/test_smac.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 import shutil
 import subprocess
 from pathlib import Path
@@ -27,3 +28,46 @@ def test_smac_branin_example():
     assert not runhistory.empty, "Run history is empty"
     assert len(runhistory) == 5, "Run history should contain 5 entries"
     shutil.rmtree(Path("branin_smac"))
+
+
+def test_smac_hyperband_termination():
+    if Path("mlp_smac_hyperband").exists():
+        shutil.rmtree(Path("mlp_smac_hyperband"))
+    process_logs = subprocess.check_output(
+        [
+            "python",
+            "examples/mlp.py",
+            "--config-name=mlp_smac",
+            "-m",
+            "hydra.sweeper.n_trials=5",
+            "hydra.run.dir=mlp_smac_hyperband",
+            "hydra.sweep.dir=mlp_smac_hyperband",
+            "hydra.sweeper.n_trials=24",
+            "+hydra.sweeper.sweeper_kwargs.max_parallelization=1",
+        ]
+    ).decode("utf-8")
+    assert Path("mlp_smac_hyperband").exists(), "Run directory not created"
+
+    # Goal: first launch 13 jobs in bracket 1, then 6 in bracket 2, then 3 in bracket 3, then 2 and terminate
+    keyword = "Launching "
+    all_keyword_indices = [m.start() for m in re.finditer(keyword, process_logs)]
+    n_jobs_spawned = [int(process_logs[occ + len(keyword)]) for occ in all_keyword_indices]
+    total_jobs_spawned = sum(n_jobs_spawned)
+
+    assert total_jobs_spawned == 24, (
+        f"Total number of spawned jobs doesn't match the overall budget. Used: {total_jobs_spawned}, expected: 24."
+    )
+    assert n_jobs_spawned[0] == 13, (
+        f"Number of spawned jobs in bracket 1 is incorrect. Used: {n_jobs_spawned[0]}, expected: 13."
+    )
+    assert n_jobs_spawned[1] == 6, (
+        f"Number of spawned jobs in bracket 2 is incorrect. Used: {n_jobs_spawned[1]}, expected: 6."
+    )
+    assert n_jobs_spawned[2] == 3, (
+        f"Number of spawned jobs in bracket 3 is incorrect. Used: {n_jobs_spawned[2]}, expected: 3."
+    )
+    assert n_jobs_spawned[3] == 2, (
+        f"Number of spawned jobs in bracket 4 before termination is incorrect. Used: {n_jobs_spawned[3]}, expected: 2."
+    )
+
+    shutil.rmtree(Path("mlp_smac_hyperband"))

tests/test_termination_conditions.py

@@ -6,9 +6,10 @@
 import re
 import shutil
 import subprocess
-import pytest
 from pathlib import Path
 
+import pytest
+
 
 @pytest.mark.parametrize("n_trials", [1, 7, 10])
 def test_terminate_n_trials(n_trials):
@@ -28,12 +29,15 @@ def test_terminate_n_trials(n_trials):
     assert Path("branin_trial_termination").exists(), "Run directory not created"
     keyword = "Launching "
     all_keyword_indices = [m.start() for m in re.finditer(keyword, process_logs)]
-    n_jobs_spawned = [int(process_logs[occ + len(keyword)]) for occ in all_keyword_indices] 
+    n_jobs_spawned = [int(process_logs[occ + len(keyword)]) for occ in all_keyword_indices]
     total_jobs_spawned = sum(n_jobs_spawned)
 
-    assert total_jobs_spawned == n_trials, f"Total number of spawned jobs doesn't match the n_trials. Used: {total_jobs_spawned}, expected: {n_trials}."
+    assert total_jobs_spawned == n_trials, (
+        f"Total number of spawned jobs doesn't match the n_trials. Used: {total_jobs_spawned}, expected: {n_trials}."
+    )
     shutil.rmtree(Path("branin_trial_termination"))
 
+
 @pytest.mark.parametrize("budget", [5, 10, 25])
 def test_terminate_budget(budget):
     if Path("mlp_budget_termination").exists():
@@ -52,6 +56,8 @@ def test_terminate_budget(budget):
     assert Path("mlp_budget_termination").exists(), "Run directory not created"
     keyword = "epochs="
     all_keyword_indices = [m.start() for m in re.finditer(keyword, process_logs)]
-    all_budgets = [int(process_logs[occ + len(keyword)]) for occ in all_keyword_indices] 
-    assert sum(all_budgets) <= budget, f"Total budget used exceeds the maximum budget. Used: {sum(all_budgets)}, budget: {budget}."
-    shutil.rmtree(Path("mlp_budget_termination"))
+    all_budgets = [int(process_logs[occ + len(keyword)]) for occ in all_keyword_indices]
+    assert sum(all_budgets) <= budget, (
+        f"Total budget used exceeds the maximum budget. Used: {sum(all_budgets)}, budget: {budget}."
+    )
+    shutil.rmtree(Path("mlp_budget_termination"))