diff --git a/.github/workflows/test_cw_engines.yml b/.github/workflows/test_cw_engines.yml new file mode 100644 index 000000000..0f06d941a --- /dev/null +++ b/.github/workflows/test_cw_engines.yml @@ -0,0 +1,129 @@ +name: Test compute worker container engines + +on: + pull_request: + paths: + - .github/workflows/test_cw_engines.yml + - packaging/container/Containerfile.compute_worker + - compute_worker/** + - src/settings/logs_loguru.py + push: + paths: + - .github/workflows/test_cw_engines.yml + - packaging/container/Containerfile.compute_worker + - compute_worker/** + - src/settings/logs_loguru.py + workflow_dispatch: + +permissions: + contents: read + +env: + TEST_IMAGE: localhost/codabench-compute-worker:ci + PYTHONDONTWRITEBYTECODE: "1" + +jobs: + build: + name: Build worker image + outputs: + image_id: ${{ steps.image.outputs.image_id }} + runs-on: ubuntu-22.04 + timeout-minutes: 30 + steps: + - uses: actions/checkout@v5 + - name: Build and export image + id: image + run: | + docker build --tag "$TEST_IMAGE" \ + --file packaging/container/Containerfile.compute_worker . + docker save --output worker.tar "$TEST_IMAGE" + docker image inspect --format 'image_id={{.Id}}' "$TEST_IMAGE" >> "$GITHUB_OUTPUT" + - uses: actions/upload-artifact@v4 + with: + name: compute-worker-image + path: worker.tar + compression-level: 0 + retention-days: 1 + + docker: + name: Test Docker + needs: build + runs-on: ubuntu-22.04 + timeout-minutes: 15 + env: + CONTAINER_ENGINE_EXECUTABLE: docker + CONTAINER_SOCKET: unix:///var/run/docker.sock + steps: + - uses: actions/checkout@v5 + - uses: actions/download-artifact@v4 + with: + name: compute-worker-image + - uses: actions/setup-python@v5 + with: + python-version: "3.13" + - name: Install test dependency and load image + run: | + python -m pip install 'docker>=7.1.0,<8' + docker load --input worker.tar + - name: Run shared tests + run: python -B -m unittest discover -s compute_worker/test -v + + podman: + name: Test Podman + needs: build + runs-on: ubuntu-22.04 + timeout-minutes: 15 + env: + CONTAINER_ENGINE_EXECUTABLE: podman + steps: + - uses: actions/checkout@v5 + - uses: actions/download-artifact@v4 + with: + name: compute-worker-image + - uses: actions/setup-python@v5 + with: + python-version: "3.13" + - name: Install Podman and test dependency + run: | + sudo apt-get update + sudo apt-get install --yes podman + python -m pip install 'docker>=7.1.0,<8' + - name: Load image and start rootless API socket + run: | + podman load --input worker.tar + podman tag "${{ needs.build.outputs.image_id }}" "$TEST_IMAGE" + podman image inspect "$TEST_IMAGE" + systemctl --user start podman.socket + - name: Run shared tests + run: | + export CONTAINER_SOCKET="unix://${XDG_RUNTIME_DIR}/podman/podman.sock" + python -B -m unittest discover -s compute_worker/test -v + + apptainer: + name: Test Apptainer + needs: build + runs-on: ubuntu-22.04 + timeout-minutes: 15 + env: + CONTAINER_ENGINE_EXECUTABLE: apptainer + TEST_SIF: worker.sif + steps: + - uses: actions/checkout@v5 + - uses: actions/download-artifact@v4 + with: + name: compute-worker-image + - uses: actions/setup-python@v5 + with: + python-version: "3.13" + - name: Install Apptainer and test dependency + run: | + sudo add-apt-repository --yes ppa:apptainer/ppa + sudo apt-get update + sudo apt-get install --yes apptainer + python -m pip install 'docker>=7.1.0,<8' + - name: Convert image to SIF + run: apptainer build worker.sif docker-archive://worker.tar + - name: Run adapter self-checks + run: python -B compute_worker/apptainer_backend.py + - name: Run shared tests + run: python -B -m unittest discover -s compute_worker/test -v diff --git a/compute_worker/apptainer_backend.py b/compute_worker/apptainer_backend.py new file mode 100644 index 000000000..8a972a7ea --- /dev/null +++ b/compute_worker/apptainer_backend.py @@ -0,0 +1,336 @@ +import codecs +import hashlib +import json +import os +from pathlib import Path +import selectors +import shlex +import signal +import subprocess +import tempfile +import uuid + + +class APIClient: + """Map the Docker API methods used by the Codabench compute worker to Apptainer. + + Only the methods and options required by the worker are implemented. + This is a partial compatibility layer, not a complete Docker API implementation. + """ + + def __init__(self): + # Explicit mappings take precedence over downloaded images. + self.images = json.loads(os.environ.get("APPTAINER_IMAGE_MAP", "{}")) + self.image_cache = Path(os.environ.get( + "APPTAINER_CACHEDIR", str(Path.home() / ".apptainer" / "cache") + )) / "codabench" + self.containers = {} + + def _image_path(self, image): + """Resolve a mapped SIF or the cache path for an image reference.""" + mapped = self.images.get(image.removeprefix("docker://")) + if mapped: + return Path(mapped) + reference = image if "://" in image else f"docker://{image}" + digest = hashlib.sha256(reference.encode()).hexdigest() + return self.image_cache / f"{digest}.sif" + + def inspect_image(self, image): + """Return the local SIF path without downloading an image.""" + path = self._image_path(image) + if not path.is_file(): + raise FileNotFoundError(f"Local Apptainer image not found: {image}") + return {"Id": str(path)} + + def pull(self, image, **kwargs): + """Use an explicit SIF mapping or pull an image into the local cache. + + Bare image references use the Docker transport. Yield status dictionaries + compatible with the compute worker's streaming pull call. + """ + if self.images.get(image.removeprefix("docker://")): + self.inspect_image(image) + yield {"status": "Status: Image is up to date (mapped SIF)"} + return + reference = image if "://" in image else f"docker://{image}" + path = self._image_path(image) + path.parent.mkdir(parents=True, exist_ok=True) + yield {"id": image, "status": "Downloading"} + # Publish only a completed SIF; a failed pull leaves any cached image intact. + with tempfile.TemporaryDirectory(dir=path.parent, prefix="pull-") as directory: + temporary = Path(directory) / "image.sif" + result = subprocess.run( + ["/usr/bin/apptainer", "pull", str(temporary), reference], + capture_output=True, text=True, + ) + if result.returncode: + raise RuntimeError( + f"Apptainer pull failed for {image}: {result.stderr.strip()}" + ) + temporary.replace(path) + yield {"id": image, "status": "Pull complete"} + + def create_host_config(self, **kwargs): + return kwargs + + def create_container(self, image, *, name, host_config, command, + working_dir, environment, network_disabled=False, **kwargs): + """Prepare an Apptainer command and return its Docker-compatible ID. + + Register the command without starting it. The image must be mapped to + a local SIF or downloaded by pull(). Bind sources stay inside /codabench; + temporary directories and the execution environment are isolated. + """ + # Resolve the local image and reject a duplicate container name. + sif = self.inspect_image(image)["Id"] + if name in self.containers: + raise RuntimeError(f"Container already exists: {name}") + # Retain the TemporaryDirectory object until remove_container cleans it. + scratch = tempfile.TemporaryDirectory(prefix="evaluation-") + root = Path(scratch.name) + # Execute the program in the SIF, rather than its image runscript. + args = [ + "/usr/bin/apptainer", "exec", + "--userns", # Use an unprivileged user namespace. + "--containall", # Isolate filesystems, PID, IPC and environment. + "--cleanenv", # Exclude ordinary inherited environment values. + "--no-eval", # Disable shell evaluation of environment/OCI args. + "--no-privs", # Drop root privileges inside the container. + "--writable-tmpfs", # Add a temporary overlay; leave the SIF unchanged. + "--no-mount", "hostfs,cwd,bind-paths", # Disable automatic host/admin binds. + "--pwd", working_dir, # Set the program's working directory. + ] + # Give this evaluation its own writable home and temporary files. + # Let --containall supply container-local /dev/shm: binding a scratch + # directory there can leave it mounted and make cleanup fail with EBUSY. + for part in ("home", "tmp"): + (root / part).mkdir() + args += ["--home", f"{root}/home:/home/competition", + "--bind", f"{root}/tmp:/tmp", + "--bind", f"{root}/tmp:/var/tmp"] + # Translate Codabench volume mappings to src:destination:mode binds. + for source, mount in host_config["binds"].items(): + # Resolve symlinks before checking the permitted source directory. + resolved = Path(source).resolve() + if not resolved.is_relative_to(Path("/codabench").resolve()): + scratch.cleanup() + raise ValueError(f"Bind source outside worker directory: {source}") + destination = mount["bind"] + # Commas/colons delimit Apptainer bind specifications, so reject them + # in paths instead of allowing an additional mount to be parsed. + if any(c in source + destination for c in ",:\n"): + scratch.cleanup() + raise ValueError("Unsupported character in bind path") + # Translate read/write access; Docker's SELinux label option is omitted. + mode = "ro" if "ro" in mount.get("mode", "").split(",") else "rw" + args += ["--bind", f"{source}:{destination}:{mode}"] + # Enable NVIDIA support when the worker requests GPU access. + if host_config.get("device_requests"): + args += ["--nv"] + # A separate network namespace with "none" removes external networking + # when the upstream worker requests it; otherwise keep its network policy. + if network_disabled: + args += ["--net", "--network", "none"] + # Match Docker's command splitting without adding a shell invocation. + argv = shlex.split(command) if isinstance(command, str) else list(command) + if not argv: + scratch.cleanup() + raise ValueError("Empty command") + args += [sif, *argv] + # Pass only the required environment to the container process. + env = {"PATH": "/usr/bin:/bin", "HOME": "/tmp", + "APPTAINER_TMPDIR": str(root / "tmp")} + # APPTAINERENV_ explicitly injects the selected values into the payload + # despite --cleanenv; values baked into the image remain available. + values = dict(item.split("=", 1) for item in environment) + # HOME is controlled by --home; Apptainer rejects APPTAINERENV_HOME. + values.pop("HOME", None) + # Preserve a host GPU visibility mask when one is configured. + if host_config.get("device_requests") and "CUDA_VISIBLE_DEVICES" in os.environ: + values["CUDA_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"] + for key, value in values.items(): + env["APPTAINERENV_" + key] = value + # Store the prepared descriptor for start/attach/wait/remove; no process + # exists until start is called. + identifier = uuid.uuid4().hex + self.containers[identifier] = {"name": name, "args": args, "env": env, + "scratch": scratch, "process": None} + return {"Id": identifier} + + def _get(self, container): + key = container["Id"] if isinstance(container, dict) else container + if key not in self.containers: + key = next(k for k, v in self.containers.items() if v["name"] == key) + return key, self.containers[key] + + def start(self, container): + """Print the prepared command and start it in its own process group.""" + _, state = self._get(container) + # Log argv with shell quoting for readability; execution still uses argv. + print("Running command:", shlex.join(state["args"]), flush=True) + # Separate the group for cleanup; capture each output stream independently. + state["process"] = subprocess.Popen( + state["args"], env=state["env"], stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, start_new_session=True) + + def attach(self, container, **kwargs): + """Yield demultiplexed UTF-8 stdout/stderr chunks until both pipes close. + + Each item is (stdout_bytes, None) or (None, stderr_bytes), as required + by the compute worker's streaming attach call. Drain buffered output + even if the process has already exited. Incremental decoding preserves + characters split across reads and replaces invalid UTF-8 input. + Docker keyword arguments are accepted for call compatibility. + """ + _, state = self._get(container) + process = state["process"] + # Keep a decoder per stream so a partial character is completed on the + # next read before the upstream worker decodes the yielded bytes. + decoders = [codecs.getincrementaldecoder("utf-8")(errors="replace") for _ in range(2)] + with selectors.DefaultSelector() as selector: + # Register non-blocking pipes: 0 identifies stdout, 1 stderr. + for index, stream in enumerate((process.stdout, process.stderr)): + os.set_blocking(stream.fileno(), False) + selector.register(stream, selectors.EVENT_READ, index) + while selector.get_map(): + for key, _ in selector.select(timeout=1): + chunk = os.read(key.fileobj.fileno(), 65536) + payload = decoders[key.data].decode(chunk, final=not chunk).encode("utf-8") + # EOF flushes any decoder remainder and removes that pipe. + if not chunk: + selector.unregister(key.fileobj) + if payload: + yield (payload, None) if key.data == 0 else (None, payload) + + def inspect_container(self, container): + _, state = self._get(container) + process = state["process"] + running = process is not None and process.poll() is None + return {"State": {"Status": "running" if running else "exited"}} + + def wait(self, container, **kwargs): + _, state = self._get(container) + code = state["process"].wait() + return {"StatusCode": code if code >= 0 else 128 - code} + + def kill(self, container, **kwargs): + _, state = self._get(container) + process = state["process"] + if process is not None: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + + def remove_container(self, container, force=False, **kwargs): + try: + key, state = self._get(container) + except (StopIteration, KeyError): + return + if force: + self.kill(container) + process = state["process"] + if process is not None: + process.wait() + process.stdout.close() + process.stderr.close() + state["scratch"].cleanup() + del self.containers[key] + + +def self_check(): + """Check container lifecycle and output handling with local Python processes.""" + import sys + # 1. Register a local SIF mapping to check image lookup and command construction. + client = APIClient() + with tempfile.TemporaryDirectory() as directory: + image = Path(directory) / "image.sif" + image.touch() + client.images = {"test:1": str(image)} + # 2. Check the command descriptor and explicitly selected environment. + container = client.create_container( + "test:1", name="test", host_config={"binds": {}}, + command=["python", "-c", "pass"], working_dir="/app/program", + environment=["HOME=/ignored", "TEST_PAYLOAD=preserved"]) + state = client.containers[container["Id"]] + assert not any("BROKER" in key for key in state["env"]) + assert "APPTAINERENV_HOME" not in state["env"] + assert state["env"]["APPTAINERENV_TEST_PAYLOAD"] == "preserved" + assert state["args"][state["args"].index("--home") + 1].endswith(":/home/competition") + assert "--containall" in state["args"] + assert not any(":/dev/shm" in arg for arg in state["args"]) + scratch_root = Path(state["scratch"].name) + assert not (scratch_root / "shm").exists() + # 3. Substitute a local Python command to check output and exit status. + state["args"] = [sys.executable, "-c", + "import sys; print('out'); print('err', file=sys.stderr); sys.exit(23)"] + state["env"] = os.environ.copy() + client.start(container) + # Deliberately wait before attach: fast exits must retain their logs. + assert client.wait(container)["StatusCode"] == 23 + chunks = list(client.attach(container)) + assert b"".join(c[0] or b"" for c in chunks) == b"out\n" + assert b"".join(c[1] or b"" for c in chunks) == b"err\n" + # 4. Check repeated removal after process and temporary directory cleanup. + client.remove_container(container, force=True) + client.remove_container(container, force=True) + assert not scratch_root.exists() + # 5. Inspection must not download an image that is unavailable locally. + try: + client.inspect_image("unapproved:1") + except FileNotFoundError: + pass + else: + raise AssertionError("Missing local image accepted") + # 6. Check image pulls without invoking Apptainer or contacting a registry. + from unittest.mock import patch + client.image_cache = Path(directory) / "cache" + + def fake_pull(args, **kwargs): + Path(args[2]).write_bytes(b"SIF") + return subprocess.CompletedProcess(args, 0, "", "") + + with patch.object(subprocess, "run", side_effect=fake_pull) as run: + assert list(client.pull("test:1")) + run.assert_not_called() + assert list(client.pull("download:1"))[-1]["status"] == "Pull complete" + assert run.call_args.args[0][-1] == "docker://download:1" + cached = Path(client.inspect_image("docker://download:1")["Id"]) + assert cached.read_bytes() == b"SIF" + with patch.object(subprocess, "run", return_value=subprocess.CompletedProcess( + [], 1, "", "pull failed" + )): + try: + list(client.pull("download:1")) + except RuntimeError as error: + assert "pull failed" in str(error) + else: + raise AssertionError("Failed image pull accepted") + assert cached.read_bytes() == b"SIF" + assert list(client.image_cache.iterdir()) == [cached] + # 7. Start a parent and child in the same group. If the child survives, + # it creates this sentinel after forced removal of the group. + sentinel = Path(directory) / "orphan" + child_code = f"import time,pathlib; time.sleep(0.3); pathlib.Path({str(sentinel)!r}).touch(); time.sleep(60)" + parent_code = ("import subprocess,time,sys; " + f"subprocess.Popen([sys.executable,'-c',{child_code!r}]); " + "print('ready',flush=True); time.sleep(60)") + process = subprocess.Popen( + [sys.executable, "-c", parent_code], start_new_session=True, + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + # Wait until the parent has spawned the child before testing termination. + assert process.stdout.readline() == b"ready\n" + scratch = tempfile.TemporaryDirectory() + client.containers["timeout"] = {"name": "timeout", "process": process, + "scratch": scratch} + client.remove_container("timeout", force=True) + assert process.returncode == -signal.SIGKILL + import time + # Allow longer than the child's 0.3-second delay, then check no file exists. + time.sleep(0.5) + assert not sentinel.exists(), "Subprocess survived group termination" + print("Apptainer adapter local checks passed") + + +if __name__ == "__main__": + self_check() diff --git a/compute_worker/compute_worker.py b/compute_worker/compute_worker.py index 3b5828b76..673fad37a 100644 --- a/compute_worker/compute_worker.py +++ b/compute_worker/compute_worker.py @@ -169,7 +169,7 @@ class SubmissionStatus: ) # ----------------------------------------------- -# Initialize Docker or Podman depending on .env +# Initialize the configured Docker, Podman or Apptainer backend # ----------------------------------------------- logger.info( f"Using {Settings.CONTAINER_ENGINE_EXECUTABLE} " @@ -181,12 +181,15 @@ class SubmissionStatus: Settings.COMPUTE_WORKER_NO_CLEANUP = True -# Intializing client -# NOTE: CONTAINER_SOCKET is set in Settings based on CONTAINER_ENGINE_EXECUTABLE which must has either podman or docker -client = docker.APIClient( - base_url=Settings.CONTAINER_SOCKET, - version="auto", -) +# Docker/Podman use their configured socket; Apptainer executes local processes. +if Settings.CONTAINER_ENGINE_EXECUTABLE == "apptainer": + from apptainer_backend import APIClient + client = APIClient() +else: + client = docker.APIClient( + base_url=Settings.CONTAINER_SOCKET, + version="auto", + ) # ----------------------------------------------- diff --git a/compute_worker/test/test_container_engines.py b/compute_worker/test/test_container_engines.py new file mode 100644 index 000000000..c8c35b16b --- /dev/null +++ b/compute_worker/test/test_container_engines.py @@ -0,0 +1,99 @@ +"""Run shared container API checks against Docker, Podman, or Apptainer.""" + +import os +from pathlib import Path +import sys +import unittest +import uuid + +import docker + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from apptainer_backend import APIClient + + +class ContainerEngineTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.engine = os.environ["CONTAINER_ENGINE_EXECUTABLE"] + cls.image = os.environ.get("TEST_IMAGE", "codabench-compute-worker:ci") + if cls.engine == "apptainer": + cls.client = APIClient() + cls.client.images = { + cls.image: str(Path(os.environ["TEST_SIF"]).resolve()) + } + elif cls.engine in ("docker", "podman"): + cls.client = docker.APIClient( + base_url=os.environ["CONTAINER_SOCKET"], version="auto", timeout=30 + ) + else: + raise ValueError(f"Unsupported container engine: {cls.engine}") + + @classmethod + def tearDownClass(cls): + if cls.engine != "apptainer": + cls.client.close() + + def create_container(self, code): + command = ["/.venv/bin/python", "-c", code] + entrypoint = [] + if self.engine == "podman": + # Explicitly replace the image entrypoint for the Podman API. + entrypoint = [command.pop(0)] + container = self.client.create_container( + self.image, + name=f"worker-test-{uuid.uuid4().hex}", + host_config=self.client.create_host_config(binds={}, auto_remove=False), + command=command, + entrypoint=entrypoint, + working_dir="/tmp", + environment=["WORKER_TEST_VALUE=configured", "PYTHONUNBUFFERED=1"], + ) + self.addCleanup(self.client.remove_container, container, v=True, force=True) + return container + + # 1. Check the image loaded by the runtime job. + def test_01_inspect_image(self): + self.assertTrue(self.client.inspect_image(self.image)["Id"]) + + # 2. Check the worker environment, output streams, and exit status. + def test_02_execution(self): + container = self.create_container( + "import os, sys, time; " + "import celery, docker, requests, websockets; " + "assert os.getcwd() == '/tmp'; " + "assert os.environ['WORKER_TEST_VALUE'] == 'configured'; " + "print('worker stdout'); " + "print('worker stderr', file=sys.stderr); " + "time.sleep(1); sys.exit(23)" + ) + self.client.start(container) + chunks = list(self.client.attach( + container, demux=True, stream=True, logs=True + )) + self.assertEqual( + b"".join(chunk[0] or b"" for chunk in chunks), b"worker stdout\n" + ) + self.assertEqual( + b"".join(chunk[1] or b"" for chunk in chunks), b"worker stderr\n" + ) + self.assertEqual(self.client.wait(container)["StatusCode"], 23) + self.assertNotEqual( + self.client.inspect_container(container)["State"]["Status"].lower(), + "running", + ) + + # 3. Check explicit termination and the resulting exit status. + def test_03_kill(self): + container = self.create_container("import time; time.sleep(300)") + self.client.start(container) + self.assertEqual( + self.client.inspect_container(container)["State"]["Status"].lower(), + "running", + ) + self.client.kill(container) + self.assertEqual(self.client.wait(container)["StatusCode"], 137) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/packaging/container/Containerfile.compute_worker b/packaging/container/Containerfile.compute_worker index 4f9728c34..67a14b04f 100644 --- a/packaging/container/Containerfile.compute_worker +++ b/packaging/container/Containerfile.compute_worker @@ -9,7 +9,7 @@ COPY --from=builder / / # This makes output not buffer and return immediately, nice for seeing results in stdout ENV PYTHONUNBUFFERED=1 -COPY compute_worker/pyproject.toml compute_worker/uv.lock compute_worker/celery_config.py compute_worker/compute_worker.py ./ +COPY compute_worker/pyproject.toml compute_worker/uv.lock compute_worker/celery_config.py compute_worker/compute_worker.py compute_worker/apptainer_backend.py ./ # Install UV and add paths to PATH for uv and the future .venv created by uv RUN curl -LsSf https://astral.sh/uv/install.sh | sh && microdnf remove -y tar gzip && microdnf clean all