refactor(template): move pool template rendering into K8sTemplateLoader

- Extend K8sTemplateLoader with optional pool_template and build_pool_manifest
- BatchSandboxProvider delegates _build_pool_manifest_from_template to loader
- Remove provider-level pool template / jinja env / render imports
- Migrate pool manifest tests from test_template_api.py to test_k8s_template_loader.py

Author: Generalwin

rock/sandbox/operator/k8s/provider.py

@@ -1,7 +1,6 @@
 """K8s provider implementations for managing sandbox resources."""
 
 import base64
-import copy
 import fnmatch
 import hashlib
 import json
@@ -11,7 +10,6 @@
 from dataclasses import dataclass
 from typing import Any, Protocol
 
-import jinja2
 import yaml
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from kubernetes import client
@@ -28,7 +26,6 @@
 from rock.sandbox.operator.k8s.constants import K8sConstants
 from rock.sandbox.operator.k8s.template_loader import K8sTemplateLoader
 from rock.sandbox.remote_sandbox import RemoteSandboxRuntime
-from rock.utils.jinja_render import render_node
 
 logger = init_logger(__name__)
 
@@ -312,14 +309,11 @@ def __init__(self, k8s_config: K8sConfig):
         self._nacos_provider = None
         self._image_auth_key = self._load_image_auth_key(k8s_config)
 
-        # Pool template for Template API (Warm path)
-        self._pool_template = k8s_config.pool_template
-        self._jinja_env = jinja2.Environment()
-
-        # 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_template=k8s_config.pool_template,
         )
         logger.info(f"Available K8S templates: {', '.join(self._template_loader.available_templates)}")
 
@@ -1057,54 +1051,10 @@ async def _delete_pool(self, pool_name: str) -> None:
     def _build_pool_manifest_from_template(self, pool_name: str, spec: TemplateSpec) -> dict[str, Any]:
         """Build Pool CRD manifest from config template and spec.
 
-        Renders Jinja2 variables in the pool_template config:
-        from_image, cpu_count, memory_mb, disk_gb, num_gpus, accelerator_type,
-        buffer_min, buffer_max, pool_min, pool_max
+        Delegates to the shared K8sTemplateLoader so pool and sandbox template
+        rendering live in one place.
         """
-        template = copy.deepcopy(self._pool_template)
-
-        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.buffer_min is not None:
-            ctx["buffer_min"] = spec.buffer_min
-        if spec.buffer_max is not None:
-            ctx["buffer_max"] = spec.buffer_max
-        if spec.pool_min is not None:
-            ctx["pool_min"] = spec.pool_min
-        if spec.pool_max is not None:
-            ctx["pool_max"] = spec.pool_max
-
-        rendered = render_node(template, self._jinja_env, ctx)
-
-        # Ensure capacitySpec values are integers (Jinja2 renders them as strings)
-        cap = rendered.get("capacitySpec", {})
-        if cap:
-            for k in ("bufferMin", "bufferMax", "poolMin", "poolMax"):
-                if k in cap:
-                    cap[k] = int(cap[k])
-
-        manifest = {
-            "apiVersion": K8sConstants.CRD_API_VERSION,
-            "kind": K8sConstants.CRD_KIND_POOL,
-            "metadata": {
-                "name": pool_name,
-                "namespace": self.namespace,
-                "labels": {
-                    K8sConstants.LABEL_MANAGED_BY: K8sConstants.LABEL_MANAGED_BY_TEMPLATE_API,
-                },
-            },
-            "spec": rendered,
-        }
-        return manifest
+        return self._template_loader.build_pool_manifest(pool_name, spec)
 
     def _map_pool_to_template_status(self, pool: dict) -> dict:
         """Map Pool CRD to template status dict (design doc format)."""

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,24 @@
 
 
 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_template: 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_template: Optional Pool CRD template for Template API (Warm path)
         """
         self._templates: dict[str, dict[str, Any]] = templates
         self._default_namespace = default_namespace
+        self._pool_template = pool_template
 
         if not self._templates:
             raise ValueError("No templates provided. At least one template must be defined in K8sConfig.templates.")
@@ -177,6 +184,74 @@ def build_manifest(
 
         return manifest
 
+    def build_pool_manifest(self, pool_name: str, spec: Any) -> 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, buffer_min, buffer_max, pool_min, pool_max. Optional
+        fields are only added to the context when they are not None.
+
+        After rendering, ``capacitySpec`` integer fields are converted from
+        strings back to ints (Jinja2 renders numbers as strings).
+
+        Args:
+            pool_name: Name for the Pool CRD (also the template ID).
+            spec: Template creation spec with the attributes listed above.
+
+        Returns:
+            Complete Pool CRD manifest.
+
+        Raises:
+            ValueError: If no pool template was configured.
+        """
+        if not self._pool_template:
+            raise ValueError("No pool template configured. Set k8s.pool_template in config.")
+
+        template = copy.deepcopy(self._pool_template)
+
+        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.buffer_min is not None:
+            ctx["buffer_min"] = spec.buffer_min
+        if spec.buffer_max is not None:
+            ctx["buffer_max"] = spec.buffer_max
+        if spec.pool_min is not None:
+            ctx["pool_min"] = spec.pool_min
+        if spec.pool_max is not None:
+            ctx["pool_max"] = spec.pool_max
+
+        rendered = render_node(template, self._jinja_env, ctx)
+
+        # Ensure capacitySpec values are integers (Jinja2 renders them as strings)
+        cap = rendered.get("capacitySpec", {})
+        if cap:
+            for k in ("bufferMin", "bufferMax", "poolMin", "poolMax"):
+                if k in cap:
+                    cap[k] = int(cap[k])
+
+        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."""

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,121 @@ 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": "{{ buffer_min | default(1) }}",
+            "bufferMax": "{{ buffer_max | default(3) }}",
+            "poolMin": "{{ pool_min | default(1) }}",
+            "poolMax": "{{ pool_max | default(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",
+                        },
+                    },
+                }],
+            },
+        },
+    }
+
+    @pytest.fixture
+    def pool_loader(self):
+        """Create a loader with a pool template."""
+        return K8sTemplateLoader(
+            templates={"default": {"ports": {"proxy": 8000}, "template": {"spec": {}}}},
+            default_namespace="rock-test",
+            pool_template=self.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 and is converted to integers."""
+        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_custom_capacity(self, pool_loader):
+        """Custom capacity values are rendered and converted to integers."""
+        spec = TemplateSpec(
+            from_image="python:3.11",
+            cpu_count=2,
+            memory_mb=2048,
+            buffer_min=2,
+            buffer_max=5,
+            pool_min=3,
+            pool_max=20,
+        )
+        manifest = pool_loader.build_pool_manifest("tpl-abc123", spec)
+
+        cap = manifest["spec"]["capacitySpec"]
+        assert cap["bufferMin"] == 2
+        assert cap["bufferMax"] == 5
+        assert cap["poolMin"] == 3
+        assert cap["poolMax"] == 20
+
+    def test_build_pool_manifest_without_pool_template(self):
+        """Calling build_pool_manifest without a pool template 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)