diff --git a/agentlightning/config/controller.yaml b/agentlightning/config/controller.yaml index d92a42a1b..590a38657 100644 --- a/agentlightning/config/controller.yaml +++ b/agentlightning/config/controller.yaml @@ -13,6 +13,10 @@ k8s_runner: ttl_after_finished: 1200 max_jobs_per_minute: 100 poll_interval: 5 + image_readiness: + enabled: true + heartbeat_seconds: 5 + lease_seconds: 30 local_runner: maximum_size: 50 diff --git a/agentlightning/controller/k8s_reconciler.py b/agentlightning/controller/k8s_reconciler.py index 7cf5f7c07..df6ff4d21 100644 --- a/agentlightning/controller/k8s_reconciler.py +++ b/agentlightning/controller/k8s_reconciler.py @@ -12,22 +12,28 @@ from __future__ import annotations import asyncio -import json import time from collections import deque +from collections.abc import Mapping, Sequence from typing import Any, cast import httpx import kr8s import kr8s.asyncio import structlog -import yaml -from jinja2 import Environment from kr8s.asyncio import objects as k8s_objects from omegaconf import DictConfig from agentlightning.client import AgentLightningAsyncClient -from agentlightning.schemas import DEFAULT_ATTEMPT_ID, Rollout, RolloutPatch, RolloutState, RolloutStatusPatch +from agentlightning.k8s import extract_pod_images, normalize_image_reference, render_job_template +from agentlightning.schemas import ( + DEFAULT_ATTEMPT_ID, + K8sImageReadinessReport, + Rollout, + RolloutPatch, + RolloutState, + RolloutStatusPatch, +) log = structlog.get_logger() @@ -40,25 +46,49 @@ def build_job_name(rollout_id: str) -> str: return f"agl-rollout-{rollout_id}" +def images_available_on_all_ready_nodes( + nodes: Sequence[Mapping[str, Any]], +) -> tuple[frozenset[str], int]: + """Return the image-name intersection for Ready, schedulable nodes.""" + eligible: list[Mapping[str, Any]] = [] + for node in nodes: + if node.get("spec", {}).get("unschedulable", False): + continue + conditions = node.get("status", {}).get("conditions", []) + if not any(condition.get("type") == "Ready" and condition.get("status") == "True" for condition in conditions): + continue + eligible.append(node) + + if not eligible: + raise RuntimeError("no Ready schedulable Kubernetes nodes found") + + per_node: list[set[str]] = [] + for node in eligible: + names = { + normalize_image_reference(name) + for image in node.get("status", {}).get("images", []) + for name in image.get("names", []) + if isinstance(name, str) and name.strip() + } + per_node.append(names) + + common = set(per_node[0]) + for names in per_node[1:]: + common.intersection_update(names) + return frozenset(common), len(eligible) + + def build_job_spec(rollout: Rollout, controller_config: DictConfig) -> dict[str, Any]: """Build a K8s Job manifest from the rollout's complete Jinja2 Job template.""" template = rollout.config.k8s.job_template if rollout.config.k8s else None if not template: raise ValueError("invalid rollout config: missing config.k8s.job_template") - env = Environment() - env.filters["yaml_escape"] = lambda value: json.dumps(str(value), ensure_ascii=True) - rendered = env.from_string(template).render( + job = render_job_template( + template, job_name=build_job_name(rollout.rollout_id), - input=rollout.input, + input_data=rollout.input, ) - docs = [doc for doc in yaml.safe_load_all(rendered) if doc is not None] - if len(docs) != 1: - raise ValueError("invalid rollout config: config.k8s.job_template must render exactly one YAML document") - - job = docs[0] - if not isinstance(job, dict) or job.get("kind") != "Job": - raise ValueError("invalid rollout config: config.k8s.job_template must render a Kubernetes Job") metadata = job.setdefault("metadata", {}) metadata["name"] = build_job_name(rollout.rollout_id) @@ -117,6 +147,15 @@ def __init__(self, api: AgentLightningAsyncClient, config: DictConfig) -> None: self._k8s_api: Any | None = None self._stop = asyncio.Event() self._job_creation_timestamps: deque[float] = deque() + self._ready_images: frozenset[str] | None = None + self._ready_images_expires_at = 0.0 + + readiness_config = self._runner_config.image_readiness + if readiness_config.enabled: + heartbeat = float(readiness_config.heartbeat_seconds) + lease = float(readiness_config.lease_seconds) + if not 0 < heartbeat < lease <= 300: + raise ValueError("k8s_runner.image_readiness requires 0 < heartbeat_seconds < lease_seconds <= 300") async def _get_k8s_api(self) -> Any: if self._k8s_api is None: @@ -131,10 +170,20 @@ async def run(self) -> None: poll_interval=self._runner_config.poll_interval, ) try: - await asyncio.gather( - self._periodic_reconcile_loop(), - self._watch_jobs_loop(), + tasks = [] + if self._runner_config.image_readiness.enabled: + try: + await self._publish_image_readiness_once() + except Exception: + log.exception("Initial K8s image readiness publish failed") + tasks.append(self._image_readiness_loop()) + tasks.extend( + [ + self._periodic_reconcile_loop(), + self._watch_jobs_loop(), + ] ) + await asyncio.gather(*tasks) except asyncio.CancelledError: log.info("Controller stopped") @@ -142,6 +191,46 @@ def stop(self) -> None: """Signal the controller to stop.""" self._stop.set() + async def _scan_preloaded_images(self) -> tuple[frozenset[str], int]: + api = await self._get_k8s_api() + nodes = [cast(k8s_objects.Node, node).raw async for node in k8s_objects.Node.async_list(api=api)] + return images_available_on_all_ready_nodes(nodes) + + async def _publish_image_readiness_once(self) -> None: + images, node_count = await self._scan_preloaded_images() + lease_seconds = float(self._runner_config.image_readiness.lease_seconds) + scanned_at = time.monotonic() + report = K8sImageReadinessReport( + images=sorted(images), + node_count=node_count, + lease_seconds=lease_seconds, + ) + response = await self._api.put( + "/api/runner-readiness/k8s", + json=report.model_dump(mode="json"), + ) + response.raise_for_status() + self._ready_images = images + self._ready_images_expires_at = scanned_at + lease_seconds + + async def _image_readiness_loop(self) -> None: + heartbeat = float(self._runner_config.image_readiness.heartbeat_seconds) + while not self._stop.is_set(): + try: + await asyncio.wait_for(self._stop.wait(), timeout=heartbeat) + return + except TimeoutError: + pass + try: + await self._publish_image_readiness_once() + except Exception: + log.exception("K8s image readiness publish failed") + + def _fresh_ready_images(self) -> frozenset[str] | None: + if self._ready_images is None or time.monotonic() >= self._ready_images_expires_at: + return None + return self._ready_images + # --- Periodic reconcile --- async def _periodic_reconcile_loop(self) -> None: @@ -242,27 +331,54 @@ async def _reconcile_once(self) -> None: async def _create_job(self, rollout: Rollout) -> None: """Create a K8s Job for a queuing rollout without changing rollout state.""" job_name = build_job_name(rollout.rollout_id) - now = time.monotonic() - window_start = now - JOB_CREATION_WINDOW_SECONDS - while self._job_creation_timestamps and self._job_creation_timestamps[0] <= window_start: - self._job_creation_timestamps.popleft() - if len(self._job_creation_timestamps) >= self._runner_config.max_jobs_per_minute: - log.info( - "Job creation rate limit reached — deferring queued rollouts", - rollout_id=rollout.rollout_id, - jobs_in_last_minute=len(self._job_creation_timestamps), - max_jobs_per_minute=self._runner_config.max_jobs_per_minute, - ) - return - try: manifest = build_job_spec(rollout, self._config) attempt_id = manifest["metadata"]["labels"]["agentlightning/attempt-id"] + requires_preloaded = bool(rollout.config.k8s and rollout.config.k8s.require_preloaded_images) + if requires_preloaded: + ready_images = self._fresh_ready_images() + if ready_images is None: + await self._patch_status( + rollout.rollout_id, + state=RolloutState.FAILED, + error_message="Fresh Kubernetes image readiness is unavailable", + ) + return + missing_images = sorted(extract_pod_images(manifest) - ready_images) + if missing_images: + await self._patch_status( + rollout.rollout_id, + state=RolloutState.FAILED, + error_message=("Required Kubernetes image(s) are not preloaded: " + ", ".join(missing_images)), + ) + return + + now = time.monotonic() + window_start = now - JOB_CREATION_WINDOW_SECONDS + while self._job_creation_timestamps and self._job_creation_timestamps[0] <= window_start: + self._job_creation_timestamps.popleft() + if len(self._job_creation_timestamps) >= self._runner_config.max_jobs_per_minute: + log.info( + "Job creation rate limit reached — deferring queued rollouts", + rollout_id=rollout.rollout_id, + jobs_in_last_minute=len(self._job_creation_timestamps), + max_jobs_per_minute=self._runner_config.max_jobs_per_minute, + ) + return + api = await self._get_k8s_api() job = k8s_objects.Job(manifest, api=api) await job.async_create() self._job_creation_timestamps.append(time.monotonic()) log.info("Job created", rollout_id=rollout.rollout_id, job_name=job_name, attempt_id=attempt_id) + except ValueError as exc: + error_str = str(exc) + log.error("Invalid Job spec — marking failed", rollout_id=rollout.rollout_id, error=error_str) + await self._patch_status( + rollout.rollout_id, + state=RolloutState.FAILED, + error_message=f"Invalid Job spec: {error_str}", + ) except Exception as exc: error_str = str(exc) lower_error = error_str.lower() diff --git a/agentlightning/k8s.py b/agentlightning/k8s.py new file mode 100644 index 000000000..4462524fc --- /dev/null +++ b/agentlightning/k8s.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Pure helpers shared by Kubernetes rollout producers and consumers.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from functools import lru_cache +from typing import Any + +import yaml +from jinja2 import Environment, Template + +__all__ = [ + "extract_pod_images", + "normalize_image_reference", + "render_job_template", +] + + +def normalize_image_reference(image: str) -> str: + """Return a canonical image reference suitable for exact comparison.""" + reference = image.strip() + if not reference: + raise ValueError("container image must be a non-empty string") + if "://" in reference: + reference = reference.split("://", 1)[1] + + if "/" not in reference: + reference = f"docker.io/library/{reference}" + else: + first, remainder = reference.split("/", 1) + if first in {"docker.io", "index.docker.io"}: + if "/" not in remainder: + remainder = f"library/{remainder}" + reference = f"docker.io/{remainder}" + elif "." not in first and ":" not in first and first != "localhost": + reference = f"docker.io/{reference}" + + last_component = reference.rsplit("/", 1)[-1] + if "@" not in reference and ":" not in last_component: + reference = f"{reference}:latest" + return reference + + +@lru_cache(maxsize=32) +def _compile_job_template(job_template: str) -> Template: + environment = Environment() + environment.filters["yaml_escape"] = lambda value: json.dumps(str(value), ensure_ascii=True) + return environment.from_string(job_template) + + +def render_job_template( + job_template: str, + *, + job_name: str, + input_data: Any, +) -> dict[str, Any]: + """Render one Kubernetes Job from the controller-compatible template.""" + rendered = _compile_job_template(job_template).render( + job_name=job_name, + input=input_data, + ) + documents = [document for document in yaml.safe_load_all(rendered) if document is not None] + if len(documents) != 1: + raise ValueError("job template must render exactly one YAML document") + job = documents[0] + if not isinstance(job, dict) or job.get("kind") != "Job": + raise ValueError("job template must render a Kubernetes Job") + return job + + +def extract_pod_images(job: Mapping[str, Any]) -> frozenset[str]: + """Extract normalized images from every container list in a Job Pod spec.""" + pod_spec = job.get("spec", {}).get("template", {}).get("spec", {}) + images: set[str] = set() + for container_key in ("initContainers", "containers", "ephemeralContainers"): + for container in pod_spec.get(container_key, []) or []: + image = container.get("image") + if not isinstance(image, str) or not image.strip(): + raise ValueError(f"{container_key} entry is missing a non-empty image") + images.add(normalize_image_reference(image)) + if not images: + raise ValueError("rendered Kubernetes Job contains no container images") + return frozenset(images) diff --git a/agentlightning/schemas.py b/agentlightning/schemas.py index cf2c5901c..232ba238a 100644 --- a/agentlightning/schemas.py +++ b/agentlightning/schemas.py @@ -102,6 +102,23 @@ class RolloutState(StrEnum): DEFAULT_ATTEMPT_ID = "0" +class K8sImageReadinessReport(BaseModel): + """Controller report of images cached on every eligible K8s node.""" + + images: list[str] = Field(default_factory=list) + node_count: int = Field(ge=1) + lease_seconds: float = Field(gt=0, le=300) + + +class K8sImageReadinessSnapshot(BaseModel): + """Server-timestamped, leased K8s image inventory.""" + + images: list[str] = Field(default_factory=list) + node_count: int = Field(ge=1) + observed_at: float + expires_at: float + + class RolloutLocalConfig(BaseModel): """Local runner config for a rollout.""" @@ -113,6 +130,7 @@ class RolloutK8sConfig(BaseModel): """K8s runner config for a rollout.""" job_template: str | None = None + require_preloaded_images: bool = False class RolloutConfig(BaseModel): diff --git a/agentlightning/server/app.py b/agentlightning/server/app.py index 97cacbb36..87612d64a 100644 --- a/agentlightning/server/app.py +++ b/agentlightning/server/app.py @@ -15,7 +15,7 @@ from omegaconf import DictConfig, OmegaConf from agentlightning.server.proxy import ProxyPauseState, ProxyRouter -from agentlightning.server.routes import events, models, proxy, rollouts +from agentlightning.server.routes import events, models, proxy, readiness, rollouts log = structlog.get_logger() @@ -83,6 +83,7 @@ async def healthz() -> dict[str, str]: app.include_router(rollouts.router, prefix="/api", dependencies=[Depends(verify_key)]) app.include_router(events.router, prefix="/api", dependencies=[Depends(verify_key)]) app.include_router(models.router, prefix="/api", dependencies=[Depends(verify_key)]) + app.include_router(readiness.router, prefix="/api", dependencies=[Depends(verify_key)]) # Proxy routes (LLM proxy + event ingestion) — require agent-facing auth. app.include_router(proxy.router, dependencies=[Depends(verify_key)]) diff --git a/agentlightning/server/routes/readiness.py b/agentlightning/server/routes/readiness.py new file mode 100644 index 000000000..fd668c1c0 --- /dev/null +++ b/agentlightning/server/routes/readiness.py @@ -0,0 +1,45 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Runner-readiness publication and lookup routes.""" + +from __future__ import annotations + +import time + +from fastapi import APIRouter, HTTPException + +from agentlightning.k8s import normalize_image_reference +from agentlightning.schemas import K8sImageReadinessReport, K8sImageReadinessSnapshot +from agentlightning.server.store import _runner_readiness + +router = APIRouter(tags=["runner-readiness"]) +_K8S_KEY = "k8s" + + +@router.put("/runner-readiness/k8s", response_model=K8sImageReadinessSnapshot) +async def publish_k8s_image_readiness(body: K8sImageReadinessReport) -> K8sImageReadinessSnapshot: + """Publish a leased image inventory using the server's clock.""" + now = time.time() + try: + images = sorted({normalize_image_reference(image) for image in body.images}) + except ValueError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc + snapshot = K8sImageReadinessSnapshot( + images=images, + node_count=body.node_count, + observed_at=now, + expires_at=now + body.lease_seconds, + ) + _runner_readiness[_K8S_KEY] = snapshot + return snapshot + + +@router.get("/runner-readiness/k8s", response_model=K8sImageReadinessSnapshot) +async def get_k8s_image_readiness() -> K8sImageReadinessSnapshot: + """Return the current snapshot, failing closed when absent or expired.""" + snapshot = _runner_readiness.get(_K8S_KEY) + if snapshot is None: + raise HTTPException(status_code=503, detail="K8s image readiness has not been published") + if snapshot.expires_at <= time.time(): + raise HTTPException(status_code=503, detail="K8s image readiness snapshot has expired") + return snapshot diff --git a/agentlightning/server/store.py b/agentlightning/server/store.py index fcdb3a9f8..dc64fcb90 100644 --- a/agentlightning/server/store.py +++ b/agentlightning/server/store.py @@ -8,11 +8,12 @@ from __future__ import annotations -from agentlightning.schemas import Event, Model, Rollout +from agentlightning.schemas import Event, K8sImageReadinessSnapshot, Model, Rollout _rollouts: dict[str, Rollout] = {} _events: dict[str, dict[str, list[Event]]] = {} _models: dict[str, dict[str, Model]] = {} +_runner_readiness: dict[str, K8sImageReadinessSnapshot] = {} # Completion-ordered ids enable cursor pagination without rescanning rollouts. _terminal_order: list[str] = [] diff --git a/agentlightning/verl/agl_rollout_manager.py b/agentlightning/verl/agl_rollout_manager.py index 7e6860746..f43b56264 100644 --- a/agentlightning/verl/agl_rollout_manager.py +++ b/agentlightning/verl/agl_rollout_manager.py @@ -255,6 +255,7 @@ def __init__( local_agent_class: str | None = None, local_env_map: dict[str, str] | None = None, k8s_job_template_path: str | None = None, + k8s_require_preloaded_images: bool = False, ) -> None: self._model = model self._step = step @@ -268,7 +269,10 @@ def __init__( "env_map": local_env_map or {}, } if k8s_job_template_path: - self._rollout_config["k8s"] = {"job_template": Path(k8s_job_template_path).read_text()} + self._rollout_config["k8s"] = { + "job_template": Path(k8s_job_template_path).read_text(), + "require_preloaded_images": k8s_require_preloaded_images, + } self.client = AgentLightningSyncClient( base_url=agl_base_url, diff --git a/agentlightning/verl/config.yaml b/agentlightning/verl/config.yaml index 7395d1d0f..75cd970b1 100644 --- a/agentlightning/verl/config.yaml +++ b/agentlightning/verl/config.yaml @@ -19,6 +19,8 @@ agentlightning: env_map: {} k8s: job_template_path: null + filter_unavailable_images: false + image_readiness_timeout_seconds: 60 reward_fillna_value: 0.0 max_ppo_update_times: null trace_aggregator: diff --git a/agentlightning/verl/entrypoint.py b/agentlightning/verl/entrypoint.py index 5c5b4031c..4e0c99d8b 100644 --- a/agentlightning/verl/entrypoint.py +++ b/agentlightning/verl/entrypoint.py @@ -22,6 +22,7 @@ from omegaconf import OmegaConf from .dataset import LoadedDataset +from .k8s_image_filter import PreparedDatasets, prepare_datasets log = logging.getLogger(__name__) @@ -34,16 +35,28 @@ def run_ppo( config: Any, train_dataset: Sequence[Any], val_dataset: Sequence[Any], + *, + max_val_instances: int | None = None, ) -> None: """Launch VERL PPO training with Agent Lightning agent orchestration. Datasets must be passed as non-empty in-memory sequences. """ - from verl.trainer.main_ppo import get_ppo_ray_runtime_env + prepared = prepare_datasets( + config, + train_dataset, + val_dataset, + max_val_instances=max_val_instances, + ) + train_dataset = prepared.train + val_dataset = prepared.val + _log_dataset_preparation(prepared) assert train_dataset is not None and len(train_dataset) > 0, "train_dataset must be non-empty" assert val_dataset is not None and len(val_dataset) > 0, "val_dataset must be non-empty" + from verl.trainer.main_ppo import get_ppo_ray_runtime_env + if not ray.is_initialized(): default_runtime_env = cast(dict[str, Any], get_ppo_ray_runtime_env()) ray_init_config = OmegaConf.to_container(config.ray_kwargs.get("ray_init", OmegaConf.create({})), resolve=True) @@ -72,6 +85,34 @@ def run_ppo( ray.get(runner.run.remote(config, train_ds, val_ds)) +def _log_dataset_preparation(prepared: PreparedDatasets) -> None: + if prepared.readiness is None: + return + print( + "K8s image readiness: " + f"{len(prepared.readiness.images)} normalized images on " + f"{prepared.readiness.node_count} eligible node(s); snapshot=fresh", + flush=True, + ) + for report in (prepared.train_report, prepared.val_report): + if report is None: + continue + suffix = f" final={len(prepared.val)}" if report.split == "val" else "" + print( + f"Image filter {report.split}: source={report.source_count} " + f"kept={report.kept_count} dropped={report.dropped_count}{suffix}", + flush=True, + ) + if report.missing_image_counts: + counts = ", ".join(f"{image}={count}" for image, count in report.missing_image_counts.items()) + print(f"Missing image counts ({report.split}): {counts}", flush=True) + if report.dropped_data_ids: + print( + f"Dropped ids ({report.split}, first 20): " + ", ".join(report.dropped_data_ids), + flush=True, + ) + + @ray.remote(num_cpus=1) class _AglTaskRunner: """TaskRunner that extends verl's TaskRunner with pre-loaded dataset support.""" diff --git a/agentlightning/verl/k8s_image_filter.py b/agentlightning/verl/k8s_image_filter.py new file mode 100644 index 000000000..4058133ff --- /dev/null +++ b/agentlightning/verl/k8s_image_filter.py @@ -0,0 +1,189 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Pre-Ray dataset filtering using CPU-side Kubernetes image readiness.""" + +from __future__ import annotations + +import math +import time +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import httpx + +from agentlightning.client import AgentLightningSyncClient +from agentlightning.k8s import extract_pod_images, normalize_image_reference, render_job_template +from agentlightning.schemas import K8sImageReadinessSnapshot + +__all__ = [ + "DatasetFilterReport", + "PreparedDatasets", + "filter_dataset_by_images", + "prepare_datasets", + "wait_for_k8s_image_readiness", +] + + +@dataclass +class DatasetFilterReport: + split: str + source_count: int + kept_count: int + dropped_count: int + dropped_data_ids: list[str] + missing_image_counts: dict[str, int] + + +@dataclass +class PreparedDatasets: + train: list[Any] + val: list[Any] + train_report: DatasetFilterReport | None + val_report: DatasetFilterReport | None + readiness: K8sImageReadinessSnapshot | None + + +def filter_dataset_by_images( + dataset: Sequence[Any], + *, + split: str, + job_template: str, + ready_images: set[str] | frozenset[str], +) -> tuple[list[Any], DatasetFilterReport]: + """Return rows whose rendered Job images are all available.""" + ready = {normalize_image_reference(image) for image in ready_images} + kept: list[Any] = [] + dropped_ids: list[str] = [] + missing_counts: Counter[str] = Counter() + + for index, row in enumerate(dataset): + job = render_job_template( + job_template, + job_name=f"agl-image-check-{index}", + input_data=row, + ) + missing = sorted(extract_pod_images(job) - ready) + if not missing: + kept.append(row) + continue + + row_id = str(row.get("data_id") or row.get("instance_id") or index) if isinstance(row, Mapping) else str(index) + if len(dropped_ids) < 20: + dropped_ids.append(row_id) + missing_counts.update(missing) + + report = DatasetFilterReport( + split=split, + source_count=len(dataset), + kept_count=len(kept), + dropped_count=len(dataset) - len(kept), + dropped_data_ids=dropped_ids, + missing_image_counts=dict(sorted(missing_counts.items())), + ) + return kept, report + + +def wait_for_k8s_image_readiness( + client: httpx.Client, + *, + timeout_seconds: float, +) -> K8sImageReadinessSnapshot: + """Wait for a fresh server-validated readiness snapshot.""" + if not math.isfinite(timeout_seconds) or timeout_seconds < 0: + raise ValueError("timeout_seconds must be a finite non-negative number") + + endpoint = "/api/runner-readiness/k8s" + deadline = time.monotonic() + timeout_seconds + last_detail = "readiness has not been published" + + while True: + remaining = deadline - time.monotonic() + try: + response = client.get(endpoint, timeout=min(5.0, max(0.0, remaining))) + except httpx.RequestError as exc: + last_detail = str(exc) + else: + if response.status_code == 200: + return K8sImageReadinessSnapshot.model_validate(response.json()) + if response.status_code != 503: + response.raise_for_status() + try: + payload = response.json() + last_detail = str(payload.get("detail", response.text)) + except ValueError: + last_detail = response.text + + remaining = deadline - time.monotonic() + if remaining <= 0: + raise RuntimeError(f"Timed out waiting for {endpoint}: {last_detail}") + time.sleep(min(1.0, remaining)) + + +def _fetch_snapshot(config: Any) -> K8sImageReadinessSnapshot: + with AgentLightningSyncClient( + base_url=str(config.agentlightning.agl_base_url), + key=str(config.agentlightning.agl_key or "") or None, + max_retries=0, + timeout=5.0, + ) as client: + return wait_for_k8s_image_readiness( + client, + timeout_seconds=float(config.agentlightning.k8s.image_readiness_timeout_seconds), + ) + + +def _require_non_empty(train_rows: Sequence[Any], val_rows: Sequence[Any]) -> None: + if not train_rows: + raise ValueError("train dataset is empty after image filtering") + if not val_rows: + raise ValueError("validation dataset is empty after image filtering") + + +def prepare_datasets( + config: Any, + train_dataset: Sequence[Any], + val_dataset: Sequence[Any], + *, + max_val_instances: int | None, +) -> PreparedDatasets: + """Prepare in-memory datasets while preserving disabled-mode behavior.""" + train_rows = list(train_dataset) + val_rows = list(val_dataset) + enabled = bool(config.agentlightning.k8s.filter_unavailable_images) + + if not enabled: + if max_val_instances: + val_rows = val_rows[:max_val_instances] + return PreparedDatasets(train_rows, val_rows, None, None, None) + + template_path = config.agentlightning.k8s.job_template_path + if not template_path: + raise ValueError("agentlightning.k8s.job_template_path is required when filter_unavailable_images=true") + + job_template = Path(str(template_path)).read_text(encoding="utf-8") + snapshot = _fetch_snapshot(config) + train_rows, train_report = filter_dataset_by_images( + train_rows, + split="train", + job_template=job_template, + ready_images=set(snapshot.images), + ) + val_rows, val_report = filter_dataset_by_images( + val_rows, + split="val", + job_template=job_template, + ready_images=set(snapshot.images), + ) + if max_val_instances: + val_rows = val_rows[:max_val_instances] + _require_non_empty(train_rows, val_rows) + return PreparedDatasets( + train=train_rows, + val=val_rows, + train_report=train_report, + val_report=val_report, + readiness=snapshot, + ) diff --git a/agentlightning/verl/trainer.py b/agentlightning/verl/trainer.py index 57575e7ba..8458bb2c4 100644 --- a/agentlightning/verl/trainer.py +++ b/agentlightning/verl/trainer.py @@ -145,6 +145,7 @@ def _make_rollout_manager(self, manager_cls: type[RolloutManagerT]) -> RolloutMa local_agent_class=al.local.agent_class, local_env_map=al.local.env_map, k8s_job_template_path=al.k8s.job_template_path, + k8s_require_preloaded_images=bool(al.k8s.filter_unavailable_images), ) def _rollout_replicas(self) -> list[Any]: diff --git a/docs/20-trainer-configuration.md b/docs/20-trainer-configuration.md index d5b24ab15..42a5269ea 100644 --- a/docs/20-trainer-configuration.md +++ b/docs/20-trainer-configuration.md @@ -18,6 +18,8 @@ agentlightning: env_map: {} k8s: job_template_path: null + filter_unavailable_images: false + image_readiness_timeout_seconds: 60 reward_fillna_value: 0.0 max_ppo_update_times: null trace_aggregator: @@ -81,6 +83,8 @@ The Controller has two execution modes: `local` and `k8s`. Configure the matchin | `agentlightning.local.agent_class` | `null` | Fully qualified Python class imported and started by the Controller in local mode. | | `agentlightning.local.env_map` | `{}` | Maps environment variable names to fields in the rollout `input`. | | `agentlightning.k8s.job_template_path` | `null` | Path to the Jinja Kubernetes Job template used by the Controller in K8s mode. | +| `agentlightning.k8s.filter_unavailable_images` | `false` | Before Ray starts, filter train and validation rows whose rendered Job images are not cached on every eligible K8s node. | +| `agentlightning.k8s.image_readiness_timeout_seconds` | `60` | Maximum time to wait for a fresh Controller readiness snapshot when filtering is enabled. | For local execution, set the agent class and map fields from each dataset row into environment variables. For example: @@ -115,6 +119,8 @@ env: The trainer reads the Jinja template and includes its text in each rollout. The Controller renders it with that rollout's `input`, then creates one Kubernetes Job per rollout. +Set `agentlightning.k8s.filter_unavailable_images=true` when repository-specific images are prepared on the Kubernetes nodes. The trainer obtains a fresh, leased inventory from the Controller, renders this same Job template for every train and validation row, and keeps a row only when all regular, init, and ephemeral container images are available. This in-memory preflight runs before Ray or GPU workers are initialized. It never modifies source dataset files, and `job_template_path` is required while the filter is enabled. Missing or expired readiness, an invalid template, or an empty filtered split fails startup instead of silently using unfiltered data. + Finally, `agentlightning.rollout_timeout_seconds` sets the maximum execution time for each rollout in both modes. The Controller uses this value and marks a rollout as failed if it does not finish within the configured number of seconds. The default is `1800`. ## Trace aggregator diff --git a/docs/30-controller-configuration.md b/docs/30-controller-configuration.md index 2af5f2093..c197a53f6 100644 --- a/docs/30-controller-configuration.md +++ b/docs/30-controller-configuration.md @@ -19,6 +19,10 @@ k8s_runner: ttl_after_finished: 1200 max_jobs_per_minute: 100 poll_interval: 5 + image_readiness: + enabled: true + heartbeat_seconds: 5 + lease_seconds: 30 local_runner: maximum_size: 50 @@ -75,6 +79,20 @@ The K8s runner provides settings that limit Job creation and clean up completed `max_jobs_per_minute` prevents the Controller from creating too many Jobs in a short period. `ttl_after_finished` prevents completed Jobs from accumulating and overloading the Kubernetes API server. +## K8s image readiness + +The K8s Controller publishes a leased inventory for trainer preflight. It reads `status.images` from Ready, schedulable Kubernetes Nodes and publishes the intersection across those nodes, so a reported image is safe regardless of which eligible node receives a Job. The Controller identity therefore needs permission to list cluster-scoped Node objects. + +Kubelet caps `Node.status.images` at 50 entries by default. Clusters that use this preflight with a larger image set should configure [`nodeStatusMaxImages: -1`](https://kubernetes.io/docs/reference/config-api/kubelet-config.v1beta1/#KubeletConfiguration) on every eligible node so the reported inventory is complete. With a capped inventory, filtering remains conservative but can exclude images that are cached but omitted from Node status. + +| Key | Default | Description | +|---|---:|---| +| `k8s_runner.image_readiness.enabled` | `true` | Publish Kubernetes node image readiness to the API Gateway. | +| `k8s_runner.image_readiness.heartbeat_seconds` | `5` | Interval between node inventory scans and publications. | +| `k8s_runner.image_readiness.lease_seconds` | `30` | How long a successful publication remains fresh. | + +The heartbeat must be shorter than the lease. A failed scan or publication does not renew stale readiness. Rollouts created by a trainer with image filtering enabled also carry an admission guard that checks the Controller's latest fresh inventory, catching image changes observed after trainer preflight. As with any preflight, an image removed only after Job submission is a residual race and the normal rollout timeout remains the fallback. + ## Local runner limits The local runner limits concurrent processes and periodically synchronizes their state: diff --git a/docs/75-example-coding-agent.md b/docs/75-example-coding-agent.md index 510416ec6..5c8a7dcee 100644 --- a/docs/75-example-coding-agent.md +++ b/docs/75-example-coding-agent.md @@ -102,10 +102,16 @@ The launcher passes additional arguments to `train_smith_agent.py`, including `v ```bash examples/swe_smith/run.sh trainer \ + --max-val-instances 1 \ trainer.total_training_steps=100 \ - actor_rollout_ref.rollout.n=4 + actor_rollout_ref.rollout.n=4 \ + agentlightning.k8s.filter_unavailable_images=true ``` +The image filter compares each row's rendered `job-template-openai.yaml` images with the fresh inventory published by the Controller before Ray allocates GPU resources. Rows whose `:openai` image was not prepared are excluded from both in-memory splits for that run; the JSONL files remain unchanged. Disable the option to retain the original unfiltered behavior. + +Because SWE-Smith uses more images than Kubelet reports by default, configure `nodeStatusMaxImages: -1` on the Kubernetes nodes as described in [K8s image readiness](30-controller-configuration.md#k8s-image-readiness). The validation cap above is applied after image filtering. + ## Preventing Reward Hacking A coding agent may obtain the reference fix without solving the task, for example by inspecting Git history, downloading upstream source code with `curl` or `wget`, installing the original package with `pip`, or using Python networking libraries such as `urllib`. diff --git a/examples/swe_smith/train_smith_agent.py b/examples/swe_smith/train_smith_agent.py index 56b139eea..252237d07 100644 --- a/examples/swe_smith/train_smith_agent.py +++ b/examples/swe_smith/train_smith_agent.py @@ -213,7 +213,7 @@ def train( raise RuntimeError("AGL_KEY is required") train_dataset = load_split_file(train_dataset_path) - val_dataset = load_split_file(val_dataset_path, max_instances=max_val_instances) + val_dataset = load_split_file(val_dataset_path) instances = train_dataset + val_dataset distinct_repos = sorted({row["repo"] for row in instances}) @@ -222,7 +222,8 @@ def train( log(f" model: {model or DEFAULT_MODEL}") log(f" train file: {train_dataset_path}") log(f" val file: {val_dataset_path}") - log(f" instances: {len(instances)} (train {len(train_dataset)} / val {len(val_dataset)})") + log(f" source instances: {len(instances)} (train {len(train_dataset)} / val {len(val_dataset)})") + log(f" validation cap: {max_val_instances if max_val_instances is not None else 'all'}") log(f" distinct repos (images to prepare): {len(distinct_repos)}") config = build_config( @@ -236,7 +237,12 @@ def train( pprint(OmegaConf.to_container(config, resolve=True)) log("\n=== Start VERL training ===") - run_ppo(config=config, train_dataset=train_dataset, val_dataset=val_dataset) + run_ppo( + config=config, + train_dataset=train_dataset, + val_dataset=val_dataset, + max_val_instances=max_val_instances, + ) def parse_args() -> tuple[argparse.Namespace, list[str]]: diff --git a/examples/swe_smith/train_smith_agent_megatron.py b/examples/swe_smith/train_smith_agent_megatron.py index fbfce8288..8ca5aa97a 100755 --- a/examples/swe_smith/train_smith_agent_megatron.py +++ b/examples/swe_smith/train_smith_agent_megatron.py @@ -205,7 +205,7 @@ def train( raise RuntimeError("AGL_KEY is required") train_dataset = load_split_file(train_dataset_path) - val_dataset = load_split_file(val_dataset_path, max_instances=max_val_instances) + val_dataset = load_split_file(val_dataset_path) instances = train_dataset + val_dataset distinct_repos = sorted({row["repo"] for row in instances}) @@ -214,7 +214,8 @@ def train( log(f" model: {model or DEFAULT_MODEL}") log(f" train file: {train_dataset_path}") log(f" val file: {val_dataset_path}") - log(f" instances: {len(instances)} (train {len(train_dataset)} / val {len(val_dataset)})") + log(f" source instances: {len(instances)} (train {len(train_dataset)} / val {len(val_dataset)})") + log(f" validation cap: {max_val_instances if max_val_instances is not None else 'all'}") log(f" distinct repos (images to prepare): {len(distinct_repos)}") config = build_config( @@ -228,7 +229,12 @@ def train( pprint(OmegaConf.to_container(config, resolve=True)) log("\n=== Start VERL training (Megatron actor, R3 router replay, vLLM rollout) ===") - run_ppo(config=config, train_dataset=train_dataset, val_dataset=val_dataset) + run_ppo( + config=config, + train_dataset=train_dataset, + val_dataset=val_dataset, + max_val_instances=max_val_instances, + ) def parse_args(): diff --git a/tests/controller/test_k8s_reconciler.py b/tests/controller/test_k8s_reconciler.py index 6fa80813f..80546b59b 100644 --- a/tests/controller/test_k8s_reconciler.py +++ b/tests/controller/test_k8s_reconciler.py @@ -1,11 +1,29 @@ # Copyright (c) Microsoft. All rights reserved. -"""Unit tests for controller manifests; no cluster or GPU is required.""" +"""Unit tests for the Kubernetes controller; no cluster or GPU is required.""" +from __future__ import annotations + +import asyncio +import time + +import pytest +from kr8s.asyncio import objects as k8s_objects from omegaconf import OmegaConf -from agentlightning.controller.k8s_reconciler import MANAGED_BY_SELECTOR, build_job_spec -from agentlightning.schemas import Rollout, RolloutConfig, RolloutK8sConfig, RolloutLifecycleStatus +from agentlightning.controller.k8s_reconciler import ( + MANAGED_BY_SELECTOR, + K8sReconciler, + build_job_spec, + images_available_on_all_ready_nodes, +) +from agentlightning.schemas import ( + Rollout, + RolloutConfig, + RolloutK8sConfig, + RolloutLifecycleStatus, + RolloutState, +) def test_build_job_spec_uses_agentlightning_labels() -> None: @@ -48,3 +66,269 @@ def test_build_job_spec_uses_agentlightning_labels() -> None: env = {item["name"]: item["value"] for item in manifest["spec"]["template"]["spec"]["containers"][0]["env"]} assert env["AGL_KEY"] == "secret" assert "/rollout/test-id/attempt/0/mode/train/" in env["AGL_OPENAI_BASE_URL"] + + +def _node( + name: str, + images: list[str], + *, + ready: bool = True, + unschedulable: bool = False, +) -> dict: + return { + "metadata": {"name": name}, + "spec": {"unschedulable": unschedulable}, + "status": { + "conditions": [{"type": "Ready", "status": "True" if ready else "False"}], + "images": [{"names": [image]} for image in images], + }, + } + + +def _controller_config(*, heartbeat: float = 5, lease: float = 30): + return OmegaConf.create( + { + "agl_server": {"url": "http://server:8080", "key": "secret"}, + "k8s_runner": { + "namespace": "default", + "ttl_after_finished": 60, + "max_jobs_per_minute": 100, + "poll_interval": 5, + "image_readiness": { + "enabled": True, + "heartbeat_seconds": heartbeat, + "lease_seconds": lease, + }, + }, + } + ) + + +def test_images_available_on_all_ready_nodes_uses_intersection() -> None: + images, node_count = images_available_on_all_ready_nodes( + [ + _node("node-a", ["docker.io/shared:v1", "docker.io/only-a:v1"]), + _node("node-b", ["docker.io/library/shared:v1", "docker.io/only-b:v1"]), + _node("not-ready", ["docker.io/not-ready:v1"], ready=False), + _node("cordoned", ["docker.io/cordoned:v1"], unschedulable=True), + ] + ) + + assert node_count == 2 + assert images == frozenset({"docker.io/library/shared:v1"}) + + +def test_images_available_on_all_ready_nodes_rejects_no_eligible_nodes() -> None: + with pytest.raises(RuntimeError, match="no Ready schedulable Kubernetes nodes"): + images_available_on_all_ready_nodes([_node("node-a", [], ready=False)]) + + +@pytest.mark.parametrize( + ("heartbeat", "lease"), + [(0, 30), (30, 30), (31, 30), (5, 301)], +) +def test_reconciler_rejects_invalid_readiness_timing(heartbeat: float, lease: float) -> None: + with pytest.raises(ValueError, match="heartbeat_seconds"): + K8sReconciler(object(), _controller_config(heartbeat=heartbeat, lease=lease)) + + +def test_publish_image_readiness_sends_normalized_snapshot(monkeypatch) -> None: + calls: list[tuple[str, dict]] = [] + + class Response: + def raise_for_status(self) -> None: + return None + + class Api: + async def put(self, path: str, *, json: dict) -> Response: + calls.append((path, json)) + return Response() + + reconciler = K8sReconciler(Api(), _controller_config()) + + async def scan() -> tuple[frozenset[str], int]: + return frozenset({"docker.io/swebench/repo:openai"}), 1 + + monkeypatch.setattr(reconciler, "_scan_preloaded_images", scan) + asyncio.run(reconciler._publish_image_readiness_once()) + + assert calls == [ + ( + "/api/runner-readiness/k8s", + { + "images": ["docker.io/swebench/repo:openai"], + "node_count": 1, + "lease_seconds": 30.0, + }, + ) + ] + assert reconciler._ready_images == frozenset({"docker.io/swebench/repo:openai"}) + + +def test_publish_failure_does_not_renew_local_cache(monkeypatch) -> None: + class Response: + def raise_for_status(self) -> None: + raise RuntimeError("server unavailable") + + class Api: + async def put(self, _path: str, *, json: dict) -> Response: + return Response() + + reconciler = K8sReconciler(Api(), _controller_config()) + reconciler._ready_images = frozenset({"docker.io/old:v1"}) + reconciler._ready_images_expires_at = 123.0 + + async def scan() -> tuple[frozenset[str], int]: + return frozenset({"docker.io/new:v1"}), 1 + + monkeypatch.setattr(reconciler, "_scan_preloaded_images", scan) + with pytest.raises(RuntimeError, match="server unavailable"): + asyncio.run(reconciler._publish_image_readiness_once()) + + assert reconciler._ready_images == frozenset({"docker.io/old:v1"}) + assert reconciler._ready_images_expires_at == 123.0 + + +def _image_rollout(*, require_preloaded_images: bool) -> Rollout: + return Rollout( + rollout_id="rollout-image-check", + input={"image_name": "swebench/missing"}, + config=RolloutConfig( + k8s=RolloutK8sConfig( + job_template=""" +apiVersion: batch/v1 +kind: Job +metadata: {} +spec: + template: + spec: + containers: + - name: agent + image: {{ (input.image_name ~ ":openai") | yaml_escape }} +""", + require_preloaded_images=require_preloaded_images, + ) + ), + status=RolloutLifecycleStatus(created_at=1.0, updated_at=1.0), + ) + + +def test_guarded_rollout_fails_before_job_creation_when_image_is_missing(monkeypatch) -> None: + reconciler = K8sReconciler(object(), _controller_config()) + reconciler._ready_images = frozenset({"docker.io/swebench/ready:openai"}) + reconciler._ready_images_expires_at = time.monotonic() + 30 + patched_status: dict = {} + + async def capture_patch(_rollout_id: str, **status) -> bool: + patched_status.update(status) + return True + + async def unexpected_k8s_api(): + raise AssertionError("guarded missing image must not create a Job") + + monkeypatch.setattr(reconciler, "_patch_status", capture_patch) + monkeypatch.setattr(reconciler, "_get_k8s_api", unexpected_k8s_api) + + asyncio.run(reconciler._create_job(_image_rollout(require_preloaded_images=True))) + + assert patched_status == { + "state": RolloutState.FAILED, + "error_message": "Required Kubernetes image(s) are not preloaded: docker.io/swebench/missing:openai", + } + + +def test_guarded_rollout_checks_images_even_when_creation_is_rate_limited(monkeypatch) -> None: + reconciler = K8sReconciler(object(), _controller_config()) + reconciler._ready_images = frozenset({"docker.io/swebench/ready:openai"}) + reconciler._ready_images_expires_at = time.monotonic() + 30 + reconciler._job_creation_timestamps.extend([time.monotonic()] * 100) + patched_status: dict = {} + + async def capture_patch(_rollout_id: str, **status) -> bool: + patched_status.update(status) + return True + + async def unexpected_k8s_api(): + raise AssertionError("guarded missing image must not create a Job") + + monkeypatch.setattr(reconciler, "_patch_status", capture_patch) + monkeypatch.setattr(reconciler, "_get_k8s_api", unexpected_k8s_api) + + asyncio.run(reconciler._create_job(_image_rollout(require_preloaded_images=True))) + + assert patched_status == { + "state": RolloutState.FAILED, + "error_message": "Required Kubernetes image(s) are not preloaded: docker.io/swebench/missing:openai", + } + + +def test_guarded_rollout_fails_when_local_readiness_is_not_fresh(monkeypatch) -> None: + reconciler = K8sReconciler(object(), _controller_config()) + reconciler._ready_images = None + reconciler._ready_images_expires_at = 0.0 + patched_status: dict = {} + + async def capture_patch(_rollout_id: str, **status) -> bool: + patched_status.update(status) + return True + + async def unexpected_k8s_api(): + raise AssertionError("guarded stale readiness must not create a Job") + + monkeypatch.setattr(reconciler, "_patch_status", capture_patch) + monkeypatch.setattr(reconciler, "_get_k8s_api", unexpected_k8s_api) + + asyncio.run(reconciler._create_job(_image_rollout(require_preloaded_images=True))) + + assert patched_status == { + "state": RolloutState.FAILED, + "error_message": "Fresh Kubernetes image readiness is unavailable", + } + + +def test_unguarded_rollout_keeps_original_job_creation_path(monkeypatch) -> None: + reconciler = K8sReconciler(object(), _controller_config()) + reconciler._ready_images = None + created: list[dict] = [] + + class FakeJob: + def __init__(self, manifest: dict, *, api) -> None: + created.append(manifest) + + async def async_create(self) -> None: + return None + + async def fake_k8s_api(): + return object() + + monkeypatch.setattr(reconciler, "_get_k8s_api", fake_k8s_api) + monkeypatch.setattr(k8s_objects, "Job", FakeJob) + + asyncio.run(reconciler._create_job(_image_rollout(require_preloaded_images=False))) + + assert len(created) == 1 + + +def test_invalid_job_template_fails_rollout_instead_of_retrying(monkeypatch) -> None: + rollout = _image_rollout(require_preloaded_images=False) + assert rollout.config.k8s is not None + rollout.config.k8s.job_template = "apiVersion: v1\nkind: Pod\nmetadata: {}\n" + reconciler = K8sReconciler(object(), _controller_config()) + patched_status: dict = {} + + async def capture_patch(_rollout_id: str, **status) -> bool: + patched_status.update(status) + return True + + async def unexpected_k8s_api(): + raise AssertionError("invalid templates must not reach the Kubernetes API") + + monkeypatch.setattr(reconciler, "_patch_status", capture_patch) + monkeypatch.setattr(reconciler, "_get_k8s_api", unexpected_k8s_api) + + asyncio.run(reconciler._create_job(rollout)) + + assert patched_status == { + "state": RolloutState.FAILED, + "error_message": "Invalid Job spec: job template must render a Kubernetes Job", + } diff --git a/tests/examples/test_swe_smith_train_config.py b/tests/examples/test_swe_smith_train_config.py new file mode 100644 index 000000000..cbeee82ea --- /dev/null +++ b/tests/examples/test_swe_smith_train_config.py @@ -0,0 +1,70 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""SWE-Smith trainer image-filter configuration behavior.""" + +from __future__ import annotations + +import importlib +import json +import sys +from pathlib import Path +from typing import Any + +import pytest + +from agentlightning.verl import entrypoint as verl_entrypoint +from examples.swe_smith import train_smith_agent +from examples.swe_smith.train_smith_agent import build_config + +SWE_SMITH_DIR = Path(__file__).resolve().parents[2] / "examples" / "swe_smith" +sys.path.insert(0, str(SWE_SMITH_DIR)) +train_smith_agent_megatron = importlib.import_module("examples.swe_smith.train_smith_agent_megatron") + + +def test_k8s_image_filter_is_disabled_by_default() -> None: + config = build_config(model="Qwen/Test-Model") + + assert config.agentlightning.k8s.filter_unavailable_images is False + assert config.agentlightning.k8s.image_readiness_timeout_seconds == 60 + + +def test_k8s_image_filter_accepts_single_hydra_override() -> None: + config = build_config( + model="Qwen/Test-Model", + config_overrides=["agentlightning.k8s.filter_unavailable_images=true"], + ) + + assert config.agentlightning.k8s.filter_unavailable_images is True + + +@pytest.mark.parametrize("trainer_module", [train_smith_agent, train_smith_agent_megatron]) +def test_train_forwards_validation_cap_after_loading_full_split( + trainer_module, + monkeypatch, + tmp_path: Path, +) -> None: + train_path = tmp_path / "train.jsonl" + val_path = tmp_path / "val.jsonl" + train_path.write_text(json.dumps({"instance_id": "train-1", "repo": "repo"}) + "\n") + val_path.write_text( + json.dumps({"instance_id": "val-1", "repo": "repo"}) + + "\n" + + json.dumps({"instance_id": "val-2", "repo": "repo"}) + + "\n" + ) + captured: dict[str, Any] = {} + + def capture_run_ppo(**kwargs) -> None: + captured.update(kwargs) + + monkeypatch.setattr(verl_entrypoint, "run_ppo", capture_run_ppo) + + trainer_module.train( + train_dataset_path=str(train_path), + val_dataset_path=str(val_path), + max_val_instances=1, + agl_key="secret", + ) + + assert len(captured["val_dataset"]) == 2 + assert captured["max_val_instances"] == 1 diff --git a/tests/server/conftest.py b/tests/server/conftest.py index 00f9e2dd9..fefa599a6 100644 --- a/tests/server/conftest.py +++ b/tests/server/conftest.py @@ -12,7 +12,7 @@ from fastapi.testclient import TestClient from agentlightning.server.app import create_app -from agentlightning.server.store import _events, _models, _rollouts, _terminal_order +from agentlightning.server.store import _events, _models, _rollouts, _runner_readiness, _terminal_order AGL_KEY = "test-secret-key" MODEL_NAME = "test-model" @@ -24,11 +24,13 @@ def clean_store() -> Iterator[None]: _events.clear() _models.clear() _terminal_order.clear() + _runner_readiness.clear() yield _rollouts.clear() _events.clear() _models.clear() _terminal_order.clear() + _runner_readiness.clear() @pytest.fixture diff --git a/tests/server/test_readiness.py b/tests/server/test_readiness.py new file mode 100644 index 000000000..048f9f754 --- /dev/null +++ b/tests/server/test_readiness.py @@ -0,0 +1,85 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +from fastapi.testclient import TestClient + + +def test_publish_and_get_k8s_image_readiness( + client: TestClient, + auth_headers: dict[str, str], + monkeypatch, +) -> None: + monkeypatch.setattr("agentlightning.server.routes.readiness.time.time", lambda: 100.0) + + published = client.put( + "/api/runner-readiness/k8s", + headers=auth_headers, + json={ + "images": ["swebench/repo:openai", "docker.io/swebench/repo:openai", "alpine"], + "node_count": 2, + "lease_seconds": 30, + }, + ) + + assert published.status_code == 200 + assert published.json() == { + "images": ["docker.io/library/alpine:latest", "docker.io/swebench/repo:openai"], + "node_count": 2, + "observed_at": 100.0, + "expires_at": 130.0, + } + fetched = client.get("/api/runner-readiness/k8s", headers=auth_headers) + assert fetched.status_code == 200 + assert fetched.json() == published.json() + + +def test_get_readiness_returns_503_when_missing(client: TestClient, auth_headers: dict[str, str]) -> None: + response = client.get("/api/runner-readiness/k8s", headers=auth_headers) + + assert response.status_code == 503 + assert "not been published" in response.json()["detail"] + + +def test_get_readiness_returns_503_when_expired( + client: TestClient, + auth_headers: dict[str, str], + monkeypatch, +) -> None: + ticks = iter([100.0, 131.0]) + monkeypatch.setattr("agentlightning.server.routes.readiness.time.time", lambda: next(ticks)) + published = client.put( + "/api/runner-readiness/k8s", + headers=auth_headers, + json={"images": [], "node_count": 1, "lease_seconds": 30}, + ) + assert published.status_code == 200 + + response = client.get("/api/runner-readiness/k8s", headers=auth_headers) + + assert response.status_code == 503 + assert "expired" in response.json()["detail"] + + +def test_readiness_endpoints_require_auth(client: TestClient) -> None: + assert client.get("/api/runner-readiness/k8s").status_code == 401 + assert ( + client.put( + "/api/runner-readiness/k8s", + json={"images": [], "node_count": 1, "lease_seconds": 30}, + ).status_code + == 401 + ) + + +def test_publish_rejects_invalid_node_count_and_lease( + client: TestClient, + auth_headers: dict[str, str], +) -> None: + response = client.put( + "/api/runner-readiness/k8s", + headers=auth_headers, + json={"images": [], "node_count": 0, "lease_seconds": 301}, + ) + + assert response.status_code == 422 diff --git a/tests/test_k8s.py b/tests/test_k8s.py new file mode 100644 index 000000000..6933ef3cb --- /dev/null +++ b/tests/test_k8s.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for Kubernetes Job rendering and image matching helpers.""" + +import pytest + +from agentlightning.k8s import extract_pod_images, normalize_image_reference, render_job_template + + +@pytest.mark.parametrize( + ("reference", "expected"), + [ + ("ubuntu", "docker.io/library/ubuntu:latest"), + ("ubuntu:24.04", "docker.io/library/ubuntu:24.04"), + ("docker.io/ubuntu:24.04", "docker.io/library/ubuntu:24.04"), + ("swebench/repo:openai", "docker.io/swebench/repo:openai"), + ("index.docker.io/swebench/repo:openai", "docker.io/swebench/repo:openai"), + ("docker.io/swebench/repo@sha256:abc", "docker.io/swebench/repo@sha256:abc"), + ("localhost:5000/repo:v1", "localhost:5000/repo:v1"), + ], +) +def test_normalize_image_reference_uses_canonical_docker_names(reference: str, expected: str) -> None: + assert normalize_image_reference(reference) == expected + + +def test_normalize_image_reference_rejects_empty_values() -> None: + with pytest.raises(ValueError, match="non-empty"): + normalize_image_reference(" ") + + +def test_render_and_extract_uses_every_pod_container_type() -> None: + template = """ +apiVersion: batch/v1 +kind: Job +metadata: {} +spec: + template: + spec: + initContainers: + - name: setup + image: busybox:1.36 + containers: + - name: agent + image: {{ (input.image_name ~ ":openai") | yaml_escape }} + ephemeralContainers: + - name: debugger + image: alpine +""" + + job = render_job_template( + template, + job_name="agl-image-check", + input_data={"image_name": "swebench/repo"}, + ) + + assert extract_pod_images(job) == frozenset( + { + "docker.io/library/busybox:1.36", + "docker.io/library/alpine:latest", + "docker.io/swebench/repo:openai", + } + ) + + +@pytest.mark.parametrize( + "job", + [ + {"kind": "Job", "spec": {"template": {"spec": {"containers": []}}}}, + {"kind": "Job", "spec": {"template": {"spec": {"containers": [{"name": "agent"}]}}}}, + ], +) +def test_extract_pod_images_rejects_missing_images(job: dict) -> None: + with pytest.raises(ValueError, match="image"): + extract_pod_images(job) + + +@pytest.mark.parametrize( + "template", + [ + "apiVersion: v1\nkind: Pod\nmetadata: {}\n", + "---\nkind: Job\n---\nkind: Job\n", + ], +) +def test_render_job_template_requires_exactly_one_job(template: str) -> None: + with pytest.raises(ValueError, match=r"Job|document"): + render_job_template(template, job_name="test", input_data={}) diff --git a/tests/verl/test_agl_rollout_manager.py b/tests/verl/test_agl_rollout_manager.py index cf13494a1..fcb88a211 100644 --- a/tests/verl/test_agl_rollout_manager.py +++ b/tests/verl/test_agl_rollout_manager.py @@ -4,6 +4,8 @@ from __future__ import annotations +from pathlib import Path + import pytest from agentlightning.schemas import Event, Rollout, RolloutConfig, RolloutLifecycleStatus, RolloutState @@ -238,3 +240,32 @@ def test_aligned_image_urls_text_only_rollout_returns_none_without_warning( # Text-only rollouts keep the exact original behavior: no alignment, no warning. assert _aligned_image_urls(raw_events, 2) is None assert capsys.readouterr().out == "" + + +def test_rollout_manager_propagates_preloaded_image_requirement(tmp_path: Path) -> None: + template_path = tmp_path / "job.yaml" + template_path.write_text( + """ +apiVersion: batch/v1 +kind: Job +metadata: {} +spec: + template: + spec: + containers: + - name: agent + image: example:v1 +""" + ) + manager = AglRolloutManagerBase( + agl_base_url="http://server:8080", + agl_key="secret", + model="test-model", + step=0, + k8s_job_template_path=str(template_path), + k8s_require_preloaded_images=True, + ) + try: + assert manager._rollout_config["k8s"]["require_preloaded_images"] is True + finally: + manager.client.close() diff --git a/tests/verl/test_entrypoint_image_filter.py b/tests/verl/test_entrypoint_image_filter.py new file mode 100644 index 000000000..da4a7a561 --- /dev/null +++ b/tests/verl/test_entrypoint_image_filter.py @@ -0,0 +1,32 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import pytest +from omegaconf import OmegaConf + +from agentlightning.verl import entrypoint + + +def test_run_ppo_preflight_happens_before_ray_init(monkeypatch) -> None: + ray_init_called = False + + def fail_preflight(*_args, **_kwargs): + raise RuntimeError("readiness unavailable") + + def ray_init(**_kwargs): + nonlocal ray_init_called + ray_init_called = True + + monkeypatch.setattr(entrypoint, "prepare_datasets", fail_preflight) + monkeypatch.setattr(entrypoint.ray, "init", ray_init) + monkeypatch.setattr(entrypoint.ray, "is_initialized", lambda: False) + + with pytest.raises(RuntimeError, match="readiness unavailable"): + entrypoint.run_ppo( + OmegaConf.create({}), + [{"instance_id": "train"}], + [{"instance_id": "val"}], + ) + + assert ray_init_called is False diff --git a/tests/verl/test_k8s_image_filter.py b/tests/verl/test_k8s_image_filter.py new file mode 100644 index 000000000..f4dacc2e7 --- /dev/null +++ b/tests/verl/test_k8s_image_filter.py @@ -0,0 +1,220 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +from pathlib import Path + +import httpx +import pytest +import yaml +from omegaconf import DictConfig, OmegaConf + +from agentlightning.schemas import K8sImageReadinessSnapshot +from agentlightning.verl.k8s_image_filter import ( + filter_dataset_by_images, + prepare_datasets, + wait_for_k8s_image_readiness, +) + +TEMPLATE = """ +apiVersion: batch/v1 +kind: Job +metadata: {} +spec: + template: + spec: + containers: + - name: agent + image: {{ (input.image_name ~ ":openai") | yaml_escape }} +""" + + +@pytest.fixture +def enabled_config(tmp_path: Path) -> DictConfig: + template_path = tmp_path / "job.yaml" + template_path.write_text(TEMPLATE) + return OmegaConf.create( + { + "agentlightning": { + "agl_base_url": "http://server:8080", + "agl_key": "secret", + "k8s": { + "job_template_path": str(template_path), + "filter_unavailable_images": True, + "image_readiness_timeout_seconds": 60, + }, + } + } + ) + + +def _snapshot(images: list[str]) -> K8sImageReadinessSnapshot: + return K8sImageReadinessSnapshot( + images=images, + node_count=1, + observed_at=100.0, + expires_at=130.0, + ) + + +def test_filter_dataset_keeps_available_rows_without_deduplicating() -> None: + rows = [ + {"instance_id": "missing", "image_name": "swebench/missing"}, + {"instance_id": "ready", "image_name": "swebench/ready"}, + {"instance_id": "ready", "image_name": "swebench/ready"}, + ] + + kept, report = filter_dataset_by_images( + rows, + split="train", + job_template=TEMPLATE, + ready_images={"docker.io/swebench/ready:openai"}, + ) + + assert kept == [rows[1], rows[2]] + assert report.source_count == 3 + assert report.kept_count == 2 + assert report.dropped_count == 1 + assert report.dropped_data_ids == ["missing"] + assert report.missing_image_counts == {"docker.io/swebench/missing:openai": 1} + + +def test_prepare_datasets_filters_before_validation_cap( + enabled_config: DictConfig, + monkeypatch, +) -> None: + monkeypatch.setattr( + "agentlightning.verl.k8s_image_filter._fetch_snapshot", + lambda _config: _snapshot(["docker.io/swebench/ready:openai"]), + ) + val = [ + {"instance_id": "missing", "image_name": "swebench/missing"}, + {"instance_id": "ready-1", "image_name": "swebench/ready"}, + {"instance_id": "ready-2", "image_name": "swebench/ready"}, + ] + + prepared = prepare_datasets( + enabled_config, + [{"instance_id": "train", "image_name": "swebench/ready"}], + val, + max_val_instances=1, + ) + + assert [row["instance_id"] for row in prepared.val] == ["ready-1"] + assert prepared.val_report is not None + assert prepared.val_report.source_count == 3 + assert prepared.val_report.kept_count == 2 + assert prepared.val_report.dropped_count == 1 + + +def test_prepare_datasets_disabled_bypasses_readiness_and_template(monkeypatch) -> None: + monkeypatch.setattr( + "agentlightning.verl.k8s_image_filter._fetch_snapshot", + lambda _config: (_ for _ in ()).throw(AssertionError("must not fetch")), + ) + config = OmegaConf.create( + { + "agentlightning": { + "k8s": { + "job_template_path": "/does/not/exist.yaml", + "filter_unavailable_images": False, + "image_readiness_timeout_seconds": 60, + } + } + } + ) + train = [{"instance_id": "train"}] + val = [{"instance_id": "val-1"}, {"instance_id": "val-2"}] + + prepared = prepare_datasets(config, train, val, max_val_instances=1) + + assert prepared.train == train + assert prepared.val == val[:1] + assert prepared.train_report is None + assert prepared.val_report is None + assert prepared.readiness is None + + +def test_enabled_filter_rejects_missing_job_template(enabled_config: DictConfig) -> None: + enabled_config.agentlightning.k8s.job_template_path = None + + with pytest.raises(ValueError, match="job_template_path"): + prepare_datasets(enabled_config, [{"id": 1}], [{"id": 2}], max_val_instances=None) + + +def test_enabled_filter_rejects_empty_train_after_filter( + enabled_config: DictConfig, + monkeypatch, +) -> None: + monkeypatch.setattr("agentlightning.verl.k8s_image_filter._fetch_snapshot", lambda _config: _snapshot([])) + + with pytest.raises(ValueError, match=r"train dataset.*empty"): + prepare_datasets( + enabled_config, + [{"instance_id": "missing", "image_name": "swebench/missing"}], + [{"instance_id": "missing-val", "image_name": "swebench/missing"}], + max_val_instances=None, + ) + + +def test_enabled_filter_surfaces_malformed_job_yaml( + enabled_config: DictConfig, + monkeypatch, +) -> None: + Path(str(enabled_config.agentlightning.k8s.job_template_path)).write_text("apiVersion: batch/v1\nkind: [") + monkeypatch.setattr( + "agentlightning.verl.k8s_image_filter._fetch_snapshot", + lambda _config: _snapshot(["docker.io/swebench/ready:openai"]), + ) + + with pytest.raises(yaml.YAMLError): + prepare_datasets( + enabled_config, + [{"instance_id": "train", "image_name": "swebench/ready"}], + [{"instance_id": "val", "image_name": "swebench/ready"}], + max_val_instances=None, + ) + + +def test_wait_for_readiness_times_out_with_last_503_detail() -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 503, + request=request, + json={"detail": "K8s image readiness snapshot has expired"}, + ) + + with ( + httpx.Client(base_url="http://server:8080", transport=httpx.MockTransport(handler)) as client, + pytest.raises(RuntimeError, match="snapshot has expired"), + ): + wait_for_k8s_image_readiness(client, timeout_seconds=0) + + +def test_wait_for_readiness_does_not_retry_auth_failure() -> None: + requests = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal requests + requests += 1 + return httpx.Response(401, request=request, json={"detail": "invalid key"}) + + with ( + httpx.Client(base_url="http://server:8080", transport=httpx.MockTransport(handler)) as client, + pytest.raises(httpx.HTTPStatusError), + ): + wait_for_k8s_image_readiness(client, timeout_seconds=60) + + assert requests == 1 + + +@pytest.mark.parametrize("timeout_seconds", [-1.0, float("nan"), float("inf")]) +def test_wait_for_readiness_rejects_invalid_timeout(timeout_seconds: float) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, request=request, json={"detail": "not ready"}) + + with ( + httpx.Client(base_url="http://server:8080", transport=httpx.MockTransport(handler)) as client, + pytest.raises(ValueError, match="finite non-negative"), + ): + wait_for_k8s_image_readiness(client, timeout_seconds=timeout_seconds)