feat(e2b): support ACR image resolution and cold-start overrides

docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.11.x/User Guides/e2b_sdk_compatibility.md

@@ -106,6 +106,28 @@ ROCK 期望部署网关对两个平面都进行身份验证。应用将传入的
 
 请求模型和创建映射位于 `rock/admin/proto/request.py` 和 `rock/admin/service/e2b_service.py`。
 
+#### 冷启动 metadata
+
+没有就绪模板、需要冷启动时，可以在 `metadata` 中指定以下可选字段，值均须为字符串。找到就绪模板时，这些字段会被忽略，资源使用模板配置。
+
+| 字段 | 用法 | 省略时的默认值 |
+|---|---|---|
+| `cpuCount` | CPU 核数，正整数，例如 `"4"`。 | 2 核 |
+| `memoryMB` | 内存大小，单位 MB，正整数，例如 `"16384"`（16 GiB）。 | 8192 MB |
+| `startup_timeout` | 提交后等待沙箱就绪的最长时间，单位秒，有限正数，例如 `"120"`。 | 85 秒 |
+
+例如，将以下字典作为 `POST /sandboxes` 请求中的 `metadata` 值，或作为 `metadata` 参数传给 `Sandbox.create()`：
+
+```json
+{
+  "cpuCount": "4",
+  "memoryMB": "16384",
+  "startup_timeout": "120"
+}
+```
+
+冷启动时，无效参数会返回 HTTP 400；CPU、内存申请仍受部署配置的资源上限约束。`startup_timeout` 包含提交后的就绪轮询与状态查询等待，不包含准备及提交耗时，也不会调整客户端或网关的 HTTP 超时。请求顶层的 `timeout` 仍控制沙箱存活时间。
+
 ### 获取沙箱信息
 
 支持 `sandbox.get_info()` 和 `Sandbox.get_info(id, ...)`。返回的 `SandboxInfo` 包含沙箱 ID、模板/镜像、元数据、开始/结束时间、状态、CPU 数量、内存和 `envdVersion`。ROCK 还在响应中报告磁盘大小，并向 `metadata` 添加 `e2b.agents.kruise.io/sandbox-ip`。详细信息中的 `template_id` 从当前沙箱状态派生，可能是解析后的镜像，而不是创建响应原样返回的 `template` 字符串。

docs/versioned_docs/version-1.11.x/User Guides/e2b_sdk_compatibility.md

