diff --git a/launch/core/platforms/base.py b/launch/core/platforms/base.py index da4991e..6ff47a8 100644 --- a/launch/core/platforms/base.py +++ b/launch/core/platforms/base.py @@ -4,6 +4,7 @@ import json import os import re +import subprocess import tarfile from dataclasses import dataclass from pathlib import Path @@ -258,29 +259,62 @@ def commit(self, image_name: str, tag: str = "latest", push: bool = False) -> st # currently restart from stopped container is not enabled. you need to launch a new instance from the new image you committed. return f"{image_name}:{tag}" + @staticmethod + def _image_pull_timeout_seconds() -> int: + """Return a bounded image-pull deadline from the environment.""" + value = os.getenv("REPOLAUNCH_IMAGE_PULL_TIMEOUT_SECONDS", "1800") + try: + timeout = int(value) + except ValueError as exc: + raise ValueError( + "REPOLAUNCH_IMAGE_PULL_TIMEOUT_SECONDS must be a positive integer" + ) from exc + if timeout <= 0: + raise ValueError( + "REPOLAUNCH_IMAGE_PULL_TIMEOUT_SECONDS must be a positive integer" + ) + return timeout + @staticmethod def pull_image(image_name: str) -> bool: - """ - Pull Docker image from registry. - - Args: - image_name (str): Name of the Docker image to pull - - Returns: - bool: True if successful, False if image not found - """ - client = docker.from_env(timeout=3600) # pull should finish in 1 hour + """Ensure an image is local, with a terminal deadline for registry pulls.""" + client = docker.from_env(timeout=30) try: - # Check if image exists locally client.images.get(image_name) return True except docker.errors.ImageNotFound: - # Image doesn't exist locally, try to pull it - try: - client.images.pull(image_name) - return True - except docker.errors.ImageNotFound: - raise ValueError(f"Image {image_name} not found in registry") + pass + + timeout = BaseRuntime._image_pull_timeout_seconds() + command = ["docker", "pull", image_name] + try: + completed = subprocess.run( + command, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + except subprocess.TimeoutExpired as exc: + raise RuntimeError( + f"Docker image pull for {image_name} timed out after {timeout} seconds" + ) from exc + + if completed.returncode != 0: + detail = (completed.stderr or completed.stdout).strip() + suffix = f": {detail}" if detail else "" + raise RuntimeError( + f"Docker image pull for {image_name} failed with exit code " + f"{completed.returncode}{suffix}" + ) + + try: + client.images.get(image_name) + except docker.errors.ImageNotFound as exc: + raise RuntimeError( + f"Docker image pull for {image_name} completed without a local image" + ) from exc + return True def cleanup(self, prune_dangling: bool = True) -> None: if self.stopped: diff --git a/tests/image_pull_deadline_test.py b/tests/image_pull_deadline_test.py new file mode 100644 index 0000000..de19286 --- /dev/null +++ b/tests/image_pull_deadline_test.py @@ -0,0 +1,78 @@ +import os +import subprocess +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import docker +from launch.core.platforms.base import BaseRuntime + + +class _Images: + def __init__(self): + self.get_calls = 0 + + def get(self, image_name): + self.get_calls += 1 + if self.get_calls == 1: + raise docker.errors.ImageNotFound("not local") + return object() + + +class _Client: + def __init__(self): + self.images = _Images() + + +class ImagePullDeadlineTests(unittest.TestCase): + def test_missing_image_uses_bounded_docker_cli_pull(self): + client = _Client() + completed = SimpleNamespace(returncode=0, stdout="pulled", stderr="") + with patch.dict(os.environ, {"REPOLAUNCH_IMAGE_PULL_TIMEOUT_SECONDS": "7"}, clear=False), \ + patch("launch.core.platforms.base.docker.from_env", return_value=client), \ + patch("launch.core.platforms.base.subprocess.run", return_value=completed) as run: + self.assertTrue(BaseRuntime.pull_image("example/image")) + + run.assert_called_once_with( + ["docker", "pull", "example/image"], + capture_output=True, + text=True, + timeout=7, + check=False, + ) + self.assertEqual(client.images.get_calls, 2) + + def test_pull_timeout_has_a_terminal_error(self): + client = _Client() + timed_out = subprocess.TimeoutExpired( + cmd=["docker", "pull", "example/image"], timeout=3, + ) + with patch.dict(os.environ, {"REPOLAUNCH_IMAGE_PULL_TIMEOUT_SECONDS": "3"}, clear=False), \ + patch("launch.core.platforms.base.docker.from_env", return_value=client), \ + patch("launch.core.platforms.base.subprocess.run", side_effect=timed_out): + with self.assertRaisesRegex(RuntimeError, "timed out after 3 seconds"): + BaseRuntime.pull_image("example/image") + + def test_nonzero_pull_exposes_the_docker_error(self): + client = _Client() + completed = SimpleNamespace(returncode=1, stdout="", stderr="registry unavailable") + with patch.dict(os.environ, {"REPOLAUNCH_IMAGE_PULL_TIMEOUT_SECONDS": "3"}, clear=False), \ + patch("launch.core.platforms.base.docker.from_env", return_value=client), \ + patch("launch.core.platforms.base.subprocess.run", return_value=completed): + with self.assertRaisesRegex(RuntimeError, "registry unavailable"): + BaseRuntime.pull_image("example/image") + + def test_invalid_timeout_configuration_is_rejected(self): + with patch.dict(os.environ, {"REPOLAUNCH_IMAGE_PULL_TIMEOUT_SECONDS": "zero"}, clear=False): + with self.assertRaisesRegex(ValueError, "positive integer"): + BaseRuntime._image_pull_timeout_seconds() + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file