Simplify and test container image names

amlb/runners/container.py

@@ -10,9 +10,11 @@
 from abc import abstractmethod
 import logging
 import re
+from typing import cast
 
 from ..benchmark import Benchmark, SetupMode
 from ..errors import InvalidStateError
+from ..frameworks.definitions import Framework
 from ..job import Job
 from ..resources import config as rconfig, get as rget
 from ..__version__ import __version__, _dev_version as dev
@@ -29,7 +31,8 @@ class ContainerBenchmark(Benchmark):
     framework_install_required = False
 
     @classmethod
-    def image_name(cls, framework_def, label=None, **kwargs):
+    def image_name(cls, framework_def: Framework, label: str | None = None) -> str:
+        """Determines the image name based on configuration data."""
         if label is None:
             label = rget().project_info.branch
 
@@ -61,8 +64,10 @@ def __init__(self, framework_name, benchmark_name, constraint_name):
         self.custom_commands = ""
         self.image = None
 
-    def _container_image_name(self, label=None):
-        return self.image_name(self.framework_def, label)
+    def _container_image_name(self, label: str | None = None) -> str:
+        return self.image_name(
+            cast(Framework, self.framework_def), label
+        )  # framework only None in AWSBenchmark
 
     def _validate(self):
         max_parallel_jobs = rconfig().job_scheduler.max_parallel_jobs

amlb/runners/singularity.py

@@ -6,12 +6,15 @@
 The image is pulled form an existing docker, yet executed in singularity framework
 """
 
+from __future__ import annotations
+
 import logging
 import os
-import re
+from typing import cast
 
 from ..benchmark import _setup_dir_
-from ..resources import config as rconfig, get as rget
+from ..frameworks.definitions import Framework
+from ..resources import config as rconfig
 from ..utils import dir_of, run_cmd, touch
 from .container import ContainerBenchmark
 
@@ -24,27 +27,6 @@ class SingularityBenchmark(ContainerBenchmark):
     an extension of ContainerBenchmark to run benchmarks inside Singularity.
     """
 
