From 1b2818da3944786c4c51a943ec6be64851544b09 Mon Sep 17 00:00:00 2001 From: PGijsbers Date: Fri, 29 Nov 2024 11:42:55 +0200 Subject: [PATCH] Simplify code Singularity custom names --- amlb/runners/container.py | 11 +- amlb/runners/singularity.py | 123 ++++++++++--------- tests/unit/amlb/benchmarks/test_benchmark.py | 25 ++-- 3 files changed, 85 insertions(+), 74 deletions(-) diff --git a/amlb/runners/container.py b/amlb/runners/container.py index dcbac638f..9a7e1a798 100644 --- a/amlb/runners/container.py +++ b/amlb/runners/container.py @@ -10,10 +10,11 @@ from abc import abstractmethod import logging import re -from typing import List, Union +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 @@ -30,7 +31,7 @@ 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 @@ -63,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 diff --git a/amlb/runners/singularity.py b/amlb/runners/singularity.py index 2bb95963b..2af2470fb 100644 --- a/amlb/runners/singularity.py +++ b/amlb/runners/singularity.py @@ -5,12 +5,16 @@ providing the same parameters and features allowing to import config and export results through mounted folders. 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 @@ -23,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): - """ - 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): """ @@ -54,28 +37,35 @@ def __init__(self, framework_name, benchmark_name, constraint_name): super().__init__(framework_name, benchmark_name, constraint_name) self._custom_image_name = rconfig().singularity.image self.minimize_instances = rconfig().singularity.minimize_instances - self.container_name = 'singularity' + self.container_name = "singularity" self.force_branch = rconfig().singularity.force_branch - self.custom_commands = self.framework_module.singularity_commands( - self.framework_def.setup_args, - setup_cmd=self.framework_def._setup_cmd - ) if hasattr(self.framework_module, 'singularity_commands') else "" + self.custom_commands = ( + self.framework_module.singularity_commands( + self.framework_def.setup_args, setup_cmd=self.framework_def._setup_cmd + ) + if hasattr(self.framework_module, "singularity_commands") + 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 + image_name = self.image_name( + 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): - return os.path.join(self._framework_dir, _setup_dir_, 'Singularityfile') + return os.path.join(self._framework_dir, _setup_dir_, "Singularityfile") def _start_container(self, script_params=""): """Implementes the container run method""" @@ -88,7 +78,7 @@ def _start_container(self, script_params=""): cmd = ( "singularity run --pwd /bench {options} " "-B {input}:/input -B {output}:/output -B {custom}:/custom " - "{image} \"{params} -i /input -o /output -u /custom -s skip -Xrun_mode=singularity {extra_params}\"" + '{image} "{params} -i /input -o /output -u /custom -s skip -Xrun_mode=singularity {extra_params}"' ).format( options=rconfig().singularity.run_extra_options, input=in_dir, @@ -102,13 +92,17 @@ def _start_container(self, script_params=""): log.info("Datasets are loaded by default from folder %s.", in_dir) log.info("Generated files will be available in folder %s.", out_dir) try: - run_cmd(cmd, _capture_error_=False) # console logs are written on stderr by default: not capturing allows live display - except: + run_cmd( + cmd, _capture_error_=False + ) # console logs are written on stderr by default: not capturing allows live display + except Exception as e: # also want to handle KeyboardInterrupt # In the foreground run mode, the user has to kill the process # There is yet no docker kill command. User has to kill PID manually - log.warning(f"Container {inst_name} may still be running, please verify and kill it manually.") - raise Exception + log.warning( + "Container may still be running, please verify and kill it manually." + ) + raise e def _image_exists(self, image): """Implements a method to see if the container image is available""" @@ -117,19 +111,25 @@ def _image_exists(self, image): return True try: # We pull from docker as there are not yet singularity org accounts - run_cmd("singularity pull {output_file} docker://{image}".format( - image=self._container_image_name(as_docker_image=True), - output_file=image, - ), _live_output_=True) + run_cmd( + "singularity pull {output_file} docker://{image}".format( + image=self._container_image_name(as_docker_image=True), + output_file=image, + ), + _live_output_=True, + ) return True except Exception: try: # If no docker image, pull from singularity hub - run_cmd("singularity pull {output_file} library://{library}/{image}".format( - image=self._container_image_name(as_docker_image=True), - output_file=image, - library=rconfig().singularity.library - ), _live_output_=True) + run_cmd( + "singularity pull {output_file} library://{library}/{image}".format( + image=self._container_image_name(as_docker_image=True), + output_file=image, + library=rconfig().singularity.library, + ), + _live_output_=True, + ) return True except Exception: pass @@ -137,18 +137,23 @@ def _image_exists(self, image): def _run_container_build_command(self, image, cache): log.info(f"Building singularity image {image}.") - run_cmd("sudo singularity build {options} {container} {script}".format( - options="" if cache else "--disable-cache", - container=image, - script=self._script, - ), _live_output_=True) + run_cmd( + "sudo singularity build {options} {container} {script}".format( + options="" if cache else "--disable-cache", + container=image, + script=self._script, + ), + _live_output_=True, + ) log.info(f"Successfully built singularity image {image}.") def _upload_image(self, image): library = rconfig().singularity.library name = self._container_image_name(as_docker_image=True) log.info(f"Publishing Singularity image {image}.") - run_cmd(f"singularity login && singularity push -U {image} library://{library}/{name}") + run_cmd( + f"singularity login && singularity push -U {image} library://{library}/{name}" + ) log.info(f"Successfully published singularity image {image}.") def _generate_script(self, custom_commands): @@ -221,17 +226,17 @@ def _generate_script(self, custom_commands): custom_commands=custom_commands.format( setup=dir_of( os.path.join(self._framework_dir, "setup/"), - rel_to_project_root=True + rel_to_project_root=True, ), pip="$PIP", - py="$PY" + py="$PY", ), - framework=self._forward_params['framework_name'], + framework=self._forward_params["framework_name"], pyv=rconfig().versions.python, pipv=rconfig().versions.pip, script=rconfig().script, ) touch(self._script) - with open(self._script, 'w') as file: + with open(self._script, "w") as file: file.write(singularity_content) diff --git a/tests/unit/amlb/benchmarks/test_benchmark.py b/tests/unit/amlb/benchmarks/test_benchmark.py index ba5bd5cd7..866deae50 100644 --- a/tests/unit/amlb/benchmarks/test_benchmark.py +++ b/tests/unit/amlb/benchmarks/test_benchmark.py @@ -1,3 +1,5 @@ +from pathlib import Path + import pytest from amlb import Benchmark, SetupMode, resources, DockerBenchmark, SingularityBenchmark @@ -78,15 +80,16 @@ def test_docker_image_name_uses_label(label, mocker, load_default_resources) -> def test_singularity_image_name( framework_name, tag, expected, load_default_resources ) -> None: - framework_def, _ = resources.get().framework_definition( - framework_name, - tag=tag, + benchmark = SingularityBenchmark( + framework_name=f"{framework_name}:{tag}", + benchmark_name="test", + constraint_name="test", ) - result = SingularityBenchmark.image_name( - framework_def, + image_path = benchmark._container_image_name( as_docker_image=False, ) - assert result == expected + image_name = Path(image_path).stem + assert image_name == expected @pytest.mark.parametrize( @@ -100,12 +103,12 @@ def test_singularity_image_name( def test_singularity_image_name_as_docker( framework_name, tag, expected, load_default_resources ) -> None: - framework_def, _ = resources.get().framework_definition( - framework_name, - tag=tag, + benchmark = SingularityBenchmark( + framework_name=f"{framework_name}:{tag}", + benchmark_name="test", + constraint_name="test", ) - result = SingularityBenchmark.image_name( - framework_def, + result = benchmark._container_image_name( as_docker_image=True, ) assert result == expected