diff --git a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.11.x/User Guides/e2b_sdk_compatibility.md b/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.11.x/User Guides/e2b_sdk_compatibility.md index c3df598965..4e9a51b10d 100644 --- a/docs/i18n/zh-Hans/docusaurus-plugin-content-docs/version-1.11.x/User Guides/e2b_sdk_compatibility.md +++ b/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` 字符串。 diff --git a/docs/versioned_docs/version-1.11.x/User Guides/e2b_sdk_compatibility.md b/docs/versioned_docs/version-1.11.x/User Guides/e2b_sdk_compatibility.md index 324e11b210..af8268a578 100644 --- a/docs/versioned_docs/version-1.11.x/User Guides/e2b_sdk_compatibility.md +++ b/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 diff --git a/rock/admin/core/redis_key.py b/rock/admin/core/redis_key.py index af23e8c5bb..896f4f8a12 100644 --- a/rock/admin/core/redis_key.py +++ b/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}" diff --git a/rock/admin/core/schema.py b/rock/admin/core/schema.py index 2a1b48bea6..055d9aaa81 100644 --- a/rock/admin/core/schema.py +++ b/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) diff --git a/rock/admin/core/template_table.py b/rock/admin/core/template_table.py index d6f3729010..89d1b9fe13 100644 --- a/rock/admin/core/template_table.py +++ b/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) diff --git a/rock/admin/main.py b/rock/admin/main.py index 4e42f3c9f3..d62585601f 100644 --- a/rock/admin/main.py +++ b/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) diff --git a/rock/admin/proto/request.py b/rock/admin/proto/request.py index c93ff5a281..2e709d18d6 100644 --- a/rock/admin/proto/request.py +++ b/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): diff --git a/rock/admin/service/acr_image_resolver.py b/rock/admin/service/acr_image_resolver.py new file mode 100644 index 0000000000..07274b94b9 --- /dev/null +++ b/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 diff --git a/rock/admin/service/e2b_service.py b/rock/admin/service/e2b_service.py index 5e3d09b816..2399a95d04 100644 --- a/rock/admin/service/e2b_service.py +++ b/rock/admin/service/e2b_service.py @@ -1,5 +1,9 @@ +from typing import TYPE_CHECKING + +from pydantic import ValidationError + from rock.admin.core.template_table import TemplateTable -from rock.admin.proto.request import ClusterInfo, UserInfo +from rock.admin.proto.request import ClusterInfo, E2BColdStartOptions, UserInfo from rock.admin.proto.response import E2BSandboxInfo, SandboxStartResponse, SandboxStatusResponse from rock.admin.service.e2b_sandbox_info import e2b_sandbox_info_fields from rock.deployments.config import DockerDeploymentConfig @@ -9,13 +13,22 @@ from rock.sdk.common.exceptions import BadRequestRockError, E2BSandboxNotFoundError, SandboxNotFoundRockError from rock.utils.format import megabytes_to_size +if TYPE_CHECKING: + from rock.admin.service.image_resolver import ImageResolver + logger = init_logger(__name__) class E2BService: - def __init__(self, sandbox_manager: SandboxManager, template_table: TemplateTable) -> None: + def __init__( + self, + sandbox_manager: SandboxManager, + template_table: TemplateTable, + image_resolver: "ImageResolver | None" = None, + ) -> None: self._sandbox_manager = sandbox_manager self._template_table = template_table + self._image_resolver = image_resolver async def start( self, @@ -24,11 +37,23 @@ async def start( cluster_info: ClusterInfo = {}, ) -> SandboxStartResponse: template = await self._template_table.get_ready_template(config.image) + wait_options = {} if template is None: + try: + cold_start = E2BColdStartOptions.model_validate(config.metadata) + except ValidationError as error: + details = "; ".join(f"{item['loc'][0]}: {item['msg']}" for item in error.errors()) + raise BadRequestRockError(f"Invalid cold-start metadata: {details}") from None + if cold_start.startup_timeout is not None: + wait_options["wait_timeout"] = cold_start.startup_timeout logger.info("Template %s is not ready or does not exist; using raw manifest", config.image) template_config = config.model_copy( update={ "template_id": None, + "cpus": cold_start.cpu_count if cold_start.cpu_count is not None else config.cpus, + "memory": ( + megabytes_to_size(cold_start.memory_mb) if cold_start.memory_mb is not None else config.memory + ), "extended_params": {**config.extended_params, EXT_USE_RAW: EXT_USE_RAW_ENABLED}, } ) @@ -41,10 +66,16 @@ async def start( "disk": megabytes_to_size(template["disk_size_mb"]), } ) + if self._image_resolver is not None: + try: + template_config.image = await self._image_resolver.resolve(template_config.image) + except Exception as error: + logger.warning("Image resolution failed; keeping the original image (%s)", type(error).__name__) return await self._sandbox_manager.start_from_template( template_config, user_info=user_info, cluster_info=cluster_info, + **wait_options, ) @property diff --git a/rock/admin/service/image_resolver.py b/rock/admin/service/image_resolver.py new file mode 100644 index 0000000000..a8b762bf52 --- /dev/null +++ b/rock/admin/service/image_resolver.py @@ -0,0 +1,32 @@ +"""Optional deployment-configured image resolution for E2B sandbox starts.""" + +from inspect import signature +from typing import Protocol + +from rock.config import ImageResolverConfig +from rock.utils.http_pool import HttpPoolManager +from rock.utils.importer import safe_import_class +from rock.utils.providers.redis_provider import RedisProvider + + +class ImageResolver(Protocol): + async def resolve(self, image: str) -> str: + """Return a replacement image, or the original when resolution is unavailable.""" + ... + + +def create_image_resolver( + config: ImageResolverConfig, + http_pool_manager: HttpPoolManager, + redis_provider: RedisProvider | None = None, +) -> ImageResolver | None: + """Create one plugin, borrowing a client owned by the admin HTTP pool manager.""" + if not config.resolver_class: + return None + resolver_class = safe_import_class(config.resolver_class) + if not isinstance(resolver_class, type): + raise ValueError(f"Invalid image resolver class: {config.resolver_class}") + options = dict(config.options) + if "redis_provider" in signature(resolver_class).parameters: + options["redis_provider"] = redis_provider + return resolver_class(client=http_pool_manager.get("probe"), **options) diff --git a/rock/common/constants.py b/rock/common/constants.py index 9757fd1b8c..4a0dadc1f3 100644 --- a/rock/common/constants.py +++ b/rock/common/constants.py @@ -17,6 +17,17 @@ E2B_CLIENT_ID = "rock" E2B_ENVD_VERSION = "0.3.0" E2B_SANDBOX_IP_METADATA_KEY = "e2b.agents.kruise.io/sandbox-ip" +DADI_IMAGE_TAG_SUFFIXES = ("_accelerated", "_containerd_accelerated") +DADI_IMAGE_REDIS_CACHE_TTL_SECONDS = 30 * 24 * 60 * 60 +DADI_IMAGE_REDIS_TIMEOUT_SECONDS = 0.1 +REGISTRY_MANIFEST_ACCEPT = ", ".join( + ( + "application/vnd.oci.image.manifest.v1+json", + "application/vnd.oci.image.index.v1+json", + "application/vnd.docker.distribution.manifest.v2+json", + "application/vnd.docker.distribution.manifest.list.v2+json", + ) +) E2B_STATE_BY_ROCK_STATE: dict[str, Literal["running", "paused"]] = { "running": "running", "pending": "running", diff --git a/rock/config.py b/rock/config.py index 6d26719014..b7ea5e6a2a 100644 --- a/rock/config.py +++ b/rock/config.py @@ -323,6 +323,14 @@ class ImageRegistryMirror: password: str | None = None +@dataclass +class ImageResolverConfig: + """Optional image resolver plugin, instantiated once at admin startup.""" + + resolver_class: str | None = None + options: dict = field(default_factory=dict) + + @dataclass class PoolConfig: """Pool configuration with resource and port settings.""" @@ -602,6 +610,7 @@ class RockConfig: aes_encrypt_key: str | None = None scheduler: SchedulerConfig = field(default_factory=SchedulerConfig) database: DatabaseConfig = field(default_factory=DatabaseConfig) + e2b_image_resolver: ImageResolverConfig = field(default_factory=ImageResolverConfig) image_registry_mirrors: list[ImageRegistryMirror] = field(default_factory=list) image_mirror_lookup_allowlist: list[str] = field(default_factory=list) http_pools: dict[str, HttpPoolConfig] = field( @@ -677,6 +686,8 @@ def from_env(cls, config_path: str | None = None): kwargs["scheduler"] = SchedulerConfig(**config["scheduler"]) if "database" in config: kwargs["database"] = DatabaseConfig(**config["database"]) + if "e2b_image_resolver" in config: + kwargs["e2b_image_resolver"] = ImageResolverConfig(**(config["e2b_image_resolver"] or {})) if "http_pools" in config: raw_pools = config["http_pools"] or {} kwargs["http_pools"] = {name: HttpPoolConfig(**params) for name, params in raw_pools.items()} @@ -780,6 +791,7 @@ async def update(self): "sandbox_config": (SandboxConfig, "sandbox_config"), "proxy_service": (ProxyServiceConfig, "proxy_service"), "lifecycle": (SandboxLifecycleConfig, "lifecycle"), + "e2b_image_resolver": (ImageResolverConfig, "e2b_image_resolver"), } for key, (config_class, attr_name) in config_map.items(): diff --git a/rock/deployments/start_config.py b/rock/deployments/start_config.py index 2c9566d1cb..32fb8cf8a6 100644 --- a/rock/deployments/start_config.py +++ b/rock/deployments/start_config.py @@ -351,6 +351,8 @@ async def apply_start_config( rock_config: RockConfig, config: DockerDeploymentConfig, rock_authorization: str | None, + *, + apply_image_mirror: bool = True, ) -> None: """Apply the shared sandbox start normalization pipeline in dependency order.""" apply_auto_clear_default(rock_config, config) @@ -361,4 +363,5 @@ async def apply_start_config( await apply_timeout_defaults(rock_config, config) await apply_cpu_overcommit_default(rock_config, config, rock_authorization) await apply_disk_limits(rock_config, config) - await apply_image_registry_mirror(rock_config, config) + if apply_image_mirror: + await apply_image_registry_mirror(rock_config, config) diff --git a/rock/sandbox/operator/factory.py b/rock/sandbox/operator/factory.py index 37e64db8e3..4180354afc 100644 --- a/rock/sandbox/operator/factory.py +++ b/rock/sandbox/operator/factory.py @@ -8,7 +8,6 @@ from rock.logger import init_logger from rock.sandbox.operator.abstract import AbstractOperator from rock.sandbox.operator.k8s.operator import K8sOperator -from rock.sandbox.operator.k8s.provider import TemplateFiberPoolLookup from rock.sandbox.operator.opensandbox.operator import OpenSandboxOperator from rock.sandbox.operator.ray import RayOperator from rock.sandbox.operator.remote.operator import RemoteOperator @@ -53,7 +52,6 @@ class OperatorContext: # K8s operator dependencies k8s_config: K8sConfig | None = None nacos_provider: NacosConfigProvider | None = None - template_table: TemplateFiberPoolLookup | None = None # OpenSandbox operator dependencies opensandbox_config: OpenSandboxConfig | None = None # Remote operator dependencies @@ -98,7 +96,7 @@ def create_operator(context: OperatorContext) -> AbstractOperator: if context.k8s_config is None: raise ValueError("K8sConfig is required for K8sOperator") logger.info("Creating K8sOperator") - k8s_operator = K8sOperator(k8s_config=context.k8s_config, template_table=context.template_table) + k8s_operator = K8sOperator(k8s_config=context.k8s_config) if context.redis_provider is not None: k8s_operator.set_redis_provider(context.redis_provider) if context.nacos_provider is not None: diff --git a/rock/sandbox/operator/k8s/operator.py b/rock/sandbox/operator/k8s/operator.py index 46384d4c3e..1be3742b4e 100644 --- a/rock/sandbox/operator/k8s/operator.py +++ b/rock/sandbox/operator/k8s/operator.py @@ -9,7 +9,7 @@ from rock.logger import init_logger from rock.sandbox.operator.abstract import AbstractOperator from rock.sandbox.operator.k8s.constants import K8sConstants -from rock.sandbox.operator.k8s.provider import BatchSandboxProvider, TemplateSpec, TemplateFiberPoolLookup +from rock.sandbox.operator.k8s.provider import BatchSandboxProvider, TemplateSpec logger = init_logger(__name__) @@ -64,16 +64,14 @@ def __init__( self, k8s_config: K8sConfig, redis_provider=None, - template_table: TemplateFiberPoolLookup | None = None, ): """Initialize K8s operator. Args: k8s_config: K8sConfig object containing kubeconfig and templates redis_provider: Optional Redis provider for caching sandbox info - template_table: Optional READY-template fiber pool lookup """ - self._provider = BatchSandboxProvider(k8s_config=k8s_config, template_table=template_table) + self._provider = BatchSandboxProvider(k8s_config=k8s_config) self._redis_provider = redis_provider logger.info("Initialized K8sOperator") diff --git a/rock/sandbox/operator/k8s/provider.py b/rock/sandbox/operator/k8s/provider.py index 3bad4f92f9..2f8c60f0e6 100644 --- a/rock/sandbox/operator/k8s/provider.py +++ b/rock/sandbox/operator/k8s/provider.py @@ -72,8 +72,6 @@ def generate_template_id(spec: TemplateSpec) -> str: raw = "|".join(parts) digest = hashlib.sha256(raw.encode()).hexdigest()[:16] return f"{K8sConstants.TEMPLATE_ID_PREFIX}{digest}" -class TemplateFiberPoolLookup(Protocol): - async def get_ready_fiber_pool_id(self, template_id: str) -> str | None: ... class PoolSelector(ABC): @@ -312,12 +310,11 @@ class BatchSandboxProvider(K8sProvider): The watch task runs in the background and automatically reconnects on network failures. """ - def __init__(self, k8s_config: K8sConfig, template_table: TemplateFiberPoolLookup | None = None): + def __init__(self, k8s_config: K8sConfig): """Initialize BatchSandbox provider. Args: k8s_config: K8sConfig object containing kubeconfig and templates - template_table: Optional READY-template fiber pool lookup """ self.kubeconfig_path = k8s_config.kubeconfig_path self.namespace = k8s_config.namespace @@ -327,7 +324,6 @@ def __init__(self, k8s_config: K8sConfig, template_table: TemplateFiberPoolLooku self._pool_api: K8sApiClient | None = None self._initialized = False self._nacos_provider = None - self._template_table = template_table self._image_auth_key = self._load_image_auth_key(k8s_config) # Initialize template loader with config templates and pool template @@ -606,7 +602,7 @@ async def _get_pool_name(self, config: DockerDeploymentConfig) -> str | None: Priority: 1. Check extended_params for explicit pool name 2. Use ResourceMatchingPoolSelector to find best matching pool - 3. Look up the READY template's fiber pool in the database + 3. Use config.template_id as the pool name Args: config: Docker deployment configuration @@ -623,10 +619,10 @@ async def _get_pool_name(self, config: DockerDeploymentConfig) -> str | None: pools = await self._get_pools() logger.info(f"Available pools from Nacos: {list(pools.keys())}") pool_name = ResourceMatchingPoolSelector().select_pool(config, pools) - if pool_name or self._template_table is None: + if pool_name: return pool_name - return await self._template_table.get_ready_fiber_pool_id(config.image) + return config.template_id async def _get_template_name(self, config: DockerDeploymentConfig) -> str: """Get template name from extended_params or Nacos template_rules. diff --git a/rock/sandbox/sandbox_manager.py b/rock/sandbox/sandbox_manager.py index 9f20db20d7..75422837c3 100644 --- a/rock/sandbox/sandbox_manager.py +++ b/rock/sandbox/sandbox_manager.py @@ -154,12 +154,18 @@ async def start_async( cluster_info: ClusterInfo = {}, *, use_template_resource_spec: bool = False, + apply_image_mirror: bool = True, ) -> SandboxStartResponse: await self._check_sandbox_exists_in_redis(config) if isinstance(config, DockerDeploymentConfig): await self.rock_config.update() - await apply_start_config(self.rock_config, config, user_info.get("rock_authorization")) + await apply_start_config( + self.rock_config, + config, + user_info.get("rock_authorization"), + apply_image_mirror=apply_image_mirror, + ) with StageTimer("startup_timing", f"[{config.image}] Init config", logger): if isinstance(config, DockerDeploymentConfig): @@ -329,12 +335,16 @@ async def start_from_template( config: DeploymentConfig, user_info: UserInfo = {}, cluster_info: ClusterInfo = {}, + *, + wait_timeout: float | None = None, ) -> SandboxStartResponse: return await self._start_and_wait( config, user_info=user_info, cluster_info=cluster_info, use_template_resource_spec=True, + apply_image_mirror=False, + wait_timeout=wait_timeout, ) async def _start_and_wait( @@ -344,23 +354,31 @@ async def _start_and_wait( cluster_info: ClusterInfo, *, use_template_resource_spec: bool = False, + wait_timeout: float | None = None, + apply_image_mirror: bool = True, ) -> SandboxStartResponse: response = await self.start_async( config, user_info=user_info, cluster_info=cluster_info, use_template_resource_spec=use_template_resource_spec, + apply_image_mirror=apply_image_mirror, ) sandbox_id = response.sandbox_id - deadline = time.time() + REQUEST_TIMEOUT_SECONDS - with StageTimer("startup_timing", f"[{sandbox_id}] Wait sandbox running", logger): + wait_timeout = REQUEST_TIMEOUT_SECONDS if wait_timeout is None else wait_timeout + + async def poll_until_running(): while True: await asyncio.sleep(1) status = await self.get_status(sandbox_id) if status.is_alive and status.host_ip: - break - if time.time() >= deadline: - raise TimeoutError(f"sandbox {sandbox_id} not running after {REQUEST_TIMEOUT_SECONDS}s") + return + + with StageTimer("startup_timing", f"[{sandbox_id}] Wait sandbox running", logger): + try: + await asyncio.wait_for(poll_until_running(), timeout=wait_timeout) + except asyncio.TimeoutError as exc: + raise TimeoutError(f"sandbox {sandbox_id} not running after {wait_timeout}s") from exc return response @monitor_sandbox_operation() diff --git a/sql/template.sql b/sql/template.sql index dbeb7a6806..f4aa1bf734 100644 --- a/sql/template.sql +++ b/sql/template.sql @@ -9,7 +9,6 @@ CREATE TABLE template ( status VARCHAR(32) NOT NULL, current_step VARCHAR(32), artifact_uri VARCHAR(1024), - fiber_pool_id VARCHAR(128), execution_context JSONB, error_code VARCHAR(128), error_message TEXT, @@ -19,4 +18,4 @@ CREATE TABLE template ( ); CREATE INDEX ix_template_status ON template (status); -CREATE UNIQUE INDEX ux_template_image ON template (image); +CREATE INDEX ix_template_image ON template (image); diff --git a/tests/integration/admin/test_e2b_sdk_commands_run.py b/tests/integration/admin/test_e2b_sdk_commands_run.py deleted file mode 100644 index 4a1b5388bb..0000000000 --- a/tests/integration/admin/test_e2b_sdk_commands_run.py +++ /dev/null @@ -1,403 +0,0 @@ -import shlex -import socket -import threading -import time -from contextlib import contextmanager -from pathlib import Path -from unittest.mock import AsyncMock, MagicMock - -import pytest -import uvicorn -from e2b import ( - CommandExitException, - FileNotFoundException, - FileType, - Sandbox, - SandboxNotFoundException, - TimeoutException, -) -from fastapi import FastAPI, Request -from httpx import AsyncClient, Client -from httpx import Request as HTTPXRequest - -from rock.actions.sandbox.response import State -from rock.admin.entrypoints import e2b_api as e2b_api_module -from rock.admin.entrypoints import e2b_proxy_api as e2b_proxy_api_module -from rock.admin.entrypoints import sandbox_api as sandbox_api_module -from rock.admin.entrypoints.e2b_api import e2b_router, set_e2b_service -from rock.admin.entrypoints.e2b_proxy_api import e2b_proxy_router, set_e2b_proxy_service -from rock.admin.proto.response import SandboxStartResponse, SandboxStatusResponse -from rock.admin.service.e2b_proxy_service import E2BProxyService -from rock.admin.service.e2b_service import E2BService -from rock.deployments.constants import Port -from rock.sandbox.sandbox_manager import SandboxManager -from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService -from rock.sdk.common.exceptions import SandboxNotFoundRockError - -_MISSING = object() -pytestmark = [pytest.mark.integration] - - -@contextmanager -def _serve_tcp(app: FastAPI): - listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - listener.bind(("127.0.0.1", 0)) - listener.listen() - port = listener.getsockname()[1] - server = uvicorn.Server( - uvicorn.Config( - app, - log_level="error", - access_log=False, - lifespan="off", - ) - ) - thread = threading.Thread( - target=server.run, - kwargs={"sockets": [listener]}, - daemon=True, - ) - thread.start() - - deadline = time.monotonic() + 5 - while not server.started and thread.is_alive() and time.monotonic() < deadline: - time.sleep(0.01) - if not server.started: - server.should_exit = True - thread.join(timeout=5) - listener.close() - raise RuntimeError("test ASGI server did not start") - - try: - yield f"http://127.0.0.1:{port}" - finally: - server.should_exit = True - thread.join(timeout=5) - listener.close() - - -class _HTTPClient: - async def request(self, **kwargs): - async with AsyncClient() as client: - return await client.request(**kwargs) - - def build_request(self, **kwargs): - return HTTPXRequest(**kwargs) - - async def send(self, request, *, stream=False): - client = AsyncClient() - try: - response = await client.send(request, stream=stream) - except Exception: - await client.aclose() - raise - if not stream: - await client.aclose() - return response - - close_response = response.aclose - - async def close(): - await close_response() - await client.aclose() - - response.aclose = close - return response - - -@pytest.fixture -def e2b_command_stack(rocklet_remote_server, monkeypatch): - api_key = "e2b_0000000000000000000000000000000000000000" - control_headers = [] - data_headers = [] - - control_app = FastAPI() - - @control_app.middleware("http") - async def capture_control_api_key(request: Request, call_next): - control_headers.append((request.url.path, request.headers.get("x-api-key"))) - return await call_next(request) - - control_app.include_router(e2b_router) - control_app.include_router(sandbox_api_module.sandbox_router, prefix="/apis/envs/sandbox/v1") - - sandbox_record = { - "sandbox_id": "sandbox-123", - "state": State.RUNNING, - "host_ip": "127.0.0.1", - "port_mapping": {int(Port.PROXY): rocklet_remote_server.port}, - "env": {"BASE": "sandbox"}, - } - timeout_state = {"value": None} - meta_store = MagicMock() - - async def get_sandbox_record(sandbox_id: str, check_db: bool = False): - return sandbox_record if sandbox_id == sandbox_record["sandbox_id"] else None - - async def update_timeout(sandbox_id: str, timeout_info: dict[str, str]): - assert sandbox_id == sandbox_record["sandbox_id"] - timeout_state["value"] = timeout_info - - meta_store.get = AsyncMock(side_effect=get_sandbox_record) - meta_store.get_timeout = AsyncMock(side_effect=lambda sandbox_id: timeout_state["value"]) - meta_store.update_timeout = AsyncMock(side_effect=update_timeout) - - control_manager = SandboxManager.__new__(SandboxManager) - control_manager._meta_store = meta_store - - async def get_control_status( - sandbox_id: str, - include_all_states: bool = False, - refresh_timeout: bool = True, - ) -> SandboxStatusResponse: - if sandbox_id != sandbox_record["sandbox_id"]: - raise SandboxNotFoundRockError(f"Sandbox {sandbox_id} not found") - return SandboxStatusResponse(sandbox_id=sandbox_id, state=State.RUNNING) - - control_manager.get_status = AsyncMock(side_effect=get_control_status) - control_manager.start_from_template = AsyncMock( - return_value=SandboxStartResponse( - sandbox_id="sandbox-123", - host_name="sandbox-123", - host_ip="127.0.0.1", - ) - ) - template_table = MagicMock() - template_table.get_ready_template = AsyncMock(return_value=None) - control_service = E2BService(control_manager, template_table) - monkeypatch.setattr(sandbox_api_module, "sandbox_manager", control_manager, raising=False) - - sandbox_manager = SandboxProxyService.__new__(SandboxProxyService) - sandbox_manager.metrics_monitor = None - sandbox_manager._meta_store = meta_store - sandbox_manager._rpc_client = _HTTPClient() - sandbox_manager._proxy_client = _HTTPClient() - - proxy_service = E2BProxyService( - meta_store=meta_store, - sandbox_manager=sandbox_manager, - ) - data_app = FastAPI() - - @data_app.middleware("http") - async def capture_data_access_token(request: Request, call_next): - data_headers.append((request.url.path, request.headers.get("x-access-token"))) - return await call_next(request) - - data_app.include_router(e2b_proxy_router) - - previous_control_service = getattr(e2b_api_module, "e2b_service", _MISSING) - previous_proxy_service = getattr(e2b_proxy_api_module, "e2b_proxy_service", _MISSING) - try: - set_e2b_service(control_service) - set_e2b_proxy_service(proxy_service) - with _serve_tcp(control_app) as api_url, _serve_tcp(data_app) as sandbox_url: - yield { - "api_url": api_url, - "sandbox_url": sandbox_url, - "api_key": api_key, - "control_headers": control_headers, - "data_headers": data_headers, - "timeout_state": timeout_state, - "control_manager": control_manager, - } - finally: - if previous_control_service is _MISSING: - delattr(e2b_api_module, "e2b_service") - else: - set_e2b_service(previous_control_service) - if previous_proxy_service is _MISSING: - delattr(e2b_proxy_api_module, "e2b_proxy_service") - else: - set_e2b_proxy_service(previous_proxy_service) - - -def _create_sandbox(stack) -> Sandbox: - return Sandbox.create( - template="linux-dind", - timeout=60, - api_url=stack["api_url"], - sandbox_url=stack["sandbox_url"], - api_key=stack["api_key"], - ) - - -def test_commands_run_success_end_to_end(e2b_command_stack, tmp_path): - sandbox = _create_sandbox(e2b_command_stack) - result = sandbox.commands.run( - 'printf "%s:%s\\n" "$BASE" "$REQUEST"; pwd; printf "warning\\n" >&2', - stdin=False, - envs={"REQUEST": "request"}, - cwd=str(tmp_path), - timeout=2.5, - user="root", - ) - - stdout_lines = result.stdout.splitlines() - assert stdout_lines[0] == "sandbox:request" - assert Path(stdout_lines[1]).resolve() == tmp_path.resolve() - assert result.stderr == "warning\n" - assert result.exit_code == 0 - assert ("/sandboxes", e2b_command_stack["api_key"]) in e2b_command_stack["control_headers"] - assert ("/process.Process/Start", e2b_command_stack["api_key"]) in e2b_command_stack["data_headers"] - - -@pytest.mark.parametrize("query,refresh_timeout", [({}, True), ({"refresh_timeout": "false"}, False)]) -def test_get_status_refresh_timeout_query_end_to_end(e2b_command_stack, query, refresh_timeout): - sandbox = _create_sandbox(e2b_command_stack) - - with Client(base_url=e2b_command_stack["api_url"]) as client: - response = client.get( - "/apis/envs/sandbox/v1/get_status", - params={"sandbox_id": sandbox.sandbox_id, **query}, - ) - - assert response.status_code == 200 - assert response.json()["result"]["sandbox_id"] == sandbox.sandbox_id - e2b_command_stack["control_manager"].get_status.assert_awaited_once_with( - sandbox.sandbox_id, - include_all_states=False, - refresh_timeout=refresh_timeout, - ) - - -def test_set_timeout_end_to_end(e2b_command_stack): - sandbox = _create_sandbox(e2b_command_stack) - - before = int(time.time()) - assert sandbox.set_timeout(120) is None - after = int(time.time()) - - timeout_info = e2b_command_stack["timeout_state"]["value"] - assert before + 120 <= int(timeout_info["expire_time"]) <= after + 120 - assert timeout_info["auto_clear_time"] == "2" - e2b_command_stack["control_manager"].get_status.assert_awaited_with( - "sandbox-123", - include_all_states=True, - refresh_timeout=False, - ) - assert ("/sandboxes/sandbox-123/timeout", e2b_command_stack["api_key"]) in e2b_command_stack["control_headers"] - - sandbox.commands.run("true", stdin=False, timeout=2.5) - assert e2b_command_stack["timeout_state"]["value"]["auto_clear_time"] == "2" - - assert sandbox.set_timeout(0) is None - assert int(e2b_command_stack["timeout_state"]["value"]["expire_time"]) <= int(time.time()) - assert e2b_command_stack["timeout_state"]["value"]["auto_clear_time"] == "0" - - with pytest.raises(SandboxNotFoundException, match="Sandbox missing not found"): - Sandbox.set_timeout( - "missing", - 120, - api_url=e2b_command_stack["api_url"], - api_key=e2b_command_stack["api_key"], - ) - - -def test_filesystem_operations_end_to_end(e2b_command_stack, tmp_path): - sandbox = _create_sandbox(e2b_command_stack) - root = tmp_path / "e2b-files" - nested = root / "nested" - single_file = root / "single.txt" - batch_text_file = nested / "batch.txt" - batch_bytes_file = nested / "batch.bin" - text_content = "hello, ROCK filesystem\n" - bytes_content = b"\x00\xffrock-bytes" - - assert sandbox.files.make_dir(str(nested), user="root") is True - assert sandbox.files.make_dir(str(nested), user="root") is False - - written = sandbox.files.write(str(single_file), text_content) - assert written.path == str(single_file) - assert written.type is FileType.FILE - - batch_written = sandbox.files.write_files( - [ - {"path": str(batch_text_file), "data": "batch text\n"}, - {"path": str(batch_bytes_file), "data": bytes_content}, - ] - ) - assert [entry.path for entry in batch_written] == [str(batch_text_file), str(batch_bytes_file)] - assert [entry.type for entry in batch_written] == [FileType.FILE, FileType.FILE] - - assert sandbox.files.read(str(single_file), format="text") == text_content - assert sandbox.files.read(str(batch_bytes_file), format="bytes") == bytearray(bytes_content) - - shallow_paths = {Path(entry.path).relative_to(root).as_posix() for entry in sandbox.files.list(str(root))} - assert shallow_paths == {"nested", "single.txt"} - - deep_entries = { - Path(entry.path).relative_to(root).as_posix(): entry for entry in sandbox.files.list(str(root), depth=2) - } - assert set(deep_entries) == {"nested", "nested/batch.bin", "nested/batch.txt", "single.txt"} - assert deep_entries["nested"].type is FileType.DIR - assert deep_entries["nested/batch.bin"].type is FileType.FILE - - file_info = sandbox.files.get_info(str(batch_bytes_file)) - assert file_info.path == str(batch_bytes_file) - assert file_info.type is FileType.FILE - assert file_info.size == len(bytes_content) - assert file_info.permissions.startswith("-") - assert file_info.modified_time.tzinfo is not None - - directory_info = sandbox.files.get_info(str(nested)) - assert directory_info.path == str(nested) - assert directory_info.type is FileType.DIR - assert directory_info.permissions.startswith("d") - - canonical_file = root / "canonical.txt" - canonical_written = sandbox.files.write(str(nested / ".." / canonical_file.name), "canonical\n") - assert canonical_written.path == str(canonical_file) - assert sandbox.files.read(str(canonical_file)) == "canonical\n" - - with pytest.raises(FileNotFoundException): - sandbox.files.read(str(root / "missing.txt")) - - -def test_commands_run_nonzero_exit_end_to_end(e2b_command_stack): - sandbox = _create_sandbox(e2b_command_stack) - - with pytest.raises(CommandExitException) as error: - sandbox.commands.run( - "printf 'before exit\\n'; printf 'failure\\n' >&2; exit 23", - stdin=False, - timeout=2.5, - ) - - assert error.value.stdout == "before exit\n" - assert error.value.stderr == "failure\n" - assert error.value.exit_code == 23 - assert error.value.error == "exit status 23" - - -def test_commands_run_background_end_to_end(e2b_command_stack, tmp_path): - sandbox = _create_sandbox(e2b_command_stack) - release_file = tmp_path / "release-command" - command = f"while [ ! -f {shlex.quote(str(release_file))} ]; do sleep 0.01; done; printf background" - - handle = sandbox.commands.run( - command, - background=True, - stdin=False, - timeout=2.5, - ) - assert not release_file.exists() - release_file.touch() - result = handle.wait() - - assert handle.pid > 0 - assert result.stdout == "background" - assert result.stderr == "" - assert result.exit_code == 0 - - -def test_commands_run_timeout_end_to_end(e2b_command_stack): - sandbox = _create_sandbox(e2b_command_stack) - - with pytest.raises(TimeoutException, match="command timed out"): - sandbox.commands.run( - "sleep 1", - stdin=False, - timeout=0.01, - ) diff --git a/tests/unit/admin/core/test_template_table.py b/tests/unit/admin/core/test_template_table.py index 9405ceb1f0..9f1695071b 100644 --- a/tests/unit/admin/core/test_template_table.py +++ b/tests/unit/admin/core/test_template_table.py @@ -24,7 +24,7 @@ def insert() -> None: await db_provider.run(insert) -async def test_get_fiber_pool_id_returns_ready_template_pool(db_provider): +async def test_get_ready_template_returns_ready_template(db_provider): table = TemplateTable(db_provider) await _insert_template( db_provider, @@ -33,7 +33,6 @@ async def test_get_fiber_pool_id_returns_ready_template_pool(db_provider): image="registry.example.com/rock/template-ready:latest", os_type="linux", status="READY", - fiber_pool_id="pool-from-db", cpu_count=4, memory_mb=8192, disk_size_mb=20480, @@ -42,7 +41,6 @@ async def test_get_fiber_pool_id_returns_ready_template_pool(db_provider): ), ) - assert await table.get_ready_fiber_pool_id("template-ready") == "pool-from-db" assert await table.get_ready_template("template-ready") == { "image": "registry.example.com/rock/template-ready:latest", "cpu_count": 4, @@ -61,7 +59,6 @@ async def test_get_ready_template_accepts_image_as_identifier(db_provider): image=image, os_type="linux", status="READY", - fiber_pool_id="pool-by-image", cpu_count=8, memory_mb=16384, disk_size_mb=40960, @@ -76,7 +73,6 @@ async def test_get_ready_template_accepts_image_as_identifier(db_provider): "memory_mb": 16384, "disk_size_mb": 40960, } - assert await table.get_ready_fiber_pool_id(image) == "pool-by-image" async def test_get_ready_template_supports_null_image(db_provider): @@ -130,7 +126,6 @@ async def test_get_ready_template_prefers_template_id_over_image(db_provider): TemplateRecord( template_id="shared-identifier", image="registry.example.com/rock/by-id:latest", - fiber_pool_id="pool-by-id", os_type="linux", status="READY", cpu_count=2, @@ -144,7 +139,6 @@ async def test_get_ready_template_prefers_template_id_over_image(db_provider): TemplateRecord( template_id="another-template", image="shared-identifier", - fiber_pool_id="pool-by-image", os_type="linux", status="READY", cpu_count=8, @@ -160,7 +154,6 @@ async def test_get_ready_template_prefers_template_id_over_image(db_provider): "memory_mb": 4096, "disk_size_mb": 10240, } - assert await table.get_ready_fiber_pool_id("shared-identifier") == "pool-by-id" async def test_template_image_must_be_unique(db_provider): @@ -200,7 +193,7 @@ async def test_template_image_must_be_unique(db_provider): @pytest.mark.parametrize("status", ["PENDING", "CREATING", "FAILED"]) -async def test_get_fiber_pool_id_ignores_non_ready_template(db_provider, status): +async def test_get_ready_template_ignores_non_ready_template(db_provider, status): table = TemplateTable(db_provider) await _insert_template( db_provider, @@ -209,7 +202,6 @@ async def test_get_fiber_pool_id_ignores_non_ready_template(db_provider, status) image="registry.example.com/rock/template-pending:latest", os_type="linux", status=status, - fiber_pool_id="pool-not-ready", cpu_count=4, memory_mb=8192, disk_size_mb=20480, @@ -218,5 +210,4 @@ async def test_get_fiber_pool_id_ignores_non_ready_template(db_provider, status) ), ) - assert await table.get_ready_fiber_pool_id("template-pending") is None assert await table.get_ready_template("template-pending") is None diff --git a/tests/unit/admin/entrypoints/test_e2b_api.py b/tests/unit/admin/entrypoints/test_e2b_api.py index aa33210f32..c69018c694 100644 --- a/tests/unit/admin/entrypoints/test_e2b_api.py +++ b/tests/unit/admin/entrypoints/test_e2b_api.py @@ -189,6 +189,41 @@ async def test_create_sandbox_falls_back_to_template_id_when_image_is_empty(e2b_ assert config.image == "template-without-image" +@pytest.mark.parametrize("ready_template", [False, True]) +async def test_create_sandbox_applies_metadata_overrides_only_on_cold_start(e2b_app, ready_template): + app, manager, templates = e2b_app + if not ready_template: + templates.get_ready_template.return_value = None + metadata = {"cpuCount": "3", "memoryMB": "12288", "startup_timeout": "120", "custom-key": "value"} + + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + response = await client.post( + "/sandboxes", json={"templateID": "image-or-template", "timeout": 1800, "metadata": metadata} + ) + + assert response.status_code == 201 + config = manager.start_from_template.await_args.args[0] + assert config.cpus == (4 if ready_template else 3) + assert config.memory == ("16g" if ready_template else "12g") + assert config.auto_clear_time_minutes == 30 + assert config.metadata == metadata + assert manager.start_from_template.await_args.kwargs.get("wait_timeout") == (None if ready_template else 120) + + +@pytest.mark.parametrize(("key", "value"), [("cpuCount", "nan"), ("memoryMB", "0"), ("startup_timeout", "-1")]) +async def test_create_sandbox_rejects_invalid_cold_start_metadata(e2b_app, key, value): + app, manager, templates = e2b_app + templates.get_ready_template.return_value = None + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: + response = await client.post( + "/sandboxes", json={"templateID": "raw-image", "timeout": 1800, "metadata": {key: value}} + ) + + assert response.status_code == 400 + assert key in response.json()["message"] + manager.start_from_template.assert_not_awaited() + + @pytest.mark.parametrize("timeout", [0, True, "3600"]) @pytest.mark.asyncio async def test_create_sandbox_returns_400_for_invalid_timeout(e2b_app, timeout): diff --git a/tests/unit/admin/service/test_e2b_image_resolution.py b/tests/unit/admin/service/test_e2b_image_resolution.py new file mode 100644 index 0000000000..868e5d05cf --- /dev/null +++ b/tests/unit/admin/service/test_e2b_image_resolution.py @@ -0,0 +1,431 @@ +"""Core image-resolution behavior through E2B sandbox start.""" + +import asyncio +from unittest.mock import AsyncMock, Mock + +import httpx +import pytest +from fakeredis.aioredis import FakeRedis + +from rock.admin.service.acr_image_resolver import AcrRegionImageResolver +from rock.admin.service.e2b_service import E2BService +from rock.admin.service.image_resolver import create_image_resolver +from rock.config import ImageResolverConfig +from rock.deployments.config import DockerDeploymentConfig +from rock.utils.providers.redis_provider import RedisProvider + + +def _challenge(region="cn-zhangjiakou"): + return {"www-authenticate": f'Bearer service="registry.aliyuncs.com:{region}:china:cri-test"'} + + +@pytest.fixture +async def make_service(): + clients = [] + + def build(handler, template=None, **options): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + clients.append(client) + templates = AsyncMock() + templates.get_ready_template.return_value = template + manager = AsyncMock() + options.setdefault("dadi_tag_suffixes", None) + redis_provider = options.pop("redis_provider", None) + pool = Mock() + pool.get.return_value = client + resolver = create_image_resolver( + ImageResolverConfig( + resolver_class="rock.admin.service.acr_image_resolver.AcrRegionImageResolver", + options={ + "source_registries": ["sg.example.com", "sh.example.com"], + "region_registry_mapping": { + "cn-zhangjiakou": "zjk.example.com", + "cn-shanghai": "sh-target.example.com", + "cn-beijing": "bj.example.com", + }, + **options, + }, + ), + pool, + redis_provider, + ) + return E2BService(manager, templates, image_resolver=resolver), manager + + yield build + for client in clients: + await client.aclose() + + +@pytest.fixture +async def dadi_redis(): + provider = RedisProvider(host="", port=0, password="") + provider.client = FakeRedis(decode_responses=True) + yield provider + await provider.close_pool() + + +@pytest.mark.parametrize( + ("source", "region", "target", "repository", "probe_registry"), + [ + ("sg.example.com", "cn-zhangjiakou", "zjk.example.com", "team/task:version", None), + ("sh.example.com", "cn-shanghai", "sh-target.example.com", "team/task:version", None), + ("sg.example.com", "cn-beijing", "bj.example.com", "team/task@sha256:" + "a" * 64, None), + ("sh.example.com", "cn-zhangjiakou", "zjk.example.com", "team/task:version", "sg.example.com"), + ], +) +async def test_start_maps_template_image_and_preserves_suffix( + make_service, source, region, target, repository, probe_registry +): + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(401, headers=_challenge(region)) + + image = f"{source}/{repository}" + template = {"image": image, "cpu_count": 4, "memory_mb": 8192, "disk_size_mb": 51200} + service, manager = make_service(respond, template=template, probe_registry=probe_registry) + config = DockerDeploymentConfig(image="template-id", template_id="template-id") + await service.start(config) + + passed = manager.start_from_template.call_args[0][0] + assert passed.image == f"{target}/{repository}" + assert passed.template_id == "template-id" + assert config.image == "template-id" + assert template["image"] == image + assert str(requests[0].url) == f"https://{probe_registry or source}/v2/" + assert "authorization" not in requests[0].headers + + +@pytest.mark.parametrize("case", ["unlisted", "unknown-region", "missing-challenge", "http-error", "dns-error"]) +async def test_start_keeps_original_image_when_resolution_is_unavailable(make_service, case): + requests = [] + + def respond(request): + requests.append(request) + if case == "dns-error": + raise httpx.ConnectError("DNS unavailable", request=request) + if case == "unknown-region": + return httpx.Response(401, headers=_challenge("ap-southeast-1")) + if case == "missing-challenge": + return httpx.Response(401) + return httpx.Response(500, headers=_challenge()) + + service, manager = make_service(respond, probe_registry="probe.example.com") + source = "sg.example.com.other.example" if case == "unlisted" else "sg.example.com" + image = f"{source}/team/task:version" + await service.start(DockerDeploymentConfig(image=image)) + + assert manager.start_from_template.call_args[0][0].image == image + assert len(requests) == (0 if case == "unlisted" else 1) + if requests: + assert str(requests[0].url) == "https://probe.example.com/v2/" + + +@pytest.mark.parametrize("outcome", ["success", "dns-error", "timeout"]) +async def test_concurrent_starts_share_bounded_probe_and_cache(make_service, outcome): + requests = [] + + async def respond(request): + requests.append(request) + await asyncio.sleep(0) + if outcome == "dns-error": + raise httpx.ConnectError("DNS unavailable", request=request) + if outcome == "timeout": + await asyncio.Event().wait() + return httpx.Response(401, headers=_challenge()) + + service, manager = make_service(respond, timeout_seconds=0.02, probe_registry="sg.example.com") + sources = ["sg.example.com", "sh.example.com"] * 5 + await asyncio.wait_for( + asyncio.gather( + *( + service.start(DockerDeploymentConfig(image=f"{source}/team/task:{index}")) + for index, source in enumerate(sources) + ) + ), + timeout=0.5, + ) + await service.start(DockerDeploymentConfig(image="sg.example.com/team/another:latest")) + + images = {call.args[0].image for call in manager.start_from_template.call_args_list} + expected = { + f"{'zjk.example.com' if outcome == 'success' else source}/team/task:{index}" + for index, source in enumerate(sources) + } + target = "zjk.example.com" if outcome == "success" else "sg.example.com" + assert images == expected | {f"{target}/team/another:latest"} + assert len(requests) == 1 + assert str(requests[0].url) == "https://sg.example.com/v2/" + + +@pytest.mark.parametrize("manifest_status", [200, 404, 403]) +async def test_start_selects_existing_dadi_tag_after_region_mapping(make_service, manifest_status): + requests = [] + + def respond(request): + requests.append(request) + if request.url.path == "/v2/": + return httpx.Response(401, headers=_challenge()) + if request.url.path == "/auth/token": + return httpx.Response(200, json={"token": "pull-only-token"}) + if "authorization" not in request.headers: + return httpx.Response( + 401, + headers={ + "www-authenticate": ( + 'Bearer realm="https://sg.example.com/auth/token", ' + 'service="registry.aliyuncs.com:cn-zhangjiakou:china:cri-test"' + ) + }, + ) + return httpx.Response(manifest_status) + + service, manager = make_service( + respond, probe_registry="sg.example.com", dadi_tag_suffixes=["_containerd_accelerated"] + ) + config = DockerDeploymentConfig(image="sh.example.com/team/task:version") + await service.start(config) + + expected = "zjk.example.com/team/task:version" + if manifest_status == 200: + expected = "zjk.example.com/team/task:version_containerd_accelerated" + assert manager.start_from_template.call_args.args[0].image == expected + assert config.image == "sh.example.com/team/task:version" + assert {request.url.host for request in requests} == {"sg.example.com"} + assert [(request.method, request.url.path) for request in requests] == [ + ("GET", "/v2/"), + ("HEAD", "/v2/team/task/manifests/version_containerd_accelerated"), + ("GET", "/auth/token"), + ("HEAD", "/v2/team/task/manifests/version_containerd_accelerated"), + ] + assert all("authorization" not in request.headers for request in requests[:3]) + assert requests[2].url.params["scope"] == "repository:team/task:pull" + assert requests[2].headers["accept"] == "application/json" + assert requests[3].headers["authorization"] == "Bearer pull-only-token" + assert "application/vnd.oci.image.manifest.v1+json" in requests[3].headers["accept"] + + +@pytest.mark.parametrize( + ("reference", "expected", "probe_count"), + [ + ("team/task", "team/task:latest_accelerated", 1), + ("team/task:v_containerd_accelerated", "team/task:v_containerd_accelerated", 0), + ("team/task:v@sha256:" + "a" * 64, "team/task:v@sha256:" + "a" * 64, 0), + ], +) +async def test_start_handles_dadi_image_references(make_service, reference, expected, probe_count): + probes = [] + + def respond(request): + if request.url.path == "/v2/": + return httpx.Response(401, headers=_challenge()) + probes.append(request) + assert request.url.path == "/v2/team/task/manifests/latest_accelerated" + return httpx.Response(200) + + service, manager = make_service(respond, dadi_tag_suffixes=["_accelerated"]) + await service.start(DockerDeploymentConfig(image=f"sg.example.com/{reference}")) + assert manager.start_from_template.call_args.args[0].image == f"zjk.example.com/{expected}" + assert len(probes) == probe_count + + +@pytest.mark.parametrize("outcome", [200, "timeout"]) +async def test_dadi_probes_run_concurrently_with_timeout(make_service, outcome): + active = peak = 0 + all_started = asyncio.Event() + + async def respond(request): + nonlocal active, peak + if request.url.path == "/v2/": + return httpx.Response(401, headers=_challenge()) + active += 1 + peak = max(peak, active) + try: + if active == 40: + all_started.set() + await all_started.wait() + if outcome == "timeout": + await asyncio.Event().wait() + return httpx.Response(outcome) + finally: + active -= 1 + + service, manager = make_service(respond, dadi_tag_suffixes=["_accelerated"], timeout_seconds=0.1) + await asyncio.wait_for( + asyncio.gather(*(service.start(DockerDeploymentConfig(image="sg.example.com/team/task:v")) for _ in range(40))), + timeout=0.5, + ) + expected = "zjk.example.com/team/task:v_accelerated" if outcome == 200 else "zjk.example.com/team/task:v" + assert {call.args[0].image for call in manager.start_from_template.call_args_list} == {expected} + assert peak == 40 + assert active == 0 + + +@pytest.mark.parametrize("failure", ["untrusted-realm", "token-denied", "invalid-json"]) +async def test_dadi_auth_failure_preserves_region_mapping(make_service, failure): + requests = [] + + def respond(request): + requests.append(request) + if request.url.path == "/v2/": + return httpx.Response(401, headers=_challenge()) + if request.url.path == "/auth/token": + if failure == "invalid-json": + return httpx.Response(200, text="not json") + return httpx.Response(403) + realm = ( + "https://untrusted.example.com/auth/token" + if failure == "untrusted-realm" + else "https://sg.example.com/auth/token" + ) + return httpx.Response(401, headers={"www-authenticate": f'Bearer realm="{realm}", service="acr"'}) + + service, manager = make_service(respond, dadi_tag_suffixes=["_accelerated"]) + await service.start(DockerDeploymentConfig(image="sg.example.com/team/task:v")) + assert manager.start_from_template.call_args.args[0].image == "zjk.example.com/team/task:v" + assert {request.url.host for request in requests} == {"sg.example.com"} + expected_requests = [("GET", "/v2/"), ("HEAD", "/v2/team/task/manifests/v_accelerated")] + if failure != "untrusted-realm": + expected_requests.append(("GET", "/auth/token")) + assert requests[-1].url.params["scope"] == "repository:team/task:pull" + assert [(request.method, request.url.path) for request in requests] == expected_requests + + +@pytest.mark.parametrize( + ("first_status", "second_status", "expected_tag", "expected_probes"), + [ + (200, 200, "v_accelerated", ["v_accelerated"]), + (404, 200, "v_containerd_accelerated", ["v_accelerated", "v_containerd_accelerated"]), + (404, 404, "v", ["v_accelerated", "v_containerd_accelerated"]), + (503, 200, "v_containerd_accelerated", ["v_accelerated", "v_containerd_accelerated"]), + ("timeout", 200, "v_containerd_accelerated", ["v_accelerated", "v_containerd_accelerated"]), + ], +) +async def test_start_probes_dadi_tags_in_default_priority_order( + first_status, second_status, expected_tag, expected_probes +): + probes = [] + + async def respond(request): + if request.url.path == "/v2/": + return httpx.Response(401, headers=_challenge()) + tag = request.url.path.rsplit("/", 1)[-1] + probes.append(tag) + status = first_status if tag == "v_accelerated" else second_status + if status == "timeout": + await asyncio.Event().wait() + return httpx.Response(status) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + resolver = AcrRegionImageResolver( + client, + source_registries=["sg.example.com"], + region_registry_mapping={"cn-zhangjiakou": "zjk.example.com"}, + timeout_seconds=0.03, + ) + manager, templates = AsyncMock(), AsyncMock() + templates.get_ready_template.return_value = None + service = E2BService(manager, templates, image_resolver=resolver) + for attempt in range(1, 3): + await service.start(DockerDeploymentConfig(image="sg.example.com/team/task:v")) + assert manager.start_from_template.call_args.args[0].image == f"zjk.example.com/team/task:{expected_tag}" + assert probes == expected_probes * attempt # Without Redis, neither hits nor misses are cached. + + +@pytest.mark.parametrize("redis_ttl", [None, 86400]) +async def test_dadi_positive_cache_is_shared_with_configured_ttl(make_service, dadi_redis, redis_ttl): + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(401, headers=_challenge()) if request.url.path == "/v2/" else httpx.Response(200) + + options = {"redis_provider": dadi_redis, "dadi_tag_suffixes": ["_accelerated"]} + if redis_ttl is not None: + options["dadi_redis_cache_ttl_seconds"] = redis_ttl + for index in range(2): + service, manager = make_service(respond, **options) + await service.start(DockerDeploymentConfig(image="sg.example.com/team/task:v")) + assert manager.start_from_template.call_args.args[0].image == "zjk.example.com/team/task:v_accelerated" + if index == 0: + keys = await dadi_redis.client.keys("e2b:dadi:exists:*") + assert keys == ["e2b:dadi:exists:v1:sg.example.com:zjk.example.com/team/task:v_accelerated"] + expected_ttl = 2592000 if redis_ttl is None else 86400 + assert expected_ttl - 5 <= await dadi_redis.get_ttl(keys[0]) <= expected_ttl + await dadi_redis.client.expire(keys[0], 60) + + assert [request.url.path for request in requests] == [ + "/v2/", + "/v2/team/task/manifests/v_accelerated", + "/v2/", + ] + keys = await dadi_redis.client.keys("e2b:dadi:exists:*") + assert len(keys) == 1 + assert await dadi_redis.json_get(keys[0], path=".") is True + assert 55 <= await dadi_redis.get_ttl(keys[0]) <= 60 # Shared hits do not renew the TTL. + await dadi_redis.client.delete(keys[0]) + await service.start(DockerDeploymentConfig(image="sg.example.com/team/task:v")) + assert [request.method for request in requests] == ["GET", "HEAD", "GET", "HEAD"] + + +@pytest.mark.parametrize("status", [404, 403, 500, "timeout"]) +async def test_dadi_non_success_is_not_shared_in_redis(make_service, dadi_redis, status): + probes = [] + + async def respond(request): + if request.url.path == "/v2/": + return httpx.Response(401, headers=_challenge()) + probes.append(request) + if status == "timeout": + await asyncio.Event().wait() + return httpx.Response(status) + + service, manager = make_service( + respond, redis_provider=dadi_redis, dadi_tag_suffixes=["_accelerated"], timeout_seconds=0.05 + ) + for _ in range(2): + await service.start(DockerDeploymentConfig(image="sg.example.com/team/task:v")) + assert manager.start_from_template.call_args.args[0].image == "zjk.example.com/team/task:v" + assert len(probes) == 2 + assert await dadi_redis.client.keys("e2b:dadi:exists:*") == [] + + +@pytest.mark.parametrize("failure", ["read-error", "read-timeout", "write-timeout"]) +async def test_dadi_redis_failure_does_not_discard_success(make_service, dadi_redis, monkeypatch, failure): + async def hang(*args, **kwargs): + await asyncio.Event().wait() + + if failure == "read-error": + monkeypatch.setattr(dadi_redis, "json_get", AsyncMock(side_effect=ConnectionError("Redis unavailable"))) + else: + method = "json_get" if failure == "read-timeout" else "json_set_with_ttl" + monkeypatch.setattr(dadi_redis, method, hang) + + def respond(request): + return httpx.Response(401, headers=_challenge()) if request.url.path == "/v2/" else httpx.Response(200) + + service, manager = make_service( + respond, redis_provider=dadi_redis, dadi_tag_suffixes=["_accelerated"], timeout_seconds=0.05 + ) + await asyncio.wait_for(service.start(DockerDeploymentConfig(image="sg.example.com/team/task:v")), timeout=0.5) + assert manager.start_from_template.call_args.args[0].image == "zjk.example.com/team/task:v_accelerated" + + +async def test_dadi_redis_cache_is_scoped_to_mapped_registry(make_service, dadi_redis): + for region, status, expected in [ + ("cn-zhangjiakou", 200, "zjk.example.com/team/task:v_accelerated"), + ("cn-shanghai", 404, "sh-target.example.com/team/task:v"), + ]: + probes = [] + + def respond(request): + if request.url.path == "/v2/": + return httpx.Response(401, headers=_challenge(region)) + probes.append(request) + return httpx.Response(status) + + service, manager = make_service(respond, redis_provider=dadi_redis, dadi_tag_suffixes=["_accelerated"]) + await service.start(DockerDeploymentConfig(image="sg.example.com/team/task:v")) + assert manager.start_from_template.call_args.args[0].image == expected + assert len(probes) == 1 diff --git a/tests/unit/admin/service/test_e2b_image_resolution_config.py b/tests/unit/admin/service/test_e2b_image_resolution_config.py new file mode 100644 index 0000000000..ccd9f8dedb --- /dev/null +++ b/tests/unit/admin/service/test_e2b_image_resolution_config.py @@ -0,0 +1,68 @@ +"""Verify configured image resolver plugins through E2B sandbox start.""" + +from unittest.mock import AsyncMock + +import pytest +import yaml + +from rock.admin.service.e2b_service import E2BService +from rock.admin.service.image_resolver import create_image_resolver +from rock.config import RockConfig +from rock.deployments.config import DockerDeploymentConfig +from rock.utils.providers.redis_provider import RedisProvider + + +class ExampleImageResolver: + """An external plugin with deployment-specific options.""" + + def __init__(self, *, client, target_registry): + self.client = client + self.target_registry = target_registry + + async def resolve(self, image: str) -> str: + return f"{self.target_registry}/{image.split('/', 1)[1]}" + + +@pytest.mark.parametrize("config_source", ["yaml", "nacos"]) +async def test_start_uses_configured_image_resolver(tmp_path, config_source): + rock_config = RockConfig() + redis_provider = RedisProvider(host="", port=0, password="") + manager = AsyncMock() + templates = AsyncMock() + templates.get_ready_template.return_value = { + "image": "source.example.com/team/image:latest", + "cpu_count": 2, + "memory_mb": 4096, + "disk_size_mb": 10240, + } + + async def start_image(): + service = E2BService( + manager, + templates, + image_resolver=create_image_resolver( + rock_config.e2b_image_resolver, rock_config.http_pool_manager, redis_provider + ), + ) + await service.start(DockerDeploymentConfig(image="template-id")) + return manager.start_from_template.call_args.args[0].image + + try: + assert await start_image() == "source.example.com/team/image:latest" + plugin_config = { + "e2b_image_resolver": { + "resolver_class": "tests.unit.admin.service.test_e2b_image_resolution_config.ExampleImageResolver", + "options": {"target_registry": "local.example.com"}, + } + } + if config_source == "yaml": + config_path = tmp_path / "rock.yml" + config_path.write_text(yaml.safe_dump(plugin_config)) + rock_config = RockConfig.from_env(str(config_path)) + else: + rock_config.nacos_provider = AsyncMock() + rock_config.nacos_provider.get_config.return_value = plugin_config + await rock_config.update() + assert await start_image() == "local.example.com/team/image:latest" + finally: + await rock_config.http_pool_manager.aclose_all() diff --git a/tests/unit/deployments/test_start_config_apply.py b/tests/unit/deployments/test_start_config_apply.py index c998af5c45..e604c656dd 100644 --- a/tests/unit/deployments/test_start_config_apply.py +++ b/tests/unit/deployments/test_start_config_apply.py @@ -34,7 +34,8 @@ def _make_rock_config( @pytest.mark.asyncio -async def test_apply_start_config_runs_all_steps_in_order(monkeypatch): +@pytest.mark.parametrize("apply_image_mirror", [True, False]) +async def test_apply_start_config_runs_all_steps_in_order(monkeypatch, apply_image_mirror): calls = [] def record_sync(name): @@ -59,7 +60,8 @@ async def apply(*args, **kwargs): monkeypatch.setattr(start_config, "apply_disk_limits", record_async("disk")) monkeypatch.setattr(start_config, "apply_image_registry_mirror", record_async("mirror")) - await start_config.apply_start_config(MagicMock(), DockerDeploymentConfig(), "Bearer token") + options = {} if apply_image_mirror else {"apply_image_mirror": False} + await start_config.apply_start_config(MagicMock(), DockerDeploymentConfig(), "Bearer token", **options) assert calls == [ "auto_clear", @@ -70,8 +72,7 @@ async def apply(*args, **kwargs): "timeout", "cpu", "disk", - "mirror", - ] + ] + (["mirror"] if apply_image_mirror else []) def test_auto_clear_default_keeps_e2b_explicit_timeout(): diff --git a/tests/unit/sandbox/operator/test_k8s_operator.py b/tests/unit/sandbox/operator/test_k8s_operator.py index 3907cfc00c..5649454e3a 100644 --- a/tests/unit/sandbox/operator/test_k8s_operator.py +++ b/tests/unit/sandbox/operator/test_k8s_operator.py @@ -20,13 +20,6 @@ def test_initialization(self, k8s_config): operator = K8sOperator(k8s_config=k8s_config) assert operator._provider is not None - def test_initialization_forwards_template_table(self, k8s_config): - template_table = AsyncMock() - with patch("rock.sandbox.operator.k8s.operator.BatchSandboxProvider") as provider_class: - K8sOperator(k8s_config=k8s_config, template_table=template_table) - - provider_class.assert_called_once_with(k8s_config=k8s_config, template_table=template_table) - def test_initialization_without_templates(self): """Test K8sOperator initialization fails without templates.""" config = K8sConfig(kubeconfig_path=None, templates={}) diff --git a/tests/unit/sandbox/operator/test_k8s_provider.py b/tests/unit/sandbox/operator/test_k8s_provider.py index de9cf63e4a..22006c0ba1 100644 --- a/tests/unit/sandbox/operator/test_k8s_provider.py +++ b/tests/unit/sandbox/operator/test_k8s_provider.py @@ -1,6 +1,6 @@ """Unit tests for BatchSandboxProvider helper methods.""" -from unittest.mock import AsyncMock, patch +from unittest.mock import patch import pytest @@ -26,13 +26,12 @@ } -def make_provider(template_table=None) -> BatchSandboxProvider: +def make_provider() -> BatchSandboxProvider: return BatchSandboxProvider( k8s_config=K8sConfig( kubeconfig_path=None, templates=BASIC_TEMPLATES, ), - template_table=template_table, ) @@ -216,37 +215,14 @@ async def test_uses_selector_when_no_extended_params(self): config = make_config(image="python:3.11", cpus=2, memory="4Gi") assert await provider._get_pool_name(config) == "pool_small" - async def test_returns_none_when_no_matching_pool(self): - """Return None when no matching pool.""" + @pytest.mark.parametrize("template_id", [None, "template-from-config"]) + async def test_returns_template_id_when_no_matching_pool(self, template_id): + """Use the configured template ID when no matching pool is available.""" provider = make_provider() # No nacos provider set, so pools is empty config = make_config() - assert await provider._get_pool_name(config) is None - - async def test_falls_back_to_ready_template_pool_from_database(self): - template_table = AsyncMock() - template_table.get_ready_fiber_pool_id.return_value = "pool-from-db" - provider = make_provider(template_table=template_table) - provider.set_nacos_provider(MockNacosProvider({K8sConstants.NACOS_POOLS_KEY: {}})) - - config = make_config(image="template-1") - - assert await provider._get_pool_name(config) == "pool-from-db" - template_table.get_ready_fiber_pool_id.assert_awaited_once_with("template-1") - - async def test_nacos_pool_takes_priority_over_database(self): - template_table = AsyncMock() - provider = make_provider(template_table=template_table) - provider.set_nacos_provider( - MockNacosProvider( - {K8sConstants.NACOS_POOLS_KEY: {"pool-nacos": {"image": "template-1", "cpus": 2, "memory": "4Gi"}}} - ) - ) - - config = make_config(image="template-1") - - assert await provider._get_pool_name(config) == "pool-nacos" - template_table.get_ready_fiber_pool_id.assert_not_awaited() + config.template_id = template_id + assert await provider._get_pool_name(config) == template_id # ========== _get_template_name ========== diff --git a/tests/unit/sandbox/operator/test_operator_factory.py b/tests/unit/sandbox/operator/test_operator_factory.py index d512a69bb8..3b6651a6a5 100644 --- a/tests/unit/sandbox/operator/test_operator_factory.py +++ b/tests/unit/sandbox/operator/test_operator_factory.py @@ -1,10 +1,10 @@ """Unit tests for OperatorFactory dispatch (opensandbox branch).""" -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest -from rock.config import K8sConfig, OpenSandboxConfig, RuntimeConfig +from rock.config import OpenSandboxConfig, RuntimeConfig from rock.sandbox.operator.factory import ( OperatorContext, OperatorFactory, @@ -53,21 +53,6 @@ def test_create_opensandbox_operator_requires_config(): OperatorFactory.create_operator(ctx) -def test_create_k8s_operator_receives_template_table(): - template_table = MagicMock() - k8s_config = K8sConfig() - ctx = OperatorContext( - runtime_config=_runtime("k8s"), - k8s_config=k8s_config, - template_table=template_table, - ) - - with patch("rock.sandbox.operator.factory.K8sOperator") as operator_class: - OperatorFactory.create_operator(ctx) - - operator_class.assert_called_once_with(k8s_config=k8s_config, template_table=template_table) - - def test_unsupported_operator_type_lists_opensandbox(): ctx = OperatorContext(runtime_config=_runtime("bogus")) with pytest.raises(ValueError, match="opensandbox"): diff --git a/tests/unit/sandbox/test_sandbox_transitions.py b/tests/unit/sandbox/test_sandbox_transitions.py index 428358406b..4032619081 100644 --- a/tests/unit/sandbox/test_sandbox_transitions.py +++ b/tests/unit/sandbox/test_sandbox_transitions.py @@ -4,6 +4,7 @@ each sandbox state, using lightweight mocks (no Ray / Docker). """ +import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -444,20 +445,27 @@ async def capture_create(sandbox_id, sandbox_info, timeout_info=None, deployment mock_meta_store.create.side_effect = capture_create mgr._start_and_wait = SandboxManager._start_and_wait.__get__(mgr) mgr.start = SandboxManager.start.__wrapped__.__get__(mgr) + mgr.start_from_template = SandboxManager.start_from_template.__wrapped__.__get__(mgr) mgr.get_status = AsyncMock(return_value=MagicMock(is_alive=True, state=State.RUNNING)) return mgr class TestManagerStart: @pytest.mark.asyncio - async def test_start_async_applies_shared_start_config(self, mgr_start): + @pytest.mark.parametrize("start_method", ["start_async", "start_from_template"]) + async def test_start_async_applies_shared_start_config(self, mgr_start, start_method): user_info = {"rock_authorization": "Bearer token"} config = DockerDeploymentConfig(image="python:3.11") with patch("rock.sandbox.sandbox_manager.apply_start_config", new_callable=AsyncMock) as apply: - await mgr_start.start_async(config, user_info=user_info) + await getattr(mgr_start, start_method)(config, user_info=user_info) - apply.assert_awaited_once_with(mgr_start.rock_config, config, "Bearer token") + apply.assert_awaited_once_with( + mgr_start.rock_config, + config, + "Bearer token", + apply_image_mirror=start_method != "start_from_template", + ) @pytest.mark.asyncio async def test_persists_sandbox_environment(self, mgr_start, mock_docker_config, mock_meta_store): @@ -475,7 +483,7 @@ async def test_refreshes_nacos_config_before_validating_spec(self, mgr_start, mo async def refresh_config(): call_order.append("refresh") - async def normalize_config(rock_config, config, rock_authorization): + async def normalize_config(rock_config, config, rock_authorization, *, apply_image_mirror): call_order.append("normalize") async def init_config(config, *, refresh_config): @@ -685,10 +693,38 @@ async def running_after_retries(sandbox_id): assert call_count >= 3 @pytest.mark.asyncio - async def test_start_timeout_raises(self, mgr_start): - mgr_start.get_status = AsyncMock(return_value=MagicMock(is_alive=False, state=State.PENDING)) + @pytest.mark.parametrize("wait_timeout, slow_status", [(None, False), (0.03, False), (0.03, True)]) + async def test_start_timeout_raises(self, mgr_start, mock_operator, wait_timeout, slow_status): + status_cancelled = asyncio.Event() + + async def get_status(sandbox_id): + if slow_status: + try: + await asyncio.Event().wait() + finally: + status_cancelled.set() + return MagicMock(is_alive=False, state=State.PENDING) + + mgr_start.get_status = AsyncMock(side_effect=get_status) config = MagicMock() config.image = "python:3.11" - with patch("rock.sandbox.sandbox_manager.REQUEST_TIMEOUT_SECONDS", 0): - with pytest.raises(TimeoutError, match="not running after"): - await mgr_start.start(config) + expected_timeout = 85 if wait_timeout is None else wait_timeout + real_wait_for, real_sleep = asyncio.wait_for, asyncio.sleep + + async def short_wait_for(awaitable, timeout): + return await real_wait_for(awaitable, min(timeout, 0.03)) + + async def short_sleep(delay): + await real_sleep(min(delay, 0.001)) + + with ( + patch("rock.sandbox.sandbox_manager.asyncio.wait_for", side_effect=short_wait_for) as wait_for, + patch("rock.sandbox.sandbox_manager.asyncio.sleep", side_effect=short_sleep), + ): + with pytest.raises(TimeoutError, match=f"sandbox sb-1 not running after {expected_timeout}s"): + await mgr_start.start_from_template(config, wait_timeout=wait_timeout) + + assert wait_for.call_args.kwargs["timeout"] == expected_timeout + assert mgr_start.get_status.await_count > 0 + assert status_cancelled.is_set() == slow_status + mock_operator.stop.assert_not_awaited()