diff --git a/docs/proposals/template-design.md b/docs/proposals/template-design.md new file mode 100644 index 0000000000..006ae9698a --- /dev/null +++ b/docs/proposals/template-design.md @@ -0,0 +1,171 @@ +# ROCK Template 管理设计文档 + +## 1. 概述 + +Template 管理提供 Sandbox 预热池的创建与查询能力。调用方通过 operator 方法提交镜像和资源规格创建 Template,获得稳定的 `templateID`。该 ID 对应一个 OpenSandbox Pool CRD,由 Pool Controller 自动维护热备 Pod。后续创建 Sandbox 时,`templateID` 作为 `poolRef` 被引用,实现秒级分配。 + +Template 能力以 operator 方法形式交付,不以 HTTP 接口暴露。上层应用(如 OpenSandbox)可自行包装为 HTTP 接口。 + +| Operator 方法 | 说明 | +| --- | --- | +| `operator.create_template(spec)` | 提交 Template 创建请求,同步返回 `templateID` 及状态 | +| `operator.get_template_status(template_id)` | 查询 Template 当前状态 | +| `operator.scale_template(template_id, capacity)` | 调整 Template 容量(PATCH 语义) | +| `operator.delete_template(template_id)` | 删除 Template(= 删除 Pool CRD) | + +### 设计要点 + +- **无独立存储**:K8s Pool CRD 即为持久化层,Informer 本地缓存提供高效读取,无需 DB。 +- **同步创建**:Pool CRD 创建不耗时,请求内同步完成,无需后台异步任务。 +- **天然去重**:相同 spec 生成相同 `templateID`(= Pool Name),K8s API Server 通过 409 Conflict 自动去重。容量字段不参与 `templateID` 计算。 +- **模版渲染**:Pool 的 PodTemplateSpec、capacitySpec 等通过 YAML 模版配置 + Jinja2 渲染。 + +### 概念说明 + +Template 遵循 E2B 风格概念,以 operator 方法形式交付(不以 HTTP 接口暴露)。Template 是统一抽象——描述 Sandbox Pod 如何被创建的配置。差异只在交付机制: + +| | Warm(当前实现) | Cold(未来可选,不实现) | +| --- | --- | --- | +| 底层资源 | Pool CRD | BatchSandbox template config | +| Pod 何时创建 | 预创建热备,引用时秒级分配 | Sandbox 请求时即时创建 | +| templateID 映射 | `tpl-xxx` → Pool Name | template name → config key | +| 状态查询 | 需查 Pool reconcile 进度 | 天然 `ready`(无需预热) | + +Operator 内部将 Template 概念转换为 Pool 操作: + +| 外部概念 | 内部映射 | K8s 资源 | +| --- | --- | --- | +| Template | Pool | Pool CRD | +| templateID | Pool Name | Pool.metadata.name | +| Template status | Pool status 映射 | Pool.status | + +## 2. 整体设计 + +### 2.1 链路 + +``` +创建 Template + ├─ 生成 templateID = tpl- + sha256(非空 spec 字段)[:16] + ├─ 从 pool_template 配置渲染 Pool CRD manifest + ├─ K8sApiClient(Pool).create_custom_object() + │ └─ 409 → Pool 已存在,返回已有 templateID + status + └─ 返回 templateID + status + +查询 Template + ├─ templateID = Pool Name + ├─ K8sApiClient(Pool).get_custom_object() ← Informer 本地缓存 + ├─ 映射 Pool status → template status + └─ 返回 templateID + status + timestamps + +删除 Template + ├─ templateID = Pool Name + ├─ K8sApiClient(Pool).delete_custom_object() + │ └─ 404 → not found = already deleted + └─ K8s GC 自动清理 Pool 拥有的 Pod + +扩缩容 Template + ├─ templateID = Pool Name + ├─ 校验容量字段(pool_min/pool_max/buffer_min/buffer_max) + ├─ 仅将提供的非空字段 PATCH 到 Pool.spec.capacitySpec + ├─ K8sApiClient(Pool).update_custom_object() + └─ 返回更新后的 templateID + status + capacity +``` + +### 2.2 templateID 生成 + +覆盖完整 spec 的**非空字段**,忽略 null/空值,保证前向兼容——后续新增字段传入非空值时自动参与哈希,传入 null 时不影响已有映射。 + +- 格式:`tpl-{sha256(排序后的非空字段)[:16]}` +- 前缀用连字符 `-`(符合 K8s RFC 1123 命名规范) +- 参与哈希的非空字段:`from_image`、`cpu_count`、`memory_mb`;`disk_gb`/`num_gpus`/`accelerator_type`/`os` 为 null 时不参与哈希 +- 容量字段(`pool_min`/`pool_max`/`buffer_min`/`buffer_max`)被排除在外,因为容量不影响 Template 身份 + +### 2.3 状态映射 + +Pool CRD status 字段为 `available`/`total`(非 `readyReplicas`/`replicas`): + +| 条件 | Template status | 说明 | +| --- | --- | --- | +| Pool 不存在 | 404 | templateID 无效 | +| `available > 0` | `ready` | 至少一个热备 Pod 就绪 | +| `available == 0 && total > 0` | `building` | Pod 已创建但未就绪 | +| `available == 0 && total == 0` | `building` | Controller 尚未处理或 Pool 为空 | +| conditions 中 `Ready=False` | `error` | Pool 创建失败 | + +返回的 Template status 中还包含容量信息,结构如下: + +```json +{ + "capacity": { + "spec": { + "poolMin": ..., + "poolMax": ..., + "bufferMin": ..., + "bufferMax": ... + }, + "status": { + "available": ..., + "total": ..., + "allocated": ... + } + } +} +``` + +### 2.4 Pool 标识 + +Template API 创建的 Pool 携带 Label `rock.sandbox/managed-by: template-api`,与 Nacos 配置的系统 Pool 区分。 + +### 2.5 关于 Update + +不支持修改 Template 的 spec 字段(镜像、CPU、内存等)。templateID 由这些 spec 字段的 hash 生成,改 spec 即产生新 templateID,本质是 create 而非 update。配置变更场景:创建新 Template → 更新引用方配置 → 删除旧 Template。 + +容量字段(`pool_min`/`pool_max`/`buffer_min`/`buffer_max`)可通过 `operator.scale_template()` 单独调整,采用 PATCH 语义:仅更新提供的非空字段,允许缩容到 0,provider 层校验字段合法性与 `min <= max`。 + +### 2.6 冒烟测试 + +Template 不再以 HTTP 接口暴露,原 HTTP 冒烟测试已 skip。后续以 operator 方法调用形式重写冒烟测试。 + +## 3. 实现要点 + +### 3.1 Pool 模版配置 + +`K8sConfig.pool_template` 字段(YAML 配置 + Jinja2 渲染)生成 Pool CRD spec。关键约定: + +- capacitySpec 使用 **camelCase** 键(`bufferMin`/`bufferMax`/`poolMin`/`poolMax`),匹配 Pool CRD spec +- Jinja2 变量需加**双引号**(`"{{ from_image }}"`),避免 YAML flow mapping 解析错误 +- 渲染上下文:`from_image`、`cpu_count`、`memory_mb`(必需);`disk_gb`/`num_gpus`/`accelerator_type`/`os`(可选,为 null 时不传入 ctx)。容量字段不来自 `TemplateSpec`,由 `pool_template` 默认配置提供 + +### 3.2 关键常量 + +- `TEMPLATE_ID_PREFIX = "tpl-"`(RFC 1123 连字符) +- `LABEL_MANAGED_BY = "rock.sandbox/managed-by"`,值为 `"template-api"` +- Pool CRD:plural=`pools`,kind=`Pool` + +### 3.3 BatchSandboxProvider 扩展 + +Provider 新增 Pool informer(复用同一 `ApiClient`,独立 watch `pools` CRD),并提供四组方法: + +| 外部方法(template 术语) | 内部方法(pool 术语) | 行为 | +| --- | --- | --- | +| `create_template(spec)` | `_create_pool(name, spec)` | 渲染 manifest → create_custom_object;409 时返回已有 Pool | +| `get_template_status(id)` | `_get_pool(name)` | 从 Informer 缓存读取 Pool → 映射 status;未找到返回 None | +| `scale_template(id, capacity)` | — | 校验字段 → PATCH capacitySpec → 返回更新后的 status | +| `delete_template(id)` | `_delete_pool(name)` | delete_custom_object;404 视为已删除 | + +辅助方法:`_build_pool_manifest_from_template`(渲染)、`_map_pool_to_template_status`(状态映射)。 + +## 4. 涉及文件 + +| 文件 | 变更 | +| --- | --- | +| `rock/sandbox/operator/k8s/constants.py` | 新增 Pool CRD 常量、Label 常量、`TEMPLATE_ID_PREFIX` | +| `rock/sandbox/operator/k8s/provider.py` | 新增 Pool informer、template/pool 转换方法、`TemplateSpec`(含 `os` 字段)、`scale_template` | +| `rock/sandbox/operator/k8s/template_loader.py` | 新增 `build_pool_manifest`,渲染上下文含 `os` 变量 | +| `rock/sandbox/operator/k8s/operator.py` | 新增 `create_template`/`get_template_status`/`scale_template`/`delete_template` 委托方法 | +| `rock/sandbox/operator/abstract.py` | 新增 `create_template`/`get_template_status`/`scale_template`/`delete_template` 抽象方法默认实现 | +| `rock-conf/rock-junxin.yml` | 新增 `pool_template` 配置段 | +| `tests/unit/sandbox/operator/test_k8s_template_provider.py` | 单元测试:ID 生成、Spec、状态映射、scale、常量 | +| `tests/unit/sandbox/operator/test_k8s_template_loader.py` | 单元测试:`build_pool_manifest` 渲染 | +| `tests/smoke/test_template.py` | 冒烟测试(已 skip,待以方法调用形式重写) | + diff --git a/rock/config.py b/rock/config.py index cd34fbe653..c13fd16f50 100644 --- a/rock/config.py +++ b/rock/config.py @@ -389,6 +389,12 @@ class K8sConfig: # ROCK_IMAGE_AUTH_KEY environment variable if not set here. image_auth_key: str | None = None + # Pool rendering templates for Template API (Warm path). + # Keyed by name (e.g. "default", "windows"); rendered via Jinja2 to produce + # Pool CRD spec. Selected by TemplateSpec.os at create time, falling back to + # "default". + pool_templates: dict[str, dict] = field(default_factory=dict) + # ============================================================================ # DEPRECATED: The following fields are deprecated and will be removed in a # future version. Do NOT use them in new code. diff --git a/rock/sandbox/operator/abstract.py b/rock/sandbox/operator/abstract.py index 6b7da32b92..af62f0f5fd 100644 --- a/rock/sandbox/operator/abstract.py +++ b/rock/sandbox/operator/abstract.py @@ -1,6 +1,10 @@ from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any from rock.actions.sandbox.sandbox_info import SandboxInfo + +if TYPE_CHECKING: + from rock.sandbox.operator.k8s.provider import TemplateSpec from rock.admin.core.redis_key import alive_sandbox_key from rock.common.constants import StopReason from rock.config import RuntimeConfig @@ -67,6 +71,47 @@ async def get_remote_status(self, sandbox_id: str, host_ip: str): return ServiceStatus() + # ======================================================================== + # Template API (Warm path) — default: raise NotImplementedError + # ======================================================================== + + async def create_template(self, spec: Any) -> dict: + """Create or reuse a template (Pool CRD). + + Only K8sOperator supports this; other operators raise BadRequestRockError. + Returns a dict with template_id and status. + """ + from rock.sdk.common.exceptions import BadRequestRockError + + raise BadRequestRockError(f"template not supported on {type(self).__name__}") + + async def get_template_status(self, template_id: str) -> dict | None: + """Get template (Pool) status. + + Only K8sOperator supports this; other operators raise BadRequestRockError. + """ + from rock.sdk.common.exceptions import BadRequestRockError + + raise BadRequestRockError(f"template not supported on {type(self).__name__}") + + async def delete_template(self, template_id: str) -> bool: + """Delete template (Pool CRD). + + Only K8sOperator supports this; other operators raise BadRequestRockError. + """ + from rock.sdk.common.exceptions import BadRequestRockError + + raise BadRequestRockError(f"template not supported on {type(self).__name__}") + + async def scale_template(self, template_id: str, capacity: dict[str, Any]) -> dict: + """Scale a template's Pool capacity. + + Only K8sOperator supports this; other operators raise BadRequestRockError. + """ + from rock.sdk.common.exceptions import BadRequestRockError + + raise BadRequestRockError(f"template not supported on {type(self).__name__}") + def set_redis_provider(self, redis_provider: RedisProvider): self._redis_provider = redis_provider diff --git a/rock/sandbox/operator/k8s/constants.py b/rock/sandbox/operator/k8s/constants.py index 99ab409176..0fb840a9c6 100644 --- a/rock/sandbox/operator/k8s/constants.py +++ b/rock/sandbox/operator/k8s/constants.py @@ -31,4 +31,15 @@ class K8sConstants: # Nacos config keys NACOS_POOLS_KEY = "pools" NACOS_TEMPLATE_RULES_KEY = "template_rules" + + # Pool CRD (Warm path: Template API creates Pool CRD) + CRD_PLURAL_POOL = "pools" + CRD_KIND_POOL = "Pool" + + # Label: distinguish Template API created Pools from system-configured Pools + LABEL_MANAGED_BY = "rock.sandbox/managed-by" + LABEL_MANAGED_BY_TEMPLATE_API = "template-api" + + # templateID prefix + TEMPLATE_ID_PREFIX = "tpl-" K8S_ALIVE_CHECK_SWITCH = "k8s_alive_check_enabled" diff --git a/rock/sandbox/operator/k8s/operator.py b/rock/sandbox/operator/k8s/operator.py index 7a74bcd2a2..46384d4c3e 100644 --- a/rock/sandbox/operator/k8s/operator.py +++ b/rock/sandbox/operator/k8s/operator.py @@ -1,5 +1,7 @@ """K8s Operator implementation for managing sandboxes via Kubernetes.""" +from typing import Any + from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.common.constants import StopReason from rock.config import K8sConfig @@ -7,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, TemplateFiberPoolLookup +from rock.sandbox.operator.k8s.provider import BatchSandboxProvider, TemplateSpec, TemplateFiberPoolLookup logger = init_logger(__name__) @@ -149,3 +151,48 @@ async def stop(self, sandbox_id: str, reason: StopReason = StopReason.MANUAL) -> async def delete(self, config: DockerDeploymentConfig, host_ip: str | None = None) -> bool: """Treat delete as successful because stop already removes the K8s resource.""" return True + + # ======================================================================== + # Template API (Warm path) + # ======================================================================== + + async def create_template(self, spec: TemplateSpec) -> dict: + """Create or reuse a template (Pool CRD). + + Returns a dict with template_id and status. + """ + return await self._provider.create_template(spec) + + async def get_template_status(self, template_id: str) -> dict | None: + """Get template (Pool) status. + + Args: + template_id: Template identifier (= Pool name) + + Returns: + Status dict or None if not found + """ + return await self._provider.get_template_status(template_id) + + async def delete_template(self, template_id: str) -> bool: + """Delete template (Pool CRD). + + Args: + template_id: Template identifier (= Pool name) + + Returns: + True if deleted or not found, False if in use + """ + return await self._provider.delete_template(template_id) + + async def scale_template(self, template_id: str, capacity: dict[str, Any]) -> dict: + """Scale a template's Pool capacity. + + Args: + template_id: Template identifier (= Pool name). + capacity: Snake-case capacity fields to update. + + Returns: + Updated template status dict. + """ + return await self._provider.scale_template(template_id, capacity) diff --git a/rock/sandbox/operator/k8s/provider.py b/rock/sandbox/operator/k8s/provider.py index d6c3d75e53..0a298c2ce6 100644 --- a/rock/sandbox/operator/k8s/provider.py +++ b/rock/sandbox/operator/k8s/provider.py @@ -2,17 +2,21 @@ import base64 import fnmatch +import hashlib import json import os import re from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import Any, Protocol from cryptography.hazmat.primitives.ciphers.aead import AESGCM from kubernetes import client from kubernetes import config as k8s_config +from rock import InternalServerRockError from rock.actions.sandbox.config import RemoteSandboxRuntimeConfig +from rock.sdk.common.exceptions import BadRequestRockError from rock.actions.sandbox.sandbox_info import SandboxInfo from rock.config import K8sConfig, PoolConfig, TemplateSelectorRule from rock.deployments.config import DockerDeploymentConfig @@ -26,6 +30,48 @@ logger = init_logger(__name__) +@dataclass +class TemplateSpec: + """Template creation spec, maps to API camelCase fields. + + Capacity (poolMin/poolMax/bufferMin/bufferMax) is intentionally excluded + because it does not affect template identity. Use the Scale API to adjust + capacity after creation. + """ + + from_image: str + cpu_count: int + memory_mb: int + disk_gb: int | None = None + num_gpus: float | None = None + accelerator_type: str | None = None + os: str | None = None + + +def generate_template_id(spec: TemplateSpec) -> str: + """Generate template ID from all non-null TemplateSpec fields. + + The ID is a hash of every provided field, so any difference in + from_image, cpu_count, memory_mb, disk_gb, num_gpus, accelerator_type, + or os produces a distinct templateID. Capacity is excluded from identity. + """ + parts: list[str] = [ + f"from_image={spec.from_image}", + f"cpu_count={spec.cpu_count}", + f"memory_mb={spec.memory_mb}", + ] + if spec.disk_gb is not None: + parts.append(f"disk_gb={spec.disk_gb}") + if spec.num_gpus is not None: + parts.append(f"num_gpus={spec.num_gpus}") + if spec.accelerator_type is not None: + parts.append(f"accelerator_type={spec.accelerator_type}") + if spec.os is not None: + parts.append(f"os={spec.os}") + + 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: ... @@ -278,15 +324,17 @@ def __init__(self, k8s_config: K8sConfig, template_table: TemplateFiberPoolLooku self._k8s_config = k8s_config self._api_client = None self._k8s_api: K8sApiClient | None = None + 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 + # Initialize template loader with config templates and pool template self._template_loader = K8sTemplateLoader( templates=k8s_config.templates, default_namespace=k8s_config.namespace, + pool_templates=k8s_config.pool_templates, ) logger.info(f"Available K8S templates: {', '.join(self._template_loader.available_templates)}") @@ -534,6 +582,19 @@ async def _ensure_initialized(self): resync_period=self._k8s_config.resync_period, ) await self._k8s_api.start() + + # Pool CRD informer for Template API (Warm path) + self._pool_api = K8sApiClient( + api_client=self._api_client, + group=K8sConstants.CRD_GROUP, + version=K8sConstants.CRD_VERSION, + plural=K8sConstants.CRD_PLURAL_POOL, + namespace=self.namespace, + qps=self._k8s_config.api_qps, + resync_period=self._k8s_config.resync_period, + ) + await self._pool_api.start() + self._initialized = True logger.info("Initialized K8s provider with informer") @@ -885,3 +946,245 @@ def _build_runtime(self, host_ip: str, port_mapping: dict[int, int]) -> RemoteSa port=proxy_port, ) return RemoteSandboxRuntime.from_config(runtime_config) + + # ======================================================================== + # Template API (Warm path) — create / get / delete Pool CRD + # ======================================================================== + + async def create_template(self, spec: TemplateSpec) -> dict: + """Create or reuse a template (Pool CRD). + + Idempotent: same non-null TemplateSpec fields produce the same templateID. + Capacity is excluded from identity. Returns a dict with template_id and status. + """ + await self._ensure_initialized() + template_id = generate_template_id(spec) + + try: + # _create_pool handles creation and 409 (already exists), returning + # the created or existing Pool object. Map it to template status. + pool = await self._create_pool(template_id, spec) + return self._map_pool_to_template_status(pool) + except InternalServerRockError: + raise + except Exception as e: + logger.error(f"Failed to create template {template_id}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to create template {template_id}: {e}") from e + + async def get_template_status(self, template_id: str) -> dict | None: + """Get template (Pool) status. + + Args: + template_id: Template identifier (= Pool name) + + Returns: + Status dict or None if not found + """ + await self._ensure_initialized() + try: + pool = await self._get_pool(template_id) + if pool is None: + return None + return self._map_pool_to_template_status(pool) + except InternalServerRockError: + raise + except Exception as e: + logger.error(f"Failed to get template status {template_id}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to get template status {template_id}: {e}") from e + + async def delete_template(self, template_id: str) -> bool: + """Delete template (Pool CRD). + + Args: + template_id: Template identifier (= Pool name) + + Returns: + True if deleted or not found. Raises on K8s error. + + Note: + We do NOT check whether the pool is referenced by running BatchSandboxes. + The Pool controller / K8s finalizer is responsible for protecting an + in-use pool. Listing all BatchSandboxes here is too expensive and couples + Template API to sandbox lifecycle details. + """ + await self._ensure_initialized() + + try: + # Check if pool exists + pool = await self._get_pool(template_id) + if pool is None: + logger.info(f"Template {template_id} not found, already deleted") + return True + + await self._delete_pool(template_id) + return True + except InternalServerRockError: + raise + except Exception as e: + logger.error(f"Failed to delete template {template_id}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to delete template {template_id}: {e}") from e + + async def scale_template(self, template_id: str, capacity: dict[str, Any]) -> dict: + """Scale a template's Pool capacity. + + Args: + template_id: Template identifier (= Pool name). + capacity: Snake-case capacity fields to update. Only provided keys + are sent to K8s (PATCH semantics). Supported keys: + pool_min, pool_max, buffer_min, buffer_max. + + Returns: + Updated template status dict. + + Raises: + BadRequestRockError: If the pool does not exist or capacity is invalid. + InternalServerRockError: On unexpected K8s errors. + """ + await self._ensure_initialized() + + if not capacity: + raise BadRequestRockError("No capacity fields provided") + + valid_keys = {"pool_min", "pool_max", "buffer_min", "buffer_max"} + invalid_keys = set(capacity.keys()) - valid_keys + if invalid_keys: + raise BadRequestRockError(f"Invalid capacity fields: {sorted(invalid_keys)}") + + key_map = { + "pool_min": "poolMin", + "pool_max": "poolMax", + "buffer_min": "bufferMin", + "buffer_max": "bufferMax", + } + capacity_spec_patch: dict[str, Any] = {} + for key, value in capacity.items(): + if value is not None: + capacity_spec_patch[key_map[key]] = int(value) + + try: + pool = await self._get_pool(template_id) + if pool is None: + raise BadRequestRockError(f"Template {template_id} not found") + + patch_body = {"spec": {"capacitySpec": capacity_spec_patch}} + updated = await self._pool_api.update_custom_object( + name=template_id, + body=patch_body, + ) + return self._map_pool_to_template_status(updated) + except BadRequestRockError: + raise + except InternalServerRockError: + raise + except Exception as e: + logger.error(f"Failed to scale template {template_id}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to scale template {template_id}: {e}") from e + + async def _create_pool(self, pool_name: str, spec: TemplateSpec) -> dict[str, Any]: + """Create Pool CRD from template. + + Returns: + The created Pool object, or the existing Pool object if it already + existed (409). + """ + manifest = self._template_loader.build_pool_manifest(pool_name, spec, template_name="default") + try: + pool = await self._pool_api.create_custom_object(body=manifest) + logger.info(f"Created Pool {pool_name} for template") + return pool + except client.exceptions.ApiException as e: + if e.status == 409: + # Idempotent: another request created it concurrently + logger.info(f"Pool {pool_name} already exists, reusing") + existing = await self._get_pool(pool_name) + if existing is None: + raise InternalServerRockError( + f"Pool {pool_name} returned 409 but could not be fetched" + ) from e + return existing + logger.error(f"Failed to create Pool {pool_name}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to create Pool {pool_name}: {e.reason}") from e + except Exception as e: + logger.error(f"Unexpected error creating Pool {pool_name}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to create Pool {pool_name}: {e}") from e + + async def _get_pool(self, pool_name: str) -> dict | None: + """Get Pool CRD from cache. + + Returns None if not found. + """ + try: + return await self._pool_api.get_custom_object(name=pool_name) + except client.exceptions.ApiException as e: + if e.status == 404: + return None + logger.error(f"Failed to get Pool {pool_name}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to get Pool {pool_name}: {e.reason}") from e + except Exception as e: + logger.error(f"Failed to get Pool {pool_name}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to get Pool {pool_name}: {e}") from e + + async def _delete_pool(self, pool_name: str) -> None: + """Delete Pool CRD.""" + try: + await self._pool_api.delete_custom_object(name=pool_name) + logger.info(f"Deleted Pool {pool_name}") + except client.exceptions.ApiException as e: + if e.status == 404: + logger.info(f"Pool {pool_name} not found, treating as deleted") + return + logger.error(f"Failed to delete Pool {pool_name}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to delete Pool {pool_name}: {e.reason}") from e + except Exception as e: + logger.error(f"Unexpected error deleting Pool {pool_name}: {e}", exc_info=True) + raise InternalServerRockError(f"Failed to delete Pool {pool_name}: {e}") from e + + def _map_pool_to_template_status(self, pool: dict) -> dict: + """Map Pool CRD to template status dict (design doc format).""" + pool_name = pool.get("metadata", {}).get("name", "") + status = pool.get("status", {}) + metadata = pool.get("metadata", {}) + + ready_replicas = status.get("available", 0) + total_replicas = status.get("total", 0) + + # If the pool is configured with zero capacity, it is immediately ready + # because no standby pods are required. + capacity_spec = pool.get("spec", {}).get("capacitySpec", {}) + pool_min = capacity_spec.get("poolMin", 1) + buffer_min = capacity_spec.get("bufferMin", 1) + zero_capacity = pool_min == 0 and buffer_min == 0 + + # Determine template status + if ready_replicas > 0 or zero_capacity: + template_status = "ready" + elif total_replicas > 0: + template_status = "building" + else: + template_status = "building" + + reason = None + + created_at = metadata.get("creationTimestamp", "") + updated_at = created_at + + return { + "template_id": pool_name, + "status": template_status, + "reason": reason, + "created_at": created_at, + "updated_at": updated_at, + "capacity": { + "spec": { + "pool_min": capacity_spec.get("poolMin"), + "pool_max": capacity_spec.get("poolMax"), + "buffer_min": capacity_spec.get("bufferMin"), + "buffer_max": capacity_spec.get("bufferMax"), + }, + "status": { + "available": status.get("available"), + "total": status.get("total"), + "allocated": status.get("allocated"), + }, + }, + } diff --git a/rock/sandbox/operator/k8s/template_loader.py b/rock/sandbox/operator/k8s/template_loader.py index 97f8cce810..4f970c1cd3 100644 --- a/rock/sandbox/operator/k8s/template_loader.py +++ b/rock/sandbox/operator/k8s/template_loader.py @@ -1,4 +1,4 @@ -"""K8S template loader for BatchSandbox manifests.""" +"""K8S template loader for BatchSandbox and Pool manifests.""" import copy import json @@ -14,17 +14,25 @@ class K8sTemplateLoader: - """Loader for K8S BatchSandbox templates.""" + """Loader for K8S BatchSandbox and Pool CRD manifests.""" - def __init__(self, templates: dict[str, dict[str, Any]], default_namespace: str = "rock"): + def __init__( + self, + templates: dict[str, dict[str, Any]], + default_namespace: str = "rock", + pool_templates: dict[str, dict[str, Any]] | None = None, + ): """Initialize template loader. Args: - templates: Dictionary of template configurations from K8sConfig + templates: Dictionary of BatchSandbox template configurations from K8sConfig default_namespace: Default namespace if template doesn't specify one + pool_templates: Optional Pool CRD templates for Template API (Warm path). + Keyed by name (e.g. "default", "windows"); selected by TemplateSpec.os. """ self._templates: dict[str, dict[str, Any]] = templates self._default_namespace = default_namespace + self._pool_templates: dict[str, dict[str, Any]] = pool_templates or {} if not self._templates: raise ValueError("No templates provided. At least one template must be defined in K8sConfig.templates.") @@ -182,6 +190,67 @@ def build_manifest( return manifest + def build_pool_manifest(self, pool_name: str, spec: Any, template_name: str = "default") -> dict[str, Any]: + """Build a complete Pool CRD manifest from the pool template and spec. + + The pool template is rendered with Jinja2 against a context built from + ``spec``: from_image, cpu_count, memory_mb, disk_gb, num_gpus, + accelerator_type, os. Capacity fields (poolMin/poolMax/bufferMin/bufferMax) + are intentionally excluded from ``spec``; they are supplied by the pool + template defaults and can be adjusted later via the Scale API. + + Args: + pool_name: Name for the Pool CRD (also the template ID). + spec: Template creation spec with the attributes listed above. + template_name: Name of the pool template to use (mirrors + ``build_manifest``'s ``template_name``). + + Returns: + Complete Pool CRD manifest. + + Raises: + ValueError: If no pool template was configured or the named + template is not found. + """ + if not self._pool_templates: + raise ValueError("No pool template configured. Set k8s.pool_templates in config.") + + if template_name not in self._pool_templates: + available = ", ".join(self._pool_templates.keys()) + raise ValueError( + f"Pool template '{template_name}' not found. Available: {available}" + ) + template = copy.deepcopy(self._pool_templates[template_name]) + + ctx: dict[str, Any] = { + "from_image": spec.from_image, + "cpu_count": spec.cpu_count, + "memory_mb": spec.memory_mb, + } + if spec.disk_gb is not None: + ctx["disk_gb"] = spec.disk_gb + if spec.num_gpus is not None: + ctx["num_gpus"] = spec.num_gpus + if spec.accelerator_type is not None: + ctx["accelerator_type"] = spec.accelerator_type + if spec.os is not None: + ctx["os"] = spec.os + + rendered = render_node(template, self._jinja_env, ctx) + + return { + "apiVersion": K8sConstants.CRD_API_VERSION, + "kind": K8sConstants.CRD_KIND_POOL, + "metadata": { + "name": pool_name, + "namespace": self._default_namespace, + "labels": { + K8sConstants.LABEL_MANAGED_BY: K8sConstants.LABEL_MANAGED_BY_TEMPLATE_API, + }, + }, + "spec": rendered, + } + @property def available_templates(self) -> list[str]: """Get list of available template names.""" diff --git a/tests/unit/sandbox/operator/test_k8s_template_loader.py b/tests/unit/sandbox/operator/test_k8s_template_loader.py index 8a05f58bb2..9586a6449d 100644 --- a/tests/unit/sandbox/operator/test_k8s_template_loader.py +++ b/tests/unit/sandbox/operator/test_k8s_template_loader.py @@ -3,6 +3,7 @@ import pytest from rock.sandbox.operator.k8s.constants import K8sConstants +from rock.sandbox.operator.k8s.provider import TemplateSpec from rock.sandbox.operator.k8s.template_loader import K8sTemplateLoader @@ -377,3 +378,171 @@ def test_build_manifest_passes_encrypted_image_auth(self): annotations = manifest["spec"]["template"]["metadata"]["annotations"] assert annotations["example.com/image-auth"] == "dGVzdC1lbmNyeXB0ZWQ=" + + +class TestBuildPoolManifest: + """Tests for Pool CRD manifest building via K8sTemplateLoader.""" + + POOL_TEMPLATE = { + "capacitySpec": { + "bufferMin": 1, + "bufferMax": 3, + "poolMin": 1, + "poolMax": 10, + }, + "template": { + "metadata": {"labels": {"app": "rock-pool"}}, + "spec": { + "tolerations": [{"operator": "Exists"}], + "containers": [{ + "name": "main", + "image": "{{ from_image }}", + "resources": { + "limits": { + "cpu": "{{ cpu_count }}", + "memory": "{{ memory_mb }}Mi", + }, + "requests": { + "cpu": "{{ cpu_count }}", + "memory": "{{ memory_mb }}Mi", + }, + }, + }], + }, + }, + } + + WINDOWS_POOL_TEMPLATE = { + "capacitySpec": { + "bufferMin": 0, + "bufferMax": 2, + "poolMin": 0, + "poolMax": 5, + }, + "template": { + "metadata": {"labels": {"app": "rock-pool-windows"}}, + "spec": { + "tolerations": [{"operator": "Exists"}], + "nodeSelector": {"kubernetes.io/os": "windows"}, + "containers": [{ + "name": "main", + "image": "{{ from_image }}", + "resources": { + "limits": { + "cpu": "{{ cpu_count }}", + "memory": "{{ memory_mb }}Mi", + }, + "requests": { + "cpu": "{{ cpu_count }}", + "memory": "{{ memory_mb }}Mi", + }, + }, + }], + }, + }, + } + + @pytest.fixture + def pool_loader(self): + """Create a loader with a single default pool template.""" + return K8sTemplateLoader( + templates={"default": {"ports": {"proxy": 8000}, "template": {"spec": {}}}}, + default_namespace="rock-test", + pool_templates={"default": self.POOL_TEMPLATE}, + ) + + @pytest.fixture + def multi_pool_loader(self): + """Create a loader with default + windows pool templates.""" + return K8sTemplateLoader( + templates={"default": {"ports": {"proxy": 8000}, "template": {"spec": {}}}}, + default_namespace="rock-test", + pool_templates={ + "default": self.POOL_TEMPLATE, + "windows": self.WINDOWS_POOL_TEMPLATE, + }, + ) + + def test_build_pool_manifest_basic(self, pool_loader): + """Pool CRD wrapper is assembled correctly.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + manifest = pool_loader.build_pool_manifest("tpl-abc123", spec) + + assert manifest["apiVersion"] == K8sConstants.CRD_API_VERSION + assert manifest["kind"] == K8sConstants.CRD_KIND_POOL + assert manifest["metadata"]["name"] == "tpl-abc123" + assert manifest["metadata"]["namespace"] == "rock-test" + assert manifest["metadata"]["labels"][K8sConstants.LABEL_MANAGED_BY] == K8sConstants.LABEL_MANAGED_BY_TEMPLATE_API + + def test_build_pool_manifest_renders_image(self, pool_loader): + """Jinja2 renders from_image variable.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + manifest = pool_loader.build_pool_manifest("tpl-abc123", spec) + + container = manifest["spec"]["template"]["spec"]["containers"][0] + assert container["image"] == "python:3.11" + + def test_build_pool_manifest_renders_cpu(self, pool_loader): + """Jinja2 renders cpu_count variable.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=4, memory_mb=2048) + manifest = pool_loader.build_pool_manifest("tpl-abc123", spec) + + container = manifest["spec"]["template"]["spec"]["containers"][0] + assert container["resources"]["limits"]["cpu"] == "4" + + def test_build_pool_manifest_renders_memory(self, pool_loader): + """Jinja2 renders memory_mb variable.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=4096) + manifest = pool_loader.build_pool_manifest("tpl-abc123", spec) + + container = manifest["spec"]["template"]["spec"]["containers"][0] + assert container["resources"]["limits"]["memory"] == "4096Mi" + + def test_build_pool_manifest_capacity_defaults(self, pool_loader): + """Capacity uses defaults from the pool template.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + manifest = pool_loader.build_pool_manifest("tpl-abc123", spec) + + cap = manifest["spec"]["capacitySpec"] + assert cap["bufferMin"] == 1 + assert cap["bufferMax"] == 3 + assert cap["poolMin"] == 1 + assert cap["poolMax"] == 10 + + def test_build_pool_manifest_without_pool_templates(self): + """Calling build_pool_manifest without pool templates raises ValueError.""" + loader = K8sTemplateLoader( + templates={"default": {"ports": {"proxy": 8000}, "template": {"spec": {}}}}, + default_namespace="rock-test", + ) + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + + with pytest.raises(ValueError, match="No pool template configured"): + loader.build_pool_manifest("tpl-abc123", spec) + + def test_build_pool_manifest_selects_by_template_name(self, multi_pool_loader): + """template_name selects the matching pool template.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + manifest = multi_pool_loader.build_pool_manifest("tpl-abc123", spec, template_name="windows") + + # windows template has poolMax=5 and nodeSelector + cap = manifest["spec"]["capacitySpec"] + assert cap["poolMax"] == 5 + assert cap["bufferMax"] == 2 + pod_spec = manifest["spec"]["template"]["spec"] + assert pod_spec["nodeSelector"]["kubernetes.io/os"] == "windows" + + def test_build_pool_manifest_default_template_name(self, multi_pool_loader): + """Omitting template_name uses the default pool template.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + manifest = multi_pool_loader.build_pool_manifest("tpl-abc123", spec) + + cap = manifest["spec"]["capacitySpec"] + assert cap["poolMax"] == 10 + + def test_build_pool_manifest_unknown_template_name_raises(self, multi_pool_loader): + """Unknown template_name raises ValueError.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + + with pytest.raises(ValueError, match="Pool template 'macos' not found"): + multi_pool_loader.build_pool_manifest("tpl-abc123", spec, template_name="macos") diff --git a/tests/unit/sandbox/operator/test_k8s_template_provider.py b/tests/unit/sandbox/operator/test_k8s_template_provider.py new file mode 100644 index 0000000000..2d2e2fd66e --- /dev/null +++ b/tests/unit/sandbox/operator/test_k8s_template_provider.py @@ -0,0 +1,285 @@ +"""Unit tests for BatchSandboxProvider template (Pool warm path) methods. + +Tests the operator/provider layer directly: +- generate_template_id +- TemplateSpec dataclass +- _map_pool_to_template_status +- scale_template +- K8sConstants for Pool CRD +""" + +import pytest + +from rock.sandbox.operator.k8s.constants import K8sConstants +from rock.sandbox.operator.k8s.provider import BatchSandboxProvider, TemplateSpec, generate_template_id +from rock.sdk.common.exceptions import BadRequestRockError + + +class TestGenerateTemplateId: + """Tests for generate_template_id function.""" + + def test_same_spec_same_id(self): + """Same (fromImage, cpuCount, memoryMB) produce the same template ID.""" + spec1 = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + spec2 = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + assert generate_template_id(spec1) == generate_template_id(spec2) + + def test_different_image_different_id(self): + """Different fromImage produces different template IDs.""" + spec1 = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + spec2 = TemplateSpec(from_image="python:3.12", cpu_count=2, memory_mb=2048) + assert generate_template_id(spec1) != generate_template_id(spec2) + + def test_different_cpu_different_id(self): + """Different cpuCount produces different template IDs.""" + spec1 = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + spec2 = TemplateSpec(from_image="python:3.11", cpu_count=4, memory_mb=2048) + assert generate_template_id(spec1) != generate_template_id(spec2) + + def test_different_memory_different_id(self): + """Different memoryMB produces different template IDs.""" + spec1 = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + spec2 = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=4096) + assert generate_template_id(spec1) != generate_template_id(spec2) + + def test_optional_fields_affect_id(self): + """All non-null optional fields participate in template ID.""" + base = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + with_disk = TemplateSpec( + from_image="python:3.11", cpu_count=2, memory_mb=2048, disk_gb=40, + ) + with_gpu = TemplateSpec( + from_image="python:3.11", cpu_count=2, memory_mb=2048, num_gpus=2, + ) + with_accelerator = TemplateSpec( + from_image="python:3.11", cpu_count=2, memory_mb=2048, accelerator_type="A100", + ) + with_os = TemplateSpec( + from_image="python:3.11", cpu_count=2, memory_mb=2048, os="linux", + ) + assert generate_template_id(base) != generate_template_id(with_disk) + assert generate_template_id(base) != generate_template_id(with_gpu) + assert generate_template_id(base) != generate_template_id(with_accelerator) + assert generate_template_id(base) != generate_template_id(with_os) + + def test_different_os_different_id(self): + """Different os produces different template IDs.""" + spec1 = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048, os="linux") + spec2 = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048, os="windows") + assert generate_template_id(spec1) != generate_template_id(spec2) + + def test_id_has_prefix(self): + """Template ID starts with the configured prefix.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + tid = generate_template_id(spec) + assert tid.startswith(K8sConstants.TEMPLATE_ID_PREFIX) + + +class TestTemplateSpec: + """Tests for TemplateSpec dataclass.""" + + def test_required_fields(self): + """Required fields are set correctly.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + assert spec.from_image == "python:3.11" + assert spec.cpu_count == 2 + assert spec.memory_mb == 2048 + + def test_optional_fields_default_none(self): + """Optional fields default to None.""" + spec = TemplateSpec(from_image="python:3.11", cpu_count=2, memory_mb=2048) + assert spec.disk_gb is None + assert spec.num_gpus is None + assert spec.accelerator_type is None + assert spec.os is None + + def test_capacity_fields_removed(self): + """Capacity fields are excluded from TemplateSpec identity.""" + with pytest.raises(TypeError): + TemplateSpec( + from_image="python:3.11", + cpu_count=2, + memory_mb=2048, + buffer_min=1, + ) + + +def _make_provider(): + """Create a BatchSandboxProvider without calling __init__.""" + return object.__new__(BatchSandboxProvider) + + +class TestMapPoolToTemplateStatus: + """Tests for Pool CRD to template status mapping (design doc format).""" + + def test_ready_pool(self): + """Pool with available > 0 maps to 'ready'.""" + provider = _make_provider() + mock_pool = { + "metadata": {"name": "tpl-abc123", "creationTimestamp": "2026-08-03T06:00:00Z"}, + "spec": { + "capacitySpec": {"poolMin": 1, "poolMax": 10, "bufferMin": 1, "bufferMax": 3}, + }, + "status": {"available": 2, "total": 3, "allocated": 1}, + } + result = provider._map_pool_to_template_status(mock_pool) + assert result["template_id"] == "tpl-abc123" + assert result["status"] == "ready" + assert result["reason"] is None + assert result["created_at"] == "2026-08-03T06:00:00Z" + assert result["capacity"]["spec"]["pool_min"] == 1 + assert result["capacity"]["status"]["available"] == 2 + + def test_building_pool(self): + """Pool with total but no available maps to 'building'.""" + provider = _make_provider() + mock_pool = { + "metadata": {"name": "tpl-abc123", "creationTimestamp": "2026-08-03T06:00:00Z"}, + "spec": { + "capacitySpec": {"poolMin": 1, "poolMax": 10, "bufferMin": 1, "bufferMax": 3}, + }, + "status": {"available": 0, "total": 3, "allocated": 0}, + } + result = provider._map_pool_to_template_status(mock_pool) + assert result["status"] == "building" + assert result["capacity"]["spec"]["pool_max"] == 10 + assert result["capacity"]["status"]["total"] == 3 + + def test_new_pool_no_status(self): + """Pool with no status section maps to 'building'.""" + provider = _make_provider() + mock_pool = { + "metadata": {"name": "tpl-abc123", "creationTimestamp": "2026-08-03T06:00:00Z"}, + "spec": { + "capacitySpec": {"poolMin": 1, "poolMax": 10, "bufferMin": 1, "bufferMax": 3}, + }, + } + result = provider._map_pool_to_template_status(mock_pool) + assert result["status"] == "building" + assert result["reason"] is None + assert result["capacity"]["spec"]["buffer_min"] == 1 + assert result["capacity"]["status"]["available"] is None + + def test_zero_capacity_pool_ready(self): + """Pool with poolMin=0 and bufferMin=0 maps to 'ready' without replicas.""" + provider = _make_provider() + mock_pool = { + "metadata": {"name": "tpl-abc123", "creationTimestamp": "2026-08-03T06:00:00Z"}, + "spec": { + "capacitySpec": {"poolMin": 0, "bufferMin": 0, "poolMax": 0, "bufferMax": 0}, + }, + "status": {"available": 0, "total": 0}, + } + result = provider._map_pool_to_template_status(mock_pool) + assert result["status"] == "ready" + assert result["reason"] is None + assert result["capacity"]["spec"]["pool_min"] == 0 + assert result["capacity"]["status"]["total"] == 0 + + def test_updated_at_fallback(self): + """updated_at falls back to creationTimestamp when status has no timestamp.""" + provider = _make_provider() + mock_pool = { + "metadata": {"name": "tpl-abc123", "creationTimestamp": "2026-08-03T06:00:00Z"}, + "spec": { + "capacitySpec": {"poolMin": 1, "poolMax": 10, "bufferMin": 1, "bufferMax": 3}, + }, + "status": {"available": 2, "total": 3}, + } + result = provider._map_pool_to_template_status(mock_pool) + assert result["updated_at"] == "2026-08-03T06:00:00Z" + + +@pytest.fixture +def anyio_backend(): + """Run anyio async tests on asyncio backend only.""" + return "asyncio" + + +class MockPoolApi: + """Minimal mock for provider scale tests.""" + + def __init__(self): + self.existing_pool = None + self.last_patch = None + + async def get_custom_object(self, name: str): + return self.existing_pool + + async def update_custom_object(self, name: str, body: dict): + self.last_patch = body + return self.existing_pool + + +class TestScaleTemplate: + """Tests for BatchSandboxProvider.scale_template.""" + + def _make_provider(self): + """Create a provider with mocked _pool_api and initialized flag.""" + provider = object.__new__(BatchSandboxProvider) + provider._initialized = True + provider._pool_api = MockPoolApi() + return provider + + @pytest.mark.anyio + async def test_scale_updates_capacity(self): + """Scale patches capacitySpec and returns mapped status.""" + provider = self._make_provider() + provider._pool_api.existing_pool = { + "metadata": {"name": "tpl-abc", "creationTimestamp": "2026-08-03T06:00:00Z"}, + "spec": {"capacitySpec": {"poolMin": 1, "poolMax": 10, "bufferMin": 1, "bufferMax": 3}}, + "status": {"available": 2, "total": 3}, + } + + result = await provider.scale_template("tpl-abc", {"pool_min": 2, "pool_max": 20}) + + assert result["template_id"] == "tpl-abc" + assert result["status"] == "ready" + assert result["capacity"]["spec"]["pool_min"] == 1 + assert result["capacity"]["status"]["available"] == 2 + assert provider._pool_api.last_patch == { + "spec": {"capacitySpec": {"poolMin": 2, "poolMax": 20}} + } + + @pytest.mark.anyio + async def test_scale_not_found(self): + """Scaling a non-existent template raises BadRequestRockError.""" + provider = self._make_provider() + provider._pool_api.existing_pool = None + + with pytest.raises(BadRequestRockError, match="Template tpl-missing not found"): + await provider.scale_template("tpl-missing", {"pool_min": 1}) + + @pytest.mark.anyio + async def test_scale_empty_capacity(self): + """Empty capacity dict raises BadRequestRockError.""" + provider = self._make_provider() + + with pytest.raises(BadRequestRockError, match="No capacity fields provided"): + await provider.scale_template("tpl-abc", {}) + + @pytest.mark.anyio + async def test_scale_invalid_field(self): + """Invalid capacity field raises BadRequestRockError.""" + provider = self._make_provider() + + with pytest.raises(BadRequestRockError, match="Invalid capacity fields"): + await provider.scale_template("tpl-abc", {"pool_min": 1, "unknown": 5}) + + +class TestK8sConstants: + """Tests for new K8sConstants entries.""" + + def test_pool_crd_constants(self): + """Pool CRD constants are set correctly.""" + assert K8sConstants.CRD_PLURAL_POOL == "pools" + assert K8sConstants.CRD_KIND_POOL == "Pool" + + def test_label_constants(self): + """Label constants are set correctly.""" + assert K8sConstants.LABEL_MANAGED_BY == "rock.sandbox/managed-by" + assert K8sConstants.LABEL_MANAGED_BY_TEMPLATE_API == "template-api" + + def test_template_id_prefix(self): + """Template ID prefix is set correctly.""" + assert K8sConstants.TEMPLATE_ID_PREFIX == "tpl-"