Skip to content

Commit

Permalink
Simplify code Singularity custom names
Browse files Browse the repository at this point in the history
  • Loading branch information
@PGijsbers
PGijsbers committed Nov 29, 2024
1 parent 7c6d2f0 commit 1b2818d
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 74 deletions.
11 changes: 7 additions & 4 deletions amlb/runners/container.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
123 changes: 64 additions & 59 deletions amlb/runners/singularity.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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):
"""
Expand All @@ -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"""
Expand All @@ -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,
Expand All @@ -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"""
Expand All @@ -117,38 +111,49 @@ 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
return False

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):
Expand Down Expand Up @@ -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)
25 changes: 14 additions & 11 deletions tests/unit/amlb/benchmarks/test_benchmark.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from pathlib import Path

import pytest

from amlb import Benchmark, SetupMode, resources, DockerBenchmark, SingularityBenchmark
Expand Down Expand Up @@ -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(
Expand All @@ -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

0 comments on commit 1b2818d

Please sign in to comment.