@@ -129,6 +129,32 @@ The create response contains `sandboxID`, requested `templateID`, `clientID="roc
 The request model and create mapping are in `rock/admin/proto/request.py` and
 `rock/admin/service/e2b_service.py`.
 
+#### Cold-start metadata
+
+When no ready template is available, ROCK accepts the following optional fields in `metadata`. All values must
+be strings. If a ready template is found, these fields are ignored and the template's resources are used.
+
+| Field | Value | Default when omitted |
+|---|---|---|
+| `cpuCount` | Positive integer number of CPU cores, such as `"4"`. | 2 CPU cores |
+| `memoryMB` | Positive integer memory size in MB, such as `"16384"` (16 GiB). | 8192 MB |
+| `startup_timeout` | Positive finite number of seconds to wait for the sandbox to become ready after submission, such as `"120"`. | 85 seconds |
+
+For example, use the following object as the `metadata` value in a `POST /sandboxes` request or `Sandbox.create()` call:
+
+```json
+{
+  "cpuCount": "4",
+  "memoryMB": "16384",
+  "startup_timeout": "120"
+}
+```
+
+Invalid values are rejected with HTTP 400 on cold starts. Resource requests remain subject to the deployment's
+CPU and memory limits. `startup_timeout` covers readiness polling and status queries after submission; it does not
+include preparation/submission time or change client/gateway HTTP timeouts. The top-level `timeout` continues to
+control sandbox lifetime.
+
 ### Get sandbox information
 
 Both `sandbox.get_info()` and `Sandbox.get_info(id, ...)` are supported. The returned `SandboxInfo` includes sandbox

rock/admin/core/redis_key.py

@@ -1,6 +1,7 @@
 ALIVE_PREFIX = "alive:"
 TIMEOUT_PREFIX = "timeout:"
 OPENSANDBOX_SESSIONS_PREFIX = "opensandbox:sessions:"
+DADI_IMAGE_CACHE_PREFIX = "e2b:dadi:exists:v1:"
 
 
 def alive_sandbox_key(sandbox_id: str) -> str:
@@ -13,3 +14,7 @@ def timeout_sandbox_key(sandbox_id: str) -> str:
 
 def opensandbox_sessions_key(sandbox_id: str) -> str:
     return f"{OPENSANDBOX_SESSIONS_PREFIX}{sandbox_id}"
+
+
+def dadi_image_cache_key(probe_registry: str, mapped_image: str) -> str:
+    return f"{DADI_IMAGE_CACHE_PREFIX}{probe_registry}:{mapped_image}"

rock/admin/core/schema.py

@@ -148,7 +148,6 @@ class TemplateRecord(Base):
     status = Column(String(32), nullable=False)
     current_step = Column(String(32), nullable=True)
     artifact_uri = Column(String(1024), nullable=True)
-    fiber_pool_id = Column(String(128), nullable=True)
     execution_context = Column(_JSONB_VARIANT, nullable=True)
     error_code = Column(String(128), nullable=True)
     error_message = Column(Text, nullable=True)

rock/admin/core/template_table.py

@@ -21,23 +21,6 @@ class TemplateTable:
     def __init__(self, db_provider: DatabaseProvider) -> None:
         self._db = db_provider
 
-    @retry_on_disconnect
-    async def get_ready_fiber_pool_id(self, template_id: str) -> str | None:
-        return await self._db.run(self._get_ready_fiber_pool_id_sync, template_id)
-
-    def _get_ready_fiber_pool_id_sync(self, template_id: str) -> str | None:
-        with self._db.session_factory() as session:
-            stmt = (
-                select(TemplateRecord.fiber_pool_id)
-                .where(
-                    or_(TemplateRecord.template_id == template_id, TemplateRecord.image == template_id),
-                    TemplateRecord.status == _READY_STATUS,
-                )
-                .order_by((TemplateRecord.template_id == template_id).desc())
-                .limit(1)
-            )
-            return session.execute(stmt).scalar_one_or_none()
-
     @retry_on_disconnect
     async def get_ready_template(self, template_id: str) -> ReadyTemplate | None:
         return await self._db.run(self._get_ready_template_sync, template_id)

rock/admin/main.py

@@ -47,6 +47,7 @@
 )
 from rock.admin.service.e2b_proxy_service import E2BProxyService
 from rock.admin.service.e2b_service import E2BService
+from rock.admin.service.image_resolver import create_image_resolver
 from rock.admin.service.ops_service import OpsService
 from rock.common.exception import request_validation_exception_handler
 from rock.config import DatabaseConfig, RockConfig, SchedulerConfig
@@ -139,7 +140,7 @@ async def lifespan(app: FastAPI):
     rock_config = RockConfig.from_env(config_file_path)
     _apply_logging_config(rock_config)
 
-    # Override config from Nacos if available (sandbox_config, proxy_service, lifecycle via update(); scheduler separately)
+    # Override config from Nacos before creating services and startup-scoped plugins.
     if rock_config.nacos_provider:
         await rock_config.update()
         nacos_config = await rock_config.nacos_provider.get_config()
@@ -212,7 +213,6 @@ async def lifespan(app: FastAPI):
             ray_service=ray_service,
             redis_provider=redis_provider,
             nacos_provider=rock_config.nacos_provider,
-            template_table=template_table,
             k8s_config=rock_config.k8s,
             opensandbox_config=rock_config.opensandbox,
             remote_config=rock_config.remote,
@@ -239,7 +239,15 @@ async def lifespan(app: FastAPI):
                 meta_store=meta_store,
             )
         set_sandbox_manager(sandbox_manager)
-        set_e2b_sandbox_manager(E2BService(sandbox_manager, template_table))
+        set_e2b_sandbox_manager(
+            E2BService(
+                sandbox_manager,
+                template_table,
+                image_resolver=create_image_resolver(
+                    rock_config.e2b_image_resolver, rock_config.http_pool_manager, redis_provider
+                ),
+            )
+        )
         warmup_service = WarmupService(rock_config.warmup)
         await warmup_service.init()
         set_warmup_service(warmup_service)

rock/admin/proto/request.py

@@ -33,6 +33,14 @@ class E2BSetSandboxTimeoutRequest(BaseModel):
     model_config = ConfigDict(extra="forbid", strict=True)
 
     timeout: int = Field(ge=0, le=2_147_483_647)
+class E2BColdStartOptions(BaseModel):
+    """ROCK-specific metadata overrides, read only when no ready template exists."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    cpu_count: int | None = Field(default=None, alias="cpuCount", gt=0)
+    memory_mb: int | None = Field(default=None, alias="memoryMB", gt=0)
+    startup_timeout: float | None = Field(default=None, gt=0, allow_inf_nan=False)
 
 
 class E2BFilePathRequest(BaseModel):

rock/admin/service/acr_image_resolver.py

@@ -0,0 +1,206 @@
+"""Opt-in image registry mapping based on the responding ACR instance."""
+
+import asyncio
+from time import monotonic
+
+import httpx
+from requests.utils import parse_dict_header
+
+from rock.admin.core.redis_key import dadi_image_cache_key
+from rock.common.constants import (
+    DADI_IMAGE_REDIS_CACHE_TTL_SECONDS,
+    DADI_IMAGE_REDIS_TIMEOUT_SECONDS,
+    DADI_IMAGE_TAG_SUFFIXES,
+    REGISTRY_MANIFEST_ACCEPT,
+)
+from rock.logger import init_logger
+from rock.utils.docker import ImageUtil
+from rock.utils.providers.redis_provider import RedisProvider
+
+logger = init_logger(__name__)
+
+
+class AcrRegionImageResolver:
+    def __init__(
+        self,
+        client: httpx.AsyncClient,
+        source_registries: list[str] | None = None,
+        region_registry_mapping: dict[str, str] | None = None,
+        timeout_seconds: float = 1.0,
+        cache_ttl_seconds: float = 300.0,
+        failure_cache_ttl_seconds: float = 5.0,
+        probe_registry: str | None = None,
+        dadi_tag_suffixes: list[str] | tuple[str, ...] | None = DADI_IMAGE_TAG_SUFFIXES,
+        redis_provider: RedisProvider | None = None,
+        dadi_redis_cache_ttl_seconds: int = DADI_IMAGE_REDIS_CACHE_TTL_SECONDS,
+    ) -> None:
+        self._client = client
+        self._source_registries = frozenset(source_registries or [])
+        self._probe_registry = probe_registry
+        self._region_registry_mapping = dict(region_registry_mapping or {})
+        self._timeout_seconds = timeout_seconds
+        self._cache_ttl_seconds = cache_ttl_seconds
+        self._failure_cache_ttl_seconds = failure_cache_ttl_seconds
+        self._cache: dict[str, tuple[str | None, float]] = {}
+        self._locks = {
+            registry: asyncio.Lock() for registry in {probe_registry or source for source in self._source_registries}
+        }
+        self._dadi_tag_suffixes = tuple(dadi_tag_suffixes or ())
+        self._redis_provider = redis_provider
+        self._dadi_redis_cache_ttl_seconds = dadi_redis_cache_ttl_seconds
+
+    async def resolve(self, image: str) -> str:
+        registry, repository = ImageUtil.parse_registry_and_others(image)
+        if registry not in self._source_registries or not self._region_registry_mapping:
+            return image
+
+        registry = self._probe_registry or registry
+        cached = self._cache.get(registry)
+        if cached is not None and cached[1] > monotonic():
+            target = cached[0]
+        else:
+            try:
+                target = await asyncio.wait_for(self._resolve_target(registry), timeout=self._timeout_seconds)
+            except asyncio.TimeoutError:
+                return image
+        if not target:
+            return image
+        mapped_image = f"{target}/{repository}"
+        return await self._resolve_dadi_image(mapped_image, probe_registry=registry)
+
+    async def _resolve_target(self, registry: str) -> str | None:
+        async with self._locks[registry]:
+            cached = self._cache.get(registry)
+            if cached is not None and cached[1] > monotonic():
+                return cached[0]
+
+            target = None
+            try:
+                region = await self._probe_region(registry)
+                target = self._region_registry_mapping.get(region)
+            except Exception as error:
+                logger.warning("ACR region probe failed for %s (%s)", registry, type(error).__name__)
+            finally:
+                # Cancellation also leaves a short negative cache entry, so a
+                # timed-out leader does not trigger a new probe for each waiter.
+                ttl = self._cache_ttl_seconds if target else self._failure_cache_ttl_seconds
+                self._cache[registry] = (target, monotonic() + ttl)
+            return target
+
+    async def _probe_region(self, registry: str) -> str | None:
+        async with self._client.stream(
+            "GET",
+            f"https://{registry}/v2/",
+            auth=None,
+            follow_redirects=False,
+            timeout=self._timeout_seconds,
+        ) as response:
+            if response.status_code not in (200, 401):
+                return None
+            challenge = response.headers.get("www-authenticate", "")
+            scheme, _, parameters = challenge.partition(" ")
+            if scheme.lower() != "bearer":
+                return None
+            challenge_params = {key.strip().lower(): value for key, value in parse_dict_header(parameters).items()}
+            service = (challenge_params.get("service") or "").strip().strip('"').split(":")
+            if len(service) < 4 or service[0] != "registry.aliyuncs.com" or not service[3].startswith("cri-"):
+                return None
+            return service[1] or None
+
+    async def _resolve_dadi_image(self, image: str, *, probe_registry: str) -> str:
+        registry, reference = ImageUtil.parse_registry_and_others(image)
+        if not self._dadi_tag_suffixes or "@" in reference:
+            return image
+        repository, separator, tag = reference.rpartition(":")
+        if not separator:
+            repository, tag = reference, ImageUtil.DEFAULT_TAG
+        if not tag or tag.endswith(("_accelerated", *self._dadi_tag_suffixes)):
+            return image
+        for suffix in self._dadi_tag_suffixes:
+            candidate_tag = tag + suffix
+            if len(candidate_tag) > 128:
+                continue
+            candidate = f"{registry}/{repository}:{candidate_tag}"
+            key = (probe_registry, candidate)
+            if await self._get_dadi_redis_cache(key):
+                return candidate
+            try:
+                exists = await asyncio.wait_for(
+                    self._dadi_manifest_exists(probe_registry, repository, candidate_tag),
+                    timeout=self._timeout_seconds,
+                )
+            except Exception as error:
+                logger.warning("DADI image probe failed for %s (%s)", probe_registry, type(error).__name__)
+                continue
+            if exists:
+                # Persistence has its own timeout; it cannot discard a confirmed manifest.
+                await self._set_dadi_redis_cache(key)
+                return candidate
+        return image
+
+    async def _get_dadi_redis_cache(self, key: tuple[str, str]) -> bool:
+        if self._redis_provider is None:
+            return False
+        try:
+            cached = await asyncio.wait_for(
+                self._redis_provider.json_get(dadi_image_cache_key(*key), path="."),
+                timeout=DADI_IMAGE_REDIS_TIMEOUT_SECONDS,
+            )
+            return cached is True
+        except Exception as error:
+            logger.warning("DADI Redis cache read failed (%s)", type(error).__name__)
+            return False
+
+    async def _set_dadi_redis_cache(self, key: tuple[str, str]) -> None:
+        if self._redis_provider is None:
+            return
+        try:
+            await asyncio.wait_for(
+                self._redis_provider.json_set_with_ttl(
+                    dadi_image_cache_key(*key), "$", True, self._dadi_redis_cache_ttl_seconds
+                ),
+                timeout=DADI_IMAGE_REDIS_TIMEOUT_SECONDS,
+            )
+        except Exception as error:
+            logger.warning("DADI Redis cache write failed (%s)", type(error).__name__)
+
+    async def _dadi_manifest_exists(self, registry: str, repository: str, tag: str) -> bool:
+        origin = httpx.URL(f"https://{registry}")
+        manifest_url = f"{origin}/v2/{repository}/manifests/{tag}"
+        options = {"auth": None, "follow_redirects": False, "timeout": self._timeout_seconds}
+        headers = {"Accept": REGISTRY_MANIFEST_ACCEPT}
+        response = await self._client.head(manifest_url, headers=headers, **options)
+        if response.status_code != 401:
+            return response.status_code == 200
+
+        scheme, _, parameters = response.headers.get("www-authenticate", "").partition(" ")
+        if scheme.lower() != "bearer":
+            return False
+        challenge = {key.strip().lower(): value for key, value in parse_dict_header(parameters).items()}
+        realm = httpx.URL(challenge.get("realm") or "")
+        # Only ask the responding registry for an anonymous, pull-scoped token.
+        if (
+            realm.scheme != "https"
+            or realm.host != origin.host
+            or realm.port != origin.port
+            or realm.username
+            or realm.password
+            or not challenge.get("service")
+        ):
+            return False
+        response = await self._client.get(
+            realm,
+            params={"service": challenge["service"], "scope": f"repository:{repository}:pull"},
+            headers={"Accept": "application/json"},
+            **options,
+        )
+        if response.status_code != 200:
+            return False
+        payload = response.json()
+        token = payload.get("token") or payload.get("access_token")
+        if not isinstance(token, str) or not token:
+            return False
+        response = await self._client.head(
+            manifest_url, headers={**headers, "Authorization": f"Bearer {token}"}, **options
+        )
+        return response.status_code == 200