feat: base on_startup auto-loads AGL_POD_SPEC_TEMPLATE; AGL_JOB_MANIFEST_TEMPLATE default
ControllerSettings.job_manifest_template: - Now has a default of "deploy/controller/job-template.yaml.j2" - cli.py --job-manifest-template option no longer required (same default) DeploySettings: - Add pod_spec_template: str | None = None (AGL_POD_SPEC_TEMPLATE) User pod spec template — plain YAML pod spec fragment loaded by the base RolloutHooks.on_startup into self._pod_spec at server startup RolloutHooks.on_startup (base): - Reads AGL_POD_SPEC_TEMPLATE from env; if set, yaml.safe_load into self._pod_spec - No-op when env var is absent (no FileNotFoundError, no KeyError) - Docstring explains the two distinct files: AGL_JOB_MANIFEST_TEMPLATE (Jinja2 scaffold used by controller) vs AGL_POD_SPEC_TEMPLATE (plain YAML used by hooks) - Subclasses that only need per-sample customisation no longer need on_startup examples/swe_bench/hooks.py: - Remove on_startup override — base handles file loading now - Remove SWE_POD_SPEC_TEMPLATE; use AGL_POD_SPEC_TEMPLATE instead deploy/agl-lite.env.example: - Expand controller section: document AGL_JOB_MANIFEST_TEMPLATE with its default, and AGL_POD_SPEC_TEMPLATE with description and example path
This commit is contained in:
yuqyang
+4
-1
@@ -40,7 +40,10 @@ def controller(
|
||||
agl_base_url: str = typer.Option("http://localhost:8000", help="agl-lite server URL"),
|
||||
namespace: str = typer.Option("default", help="K8s namespace for agent Jobs"),
|
||||
secret_name: str = typer.Option("agl-api-keys", help="K8s Secret name with API keys"),
|
||||
job_manifest_template: str = typer.Option(..., help="Path to Jinja2 job manifest template file"),
|
||||
job_manifest_template: str = typer.Option(
|
||||
"deploy/controller/job-template.yaml.j2",
|
||||
help="Path to Jinja2 job manifest template (AGL_JOB_MANIFEST_TEMPLATE)",
|
||||
),
|
||||
) -> None:
|
||||
"""Start the K8s controller (reconcile loop)."""
|
||||
import asyncio
|
||||
|
||||
@@ -28,4 +28,4 @@ class ControllerSettings(BaseSettings):
|
||||
# Job defaults (can be overridden by resources snapshot).
|
||||
ttl_after_finished: int = 3600 # ttlSecondsAfterFinished on Jobs (pod GC safety)
|
||||
|
||||
job_manifest_template: str # path to Jinja2 job manifest template; always required
|
||||
job_manifest_template: str = "deploy/controller/job-template.yaml.j2" # Jinja2 job scaffold; AGL_JOB_MANIFEST_TEMPLATE
|
||||
|
||||
@@ -47,6 +47,12 @@ class DeploySettings(BaseSettings):
|
||||
host_port: int = 8080
|
||||
|
||||
job_manifest_template: str | None = None
|
||||
|
||||
# User pod spec template — plain YAML file loaded by the base RolloutHooks.on_startup
|
||||
# into self._pod_spec (AGL_POD_SPEC_TEMPLATE). Often the job-template.yaml in the
|
||||
# example folder. When set, hooks get self._pod_spec for free without overriding
|
||||
# on_startup. Defaults to None (no pod spec loaded by base).
|
||||
pod_spec_template: str | None = None
|
||||
gateway_config: str | None = None
|
||||
hooks: str | None = None
|
||||
artifact_dir: str | None = None
|
||||
|
||||
+26
-13
@@ -10,14 +10,12 @@ loads the module at startup via ``--hooks path/to/hooks.py``.
|
||||
Typical pattern::
|
||||
|
||||
class MyHooks(RolloutHooks):
|
||||
def on_startup(self, store: InMemoryStore) -> None:
|
||||
# Hook-specific config from env vars — document what your hook expects.
|
||||
self._pod_spec = yaml.safe_load(
|
||||
Path(os.environ["MY_POD_SPEC_TEMPLATE"]).read_text()
|
||||
)
|
||||
# on_startup is optional — if AGL_POD_SPEC_TEMPLATE is set in the
|
||||
# environment the base implementation loads the pod spec automatically.
|
||||
# Override only when you need extra setup beyond file loading.
|
||||
|
||||
def on_enqueue(self, request: EnqueueRolloutRequest) -> EnqueueRolloutRequest:
|
||||
pod_spec = self.copy_pod_spec()
|
||||
pod_spec = self.copy_pod_spec() # deep copy of the loaded template
|
||||
agent = self.get_container(pod_spec, "agent")
|
||||
agent["image"] = f"my-image:{request.input['version']}"
|
||||
request.config.pod_spec = pod_spec
|
||||
@@ -27,8 +25,12 @@ Typical pattern::
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import yaml
|
||||
|
||||
from agl_lite.schemas.api import EnqueueRolloutRequest
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -57,17 +59,28 @@ class RolloutHooks:
|
||||
def on_startup(self, store: InMemoryStore) -> None:
|
||||
"""Called once by the server after startup and store initialisation.
|
||||
|
||||
Override to load per-dataset resources (pod spec templates, eval configs,
|
||||
etc.) that are needed for every request. Hook-specific config should come
|
||||
from environment variables — document what your hook expects.
|
||||
The base implementation reads ``AGL_POD_SPEC_TEMPLATE`` from the
|
||||
environment and, if set, loads that YAML file into ``self._pod_spec``.
|
||||
This covers the common case where all instances use the same container
|
||||
image base and only differ in per-sample env vars.
|
||||
|
||||
Example::
|
||||
``AGL_POD_SPEC_TEMPLATE`` — path to a plain YAML file that describes the
|
||||
pod spec fragment: ``containers``, optional ``volumes``, ``nodeSelector``,
|
||||
``tolerations``, ``activeDeadlineSeconds``, etc. Typically
|
||||
``examples/<project>/job-template.yaml``.
|
||||
|
||||
Override only when you need setup beyond file loading, e.g. loading a
|
||||
dataset index or connecting to an external registry. When overriding,
|
||||
call ``super().on_startup(store)`` first so the base pod spec load still
|
||||
happens::
|
||||
|
||||
def on_startup(self, store: InMemoryStore) -> None:
|
||||
self._pod_spec = yaml.safe_load(
|
||||
Path(os.environ["MY_POD_SPEC_TEMPLATE"]).read_text()
|
||||
)
|
||||
super().on_startup(store) # loads AGL_POD_SPEC_TEMPLATE
|
||||
self._index = load_index(os.environ["MY_INDEX"])
|
||||
"""
|
||||
template_path = os.environ.get("AGL_POD_SPEC_TEMPLATE")
|
||||
if template_path:
|
||||
self._pod_spec = yaml.safe_load(Path(template_path).read_text())
|
||||
|
||||
def copy_pod_spec(self) -> dict[str, Any]:
|
||||
"""Return a deep copy of the stored pod spec template.
|
||||
|
||||
@@ -27,9 +27,16 @@ AGL_HOST_IP_BIND=0.0.0.0
|
||||
AGL_HOST_PORT=8080
|
||||
|
||||
# ── Controller ────────────────────────────────────────────────────
|
||||
# Optional: path to a custom Jinja2 job manifest template.
|
||||
# Defaults to deploy/controller/job-template.yaml.j2 when unset.
|
||||
# AGL_JOB_MANIFEST_TEMPLATE=/path/to/job-template.yaml.j2
|
||||
# Jinja2 job manifest template — two YAML docs: Job scaffold + PodPatcher.
|
||||
# Defaults to deploy/controller/job-template.yaml.j2 (bundled with agl-lite).
|
||||
# AGL_JOB_MANIFEST_TEMPLATE=deploy/controller/job-template.yaml.j2
|
||||
|
||||
# User pod spec template — plain YAML pod spec fragment loaded into the hook’s
|
||||
# self._pod_spec at startup (containers, volumes, nodeSelector, tolerations,
|
||||
# activeDeadlineSeconds, …). Often examples/<project>/job-template.yaml.
|
||||
# When set, the base RolloutHooks.on_startup handles the load automatically;
|
||||
# hooks only need to override on_enqueue for per-sample customisation.
|
||||
# AGL_POD_SPEC_TEMPLATE=examples/math-poc/job-template.yaml
|
||||
|
||||
# ── Server runtime ────────────────────────────────────────────────
|
||||
# AGL_GATEWAY_CONFIG=examples/swe_bench/gateway-config.yaml
|
||||
|
||||
@@ -1,25 +1,20 @@
|
||||
"""SWE-bench hooks — task-specific logic for SWE-bench rollouts.
|
||||
|
||||
on_startup: load pod spec template from SWE_POD_SPEC_TEMPLATE env var.
|
||||
on_enqueue: deep copy template, set per-instance image, inject env vars
|
||||
(AGL_TASK_INPUT, AGL_EVAL_SCRIPT, AGL_EVAL_META, etc.),
|
||||
set config.timeout from template.
|
||||
|
||||
on_succeeded / on_failed: post zero-reward fallback if container didn't post one.
|
||||
Grading is done in the container using official swebench tools.
|
||||
The base RolloutHooks.on_startup automatically loads the pod spec from
|
||||
AGL_POD_SPEC_TEMPLATE (set in the project .env file). This hook only needs
|
||||
to customise on_enqueue (per-instance image + env vars) and on_succeeded /
|
||||
on_failed (reward posting).
|
||||
|
||||
Required env vars:
|
||||
SWE_POD_SPEC_TEMPLATE path to the pod spec YAML (e.g. examples/swe_bench/job-template.yaml)
|
||||
AGL_POD_SPEC_TEMPLATE path to the pod spec YAML (e.g. examples/swe_bench/job-template.yaml)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
from swebench.harness.test_spec.test_spec import make_test_spec
|
||||
|
||||
from agl_lite.hooks import RolloutHooks
|
||||
@@ -30,11 +25,6 @@ from agl_lite.store.memory import InMemoryStore
|
||||
|
||||
class SWEBenchHooks(RolloutHooks):
|
||||
|
||||
def on_startup(self, store: InMemoryStore) -> None:
|
||||
"""Load pod spec template once from SWE_POD_SPEC_TEMPLATE env var."""
|
||||
path = os.environ["SWE_POD_SPEC_TEMPLATE"]
|
||||
self._pod_spec = yaml.safe_load(Path(path).read_text())
|
||||
|
||||
def on_enqueue(self, request: EnqueueRolloutRequest) -> EnqueueRolloutRequest:
|
||||
instance = request.input
|
||||
if not isinstance(instance, dict) or "instance_id" not in instance:
|
||||
|
||||
+40
-11
@@ -2,10 +2,12 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import textwrap
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -79,21 +81,48 @@ class ErrorHooks(RolloutHooks):
|
||||
|
||||
|
||||
class TestOnStartup:
|
||||
def test_on_startup_called_with_store(self) -> None:
|
||||
"""on_startup receives the store and can load resources into self."""
|
||||
class StartupHooks(RolloutHooks):
|
||||
started: bool = False
|
||||
def test_base_loads_pod_spec_from_env(self, tmp_path) -> None:
|
||||
"""Base on_startup reads AGL_POD_SPEC_TEMPLATE and populates self._pod_spec."""
|
||||
import yaml
|
||||
pod_spec = {"containers": [{"name": "agent", "image": "auto:v1"}]}
|
||||
f = tmp_path / "pod-spec.yaml"
|
||||
f.write_text(yaml.dump(pod_spec))
|
||||
|
||||
hooks = RolloutHooks()
|
||||
with patch.dict(os.environ, {"AGL_POD_SPEC_TEMPLATE": str(f)}):
|
||||
hooks.on_startup(InMemoryStore())
|
||||
|
||||
assert hooks._pod_spec is not None
|
||||
assert hooks._pod_spec["containers"][0]["image"] == "auto:v1"
|
||||
|
||||
def test_base_no_op_when_env_unset(self) -> None:
|
||||
"""Base on_startup is a no-op when AGL_POD_SPEC_TEMPLATE is not set."""
|
||||
env = {k: v for k, v in os.environ.items() if k != "AGL_POD_SPEC_TEMPLATE"}
|
||||
hooks = RolloutHooks()
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
hooks.on_startup(InMemoryStore()) # must not raise
|
||||
assert hooks._pod_spec is None
|
||||
|
||||
def test_subclass_can_call_super(self, tmp_path) -> None:
|
||||
"""Subclass calling super().on_startup gets the env-loaded pod spec."""
|
||||
import yaml
|
||||
pod_spec = {"containers": [{"name": "agent", "image": "base:v1"}]}
|
||||
f = tmp_path / "pod-spec.yaml"
|
||||
f.write_text(yaml.dump(pod_spec))
|
||||
|
||||
class MyHooks(RolloutHooks):
|
||||
index_loaded: bool = False
|
||||
|
||||
def on_startup(self, store: InMemoryStore) -> None:
|
||||
self.started = True
|
||||
self._pod_spec = {"containers": [{"name": "agent", "image": "startup:v1"}]}
|
||||
super().on_startup(store)
|
||||
self.index_loaded = True # simulate extra setup
|
||||
|
||||
hooks = StartupHooks()
|
||||
store = InMemoryStore()
|
||||
hooks.on_startup(store)
|
||||
h = MyHooks()
|
||||
with patch.dict(os.environ, {"AGL_POD_SPEC_TEMPLATE": str(f)}):
|
||||
h.on_startup(InMemoryStore())
|
||||
|
||||
assert hooks.started is True
|
||||
assert hooks._pod_spec is not None
|
||||
assert h._pod_spec is not None
|
||||
assert h.index_loaded is True
|
||||
|
||||
def test_copy_pod_spec_deep_copies(self) -> None:
|
||||
hooks = RolloutHooks()
|
||||
|
||||
@@ -74,27 +74,28 @@ def swe_instance() -> dict:
|
||||
|
||||
|
||||
class TestOnStartup:
|
||||
def test_loads_pod_spec_from_env_var(self, hooks, tmp_path) -> None:
|
||||
"""on_startup reads SWE_POD_SPEC_TEMPLATE env var and loads pod spec."""
|
||||
def test_loads_pod_spec_from_agl_env_var(self, hooks, tmp_path) -> None:
|
||||
"""Base on_startup reads AGL_POD_SPEC_TEMPLATE; SWEBenchHooks inherits this."""
|
||||
import yaml
|
||||
pod_spec = {"containers": [{"name": "agent", "image": "test:v1"}]}
|
||||
template_file = tmp_path / "job-template.yaml"
|
||||
template_file.write_text(yaml.dump(pod_spec))
|
||||
|
||||
h = hooks.SWEBenchHooks()
|
||||
with patch.dict(os.environ, {"SWE_POD_SPEC_TEMPLATE": str(template_file)}):
|
||||
with patch.dict(os.environ, {"AGL_POD_SPEC_TEMPLATE": str(template_file)}):
|
||||
from agl_lite.store.memory import InMemoryStore
|
||||
h.on_startup(InMemoryStore())
|
||||
|
||||
assert h._pod_spec is not None
|
||||
assert h._pod_spec["containers"][0]["name"] == "agent"
|
||||
|
||||
def test_missing_env_var_raises(self, hooks) -> None:
|
||||
def test_missing_env_var_leaves_pod_spec_none(self, hooks) -> None:
|
||||
"""No AGL_POD_SPEC_TEMPLATE → _pod_spec stays None (copy_pod_spec raises later)."""
|
||||
h = hooks.SWEBenchHooks()
|
||||
env = {k: v for k, v in os.environ.items() if k != "SWE_POD_SPEC_TEMPLATE"}
|
||||
env = {k: v for k, v in os.environ.items() if k != "AGL_POD_SPEC_TEMPLATE"}
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
with pytest.raises(KeyError):
|
||||
h.on_startup(InMemoryStore())
|
||||
h.on_startup(InMemoryStore()) # must not raise
|
||||
assert h._pod_spec is None
|
||||
|
||||
|
||||
class TestOnEnqueue:
|
||||
|
||||
Reference in New Issue
Block a user