support docker login

Author: zhongwen666

rock/admin/proto/request.py

@@ -27,6 +27,10 @@ class SandboxStartRequest(BaseModel):
     """The amount of CPUs to allocate for the container."""
     sandbox_id: str | None = Field(default=None)
     """The id of the sandbox."""
+    registry_username: str | None = None
+    """Username for Docker registry authentication. When both username and password are provided, docker login will be performed before pulling the image."""
+    registry_password: str | None = None
+    """Password for Docker registry authentication. When both username and password are provided, docker login will be performed before pulling the image."""
 
 
 class SandboxCommand(Command):

rock/deployments/config.py

@@ -109,6 +109,12 @@ class DockerDeploymentConfig(DeploymentConfig):
     actor_resource_num: float = 1
     """Number of actor resources to allocate (to be refined)."""
 
+    registry_username: str | None = None
+    """Username for Docker registry authentication. When both username and password are provided, docker login will be performed before pulling the image."""
+
+    registry_password: str | None = None
+    """Password for Docker registry authentication. When both username and password are provided, docker login will be performed before pulling the image."""
+
     runtime_config: RuntimeConfig = Field(default_factory=RuntimeConfig)
     """Runtime configuration settings."""
 

rock/deployments/docker.py

@@ -18,6 +18,7 @@
 from rock.deployments.config import DockerDeploymentConfig
 from rock.deployments.constants import Port, Status
 from rock.deployments.hooks.abstract import CombinedDeploymentHook, DeploymentHook
+from rock.deployments.hooks.docker_login import DockerLoginHook
 from rock.deployments.runtime_env import DockerRuntimeEnv, LocalRuntimeEnv, PipRuntimeEnv, UvRuntimeEnv
 from rock.deployments.sandbox_validator import DockerSandboxValidator
 from rock.deployments.status import PersistedServiceStatus, ServiceStatus
@@ -76,6 +77,11 @@ def __init__(
         else:
             raise Exception(f"Invalid ROCK_WORKER_ENV_TYPE: {env_vars.ROCK_WORKER_ENV_TYPE}")
 
+        if self._config.registry_username is not None and self._config.registry_password is not None:
+            self.add_hook(
+                DockerLoginHook(self._config.image, self._config.registry_username, self._config.registry_password)
+            )
+
         self.sandbox_validator: DockerSandboxValidator | None = DockerSandboxValidator()
 
     def add_hook(self, hook: DeploymentHook):

rock/deployments/hooks/docker_login.py

@@ -0,0 +1,37 @@
+import asyncio
+
+from rock.deployments.hooks.abstract import DeploymentHook
+from rock.logger import init_logger
+from rock.utils import DockerUtil, ImageUtil
+
+logger = init_logger(__name__)
+
+
+class DockerLoginHook(DeploymentHook):
+    """Hook that performs Docker registry authentication before pulling images.
+
+    When triggered by the "Pulling docker image" step, this hook parses the
+    registry from the image name and logs in using the provided credentials.
+    """
+
+    _PULL_STEP_MESSAGE = "Pulling docker image"
+
+    def __init__(self, image: str, username: str, password: str):
+        self._image = image
+        self._username = username
+        self._password = password
+
+    def on_custom_step(self, message: str):
+        if message != self._PULL_STEP_MESSAGE:
+            return
+
+        loop = asyncio.new_event_loop()
+        try:
+            registry, _ = loop.run_until_complete(ImageUtil.parse_registry_and_others(self._image))
+        finally:
+            loop.close()
+        if registry:
+            logger.info(f"Authenticating to registry {registry!r} before pulling image")
+            DockerUtil.login(registry, self._username, self._password)
+        else:
+            logger.warning(f"No registry found in image name {self._image!r}, skipping docker login")

rock/sdk/sandbox/client.py

@@ -171,6 +171,8 @@ async def start(self):
             "startup_timeout": self.config.startup_timeout,
             "memory": self.config.memory,
             "cpus": self.config.cpus,
+            "registry_username": self.config.registry_username,
+            "registry_password": self.config.registry_password,
         }
         try:
             response = await HttpUtils.post(url, headers, data)

rock/sdk/sandbox/config.py

@@ -39,6 +39,8 @@ class SandboxConfig(BaseConfig):
     experiment_id: str | None = None
     cluster: str = "zb"
     namespace: str | None = None
+    registry_username: str | None = None
+    registry_password: str | None = None
 
 
 class SandboxGroupConfig(SandboxConfig):

rock/utils/docker.py

@@ -36,6 +36,39 @@ def pull_image(cls, image: str) -> bytes:
             # e.stderr contains the error message as bytes
             raise subprocess.CalledProcessError(e.returncode, e.cmd, e.output, e.stderr) from None
 
+    @classmethod
+    def login(cls, registry: str, username: str, password: str, timeout: int = 30) -> str:
+        """Login to a Docker registry
+
+        Args:
+            registry: Docker registry URL (e.g. registry.example.com)
+            username: Registry username
+            password: Registry password
+            timeout: Command timeout in seconds
+
+        Returns:
+            Command output as string on success
+
+        Raises:
+            subprocess.CalledProcessError: If login fails
+        """
+        try:
+            result = subprocess.run(
+                ["docker", "login", registry, "-u", username, "--password-stdin"],
+                input=password,
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+            )
+            if result.returncode != 0:
+                logger.error(f"Docker login to {registry} failed: {result.stderr.strip()}")
+                raise subprocess.CalledProcessError(result.returncode, result.args, result.stdout, result.stderr)
+            logger.info(f"Successfully logged in to {registry}")
+            return result.stdout.strip()
+        except subprocess.TimeoutExpired:
+            logger.error(f"Docker login to {registry} timed out after {timeout}s")
+            raise
+
     @classmethod
     def remove_image(image: str) -> bytes:
         """Remove a Docker image"""