feat(sandbox): inject INSTANCE_ROCK_REGISTRY env var into launched sandboxes (#1090)

Add RuntimeConfig.instance_registry_mirrors (configurable via YAML and
Nacos hot-reload). Docker, Ray, and K8s backends inject
INSTANCE_ROCK_REGISTRY=<comma-joined> into every launched container.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: jake11-oho

rock-conf/rock-local.yml

@@ -24,6 +24,16 @@ warmup:
 #   - "swe-bench:"               # bare image name prefix
 #   - "aaaaaa/bbb/"              # registry/namespace prefix
 
+# Inject INSTANCE_ROCK_REGISTRY into every launched sandbox so the
+# runtime can rewrite prebuilt images to a ROCK mirror first.
+# Configured via YAML or Nacos (hot-reload supported).
+# Nacos YAML format:
+#
+# runtime:
+#   instance_registry_mirrors:
+#     - "reg-a.aliyuncs.com/mirror-1"
+#     - "reg-b.aliyuncs.com/mirror-2"
+
 # Scheduler configuration
 scheduler:
   enabled: true                    # Whether to enable the scheduler

rock/config.py

@@ -274,6 +274,12 @@ class RuntimeConfig:
     sandbox_disk_limit_rootfs: str | None = None
     """Default rootfs quota per container. None means no limit. Can be overridden by nacos key 'default_disk_limit'."""
 
+    instance_registry_mirrors: list[str] = field(default_factory=list)
+    """Registry mirrors injected into each launched sandbox as
+    INSTANCE_ROCK_REGISTRY=<comma-joined>. Each entry is a host/namespace
+    string (e.g. "reg-a.aliyuncs.com/mirror-1"). When empty, the env var is
+    not set and downstream tools skip mirror rewriting."""
+
     def __post_init__(self) -> None:
         # Convert dict to StandardSpec if needed
         if isinstance(self.standard_spec, dict):
@@ -511,8 +517,14 @@ async def update(self):
         if "image_mirror_lookup_allowlist" in nacos_result:
             self.image_mirror_lookup_allowlist = list(nacos_result["image_mirror_lookup_allowlist"] or [])
 
+        if "runtime" in nacos_result and isinstance(nacos_result["runtime"], dict):
+            runtime_overrides = nacos_result["runtime"]
+            if "instance_registry_mirrors" in runtime_overrides:
+                self.runtime.instance_registry_mirrors = list(runtime_overrides["instance_registry_mirrors"] or [])
+
         logger.info(
             f"Updated config from Nacos: sandbox_config={self.sandbox_config}, proxy_service={self.proxy_service}"
             f", image_registry_mirrors={self.image_registry_mirrors}"
             f", image_mirror_lookup_allowlist={self.image_mirror_lookup_allowlist}"
+            f", instance_registry_mirrors={self.runtime.instance_registry_mirrors}"
         )

rock/deployments/docker.py

@@ -523,8 +523,6 @@ async def start(self):
         if self._config.remove_container:
             rm_arg = ["--rm"]
 
-        env_arg = []
-
         # Conditionally set up logging path mount based on ROCK_LOGGING_PATH.
         log_file_path: str | None = None
         volume_args = self._prepare_volume_mounts()
@@ -533,23 +531,17 @@ async def start(self):
             os.makedirs(log_file_path, exist_ok=True)
             os.chmod(log_file_path, 0o777)
             volume_args.extend(["-v", f"{log_file_path}:{env_vars.ROCK_LOGGING_PATH}"])
-            env_arg = [
-                "-e",
-                f"ROCK_LOGGING_PATH={env_vars.ROCK_LOGGING_PATH}",
-                "-e",
-                f"ROCK_LOGGING_LEVEL={env_vars.ROCK_LOGGING_LEVEL}",
-            ]
 
-        env_arg.extend(["-e", f"ROCK_TIME_ZONE={env_vars.ROCK_TIME_ZONE}"])
         volume_args.extend(self._prepare_timezone_mount())
 
-        # Kata DinD: prepare disk image and add volume mount + env var
+        # Kata DinD: prepare disk image and add volume mount
         if self._config.use_kata_runtime:
             with StageTimer("startup_timing", f"[{self._container_name}] Kata disk prepare", logger):
                 self._prepare_kata_disk()
             disk_path = self._get_kata_disk_image_path()
             volume_args.extend(["-v", f"{disk_path}:/docker-disk.img"])
-            env_arg.extend(["-e", "ROCK_KATA_RUNTIME=true"])
+
+        env_arg = self._build_env_args()
 
         with StageTimer("startup_timing", f"[{self._container_name}] Random sleep", logger):
             time.sleep(random.randint(0, 5))
@@ -611,6 +603,32 @@ async def start(self):
         if self._config.enable_auto_clear:
             self._check_stop_task = asyncio.create_task(self._check_stop())
 
+    def _build_env_args(self) -> list[str]:
+        """Construct `-e KEY=VALUE` argv pairs for `docker run`.
+
+        Centralizes every env var rock injects into a sandbox in one place so
+        callers (and tests) can audit them together. Volume mounts that pair
+        with these vars (ROCK_LOGGING_PATH, kata disk) are still set up
+        alongside in `start()`.
+        """
+        args: list[str] = []
+        if env_vars.ROCK_LOGGING_PATH:
+            args.extend(
+                [
+                    "-e",
+                    f"ROCK_LOGGING_PATH={env_vars.ROCK_LOGGING_PATH}",
+                    "-e",
+                    f"ROCK_LOGGING_LEVEL={env_vars.ROCK_LOGGING_LEVEL}",
+                ]
+            )
+        args.extend(["-e", f"ROCK_TIME_ZONE={env_vars.ROCK_TIME_ZONE}"])
+        if self._config.use_kata_runtime:
+            args.extend(["-e", "ROCK_KATA_RUNTIME=true"])
+        mirrors = self._config.runtime_config.instance_registry_mirrors
+        if mirrors:
+            args.extend(["-e", f"INSTANCE_ROCK_REGISTRY={','.join(mirrors)}"])
+        return args
+
     def _prepare_timezone_mount(self) -> list[str]:
         tz = env_vars.ROCK_TIME_ZONE
         localtime_src = f"/usr/share/zoneinfo/{tz}"

tests/unit/deployments/test_docker_env_injection.py

@@ -0,0 +1,101 @@
+"""Unit tests for DockerDeployment._build_env_args.
+
+Covers the env-arg list assembled for `docker run -e ...`, focusing on the
+INSTANCE_ROCK_REGISTRY injection sourced from RuntimeConfig.instance_registry_mirrors.
+Existing entries (ROCK_LOGGING_*, ROCK_TIME_ZONE, ROCK_KATA_RUNTIME) are regression-checked.
+"""
+
+from unittest.mock import patch
+
+from rock.config import RuntimeConfig
+from rock.deployments.config import DockerDeploymentConfig
+from rock.deployments.docker import DockerDeployment
+
+
+def _make_deployment(**config_kwargs) -> DockerDeployment:
+    with patch("rock.deployments.docker.DockerSandboxValidator"):
+        return DockerDeployment.from_config(DockerDeploymentConfig(**config_kwargs))
+
+
+class TestBuildEnvArgsInstanceRegistry:
+    def test_no_mirrors_means_no_registry_env(self):
+        deployment = _make_deployment()
+        with patch("rock.deployments.docker.env_vars") as mock_env:
+            mock_env.ROCK_LOGGING_PATH = ""
+            mock_env.ROCK_TIME_ZONE = "UTC"
+            args = deployment._build_env_args()
+        assert "INSTANCE_ROCK_REGISTRY" not in " ".join(args)
+
+    def test_single_mirror_emitted_unquoted(self):
+        deployment = _make_deployment(
+            runtime_config=RuntimeConfig(instance_registry_mirrors=["reg-a.example.com/ns"]),
+        )
+        with patch("rock.deployments.docker.env_vars") as mock_env:
+            mock_env.ROCK_LOGGING_PATH = ""
+            mock_env.ROCK_TIME_ZONE = "UTC"
+            args = deployment._build_env_args()
+        assert "-e" in args
+        assert "INSTANCE_ROCK_REGISTRY=reg-a.example.com/ns" in args
+
+    def test_multiple_mirrors_comma_joined(self):
+        deployment = _make_deployment(
+            runtime_config=RuntimeConfig(
+                instance_registry_mirrors=[
+                    "reg-a.example.com/ns-1",
+                    "reg-b.example.com/ns-2",
+                ],
+            ),
+        )
+        with patch("rock.deployments.docker.env_vars") as mock_env:
+            mock_env.ROCK_LOGGING_PATH = ""
+            mock_env.ROCK_TIME_ZONE = "UTC"
+            args = deployment._build_env_args()
+        assert "INSTANCE_ROCK_REGISTRY=reg-a.example.com/ns-1,reg-b.example.com/ns-2" in args
+
+
+class TestBuildEnvArgsRegression:
+    """The existing env entries must keep working after the extraction."""
+
+    def test_timezone_always_present(self):
+        deployment = _make_deployment()
+        with patch("rock.deployments.docker.env_vars") as mock_env:
+            mock_env.ROCK_LOGGING_PATH = ""
+            mock_env.ROCK_TIME_ZONE = "Asia/Shanghai"
+            args = deployment._build_env_args()
+        assert "ROCK_TIME_ZONE=Asia/Shanghai" in args
+
+    def test_logging_entries_present_when_logging_path_set(self):
+        deployment = _make_deployment()
+        with patch("rock.deployments.docker.env_vars") as mock_env:
+            mock_env.ROCK_LOGGING_PATH = "/var/log/rock"
+            mock_env.ROCK_LOGGING_LEVEL = "INFO"
+            mock_env.ROCK_TIME_ZONE = "UTC"
+            args = deployment._build_env_args()
+        assert "ROCK_LOGGING_PATH=/var/log/rock" in args
+        assert "ROCK_LOGGING_LEVEL=INFO" in args
+
+    def test_logging_entries_absent_when_logging_path_empty(self):
+        deployment = _make_deployment()
+        with patch("rock.deployments.docker.env_vars") as mock_env:
+            mock_env.ROCK_LOGGING_PATH = ""
+            mock_env.ROCK_TIME_ZONE = "UTC"
+            args = deployment._build_env_args()
+        joined = " ".join(args)
+        assert "ROCK_LOGGING_PATH" not in joined
+        assert "ROCK_LOGGING_LEVEL" not in joined
+
+    def test_kata_runtime_env_present_when_enabled(self):
+        deployment = _make_deployment(use_kata_runtime=True)
+        with patch("rock.deployments.docker.env_vars") as mock_env:
+            mock_env.ROCK_LOGGING_PATH = ""
+            mock_env.ROCK_TIME_ZONE = "UTC"
+            args = deployment._build_env_args()
+        assert "ROCK_KATA_RUNTIME=true" in args
+
+    def test_kata_runtime_env_absent_when_disabled(self):
+        deployment = _make_deployment(use_kata_runtime=False)
+        with patch("rock.deployments.docker.env_vars") as mock_env:
+            mock_env.ROCK_LOGGING_PATH = ""
+            mock_env.ROCK_TIME_ZONE = "UTC"
+            args = deployment._build_env_args()
+        assert "ROCK_KATA_RUNTIME=true" not in args

tests/unit/test_config.py

@@ -56,6 +56,42 @@ async def test_runtime_config():
     assert runtime_config.standard_spec.cpus == 2
 
 
+def test_runtime_config_instance_registry_mirrors_default_empty():
+    cfg = RuntimeConfig()
+    assert cfg.instance_registry_mirrors == []
+
+
+def test_runtime_config_instance_registry_mirrors_accepts_list():
+    mirrors = ["reg-a.example.com/ns-1", "reg-b.example.com/ns-2"]
+    cfg = RuntimeConfig(instance_registry_mirrors=mirrors)
+    assert cfg.instance_registry_mirrors == mirrors
+
+
+@pytest.mark.asyncio
+async def test_runtime_config_instance_registry_mirrors_from_yaml(tmp_path, monkeypatch):
+    yaml_path = tmp_path / "rock-test.yml"
+    yaml_path.write_text(
+        yaml.safe_dump(
+            {
+                "runtime": {
+                    "instance_registry_mirrors": [
+                        "reg-a.example.com/mirror-1",
+                        "reg-b.example.com/mirror-2",
+                    ],
+                },
+            }
+        )
+    )
+    # RuntimeConfig.__post_init__ requires ROCK_PYTHON_ENV_PATH + ROCK_ENVHUB_DB_URL.
+    monkeypatch.setenv("ROCK_PYTHON_ENV_PATH", "/usr")
+    monkeypatch.setenv("ROCK_ENVHUB_DB_URL", "sqlite:////tmp/test.db")
+    rock_config = RockConfig.from_env(str(yaml_path))
+    assert rock_config.runtime.instance_registry_mirrors == [
+        "reg-a.example.com/mirror-1",
+        "reg-b.example.com/mirror-2",
+    ]
+
+
 def test_oss_config_defaults():
     from rock.config import OssAccountConfig, OssConfig
 
@@ -214,6 +250,56 @@ async def test_nacos_update_without_mirror_keys_preserves_existing():
     assert rock_config.image_mirror_lookup_allowlist == ["*"]
 
 
+# ===== Nacos hot-reload for instance_registry_mirrors =====
+
+
+@pytest.mark.asyncio
+async def test_nacos_update_loads_instance_registry_mirrors():
+    rock_config = RockConfig()
+    rock_config.nacos_provider = MagicMock()
+    rock_config.nacos_provider.get_config = AsyncMock(
+        return_value={
+            "runtime": {
+                "instance_registry_mirrors": ["reg-a.example.com/ns-1", "reg-b.example.com/ns-2"],
+            },
+        }
+    )
+
+    await rock_config.update()
+
+    assert rock_config.runtime.instance_registry_mirrors == ["reg-a.example.com/ns-1", "reg-b.example.com/ns-2"]
+
+
+@pytest.mark.asyncio
+async def test_nacos_update_empty_instance_mirrors_clears_list():
+    rock_config = RockConfig()
+    rock_config.runtime.instance_registry_mirrors = ["old.example.com/ns"]
+    rock_config.nacos_provider = MagicMock()
+    rock_config.nacos_provider.get_config = AsyncMock(
+        return_value={
+            "runtime": {
+                "instance_registry_mirrors": [],
+            },
+        }
+    )
+
+    await rock_config.update()
+
+    assert rock_config.runtime.instance_registry_mirrors == []
+
+
+@pytest.mark.asyncio
+async def test_nacos_update_without_runtime_key_preserves_instance_mirrors():
+    rock_config = RockConfig()
+    rock_config.runtime.instance_registry_mirrors = ["keep.example.com/ns"]
+    rock_config.nacos_provider = MagicMock()
+    rock_config.nacos_provider.get_config = AsyncMock(return_value={"sandbox_config": {}})
+
+    await rock_config.update()
+
+    assert rock_config.runtime.instance_registry_mirrors == ["keep.example.com/ns"]
+
+
 # ===== _resolve_k8s_template_includes =====