Restore Python lint checks and extend them to the verl subtree (#553)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yuge Zhang
2026-08-22 10:08:06 +08:00
committed by GitHub
parent 07c4ce894d
commit c632d8dd80
46 changed files with 2967 additions and 280 deletions
+49 -10
View File
@@ -15,10 +15,54 @@ concurrency:
cancel-in-progress: true
jobs:
test:
name: Run core tests
lint:
name: Lint Python and repository files
runs-on: ubuntu-latest
timeout-minutes: 15
timeout-minutes: 10
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
with:
python-version: '3.12'
- uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
enable-cache: true
- name: Sync lint dependencies
run: uv sync --frozen --no-default-groups --extra dev --group dev
- name: Run pre-commit checks
run: uv run --locked --no-sync pre-commit run --all-files --show-diff-on-failure
- name: Run Ruff
run: uv run --locked --no-sync ruff check .
- name: Check Ruff formatting
run: uv run --locked --no-sync ruff format --check .
- name: Check Python headers
run: uv run --locked --no-sync python scripts/check_headers.py
typecheck:
name: Type-check Python
runs-on: ubuntu-latest
# Split from `lint` because the verl-cpu group pulls torch and friends.
# Pyright needs them to check agentlightning/verl and tests/verl; the fast
# checks above should not wait on that download.
timeout-minutes: 20
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
with:
python-version: '3.12'
- uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
enable-cache: true
- name: Sync type-check dependencies
run: uv sync --frozen --no-default-groups --extra dev --group dev --group verl-cpu
- name: Run Pyright
run: uv run --locked --no-sync pyright
test:
name: Run tests
runs-on: ubuntu-latest
# verl-cpu is needed for tests/verl; the rest of the suite only needs dev.
timeout-minutes: 20
steps:
- uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
- uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
@@ -28,14 +72,9 @@ jobs:
with:
enable-cache: true
- name: Sync test dependencies
run: uv sync --frozen --no-default-groups --extra dev --group dev
run: uv sync --frozen --no-default-groups --extra dev --group dev --group verl-cpu
- name: Run tests
run: >-
uv run --locked --no-sync pytest -v --durations=20
tests/server
tests/controller
tests/test_package.py
tests/examples/test_swe_smith_images.py
run: uv run --locked --no-sync pytest -v --durations=20 tests
package:
name: Build package
+16
View File
@@ -0,0 +1,16 @@
exclude: ^(\.agents/|examples/llm-in-sandbox/vendor/)
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v6.0.0
hooks:
- id: end-of-file-fixer
- id: trailing-whitespace
- id: check-yaml
exclude: ^(mkdocs\.yml|examples/calc_x/job-template\.yaml|examples/llm-in-sandbox/job-template\.yaml|examples/swe_smith/job-template-openai\.yaml)$
- id: check-toml
- id: check-added-large-files
args: ["--maxkb=1024"]
exclude: ^uv\.lock$
- id: check-shebang-scripts-are-executable
- id: detect-private-key
+1 -1
View File
@@ -7,4 +7,4 @@ default_proxy:
train:
temperature: 1
val:
temperature: 0.7
temperature: 0.7
+12 -14
View File
@@ -15,8 +15,9 @@ import asyncio
import json
import time
from collections import deque
from typing import Any
from typing import Any, cast
import httpx
import kr8s
import kr8s.asyncio
import structlog
@@ -26,7 +27,7 @@ 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
from agentlightning.schemas import DEFAULT_ATTEMPT_ID, Rollout, RolloutPatch, RolloutState, RolloutStatusPatch
log = structlog.get_logger()
@@ -77,20 +78,17 @@ def build_job_spec(rollout: Rollout, controller_config: DictConfig) -> dict[str,
mode = "train" if rollout.is_train else "val"
agent_base_url = str(
controller_config.agl_server.get("agent_url", None)
or controller_config.agl_server.url
controller_config.agl_server.get("agent_url", None) or controller_config.agl_server.url
).rstrip("/")
agl_openai_base_url = (
f"{agent_base_url}/proxy/rollout/{rollout.rollout_id}"
f"/attempt/{DEFAULT_ATTEMPT_ID}/mode/{mode}/openai/v1"
f"{agent_base_url}/proxy/rollout/{rollout.rollout_id}/attempt/{DEFAULT_ATTEMPT_ID}/mode/{mode}/openai/v1"
)
for container in pod_spec.get("containers", []):
env = container.setdefault("env", [])
for name, value in {
"AGL_OPENAI_BASE_URL": agl_openai_base_url,
"AGL_EVENT_URL": (
f"{agent_base_url}/api/rollouts/{rollout.rollout_id}"
f"/attempt/{DEFAULT_ATTEMPT_ID}/events"
f"{agent_base_url}/api/rollouts/{rollout.rollout_id}/attempt/{DEFAULT_ATTEMPT_ID}/events"
),
"AGL_KEY": str(controller_config.agl_server.key or ""),
}.items():
@@ -165,7 +163,7 @@ class K8sReconciler:
rollouts = await self._query_rollouts(state_in=[RolloutState.QUEUING, RolloutState.RUNNING], limit=500)
api = await self._get_k8s_api()
jobs = [
job.raw
cast(k8s_objects.Job, job).raw
async for job in k8s_objects.Job.async_list(
namespace=self._namespace,
label_selector=MANAGED_BY_SELECTOR,
@@ -342,15 +340,17 @@ class K8sReconciler:
state_in: list[RolloutState],
limit: int = 50,
) -> list[Rollout]:
params: list[tuple[str, str | int]] = [("state_in", state.value) for state in state_in]
params.append(("limit", limit))
params = httpx.QueryParams()
for state in state_in:
params = params.add("state_in", state.value)
params = params.add("limit", limit)
response = await self._api.get("/api/rollouts", params=params)
response.raise_for_status()
return [Rollout.model_validate(item) for item in response.json()]
async def _patch_status(self, rollout_id: str, **status: Any) -> bool:
try:
patch = RolloutPatch(status=status)
patch = RolloutPatch(status=RolloutStatusPatch.model_validate(status))
response = await self._api.patch(
f"/api/rollouts/{rollout_id}",
json=patch.model_dump(mode="json", exclude_unset=True),
@@ -360,5 +360,3 @@ class K8sReconciler:
except Exception as exc:
log.warning("Failed to patch rollout", rollout_id=rollout_id, error=str(exc))
return False
+16 -15
View File
@@ -16,11 +16,12 @@ import time
import traceback
from dataclasses import dataclass
import httpx
import structlog
from omegaconf import DictConfig
from agentlightning.client import AgentLightningAsyncClient
from agentlightning.schemas import DEFAULT_ATTEMPT_ID, Rollout, RolloutPatch, RolloutState
from agentlightning.schemas import DEFAULT_ATTEMPT_ID, Rollout, RolloutPatch, RolloutState, RolloutStatusPatch
log = structlog.get_logger()
@@ -126,11 +127,10 @@ class LocalReconciler:
pass
async def _reconcile_once(self) -> None:
params: list[tuple[str, str | int]] = [
("state_in", RolloutState.QUEUING.value),
("state_in", RolloutState.RUNNING.value),
("limit", 50),
]
params = httpx.QueryParams()
params = params.add("state_in", RolloutState.QUEUING.value)
params = params.add("state_in", RolloutState.RUNNING.value)
params = params.add("limit", 50)
response = await self._api.get("/api/rollouts", params=params)
response.raise_for_status()
rollouts = [Rollout.model_validate(item) for item in response.json()]
@@ -161,9 +161,12 @@ class LocalReconciler:
continue
rollout = rollouts_by_id.get(rollout_id)
timeout = float(rollout.config.timeout_seconds) if rollout and rollout.config.timeout_seconds else None
if timeout is not None and (now - item.spawned_at) > timeout:
if await self._kill_process_group(rollout_id, item):
await self._patch(rollout_id, RolloutState.FAILED, "local subprocess timed out")
if (
timeout is not None
and (now - item.spawned_at) > timeout
and await self._kill_process_group(rollout_id, item)
):
await self._patch(rollout_id, RolloutState.FAILED, "local subprocess timed out")
async def _finish_proc(self, rollout: Rollout, item: Proc) -> bool:
if rollout.status.state == RolloutState.QUEUING:
@@ -210,8 +213,7 @@ class LocalReconciler:
f"/attempt/{attempt_id}/mode/{mode}/openai/v1"
),
"AGL_EVENT_URL": (
f"{self._config.agl_server.url}/api/rollouts/{rollout.rollout_id}"
f"/attempt/{attempt_id}/events"
f"{self._config.agl_server.url}/api/rollouts/{rollout.rollout_id}/attempt/{attempt_id}/events"
),
}
env.update(_build_env_from_map(rollout.input, rollout.config.local.env_map))
@@ -263,11 +265,11 @@ class LocalReconciler:
*,
last_attempt_id: str | None = None,
) -> bool:
status: dict[str, object] = {"state": state}
status = RolloutStatusPatch(state=state)
if error_message is not None:
status["error_message"] = error_message
status.error_message = error_message
if last_attempt_id is not None:
status["last_attempt_id"] = last_attempt_id
status.last_attempt_id = last_attempt_id
patch = RolloutPatch(status=status)
try:
response = await self._api.patch(
@@ -279,4 +281,3 @@ class LocalReconciler:
except Exception as e:
log.warning("Failed to patch rollout", rollout_id=rollout_id, error=str(e))
return False
+2 -2
View File
@@ -6,7 +6,7 @@ from __future__ import annotations
from collections.abc import AsyncIterator, Mapping
from contextlib import asynccontextmanager
from typing import Any
from typing import Any, cast
import httpx
import structlog
@@ -24,7 +24,7 @@ def _server_config(config: Mapping[str, Any] | DictConfig | None) -> dict[str, A
if config is None:
raise ValueError("server config is required")
elif OmegaConf.is_config(config):
raw = dict(OmegaConf.to_container(config, resolve=True))
raw = dict(cast(Any, OmegaConf.to_container(config, resolve=True)))
else:
raw = dict(config)
+1 -5
View File
@@ -180,11 +180,7 @@ def _dedupe_model_requests_by_prompt_token_ids(events: list[Event]) -> list[Even
last_index_by_prompt[prompt_key] = index
last_indexes = set(last_index_by_prompt.values())
return [
event
for index, event in enumerate(events)
if event.event_type != "model_request" or index in last_indexes
]
return [event for index, event in enumerate(events) if event.event_type != "model_request" or index in last_indexes]
@router.post("/rollouts/{rollout_id}/attempt/{attempt_id}/events", response_model=Event)
-2
View File
@@ -218,5 +218,3 @@ async def delete_rollout(rollout_id: str) -> None:
"""Delete a rollout and its events. Idempotent: missing id is a no-op."""
_rollouts.pop(rollout_id, None)
_events.pop(rollout_id, None)
+2 -2
View File
@@ -11,7 +11,7 @@ from collections import defaultdict
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, cast
import numpy as np
from httpx_retries import Retry, RetryTransport
@@ -257,7 +257,7 @@ class AglRolloutManagerBase:
request = RolloutCreate(
input=_to_native(original),
is_train=is_train,
config=self._rollout_config,
config=cast(Any, self._rollout_config), # pydantic coerces the dict
metadata={},
)
if self._hooks is not None:
+2 -2
View File
@@ -2,13 +2,13 @@
# type: ignore
from typing import Any, Sequence
from collections.abc import Sequence
from typing import Any
import torch
from datasets import Dataset as HuggingFaceDataset
from verl.utils.dataset.rl_dataset import RLHFDataset
__all__ = [
"LoadedDataset",
]
-1
View File
@@ -27,7 +27,6 @@ __all__ = [
]
def run_ppo(
config: Any,
train_dataset: Sequence[Any],
+2
View File
@@ -1,3 +1,5 @@
# Copyright (c) Microsoft. All rights reserved.
"""Rollout-level mean policy loss for VERL."""
from __future__ import annotations
+5 -10
View File
@@ -7,7 +7,7 @@ from __future__ import annotations
import io
import json
import zipfile
from typing import Any
from typing import Any, cast
import numpy as np
import torch
@@ -106,9 +106,7 @@ def _build_compact_rollout_trajectory_records(
def _build_zipped_jsonl(records: list[dict[str, Any]], jsonl_name: str) -> bytes:
jsonl_text = "".join(
json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n" for record in records
)
jsonl_text = "".join(json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n" for record in records)
buffer = io.BytesIO()
with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as zip_file:
zip_file.writestr(jsonl_name, jsonl_text.encode("utf-8"))
@@ -121,7 +119,7 @@ def _upload_trace_merge_mismatches_to_wandb(rows: list[dict[str, Any]], global_s
if wandb.run is None:
return
table = wandb.Table(columns=_TRACE_MERGE_MISMATCH_COLUMNS)
table = wandb.Table(columns=cast(list[str | int], _TRACE_MERGE_MISMATCH_COLUMNS))
for row in rows:
table.add_data(*(row.get(column) for column in _TRACE_MERGE_MISMATCH_COLUMNS))
wandb.log({"training/trace_merge_mismatches": table}, step=global_steps)
@@ -156,7 +154,7 @@ def _upload_compact_rollout_trajectories_to_wandb(
trajectory_file.write(_build_zipped_jsonl(records, f"{artifact_type}.jsonl"))
run.log_artifact(artifact)
table = wandb.Table(columns=_ROLLOUT_TRAJECTORY_COLUMNS)
table = wandb.Table(columns=cast(list[str | int], _ROLLOUT_TRAJECTORY_COLUMNS))
table.add_data(global_steps, artifact_name, artifact_path, len(records))
table_key = "val/rollout_trajectories" if is_validation else "training/rollout_trajectories"
wandb.log({table_key: table}, step=global_steps)
@@ -512,10 +510,7 @@ class RolloutAdapter:
"has_reward": rollout.final_reward is not None,
}
if rollout.triplets:
response_length_list = [
len(triplet.response.get("token_ids") or [])
for triplet in rollout.triplets
]
response_length_list = [len(triplet.response.get("token_ids") or []) for triplet in rollout.triplets]
sample_stat.update(
{
"total_response_length": np.sum(response_length_list),
@@ -26,9 +26,7 @@ def compute_rollout_level_advantage(
"""Compute advantages once per rollout, then broadcast to rollout triplets."""
rollout_ids = _required_non_tensor(batch, "rollout_id_list")
if len(rollout_ids) != len(batch):
raise RuntimeError(
f"rollout_id_list length ({len(rollout_ids)}) must match batch length ({len(batch)})"
)
raise RuntimeError(f"rollout_id_list length ({len(rollout_ids)}) must match batch length ({len(batch)})")
uid_values = batch.non_tensor_batch.get("uid")
if uid_values is None:
+23 -33
View File
@@ -20,7 +20,7 @@ from tqdm import tqdm
from verl import DataProto
from verl.trainer.ppo.metric_utils import compute_data_metrics, compute_throughout_metrics, compute_timing_metrics
from verl.trainer.ppo.ray_trainer import (
AdvantageEstimator,
AdvantageEstimator, # pyright: ignore[reportPrivateImportUsage]
RayPPOTrainer,
apply_kl_penalty,
compute_advantage,
@@ -149,8 +149,8 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
def _rollout_replicas(self) -> list[Any]:
if hasattr(self, "llm_server_manager"):
return list(self.llm_server_manager.get_replicas())
return list(self.async_rollout_manager.rollout_replicas)
return list(self.llm_server_manager.get_replicas()) # pyright: ignore[reportAttributeAccessIssue]
return list(self.async_rollout_manager.rollout_replicas) # pyright: ignore[reportAttributeAccessIssue]
@auto_await
async def _abort_all_rollout_requests(self) -> None:
@@ -233,9 +233,7 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
running_at = rollout.running_at
finished_at = rollout.finished_at
queue_wait = (running_at - submitted) if running_at is not None else None
run_duration = (
finished_at - running_at if (running_at is not None and finished_at is not None) else None
)
run_duration = finished_at - running_at if (running_at is not None and finished_at is not None) else None
total = (finished_at - submitted) if finished_at is not None else None
if queue_wait is not None:
queue_waits.append(queue_wait)
@@ -365,10 +363,10 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
# verl 0.8.0 moved rollout server state behind llm_server_manager.
has_llm_server_manager = hasattr(self, "llm_server_manager")
if has_llm_server_manager:
server_addresses = list(self.llm_server_manager.get_addresses())
server_addresses = list(self.llm_server_manager.get_addresses()) # pyright: ignore[reportAttributeAccessIssue]
else:
server_addresses = list(self.async_rollout_manager.server_addresses)
self._resume_all_rollout_generation()
server_addresses = list(self.async_rollout_manager.server_addresses) # pyright: ignore[reportAttributeAccessIssue]
self._resume_all_rollout_generation() # pyright: ignore[reportUnusedCoroutine]
if self.is_async:
self._resume_gateway()
data_dict = dict(gen_batch.non_tensor_batch)
@@ -437,7 +435,7 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
print("AgentLightningRayPPOTrainer: agl gateway paused and drained.")
else:
print("AgentLightningRayPPOTrainer: aborting residual vLLM requests.")
self._abort_all_rollout_requests()
self._abort_all_rollout_requests() # pyright: ignore[reportUnusedCoroutine]
print("AgentLightningRayPPOTrainer: residual vLLM requests aborted.")
return out, metrics
@@ -463,16 +461,16 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
with marked_timer("gen", timing_raw, color="red"):
if curr_step_profile:
if has_llm_server_manager:
self.llm_server_manager.start_profile()
self.llm_server_manager.start_profile() # pyright: ignore[reportAttributeAccessIssue]
else:
self.async_rollout_manager.start_profile()
self.async_rollout_manager.start_profile() # pyright: ignore[reportAttributeAccessIssue]
gen_batch_output, agent_metrics = self._rollout(gen_batch, is_train=True)
if curr_step_profile:
if has_llm_server_manager:
self.llm_server_manager.stop_profile()
self.llm_server_manager.stop_profile() # pyright: ignore[reportAttributeAccessIssue]
else:
self.async_rollout_manager.stop_profile()
self.async_rollout_manager.stop_profile() # pyright: ignore[reportAttributeAccessIssue]
metrics.update(agent_metrics)
metrics["timing/rollout_phase_end_wall"] = time.time()
@@ -496,7 +494,7 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
if "is_drop_mask" in batch.batch:
keep = (~batch.batch["is_drop_mask"].bool()).nonzero(as_tuple=True)[0].tolist()
metrics["training/n_sample_dropped/marked"] = len(batch) - len(keep)
batch = batch[keep]
batch = batch[keep] # pyright: ignore[reportAssignmentType]
mini_bs = self.config.actor_rollout_ref.actor.ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n
n_transition = len(batch)
@@ -527,7 +525,7 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
drop_indices = same_reward_drop_set | set(random_drop_indices)
keep_indices = [sample_idx for sample_idx in range(n_transition) if sample_idx not in drop_indices]
batch = batch[keep_indices]
batch = batch[keep_indices] # pyright: ignore[reportAssignmentType]
metrics["training/n_sample_dropped/same_reward"] = n_dropped_same_reward
metrics["training/n_sample_dropped/random"] = n_dropped_random
metrics["training/n_sample_trained"] = len(batch)
@@ -536,7 +534,7 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
return None
print("AgentLightningRayPPOTrainer: sleeping rollout replicas.")
self.checkpoint_manager.sleep_replicas()
self.checkpoint_manager.sleep_replicas() # pyright: ignore[reportOptionalMemberAccess]
print("AgentLightningRayPPOTrainer: rollout replicas slept.")
if self.config.trainer.balance_batch:
@@ -578,25 +576,19 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
if self.config.algorithm.use_kl_in_reward:
batch, kl_metrics = apply_kl_penalty(
batch,
kl_ctrl=self.kl_ctrl_in_reward,
kl_ctrl=self.kl_ctrl_in_reward, # pyright: ignore[reportArgumentType]
kl_penalty=self.config.algorithm.kl_penalty,
)
metrics.update(kl_metrics)
else:
batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]
if (
rollout_corr_config is not None
and not bypass_mode
and "rollout_log_probs" in batch.batch
):
if rollout_corr_config is not None and not bypass_mode and "rollout_log_probs" in batch.batch:
from verl.trainer.ppo.rollout_corr_helper import (
compute_rollout_correction_and_add_to_batch,
)
batch, is_metrics = compute_rollout_correction_and_add_to_batch(
batch, rollout_corr_config
)
batch, is_metrics = compute_rollout_correction_and_add_to_batch(batch, rollout_corr_config)
metrics.update(is_metrics)
adv_kwargs = {
@@ -641,7 +633,7 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
metrics.update(reduce_metrics(actor_output.meta_info["metrics"]))
with marked_timer("update_weights", timing_raw, color="red"):
self.checkpoint_manager.update_weights(self.global_steps)
self.checkpoint_manager.update_weights(self.global_steps) # pyright: ignore[reportOptionalMemberAccess]
batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
# Return the batch so fit() can compute throughput after the step timer closes.
@@ -663,7 +655,7 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
self._carry_over_rollouts = []
self._load_checkpoint()
# Push loaded weights before the first rollout or validation.
self.checkpoint_manager.update_weights(self.global_steps)
self.checkpoint_manager.update_weights(self.global_steps) # pyright: ignore[reportOptionalMemberAccess]
if self.config.trainer.get("val_before_train", True):
val_metrics = self._validate()
@@ -704,9 +696,7 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
metrics.update(compute_timing_metrics(batch=step_batch, timing_raw=timing_raw))
n_gpus = self.resource_pool_manager.get_n_gpus()
if n_gpus > 0 and "step" in timing_raw:
metrics.update(
compute_throughout_metrics(batch=step_batch, timing_raw=timing_raw, n_gpus=n_gpus)
)
metrics.update(compute_throughout_metrics(batch=step_batch, timing_raw=timing_raw, n_gpus=n_gpus))
is_last_step = self.global_steps >= self.total_training_steps
@@ -740,9 +730,9 @@ class AgentLightningRayPPOTrainer(RayPPOTrainer):
# verl 0.8.0 moved rollout server state behind llm_server_manager.
has_llm_server_manager = hasattr(self, "llm_server_manager")
if has_llm_server_manager:
server_addresses = list(self.llm_server_manager.get_addresses())
server_addresses = list(self.llm_server_manager.get_addresses()) # pyright: ignore[reportAttributeAccessIssue]
else:
server_addresses = list(self.async_rollout_manager.server_addresses)
server_addresses = list(self.async_rollout_manager.server_addresses) # pyright: ignore[reportAttributeAccessIssue]
assert server_addresses, "_validate called before rollout server addresses are available"
merged_metrics: dict[str, Any] = {}
+4 -6
View File
@@ -4,7 +4,7 @@ from __future__ import annotations
import html
import os
from typing import Any, Dict
from typing import Any
def define_env(env: Any):
@@ -29,9 +29,9 @@ def define_env(env: Any):
```
"""
cfg: Dict[str, Any] = env.conf or {}
cfg: dict[str, Any] = env.conf or {}
repo_url = cfg.get("repo_url", "").rstrip("/")
extra: Dict[str, Any] = cfg.get("extra", {}) or {}
extra: dict[str, Any] = cfg.get("extra", {}) or {}
default_commit = extra.get("source_commit", "main")
project_dir = env.project_dir
@@ -61,9 +61,7 @@ def define_env(env: Any):
if not os.path.exists(abs_path):
_warn(f"Source path not found: {path}. Rendering a visible broken-link marker.")
label = html.escape(text or path)
return (
f'<span class="broken-source-link" title="Missing: {html.escape(url)}">' f"{label} (missing)" f"</span>"
)
return f'<span class="broken-source-link" title="Missing: {html.escape(url)}">{label} (missing)</span>'
if text:
return f"[{text}]({url})"
+1 -1
View File
@@ -23,4 +23,4 @@ RUN python -m compileall -q "$(python -c 'import site; print(site.getsitepackage
RUN python -c "import mcp_server_calculator" \
&& timeout 3 python -m mcp_server_calculator; \
status=$?; \
if [ "$status" -ne 0 ] && [ "$status" -ne 124 ]; then exit "$status"; fi
if [ "$status" -ne 0 ] && [ "$status" -ne 124 ]; then exit "$status"; fi
+1 -2
View File
@@ -12,7 +12,6 @@ import os
import re
import httpx
from eval_utils import scalar_are_results_same
@@ -78,4 +77,4 @@ class Agent:
if __name__ == "__main__":
asyncio.run(Agent().run())
asyncio.run(Agent().run())
+1
View File
@@ -9,6 +9,7 @@ from datasets import Dataset as HuggingFaceDataset
from hydra import compose, initialize_config_dir
from omegaconf import OmegaConf
def verl_default_config() -> dict[str, Any]:
"""VERL config overrides for Calc-X training.
-2
View File
@@ -142,5 +142,3 @@ class CompletionAgent:
prediction=decode_choice_token_ids(response.choices[0], model),
answer=answer,
)
Regular → Executable
+1 -1
View File
@@ -48,4 +48,4 @@ python train_gsm8k_agent.py \
--agl-base-url "http://localhost:$AGL_SERVER_PORT" \
--agl-key "$AGL_KEY" \
--run-name local \
"$@"
"$@"
+1 -1
View File
@@ -130,7 +130,7 @@ def build_config(
if agl_key is not None:
overrides["agentlightning"]["agl_key"] = agl_key
if run_name:
overrides["trainer"]["experiment_name"] = f'{overrides["trainer"]["experiment_name"]}_{run_name}'
overrides["trainer"]["experiment_name"] = f"{overrides['trainer']['experiment_name']}_{run_name}"
override_conf = OmegaConf.create(overrides)
cli_override_conf = OmegaConf.from_dotlist(list(config_overrides))
+1 -1
View File
@@ -5,4 +5,4 @@ mismatch_cases/
__pycache__/
*.pyc
data/
data/
+1 -1
View File
@@ -34,4 +34,4 @@ WORKDIR /testbed
RUN git init \
&& python -m compileall -q /opt/llm-in-sandbox /app
CMD ["python", "/app/runner.py"]
CMD ["python", "/app/runner.py"]
+1 -3
View File
@@ -79,9 +79,7 @@ def configure_llm_env() -> None:
model_name = os.environ.get("AGL_MODEL_NAME", "Qwen/Qwen3-4B-Instruct-2507")
os.environ["LLM_NAME"] = openai_model_name(model_name)
os.environ["LLM_BASE_URL"] = (
os.environ.get("LLM_BASE_URL")
or os.environ.get("AGL_OPENAI_BASE_URL")
or os.environ.get("OPENAI_BASE_URL")
os.environ.get("LLM_BASE_URL") or os.environ.get("AGL_OPENAI_BASE_URL") or os.environ.get("OPENAI_BASE_URL")
)
os.environ["LLM_API_KEY"] = os.environ.get("LLM_API_KEY") or os.environ.get("OPENAI_API_KEY", "dummy")
os.environ["LLM_TEMPERATURE"] = os.environ.get("LLM_TEMPERATURE") or os.environ.get("AGL_LLM_TEMPERATURE", "1.0")
+1 -1
View File
@@ -25,4 +25,4 @@ spec:
memory: "2Gi"
limits:
cpu: "4"
memory: "8Gi"
memory: "8Gi"
+1 -1
View File
@@ -53,4 +53,4 @@ python examples/llm-in-sandbox/train_llm_in_sandbox.py \
--agl-base-url "http://localhost:$AGL_SERVER_PORT" \
--agl-key "$AGL_KEY" \
--run-name minikube \
"$@"
"$@"
+3 -2
View File
@@ -10,9 +10,10 @@ import copy
import importlib.resources
import json
import os
from collections.abc import Sequence
from pathlib import Path
from pprint import pprint
from typing import Any, Sequence
from typing import Any
from hydra import compose, initialize_config_dir
from omegaconf import DictConfig, OmegaConf
@@ -192,7 +193,7 @@ def build_config(
if agl_key is not None:
overrides["agentlightning"]["agl_key"] = agl_key
if run_name:
overrides["trainer"]["experiment_name"] = f'{overrides["trainer"]["experiment_name"]}_{run_name}'
overrides["trainer"]["experiment_name"] = f"{overrides['trainer']['experiment_name']}_{run_name}"
override_conf = OmegaConf.create(overrides)
cli_override_conf = OmegaConf.from_dotlist(list(config_overrides))
+2 -3
View File
@@ -81,8 +81,7 @@ def build_dataset(
n_train = budget - n_val
if n_train <= 0:
raise ValueError(
f"task {task_name!r} has only {budget} variations after "
f"val_fraction={val_fraction}; no train rows left"
f"task {task_name!r} has only {budget} variations after val_fraction={val_fraction}; no train rows left"
)
for v in range(n_train):
train.append(_row(task_name, v, simplification))
@@ -217,7 +216,7 @@ def build_config(
if agl_key is not None:
overrides["agentlightning"]["agl_key"] = agl_key
if run_name:
overrides["trainer"]["experiment_name"] = f'{overrides["trainer"]["experiment_name"]}_{run_name}'
overrides["trainer"]["experiment_name"] = f"{overrides['trainer']['experiment_name']}_{run_name}"
override_conf = OmegaConf.create(overrides)
OmegaConf.set_struct(base_cfg, False)
-1
View File
@@ -1 +0,0 @@
-1
View File
@@ -32,4 +32,3 @@ python "$SCRIPT_DIR/retrieval_server.py" \
--port "$PORT" \
--device "$DEVICE" \
--faiss-gpu
+3 -3
View File
@@ -133,9 +133,9 @@ def build_config(
overrides["agentlightning"]["local"]["agent_class"] = (
"examples.search_r1.agents.search_r1_agent:SearchR1CompletionAgent"
)
overrides["agentlightning"]["local"]["env_map"]["SEARCH_R1_TOKENIZER_MODEL"] = overrides[
"actor_rollout_ref"
]["model"]["path"]
overrides["agentlightning"]["local"]["env_map"]["SEARCH_R1_TOKENIZER_MODEL"] = overrides["actor_rollout_ref"][
"model"
]["path"]
if agl_base_url:
overrides["agentlightning"]["agl_base_url"] = agl_base_url
if agl_key is not None:
+169 -94
View File
@@ -40,7 +40,7 @@ _ACTION_RE = re.compile(r"```(?:bash|mswea_bash_command)[^\S\n]*\n(.*?)\n?```",
# Block git use/metadata access. `git` matched only in command position (after
# ; && || | ( ` $( and VAR=val prefixes) so args/prose containing "git" pass.
_GIT_INVOKE_RE = re.compile(
r'(?:^|[\n;`(]|&&|\|\|?|\$\()\s*(?:\w+=\S+\s+)*(?:[\w./-]*/)?git(?:-[a-z]+)?(?=\s|$|;|&|\|)',
r"(?:^|[\n;`(]|&&|\|\|?|\$\()\s*(?:\w+=\S+\s+)*(?:[\w./-]*/)?git(?:-[a-z]+)?(?=\s|$|;|&|\|)",
re.I,
)
_GIT_ACCESS_RE = re.compile(r'--git-dir|--work-tree|(?:^|[\s=:"\'/])\.git(?:/|\b)', re.I)
@@ -49,30 +49,30 @@ _GIT_ACCESS_RE = re.compile(r'--git-dir|--work-tree|(?:^|[\s=:"\'/])\.git(?:/|\b
# (curl/wget, pip install, python urllib). Block these; same command-position
# anchoring as _GIT_INVOKE_RE. Code-level backstop for the egress NetworkPolicy.
_NET_FETCH_RE = re.compile(
r'(?:^|[\n;`(]|&&|\|\|?|\$\()\s*(?:\w+=\S+\s+)*(?:[\w./-]*/)?'
r'(?:curl|wget|httpie|http|https|aria2c|scp|sftp|rsync|nc|ncat|netcat|telnet)'
r'(?=\s|$|;|&|\|)',
r"(?:^|[\n;`(]|&&|\|\|?|\$\()\s*(?:\w+=\S+\s+)*(?:[\w./-]*/)?"
r"(?:curl|wget|httpie|http|https|aria2c|scp|sftp|rsync|nc|ncat|netcat|telnet)"
r"(?=\s|$|;|&|\|)",
re.I,
)
# Any install can pull the target package's correct source. Match
# pip/pip3/conda/mamba/uv/easy_install install and `python -m pip install`.
_PKG_INSTALL_RE = re.compile(
r'(?:^|[\n;`(]|&&|\|\|?|\$\()\s*(?:\w+=\S+\s+)*(?:'
r'(?:[\w./-]*/)?(?:pip|pip3|conda|mamba|easy_install|uv)\b[^\n;&|]*?\binstall\b'
r'|(?:[\w./-]*/)?python[0-9.]*\s+-m\s+pip\b[^\n;&|]*?\binstall\b)',
r"(?:^|[\n;`(]|&&|\|\|?|\$\()\s*(?:\w+=\S+\s+)*(?:"
r"(?:[\w./-]*/)?(?:pip|pip3|conda|mamba|easy_install|uv)\b[^\n;&|]*?\binstall\b"
r"|(?:[\w./-]*/)?python[0-9.]*\s+-m\s+pip\b[^\n;&|]*?\binstall\b)",
re.I,
)
# Python one-liners that reach the network -- a route around curl/wget.
_PY_NET_RE = re.compile(
r'urllib\.request|\burlopen\b|\brequests\.(?:get|post|put|head|Session)\b|'
r'\bhttpx\.|\bsocket\.(?:socket|create_connection)\b|\burllib3\b',
r"urllib\.request|\burlopen\b|\brequests\.(?:get|post|put|head|Session)\b|"
r"\bhttpx\.|\bsocket\.(?:socket|create_connection)\b|\burllib3\b",
re.I,
)
# Writing test-harness/config files (conftest, .pth, sitecustomize) can force
# PASS or patch imports, bypassing evaluate(). Match writes TO these files only.
_TEST_TAMPER_RE = re.compile(
r'(?:>>?|\btee\b(?:\s+-a)?\s+)\s*[\'"]?[^\s\'"|;&<>]*'
r'(?:conftest\.py|pytest\.ini|tox\.ini|sitecustomize\.py|usercustomize\.py|'
r"(?:conftest\.py|pytest\.ini|tox\.ini|sitecustomize\.py|usercustomize\.py|"
r'setup\.cfg|pyproject\.toml|\.pth)(?=[\'"\s;&|]|$)',
re.I,
)
@@ -263,25 +263,35 @@ def _forbidden_action(action: str) -> str | None:
backstop; the authoritative fix is a default-deny egress NetworkPolicy.
"""
if _GIT_INVOKE_RE.search(action):
return ("git is disabled in this environment; do not use it for any "
"purpose. Inspect and edit the source files under /testbed "
"directly (cat, grep, sed, python) to fix the bug.")
return (
"git is disabled in this environment; do not use it for any "
"purpose. Inspect and edit the source files under /testbed "
"directly (cat, grep, sed, python) to fix the bug."
)
if _GIT_ACCESS_RE.search(action) or _HIDDEN_GIT_DIR in action:
return ("Accessing the git metadata directory is not allowed. Work only "
"with the source files under /testbed; do not read .git.")
return (
"Accessing the git metadata directory is not allowed. Work only "
"with the source files under /testbed; do not read .git."
)
if _NET_FETCH_RE.search(action) or _PY_NET_RE.search(action):
return ("Network access is disabled. Do not fetch code from the internet "
"(curl, wget, urllib, requests, etc.); all dependencies are "
"already installed. Solve the bug using only the source files "
"already present under /testbed.")
return (
"Network access is disabled. Do not fetch code from the internet "
"(curl, wget, urllib, requests, etc.); all dependencies are "
"already installed. Solve the bug using only the source files "
"already present under /testbed."
)
if _PKG_INSTALL_RE.search(action):
return ("Installing packages is not allowed. Everything needed to run the "
"code and its tests is already installed. Fix the bug by editing "
"the source under /testbed; do not install anything.")
return (
"Installing packages is not allowed. Everything needed to run the "
"code and its tests is already installed. Fix the bug by editing "
"the source under /testbed; do not install anything."
)
if _TEST_TAMPER_RE.search(action):
return ("Modifying test-harness or config files (conftest.py, pytest.ini, "
"tox.ini, setup.cfg, pyproject.toml, sitecustomize.py, .pth) is "
"not allowed. Fix the bug in the source under /testbed instead.")
return (
"Modifying test-harness or config files (conftest.py, pytest.ini, "
"tox.ini, setup.cfg, pyproject.toml, sitecustomize.py, .pth) is "
"not allowed. Fix the bug in the source under /testbed instead."
)
return None
@@ -337,9 +347,10 @@ def _run(command: str, timeout: int) -> tuple[str, int]:
return proc.stdout + proc.stderr, proc.returncode
except subprocess.TimeoutExpired:
return f"[timed out after {timeout}s]", 124
except Exception as exc: # noqa: BLE001 — surface exec errors to the model
except Exception as exc:
return f"[failed to start: {exc}]", 1
def post_event(event_type: str, data: dict[str, Any], retry: bool = False) -> None:
event_url = os.environ.get("AGL_EVENT_URL")
if not event_url:
@@ -365,10 +376,14 @@ def post_event(event_type: str, data: dict[str, Any], retry: bool = False) -> No
backoff = min(2 ** min(attempt, 5), 30)
log.warning(
"failed to post %s event (attempt %d): %s; retrying in %ds",
event_type, attempt, exc, backoff,
event_type,
attempt,
exc,
backoff,
)
time.sleep(backoff)
def fetch_eval_meta() -> dict[str, Any]:
event_url = os.environ.get("AGL_EVENT_URL", "")
match = re.match(r"^(?P<base>.*)/api/rollouts/(?P<rid>[^/]+)/attempt/", event_url)
@@ -388,11 +403,16 @@ def fetch_eval_meta() -> dict[str, Any]:
with urllib.request.urlopen(request, timeout=30) as response:
payload = json.loads(response.read().decode("utf-8"))
break
except (urllib.error.URLError, TimeoutError, ConnectionError,
http.client.HTTPException, json.JSONDecodeError) as exc:
except (
urllib.error.URLError,
TimeoutError,
ConnectionError,
http.client.HTTPException,
json.JSONDecodeError,
) as exc:
log.warning("fetch eval meta failed (attempt %d/%d): %s", i + 1, attempts, exc)
if i < attempts - 1:
time.sleep(min(2 ** i, 8))
time.sleep(min(2**i, 8))
if payload is None:
log.error("failed to fetch eval meta from server after %d attempts", attempts)
return {}
@@ -404,6 +424,7 @@ def fetch_eval_meta() -> dict[str, Any]:
"PASS_TO_PASS": inp.get("PASS_TO_PASS", []),
}
def _git_dir() -> str:
"""Path to the live git directory (relocated once the agent loop is set up)."""
return _HIDDEN_GIT_DIR if os.path.isdir(_HIDDEN_GIT_DIR) else os.path.join(TESTBED, ".git")
@@ -432,12 +453,12 @@ def _remove_stale_git_index_lock() -> None:
except OSError as exc:
log.warning("failed to remove stale git index lock %s: %s", lock_path, exc)
def _proc_output(proc: subprocess.CompletedProcess) -> str:
return "\n".join(part for part in (proc.stdout, proc.stderr) if part).strip()
def _git_retry(
args: list[str], *, timeout: int = 120, attempts: int = 3
) -> subprocess.CompletedProcess | None:
def _git_retry(args: list[str], *, timeout: int = 120, attempts: int = 3) -> subprocess.CompletedProcess | None:
"""Run a git command in /testbed with retries on timeout/failure.
Under high pod concurrency the node is heavily oversubscribed and IO-bound
@@ -453,25 +474,24 @@ def _git_retry(
log.info("sleeping %.2fs before git checkout", delay)
time.sleep(delay)
try:
proc = subprocess.run(
[*_git_base(), *args], cwd=TESTBED, capture_output=True, text=True, timeout=timeout
)
proc = subprocess.run([*_git_base(), *args], cwd=TESTBED, capture_output=True, text=True, timeout=timeout)
if proc.returncode == 0:
return proc
last = proc
if "index.lock" in _proc_output(proc):
_remove_stale_git_index_lock()
log.warning("git %s rc=%d (attempt %d/%d): %s",
args[0], proc.returncode, i + 1, attempts, _proc_output(proc)[:200])
log.warning(
"git %s rc=%d (attempt %d/%d): %s", args[0], proc.returncode, i + 1, attempts, _proc_output(proc)[:200]
)
except subprocess.TimeoutExpired:
last = None
_remove_stale_git_index_lock()
log.warning("git %s timed out after %ds (attempt %d/%d)",
args[0], timeout, i + 1, attempts)
log.warning("git %s timed out after %ds (attempt %d/%d)", args[0], timeout, i + 1, attempts)
if i < attempts - 1:
time.sleep(min(2 ** i, 8))
time.sleep(min(2**i, 8))
return last
def checkout_bug_commit(instance_id: str) -> None:
"""Bring the injected bug into /testbed using SWE-smith's official checkout.
@@ -488,6 +508,7 @@ def checkout_bug_commit(instance_id: str) -> None:
raise SystemExit(f"git checkout {instance_id} failed: {reason}")
log.info("checked out SWE-smith bug branch in /testbed: %s", instance_id)
def relocate_git() -> None:
"""Move /testbed/.git out of the agent's working tree (call after checkout).
@@ -513,8 +534,8 @@ def relocate_git() -> None:
except OSError:
same_fs = False
shutil.move(src, _HIDDEN_GIT_DIR)
log.info("relocated .git -> %s (%s)", _HIDDEN_GIT_DIR,
"rename O(1)" if same_fs else "cross-fs copy")
log.info("relocated .git -> %s (%s)", _HIDDEN_GIT_DIR, "rename O(1)" if same_fs else "cross-fs copy")
def restore_f2p_tests(instance_id: str, test_nodes: list[str]) -> None:
"""Restore the test files deleted by `Remove F2P Tests`, keeping the agent's fix.
@@ -532,17 +553,22 @@ def restore_f2p_tests(instance_id: str, test_nodes: list[str]) -> None:
if proc is None or proc.returncode != 0:
log.warning("restore_f2p_tests failed: %s", getattr(proc, "stderr", "timeout"))
def capture_patch() -> str:
try:
proc = subprocess.run(
[*_git_base(), "-c", "core.fileMode=false", "diff", "HEAD"],
cwd=TESTBED, capture_output=True, text=True, timeout=60,
cwd=TESTBED,
capture_output=True,
text=True,
timeout=60,
)
return proc.stdout
except Exception as exc: # noqa: BLE001
except Exception as exc:
log.error("patch capture failed: %s", exc)
return ""
def parse_test_statuses(test_output: str) -> dict[str, str]:
statuses: dict[str, str] = {}
for line in test_output.splitlines():
@@ -556,6 +582,7 @@ def parse_test_statuses(test_output: str) -> dict[str, str]:
statuses[node] = status
return statuses
def evaluate(eval_meta: dict[str, Any], timeout: int, f2p_only: bool = True) -> tuple[float, bool, str, bool]:
fail_to_pass = list(eval_meta.get("FAIL_TO_PASS", []))
pass_to_pass = list(eval_meta.get("PASS_TO_PASS", []))
@@ -567,20 +594,16 @@ def evaluate(eval_meta: dict[str, Any], timeout: int, f2p_only: bool = True) ->
# unrelated cross-file P2P that cause eval timeouts). Resolve = all F2P + P2P pass.
if f2p_only:
f2p_files = sorted({t.split("::", 1)[0] for t in fail_to_pass})
pass_to_pass = [t for t in pass_to_pass
if any(t.startswith(f) for f in f2p_files)]
pass_to_pass = [t for t in pass_to_pass if any(t.startswith(f) for f in f2p_files)]
nodes = fail_to_pass + pass_to_pass
restore_f2p_tests(eval_meta.get("instance_id", ""), nodes)
# Cap pytest at 4 workers: the project's own `-n auto` reads HOST cores (dozens),
# ignores the pod cgroup, and spawns ~70 workers that OOMKill the 4Gi pod. `-n`
# needs xdist, so probe first and fall back to serial `-p no:xdist` when absent.
xdist_flag = (
"-n4" if _run("python -c 'import xdist'", 30)[1] == 0 else "-p no:xdist"
)
xdist_flag = "-n4" if _run("python -c 'import xdist'", 30)[1] == 0 else "-p no:xdist"
output, rc = _run(
"python -m pytest -rA -p no:cacheprovider " + xdist_flag + " "
+ " ".join(map(_shq, nodes)),
"python -m pytest -rA -p no:cacheprovider " + xdist_flag + " " + " ".join(map(_shq, nodes)),
timeout,
)
# rc 124 == subprocess.TimeoutExpired (see _run); pytest itself never exits 124.
@@ -589,12 +612,8 @@ def evaluate(eval_meta: dict[str, Any], timeout: int, f2p_only: bool = True) ->
f2p_pass = [t for t in fail_to_pass if statuses.get(t) in ("PASSED", "XFAIL")]
# only count P2P reporting PASSED/XFAIL — a missing status means it never ran.
p2p_ok = [t for t in pass_to_pass if statuses.get(t) in ("PASSED", "XFAIL")]
resolved = (
not timed_out
and len(f2p_pass) == len(fail_to_pass)
and len(p2p_ok) == len(pass_to_pass)
)
prefix = "EVAL TIMEOUT after {}s — ".format(timeout) if timed_out else ""
resolved = not timed_out and len(f2p_pass) == len(fail_to_pass) and len(p2p_ok) == len(pass_to_pass)
prefix = f"EVAL TIMEOUT after {timeout}s — " if timed_out else ""
suffix = " (f2p_only)" if f2p_only else ""
reason = (
f"{prefix}FAIL_TO_PASS {len(f2p_pass)}/{len(fail_to_pass)} passed, "
@@ -602,17 +621,22 @@ def evaluate(eval_meta: dict[str, Any], timeout: int, f2p_only: bool = True) ->
)
return (1.0 if resolved else 0.0), resolved, reason, timed_out
def _shq(value: str) -> str:
return "'" + value.replace("'", "'\\''") + "'"
class _ContextOverflow(Exception):
pass
class _GatewayPaused(Exception):
pass
_OVERFLOW_MARKERS = ("maximum context length", "'max_tokens' is too large")
def _is_context_overflow(exc: Exception) -> bool:
if getattr(exc, "status_code", None) != 400:
@@ -620,6 +644,7 @@ def _is_context_overflow(exc: Exception) -> bool:
message = str(getattr(exc, "message", None) or exc).lower()
return any(marker in message for marker in _OVERFLOW_MARKERS)
def _is_gateway_paused(exc: Exception) -> bool:
# async weight-sync pauses the proxy (429 gateway paused). Transient — callers
# wait it out instead of burning a turn.
@@ -627,8 +652,10 @@ def _is_gateway_paused(exc: Exception) -> bool:
return False
return "gateway paused" in str(getattr(exc, "message", None) or exc).lower()
def length_penalized_reward(reward: float, n_turns: int, max_turns: int,
*, t0: int, lam: float, is_train: bool) -> float:
def length_penalized_reward(
reward: float, n_turns: int, max_turns: int, *, t0: int, lam: float, is_train: bool
) -> float:
"""Apply the long-turn penalty (plan A) to a SOLVED *training* trajectory's reward.
The penalty applies **only** when ``is_train`` is True and ``reward >= 1.0``
@@ -650,9 +677,16 @@ def length_penalized_reward(reward: float, n_turns: int, max_turns: int,
return 1.0 - lam * frac
def prompt_length_penalty(reward: float, max_prompt_tokens: int, *,
soft_start: int, hard_cap: int, max_pen: float,
is_train: bool, solved: bool) -> float:
def prompt_length_penalty(
reward: float,
max_prompt_tokens: int,
*,
soft_start: int,
hard_cap: int,
max_pen: float,
is_train: bool,
solved: bool,
) -> float:
"""Penalizes context bloat: the longest single-turn prompt of a rollout is its
true upper bound on context pressure (each turn's prompt embeds all prior
history). Gated identically to the turn penalty only ``is_train`` and
@@ -678,9 +712,18 @@ def prompt_length_penalty(reward: float, max_prompt_tokens: int, *,
return reward - max_pen * frac
def run_agent_loop(client: Any, problem: str, *, max_turns: int, cmd_timeout: int,
obs_cap: int, max_tokens: int, max_format_errors: int = 3,
gateway_wait_s: float = 600.0, gateway_poll_s: float = 5.0) -> tuple[bool, int, int]:
def run_agent_loop(
client: Any,
problem: str,
*,
max_turns: int,
cmd_timeout: int,
obs_cap: int,
max_tokens: int,
max_format_errors: int = 3,
gateway_wait_s: float = 600.0,
gateway_poll_s: float = 5.0,
) -> tuple[bool, int, int]:
"""Drive the agent until it submits, errors out, or hits ``max_turns``.
Mirrors mini-swe-agent's control flow: each turn we query the model, parse a
@@ -718,12 +761,12 @@ def run_agent_loop(client: Any, problem: str, *, max_turns: int, cmd_timeout: in
return False, turns_used, max_prompt_tokens
except _GatewayPaused as exc:
if paused_for >= gateway_wait_s:
log.warning("turn=%d gateway paused >%.0fs, giving up: %s",
turn, gateway_wait_s, exc)
log.warning("turn=%d gateway paused >%.0fs, giving up: %s", turn, gateway_wait_s, exc)
content, finish_reason, prompt_tokens = "", "error", 0
break
log.info("turn=%d gateway paused, waiting %.0fs (waited %.0fs): %s",
turn, gateway_poll_s, paused_for, exc)
log.info(
"turn=%d gateway paused, waiting %.0fs (waited %.0fs): %s", turn, gateway_poll_s, paused_for, exc
)
time.sleep(gateway_poll_s)
paused_for += gateway_poll_s
@@ -747,12 +790,10 @@ def run_agent_loop(client: Any, problem: str, *, max_turns: int, cmd_timeout: in
except FormatError as exc:
n_format_errors += 1
if 0 < max_format_errors <= n_format_errors:
log.warning("turn=%d %d consecutive format errors, ending episode",
turn, n_format_errors)
log.warning("turn=%d %d consecutive format errors, ending episode", turn, n_format_errors)
return False, turns_used, max_prompt_tokens
log.info("turn=%d format error (%d/%d)", turn, n_format_errors, max_format_errors)
messages.append({"role": "user",
"content": format_error_message(exc.n_actions, finish_reason)})
messages.append({"role": "user", "content": format_error_message(exc.n_actions, finish_reason)})
continue
n_format_errors = 0
@@ -770,6 +811,7 @@ def run_agent_loop(client: Any, problem: str, *, max_turns: int, cmd_timeout: in
messages.append({"role": "user", "content": render_observation(rc, output, obs_cap)})
return False, turns_used, max_prompt_tokens
def _query(client: Any, messages: list[dict[str, Any]], max_tokens: int) -> tuple[str, str, int]:
"""Return ``(content, finish_reason, prompt_tokens)`` for one model call.
@@ -780,14 +822,16 @@ def _query(client: Any, messages: list[dict[str, Any]], max_tokens: int) -> tupl
"""
try:
completion = client.chat.completions.create(
model="auto", messages=messages, max_tokens=max_tokens,
model="auto",
messages=messages,
max_tokens=max_tokens,
temperature=1.0,
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
)
choice = completion.choices[0]
prompt_tokens = getattr(getattr(completion, "usage", None), "prompt_tokens", 0) or 0
return (choice.message.content or ""), (choice.finish_reason or "stop"), int(prompt_tokens)
except Exception as exc: # noqa: BLE001
except Exception as exc:
if _is_context_overflow(exc):
raise _ContextOverflow(str(exc)) from exc
if _is_gateway_paused(exc):
@@ -795,9 +839,11 @@ def _query(client: Any, messages: list[dict[str, Any]], max_tokens: int) -> tupl
log.error("LLM call failed: %s", exc)
return "", "error", 0
def main() -> int:
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s",
handlers=[logging.StreamHandler(sys.stdout)])
logging.basicConfig(
level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s", handlers=[logging.StreamHandler(sys.stdout)]
)
problem = os.environ.get("AGL_TASK_INPUT", "").strip()
if not problem:
@@ -828,8 +874,13 @@ def main() -> int:
checkout_bug_commit(instance_id)
relocate_git()
submitted, n_turns, max_prompt_tokens = run_agent_loop(
client, problem, max_turns=max_turns, cmd_timeout=cmd_timeout,
obs_cap=obs_cap, max_tokens=max_tokens, max_format_errors=max_format_errors,
client,
problem,
max_turns=max_turns,
cmd_timeout=cmd_timeout,
obs_cap=obs_cap,
max_tokens=max_tokens,
max_format_errors=max_format_errors,
gateway_wait_s=gateway_wait_s,
)
@@ -843,8 +894,7 @@ def main() -> int:
is_train = "/mode/train/" in base_url
len_pen_t0 = int(os.environ.get("SMITH_LEN_PEN_T0", "80"))
len_pen_lambda = float(os.environ.get("SMITH_LEN_PEN_LAMBDA", "0.1"))
reward = length_penalized_reward(reward, n_turns, max_turns,
t0=len_pen_t0, lam=len_pen_lambda, is_train=is_train)
reward = length_penalized_reward(reward, n_turns, max_turns, t0=len_pen_t0, lam=len_pen_lambda, is_train=is_train)
# Prompt-length penalty (plan B): stack a context-bloat penalty on the same
# SOLVED-train gating, keyed on the rollout's largest prompt_tokens.
@@ -852,23 +902,48 @@ def main() -> int:
prompt_pen_hard = int(os.environ.get("SMITH_PROMPT_PEN_HARD_CAP", "64000"))
prompt_pen_max = float(os.environ.get("SMITH_PROMPT_PEN_MAX", "0.1"))
if is_train and max_prompt_tokens >= prompt_pen_hard:
log.warning("max_prompt_tokens=%d >= hard_cap=%d (context near budget)",
max_prompt_tokens, prompt_pen_hard)
reward = prompt_length_penalty(reward, max_prompt_tokens, soft_start=prompt_pen_soft,
hard_cap=prompt_pen_hard, max_pen=prompt_pen_max,
is_train=is_train, solved=resolved)
log.warning("max_prompt_tokens=%d >= hard_cap=%d (context near budget)", max_prompt_tokens, prompt_pen_hard)
reward = prompt_length_penalty(
reward,
max_prompt_tokens,
soft_start=prompt_pen_soft,
hard_cap=prompt_pen_hard,
max_pen=prompt_pen_max,
is_train=is_train,
solved=resolved,
)
log.info("done: mode=%s submitted=%s patch=%dB reward=%.3f raw_reward=%.3f n_turns=%d "
"max_prompt_tokens=%d reason=%s",
"train" if is_train else "val", submitted, len(patch), reward, raw_reward, n_turns,
max_prompt_tokens, reason)
log.info(
"done: mode=%s submitted=%s patch=%dB reward=%.3f raw_reward=%.3f n_turns=%d max_prompt_tokens=%d reason=%s",
"train" if is_train else "val",
submitted,
len(patch),
reward,
raw_reward,
n_turns,
max_prompt_tokens,
reason,
)
event_base = {"instance_id": instance_id, "repo": eval_meta.get("repo", "")}
post_event("agent_output", {**event_base, "patch": patch, "patch_size": len(patch), "submitted": submitted})
post_event("reward", {**event_base, "value": reward, "raw_value": raw_reward, "resolved": resolved,
"reason": reason, "n_turns": n_turns, "max_prompt_tokens": max_prompt_tokens,
"eval_timeout": timed_out, "source": "agent"}, retry=True)
post_event(
"reward",
{
**event_base,
"value": reward,
"raw_value": raw_reward,
"resolved": resolved,
"reason": reason,
"n_turns": n_turns,
"max_prompt_tokens": max_prompt_tokens,
"eval_timeout": timed_out,
"source": "agent",
},
retry=True,
)
return 0
if __name__ == "__main__":
sys.exit(main())
+2 -3
View File
@@ -12,6 +12,7 @@ from the active Docker daemon.
from __future__ import annotations
import argparse
import contextlib
import json
import shlex
from collections.abc import Sequence
@@ -182,10 +183,8 @@ def prepare_one(
return ImagePreparationResult(source_image, target_image, f"{source_status},{status}")
except Exception as exc:
target = f"{source_image}:{OPENAI_IMAGE_SUFFIX}"
try:
with contextlib.suppress(Exception):
target = openai_image_name(source_image)
except Exception:
pass
return ImagePreparationResult(source_image, target, "failed", str(exc))
@@ -23,16 +23,16 @@
{%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
{%- elif message.role == "assistant" %}
{{- '<|im_start|>' + message.role + '\n' }}
{%- if enable_thinking is defined and enable_thinking is false %}
{{- '<think>\n\n</think>\n\n' }}
{%- endif %}
{{- content }}
{%- if message.tool_calls %}
{%- for tool_call in message.tool_calls %}
{%- if tool_call.function %}
{%- set tool_call = tool_call.function %}
{%- endif %}
+15 -10
View File
@@ -5,9 +5,10 @@ from __future__ import annotations
import argparse
import importlib.resources
import json
from collections.abc import Sequence
from pathlib import Path
from pprint import pprint
from typing import Any, Sequence
from typing import Any
from hydra import compose, initialize_config_dir
from omegaconf import DictConfig, OmegaConf
@@ -27,22 +28,26 @@ INSTANCE_FIELDS = (
"PASS_TO_PASS",
)
def log(message: str) -> None:
print(message, flush=True)
def _project(instance: dict[str, Any]) -> dict[str, Any]:
row = {field: instance.get(field) for field in INSTANCE_FIELDS}
row["data_source"] = DATA_SOURCE
row["data_id"] = str(instance.get("instance_id", ""))
return row
def load_split_file(
path: str,
*,
max_instances: int | None = None,
) -> list[dict[str, Any]]:
"""Load a pre-split, pre-curated JSONL dataset and project to the VERL schema."""
rows = [json.loads(line) for line in Path(path).open() if line.strip()]
with Path(path).open() as file:
rows = [json.loads(line) for line in file if line.strip()]
selected = [_project(row) for row in rows]
if max_instances:
selected = selected[:max_instances]
@@ -50,6 +55,7 @@ def load_split_file(
raise ValueError(f"No instances loaded from {path}")
return selected
def verl_default_config() -> dict[str, Any]:
return {
"algorithm": {
@@ -88,7 +94,6 @@ def verl_default_config() -> dict[str, Any]:
"val_kwargs": {"temperature": 0.7, "do_sample": True},
"enable_prefix_caching": True,
"enable_chunked_prefill": True,
"checkpoint_engine": {"update_weights_bucket_megabytes": 4096},
},
"actor": {
@@ -101,7 +106,6 @@ def verl_default_config() -> dict[str, Any]:
"entropy_coeff": 0,
"clip_ratio_low": 0.2,
"clip_ratio_high": 0.28,
"fsdp_config": {
"param_offload": True,
"optimizer_offload": True,
@@ -129,7 +133,6 @@ def verl_default_config() -> dict[str, Any]:
"project_name": "agentlightning",
"experiment_name": "swe_smith",
"nnodes": 1,
"nccl_timeout": 1800,
"test_freq": 16,
"save_freq": 16,
@@ -139,7 +142,6 @@ def verl_default_config() -> dict[str, Any]:
"agentlightning": {
"agl_base_url": "http://localhost:8080",
"agl_key": "",
"rollout_timeout_seconds": 5400,
"reward_fillna_value": 0.0,
"max_ppo_update_times": 2,
@@ -158,6 +160,7 @@ def verl_default_config() -> dict[str, Any]:
},
}
def build_config(
*,
model: str | None = None,
@@ -181,11 +184,9 @@ def build_config(
rollout_mode = overrides["actor_rollout_ref"]["rollout"]["mode"]
model_path = overrides["actor_rollout_ref"]["model"]["path"]
overrides["trainer"]["experiment_name"] = (
f"swe_smith_{rollout_mode}_{model_path.split('/')[-1]}_{TRAIN_BACKEND}"
)
overrides["trainer"]["experiment_name"] = f"swe_smith_{rollout_mode}_{model_path.split('/')[-1]}_{TRAIN_BACKEND}"
if run_name:
overrides["trainer"]["experiment_name"] = f'{overrides["trainer"]["experiment_name"]}_{run_name}'
overrides["trainer"]["experiment_name"] = f"{overrides['trainer']['experiment_name']}_{run_name}"
override_conf = OmegaConf.create(overrides)
cli_override_conf = OmegaConf.from_dotlist(list(config_overrides))
@@ -194,6 +195,7 @@ def build_config(
OmegaConf.set_struct(config, False)
return config
def train(
*,
train_dataset_path: str,
@@ -236,6 +238,7 @@ def train(
log("\n=== Start VERL training ===")
run_ppo(config=config, train_dataset=train_dataset, val_dataset=val_dataset)
def parse_args() -> tuple[argparse.Namespace, list[str]]:
parser = argparse.ArgumentParser(description="Train a SWE-smith agent with VERL/GRPO via Agent Lightning")
parser.add_argument(
@@ -263,6 +266,7 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
args, config_overrides = parser.parse_known_args()
return args, config_overrides
def main() -> None:
args, config_overrides = parse_args()
train(
@@ -276,5 +280,6 @@ def main() -> None:
config_overrides=config_overrides,
)
if __name__ == "__main__":
main()
+12 -14
View File
@@ -4,23 +4,23 @@
from __future__ import annotations
import importlib.resources
from collections.abc import Sequence
from pprint import pprint
from typing import Any, Sequence
from typing import Any
from omegaconf import DictConfig, OmegaConf
from train_smith_agent import ( # noqa: E402 — sibling module, run from example dir
from train_smith_agent import (
DEFAULT_MODEL,
EXAMPLE_DIR,
load_split_file,
log,
)
import importlib.resources # noqa: E402
CHAT_TEMPLATE_PATH = str(EXAMPLE_DIR / "swe_smith_chat_template.jinja")
TRAIN_BACKEND = "megatron"
def verl_megatron_config() -> dict[str, Any]:
return {
"algorithm": {
@@ -42,9 +42,7 @@ def verl_megatron_config() -> dict[str, Any]:
"gpu_memory_utilization": 0.7,
"max_model_len": 32768,
"enforce_eager": True,
"enable_rollout_routing_replay": True,
"calculate_log_probs": True,
"log_prob_micro_batch_size_per_gpu": 1,
"log_prob_use_dynamic_bsz": False,
@@ -75,7 +73,6 @@ def verl_megatron_config() -> dict[str, Any]:
"clip_ratio_low": 0.2,
"clip_ratio_high": 0.28,
"loss_agg_mode": "seq-mean-token-sum",
"megatron": {
"pipeline_model_parallel_size": 1,
"tensor_model_parallel_size": 2,
@@ -85,9 +82,7 @@ def verl_megatron_config() -> dict[str, Any]:
"optimizer_offload": True,
"grad_offload": True,
"use_mbridge": True,
"router_replay": {"mode": "R3"},
"override_transformer_config": {
"moe_enable_deepep": True,
"moe_token_dispatcher_type": "flex",
@@ -156,6 +151,7 @@ def verl_megatron_config() -> dict[str, Any]:
},
}
def build_config(
*,
model: str | None = None,
@@ -180,11 +176,9 @@ def build_config(
rollout_mode = overrides["actor_rollout_ref"]["rollout"]["mode"]
model_path = overrides["actor_rollout_ref"]["model"]["path"]
overrides["trainer"]["experiment_name"] = (
f"swe_smith_{rollout_mode}_{model_path.split('/')[-1]}_{TRAIN_BACKEND}"
)
overrides["trainer"]["experiment_name"] = f"swe_smith_{rollout_mode}_{model_path.split('/')[-1]}_{TRAIN_BACKEND}"
if run_name:
overrides["trainer"]["experiment_name"] = f'{overrides["trainer"]["experiment_name"]}_{run_name}'
overrides["trainer"]["experiment_name"] = f"{overrides['trainer']['experiment_name']}_{run_name}"
override_conf = OmegaConf.create(overrides)
cli_override_conf = OmegaConf.from_dotlist(list(config_overrides))
@@ -193,6 +187,7 @@ def build_config(
OmegaConf.set_struct(config, False)
return config
def train(
*,
train_dataset_path: str,
@@ -235,6 +230,7 @@ def train(
log("\n=== Start VERL training (Megatron actor, R3 router replay, vLLM rollout) ===")
run_ppo(config=config, train_dataset=train_dataset, val_dataset=val_dataset)
def parse_args():
import argparse
@@ -264,6 +260,7 @@ def parse_args():
parser.add_argument("--run-name", default=None)
return parser.parse_known_args()
def main() -> None:
args, config_overrides = parse_args()
train(
@@ -277,5 +274,6 @@ def main() -> None:
config_overrides=config_overrides,
)
if __name__ == "__main__":
main()
+38
View File
@@ -22,6 +22,8 @@ dev = [
"pytest>=8.3.0",
"pytest-asyncio>=0.25.0",
"pytest-cov>=6.0.0",
"pre-commit>=4.0.0",
"openai>=2.0.0,<3",
"ruff>=0.11.0",
"pyright>=1.1.390",
"sympy>=1.13.0",
@@ -31,9 +33,23 @@ dev = [
agl-server = "agentlightning.server.__main__:main"
agl-controller = "agentlightning.controller.__main__:main"
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu", group = "verl-cpu" },
]
[[tool.uv.index]]
name = "pypi"
url = "https://pypi.org/simple"
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
[tool.ruff]
line-length = 120
target-version = "py312"
exclude = ["examples/llm-in-sandbox/vendor"]
[tool.ruff.lint]
select = ["E", "F", "W", "I", "UP", "B", "SIM", "RUF"]
@@ -42,8 +58,11 @@ select = ["E", "F", "W", "I", "UP", "B", "SIM", "RUF"]
quote-style = "double"
[tool.pyright]
include = ["agentlightning", "tests", "docs/macros", "scripts"]
pythonVersion = "3.12"
typeCheckingMode = "standard"
venvPath = "."
venv = ".venv"
[tool.pytest.ini_options]
testpaths = ["tests"]
@@ -66,3 +85,22 @@ dev = [
"pytest-httpx>=0.36.0",
"sympy>=1.14.0",
]
# Lint/type-check environment for the VERL integration under agentlightning/verl.
# Pyright needs verl and its peers importable to check that subtree; it only ever
# reads .py/.pyi files, so torch comes from the CPU index (see [tool.uv.sources])
# and ~2.3GB of CUDA runtime is skipped. This is not a training environment --
# install verl against the GPU torch build of your choice for that.
#
# Upper bound is deliberate: verl 0.9.0 renamed main_ppo.TaskRunner to
# TaskRunnerV1 (the old name moved to main_ppo_v0) and dropped
# create_rl_sampler, both of which entrypoint.py imports. auto_await landed in
# 0.7.1, so that is the floor.
verl-cpu = [
"verl>=0.7.1,<0.9.0",
"torch",
"ray",
"tensordict",
"datasets",
"numpy",
"tqdm",
]
+79
View File
@@ -0,0 +1,79 @@
# Copyright (c) Microsoft. All rights reserved.
"""Ensure tracked Python files include the required copyright header."""
from __future__ import annotations
import subprocess
import sys
from pathlib import Path
HEADER = "# Copyright (c) Microsoft. All rights reserved."
REPO_ROOT = Path(__file__).resolve().parent.parent
SOURCE_SUFFIXES = (".py", ".pyi", ".pyw")
EXCLUDED_PREFIXES = ("examples/llm-in-sandbox/vendor/",)
def iter_source_files() -> list[Path]:
"""Return tracked and untracked Python source files."""
result = subprocess.run(
[
"git",
"ls-files",
"--cached",
"--others",
"--exclude-standard",
"--",
*(f"*{suffix}" for suffix in SOURCE_SUFFIXES),
],
capture_output=True,
text=True,
check=True,
cwd=REPO_ROOT,
)
paths = (line.strip() for line in result.stdout.splitlines())
return [
REPO_ROOT / path for path in paths if path and not any(path.startswith(prefix) for prefix in EXCLUDED_PREFIXES)
]
def main() -> int:
missing_header: list[str] = []
missing_blank_line: list[str] = []
for file_path in iter_source_files():
if not file_path.exists():
continue
try:
with file_path.open("r", encoding="utf-8") as file:
first_line = file.readline().rstrip("\r\n")
header_line = file.readline().rstrip("\r\n") if first_line.startswith("#!") else first_line
following_line = file.readline()
except OSError as exc:
print(f"Failed to read {file_path}: {exc}", file=sys.stderr)
return 1
relative_path = str(file_path.relative_to(REPO_ROOT))
if header_line != HEADER:
missing_header.append(relative_path)
continue
if following_line and following_line.strip():
missing_blank_line.append(relative_path)
if missing_header:
print("The following files are missing the required copyright header:")
for path in missing_header:
print(f" - {path}")
print(f"Add this header after any shebang:\n{HEADER}")
if missing_blank_line:
print("The following files are missing a blank line after the copyright header:")
for path in missing_blank_line:
print(f" - {path}")
return 1 if missing_header or missing_blank_line else 0
if __name__ == "__main__":
sys.exit(main())
+15 -6
View File
@@ -188,9 +188,7 @@ def test_git_retry_sleeps_before_each_checkout(monkeypatch, tmp_path) -> None:
def test_checkout_bug_commit_uses_swesmith_official_checkout(monkeypatch) -> None:
calls: list[tuple[list[str], int, int]] = []
def fake_git_retry(
args: list[str], *, timeout: int = 120, attempts: int = 3
) -> subprocess.CompletedProcess:
def fake_git_retry(args: list[str], *, timeout: int = 120, attempts: int = 3) -> subprocess.CompletedProcess:
calls.append((args, timeout, attempts))
return subprocess.CompletedProcess(["git", *args], 0, "ok", "")
@@ -298,7 +296,13 @@ def test_relocate_git_moves_worktree_git_and_routes_harness(monkeypatch, tmp_pat
assert (hidden / "HEAD").read_text() == "ref: refs/heads/main\n"
assert smith_agent._git_dir() == str(hidden)
assert smith_agent._git_base() == [
"git", "--git-dir", str(hidden), "--work-tree", str(testbed), "-c", "safe.directory=*",
"git",
"--git-dir",
str(hidden),
"--work-tree",
str(testbed),
"-c",
"safe.directory=*",
]
@@ -354,8 +358,13 @@ def test_length_penalty_guards_degenerate_span() -> None:
def _plp(reward, mpt, is_train=True, solved=True):
return smith_agent.prompt_length_penalty(
reward, mpt, soft_start=50000, hard_cap=64000, max_pen=0.1,
is_train=is_train, solved=solved,
reward,
mpt,
soft_start=50000,
hard_cap=64000,
max_pen=0.1,
is_train=is_train,
solved=solved,
)
+1 -3
View File
@@ -63,9 +63,7 @@ class _FakeClient:
def test_openai_image_name_matches_job_template_convention() -> None:
assert pull_images.openai_image_name("jyangballin/swesmith.foo") == (
"jyangballin/swesmith.foo:openai"
)
assert pull_images.openai_image_name("jyangballin/swesmith.foo") == ("jyangballin/swesmith.foo:openai")
def test_openai_image_name_rejects_already_tagged_sources() -> None:
+7 -3
View File
@@ -4,7 +4,11 @@
from __future__ import annotations
from collections.abc import Iterator
from typing import Any
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from agentlightning.server.app import create_app
@@ -15,7 +19,7 @@ MODEL_NAME = "test-model"
@pytest.fixture(autouse=True)
def clean_store():
def clean_store() -> Iterator[None]:
_rollouts.clear()
_events.clear()
_models.clear()
@@ -40,12 +44,12 @@ def server_config() -> dict:
@pytest.fixture
def app(server_config: dict[str, str]):
def app(server_config: dict[str, Any]) -> FastAPI:
return create_app(server_config)
@pytest.fixture
def client(app) -> TestClient:
def client(app: FastAPI) -> Iterator[TestClient]:
with TestClient(app) as c:
yield c
+10 -5
View File
@@ -360,16 +360,21 @@ def _terminal(client: TestClient, headers: dict[str, str], data_id: str, is_trai
)
assert created.status_code == 201
rid = created.json()[0]["rollout_id"]
assert client.patch(f"/api/rollouts/{rid}", json={"status": {"state": "running"}}, headers=headers).status_code == 200
assert client.patch(f"/api/rollouts/{rid}", json={"status": {"state": "succeeded"}}, headers=headers).status_code == 200
assert (
client.patch(f"/api/rollouts/{rid}", json={"status": {"state": "running"}}, headers=headers).status_code == 200
)
assert (
client.patch(f"/api/rollouts/{rid}", json={"status": {"state": "succeeded"}}, headers=headers).status_code
== 200
)
return rid
def test_terminal_rollouts_cursor_pagination(client: TestClient, auth_headers: dict[str, str]):
# A rollout that never reaches a terminal state must NOT appear in the log.
pending = client.post(
"/api/rollouts", json=[{"input": {"data_id": "pending"}}], headers=auth_headers
).json()[0]["rollout_id"]
pending = client.post("/api/rollouts", json=[{"input": {"data_id": "pending"}}], headers=auth_headers).json()[0][
"rollout_id"
]
# Complete three rollouts; the log is ordered by COMPLETION (append-on-terminal).
rid_a = _terminal(client, auth_headers, "a", is_train=True)
+3 -1
View File
@@ -1,3 +1,5 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import pytest
@@ -61,7 +63,7 @@ def test_policy_loss_matches_masked_sum() -> None:
log_prob = torch.zeros(2, 3)
loss, metrics = compute_policy_loss_per_rollout_mean(
old_log_prob=log_prob,
old_log_prob=log_prob, # pyright: ignore[reportCallIssue]
log_prob=log_prob,
advantages=advantages,
response_mask=response_mask,
+1 -1
View File
@@ -230,7 +230,7 @@ def test_training_step_uploads_24_compact_rollout_trajectories_to_wandb_zip(
_triplet([index, index + 100, index + 200], [index + 300]),
],
)
for index in range(1, 27)
for index in range(1, 27)
]
_adapter().get_train_data_batch(rollouts, global_steps=23)
Generated
+2458 -7
View File
File diff suppressed because it is too large Load Diff