-    @classmethod
-    def image_name(cls, framework_def, label=None, as_docker_image=False, **kwargs):
-        """
-        We prefer to pull from docker, so we have to mind the docker tag
-        When downloading from Docker, the colon is changed to underscore
-        """
-        if label is None:
-            label = rget().project_info.branch
-        di = framework_def.image
-
-        # If we want to pull from docker, the separator is a colon for tag
-        separator = "_" if not as_docker_image else ":"
-        # Also, no need for author in image name
-        author = "" if not as_docker_image else f"{di.author}/"
-        image = di.image if di.image else framework_def.name.lower()
-        tags = [di.tag if di.tag else framework_def.version.lower()]
-        if label not in rconfig().container.ignore_labels:
-            tags.append(label)
-        tag = re.sub(r"([^\w.-])", ".", "-".join(tags))
-        return f"{author}{image}{separator}{tag}"
-
     def __init__(self, framework_name, benchmark_name, constraint_name):
         """
 
@@ -65,19 +47,21 @@ def __init__(self, framework_name, benchmark_name, constraint_name):
             else ""
         )
 
-    def _container_image_name(self, label=None, as_docker_image=False):
+    def _container_image_name(
+        self, label: str | None = None, as_docker_image: bool = False
+    ) -> str:
         """
         Singularity Images would be located on the framework directory
         """
         image_name = self.image_name(
-            self.framework_def, label=label, as_docker_image=as_docker_image
-        )
-
-        # Make sure image is in the framework directory
+            cast(Framework, self.framework_def), label=label
+        )  # Framework only none in AWSBench
         if as_docker_image:
             return image_name
-        else:
-            return os.path.join(self._framework_dir, _setup_dir_, image_name + ".sif")
+
+        author, image_name = image_name.split("/", maxsplit=1)
+        image_name = image_name.replace(":", "_")
+        return os.path.join(self._framework_dir, _setup_dir_, image_name + ".sif")
 
     @property
     def _script(self):

tests/conftest.py

@@ -1,14 +1,37 @@
-from pathlib import Path
-from typing import Generator
+import os
 
 import pytest
-from amlb import Resources
-from amlb.utils import Namespace
+from amlb import Resources, resources
+from amlb.defaults import default_dirs
+from amlb.utils import Namespace, config_load
 
 
-@pytest.fixture(autouse=True)
-def tmp_output_directory(tmp_path: Path) -> Generator[Path, None, None]:
-    yield tmp_path
+@pytest.fixture
+def load_default_resources(tmp_path):
+    config_default = config_load(
+        os.path.join(default_dirs.root_dir, "resources", "config.yaml")
+    )
+    config_default_dirs = default_dirs
+    # allowing config override from user_dir: useful to define custom benchmarks and frameworks for example.
+    config_user = Namespace()
+    # config listing properties set by command line
+    config_args = Namespace.parse(
+        {"results.global_save": False},
+        output_dir=str(tmp_path),
+        script=os.path.basename(__file__),
+        run_mode="local",
+        parallel_jobs=1,
+        sid="pytest.session",
+        exit_on_error=True,
+        test_server=False,
+        tag=None,
+        command="pytest invocation",
+    )
+    config_args = Namespace({k: v for k, v in config_args if v is not None})
+    # merging all configuration files and saving to the global variable
+    resources.from_configs(
+        config_default, config_default_dirs, config_user, config_args
+    )
 
 
 @pytest.fixture

tests/unit/amlb/benchmarks/test_benchmark.py

@@ -1,37 +1,12 @@
-from amlb import Benchmark, SetupMode, resources
-import os
-
-from amlb.defaults import default_dirs
-from amlb.utils import config_load
-from amlb.utils import Namespace as ns
-from tests.conftest import tmp_output_directory
-
-
-def test_benchmark(tmp_path) -> None:
-    config_default = config_load(
-        os.path.join(default_dirs.root_dir, "resources", "config.yaml")
-    )
-    config_default_dirs = default_dirs
-    # allowing config override from user_dir: useful to define custom benchmarks and frameworks for example.
-    config_user = ns()
-    # config listing properties set by command line
-    config_args = ns.parse(
-        {"results.global_save": False},
-        output_dir=str(tmp_output_directory),
-        script=os.path.basename(__file__),
-        run_mode="local",
-        parallel_jobs=1,
-        sid="pytest.session",
-        exit_on_error=True,
-        test_server=False,
-        tag=None,
-        command="pytest invocation",
-    )
-    config_args = ns({k: v for k, v in config_args if v is not None})
-    # merging all configuration files and saving to the global variable
-    resources.from_configs(
-        config_default, config_default_dirs, config_user, config_args
-    )
+from pathlib import Path
+
+import pytest
+
+from amlb import Benchmark, SetupMode, resources, DockerBenchmark, SingularityBenchmark
+from amlb.utils import Namespace
+
+
+def test_benchmark(load_default_resources) -> None:
     benchmark = Benchmark(
         framework_name="constantpredictor",
         benchmark_name="test",
@@ -42,3 +17,98 @@ def test_benchmark(tmp_path) -> None:
     results = benchmark.run()
     assert len(results) == 6
     assert not results["result"].isna().any()
+
+
+@pytest.mark.parametrize(
+    ("framework_name", "tag", "expected"),
+    [
+        ("constantpredictor", "latest", "automlbenchmark/constantpredictor:stable"),
+        ("flaml", "2023Q2", "automlbenchmark/flaml:1.2.4"),
+        ("autosklearn", "stable", "automlbenchmark/autosklearn:stable"),
+    ],
+)
+def test_docker_image_name(
+    framework_name, tag, expected, load_default_resources
+) -> None:
+    framework_def, _ = resources.get().framework_definition(
+        framework_name,
+        tag=tag,
+    )
+    # The docker image name is based entirely on configuration (i.e. configured branch, not checked out branch).
+    # If there is a different branch checked out locally and you run the benchmark,
+    # you will get a prompt to verify you want to build anyway
+    result = DockerBenchmark.image_name(
+        framework_def,
+    )
+    assert result == expected
+
+
+@pytest.mark.parametrize("branch", ["master", "foo"])
+def test_docker_image_name_uses_branch(branch, mocker, load_default_resources) -> None:
+    mocker.patch(
+        "amlb.runners.container.rget",
+        return_value=Namespace(project_info=Namespace(branch=branch)),
+    )
+    framework_def, _ = resources.get().framework_definition("constantpredictor")
+    result = DockerBenchmark.image_name(framework_def)
+    assert result == f"automlbenchmark/constantpredictor:stable-{branch}"
+
+
+@pytest.mark.parametrize("label", [None, "master", "foo"])
+def test_docker_image_name_uses_label(label, mocker, load_default_resources) -> None:
+    branch = "bar-branch"
+    mocker.patch(
+        "amlb.runners.container.rget",
+        return_value=Namespace(project_info=Namespace(branch=branch)),
+    )
+
+    framework_def, _ = resources.get().framework_definition("constantpredictor")
+    result = DockerBenchmark.image_name(framework_def, label=label)
+
+    used_label = label or branch
+    assert result == f"automlbenchmark/constantpredictor:stable-{used_label}"
+
+
+@pytest.mark.parametrize(
+    ("framework_name", "tag", "expected"),
+    [
+        ("constantpredictor", "latest", "constantpredictor_stable"),
+        ("flaml", "2023Q2", "flaml_1.2.4"),
+        ("autosklearn", "stable", "autosklearn_stable"),
+    ],
+)
+def test_singularity_image_name(
+    framework_name, tag, expected, load_default_resources
+) -> None:
+    benchmark = SingularityBenchmark(
+        framework_name=f"{framework_name}:{tag}",
+        benchmark_name="test",
+        constraint_name="test",
+    )
+    image_path = benchmark._container_image_name(
+        as_docker_image=False,
+    )
+    image_name = Path(image_path).stem
+    assert image_name == expected
+
+
+@pytest.mark.parametrize(
+    ("framework_name", "tag", "expected"),
+    [
+        ("constantpredictor", "latest", "automlbenchmark/constantpredictor:stable"),
+        ("flaml", "2023Q2", "automlbenchmark/flaml:1.2.4"),
+        ("autosklearn", "stable", "automlbenchmark/autosklearn:stable"),
+    ],
+)
+def test_singularity_image_name_as_docker(
+    framework_name, tag, expected, load_default_resources
+) -> None:
+    benchmark = SingularityBenchmark(
+        framework_name=f"{framework_name}:{tag}",
+        benchmark_name="test",
+        constraint_name="test",
+    )
+    result = benchmark._container_image_name(
+        as_docker_image=True,
+    )
+    assert result == expected