[codex] Add Tinker checkpoint delete endpoint

.claude/docs/tinker.md

@@ -37,6 +37,7 @@ All endpoints are under `/api/v1/`. Requests are async -- submit via POST, get a
 | `/save_weights` | POST | Save full training checkpoint (weights + optimizer) |
 | `/save_weights_for_sampler` | POST | Sync weights to inference engines |
 | `/load_weights` | POST | Load a previously saved checkpoint |
+| `/training_runs/{unique_id}/checkpoints/{checkpoint_id}` | DELETE | Delete a saved checkpoint archive from disk |
 | `/retrieve_future` | POST | Long-poll for async result (300s timeout) |
 | `/healthz` | GET | Liveness check |
 
@@ -52,6 +53,7 @@ All endpoints are under `/api/v1/`. Requests are async -- submit via POST, get a
 - **Persistent**: `save_weights_for_sampler(name="...")` -- syncs to inference engines AND writes HF checkpoint to disk. Expensive.
 - **Ephemeral**: `save_weights_and_get_sampling_client(name="...")` -- syncs to inference engines only, skips disk write. Triggered when `sampling_session_seq_id` is present in the request.
 - In RL loops, always prefer ephemeral mode; reserve persistent saves for periodic checkpoints.
+- Delete persistent checkpoints with `DELETE /training_runs/{unique_id}/checkpoints/{checkpoint_id}` after they are no longer needed. This removes the saved archive from `checkpoints_base`; it does not unload the live model.
 
 ## Testing
 
@@ -74,4 +76,4 @@ uv run --extra dev --extra jax pytest tests/tinker/test_api.py -v
 
 - **Token shifting**: Tinker pre-shifts inputs/targets; SkyRL-Train shifts internally. The backend appends the last target token to reconstruct full sequences during batch conversion -- be careful if modifying `prepare_model_pass_batch`.
 - **Left-padding**: SkyRL-Train expects left-padded tensors. The backend handles this during batch prep.
-- **API models vs internal types**: `api.py` defines its own Pydantic models (e.g., `api.ForwardBackwardInput`) that mirror but differ from `types.ForwardBackwardInput`. Each API model has a `.to_types()` method for conversion. Do not confuse the two.
+- **API models vs internal types**: `api.py` defines its own Pydantic models (e.g., `api.ForwardBackwardInput`) that mirror but differ from `types.ForwardBackwardInput`. Each API model has a `.to_types()` method for conversion. Do not confuse the two.

skyrl/tinker/api.py

@@ -1,6 +1,8 @@
 import asyncio
 import os
 import random
+import re
+import shutil
 import signal
 import threading
 import time
@@ -1219,8 +1221,59 @@ async def validate_checkpoint(
     if checkpoint_db.status == CheckpointStatus.FAILED:
         raise HTTPException(status_code=500, detail=f"Checkpoint creation failed: {checkpoint_db.error_message}")
 
-    subdir = "sampler_weights" if checkpoint_type == types.CheckpointType.SAMPLER else ""
-    return request.app.state.engine_config.checkpoints_base / unique_id / subdir / f"{checkpoint_id}.tar.gz"
+    return checkpoint_file_path(request, unique_id, checkpoint_id, checkpoint_type)
+
+
+def checkpoint_file_path(
+    request: Request, unique_id: str, checkpoint_id: str, checkpoint_type: types.CheckpointType
+) -> Any:
+    checkpoint_dir = request.app.state.engine_config.checkpoints_base / unique_id
+    if checkpoint_type == types.CheckpointType.SAMPLER:
+        checkpoint_dir = checkpoint_dir / "sampler_weights"
+    return checkpoint_dir / f"{checkpoint_id}.tar.gz"
+
+
+def parse_checkpoint_delete_path(
+    checkpoint_path: str, checkpoint_type: types.CheckpointType | None
+) -> tuple[str, types.CheckpointType | None]:
+    path_kind, separator, checkpoint_id = checkpoint_path.partition("/")
+    if separator:
+        if path_kind == "weights":
+            inferred_checkpoint_type = types.CheckpointType.TRAINING
+        elif path_kind == "sampler_weights":
+            inferred_checkpoint_type = types.CheckpointType.SAMPLER
+        else:
+            raise HTTPException(status_code=400, detail=f"Invalid checkpoint path: {checkpoint_path}")
+
+        if checkpoint_type is not None and checkpoint_type != inferred_checkpoint_type:
+            raise HTTPException(status_code=400, detail="checkpoint_type does not match checkpoint path")
+    else:
+        checkpoint_id = checkpoint_path
+        inferred_checkpoint_type = checkpoint_type
+
+    if not re.fullmatch(ID_PATTERN, checkpoint_id) or len(checkpoint_id) > ID_MAX_LENGTH:
+        raise HTTPException(status_code=422, detail="Invalid checkpoint_id")
+
+    return checkpoint_id, inferred_checkpoint_type
+
+
+def delete_checkpoint_file(checkpoint_path: Any) -> None:
+    if checkpoint_path.is_dir():
+        if hasattr(checkpoint_path, "rmtree"):
+            checkpoint_path.rmtree()
+        else:
+            shutil.rmtree(checkpoint_path)
+    else:
+        try:
+            checkpoint_path.unlink()
+        except FileNotFoundError:
+            return
+
+    with suppress(OSError):
+        checkpoint_path.parent.rmdir()
+    if checkpoint_path.parent.name == "sampler_weights":
+        with suppress(OSError):
+            checkpoint_path.parent.parent.rmdir()
 
 
 @app.get("/api/v1/training_runs")
@@ -1303,6 +1356,38 @@ async def download_checkpoint_archive(
     return StreamingResponse(file_buffer, media_type="application/octet-stream", headers=headers)
 
 
+@app.delete("/api/v1/training_runs/{unique_id}/checkpoints/{checkpoint_path:path}", status_code=204)
+async def delete_checkpoint(
+    request: Request,
+    unique_id: str = fastapi.Path(..., pattern=ID_PATTERN, max_length=ID_MAX_LENGTH),
+    checkpoint_path: str = fastapi.Path(...),
+    checkpoint_type: types.CheckpointType | None = fastapi.Query(None),
+    session: AsyncSession = Depends(get_session),
+) -> None:
+    """Delete a saved checkpoint artifact and its database row."""
+    checkpoint_id, resolved_checkpoint_type = parse_checkpoint_delete_path(checkpoint_path, checkpoint_type)
+    checkpoint_db = None
+    if resolved_checkpoint_type is None:
+        resolved_checkpoint_type = types.CheckpointType.TRAINING
+        checkpoint_db = await session.get(CheckpointDB, (unique_id, checkpoint_id, resolved_checkpoint_type))
+        if checkpoint_db is None:
+            resolved_checkpoint_type = types.CheckpointType.SAMPLER
+            checkpoint_db = await session.get(CheckpointDB, (unique_id, checkpoint_id, resolved_checkpoint_type))
+    else:
+        checkpoint_db = await session.get(CheckpointDB, (unique_id, checkpoint_id, resolved_checkpoint_type))
+
+    if not checkpoint_db:
+        raise HTTPException(status_code=404, detail=f"Checkpoint not found: {unique_id}/{checkpoint_id}")
+
+    if checkpoint_db.status == CheckpointStatus.PENDING:
+        raise HTTPException(status_code=425, detail="Checkpoint is still being created")
+
+    path = checkpoint_file_path(request, unique_id, checkpoint_id, resolved_checkpoint_type)
+    await asyncio.to_thread(delete_checkpoint_file, path)
+    await session.delete(checkpoint_db)
+    await session.commit()
+
+
 @app.get("/api/v1/training_runs/{unique_id}/checkpoints")
 async def list_checkpoints(
     unique_id: str = fastapi.Path(..., pattern=ID_PATTERN, max_length=ID_MAX_LENGTH),
@@ -1379,12 +1464,19 @@ async def root():
         "name": "Tinker API Mock",
         "version": "0.0.1",
         "endpoints": {
-            "models": ["/api/v1/create_model", "/api/v1/get_info", "/api/v1/training_runs/{model_id}"],
+            "models": [
+                "/api/v1/create_model",
+                "/api/v1/get_info",
+                "/api/v1/training_runs/{model_id}",
+            ],
             "training": ["/api/v1/forward_backward", "/api/v1/optim_step"],
             "futures": ["/api/v1/retrieve_future"],
             "service": ["/api/v1/get_server_capabilities"],
             "telemetry": ["/api/v1/telemetry"],
-            "checkpoints": ["/api/v1/training_runs/{unique_id}/checkpoints"],
+            "checkpoints": [
+                "/api/v1/training_runs/{unique_id}/checkpoints",
+                "DELETE /api/v1/training_runs/{unique_id}/checkpoints/{checkpoint_id}",
+            ],
             "download": [
                 "/api/v1/training_runs/{unique_id}/checkpoints/{checkpoint_id}/archive",
                 "/api/v1/training_runs/{unique_id}/checkpoints/{checkpoint_id}/download",

tests/tinker/test_checkpoint_delete_api.py

@@ -0,0 +1,158 @@
+"""Fast tests for checkpoint deletion API behavior."""
+
+import asyncio
+from collections.abc import AsyncGenerator, Iterator
+from datetime import datetime, timezone
+from pathlib import Path
+
+import pytest
+from fastapi.testclient import TestClient
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+from sqlmodel import SQLModel
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+from skyrl.tinker import types
+from skyrl.tinker.api import app, get_session
+from skyrl.tinker.config import EngineConfig
+from skyrl.tinker.db_models import CheckpointDB, CheckpointStatus, ModelDB, SessionDB
+
+MODEL_ID = "model_fast_delete"
+
+
+@pytest.fixture
+def checkpoint_api(tmp_path: Path) -> Iterator[tuple[TestClient, Path, AsyncEngine]]:
+    db_path = tmp_path / "checkpoint_delete.db"
+    checkpoint_base = tmp_path / "checkpoints"
+    db_engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}")
+
+    async def setup_db() -> None:
+        async with db_engine.begin() as conn:
+            await conn.run_sync(SQLModel.metadata.create_all)
+
+    async def override_get_session() -> AsyncGenerator[AsyncSession, None]:
+        async with AsyncSession(db_engine) as session:
+            yield session
+
+    asyncio.run(setup_db())
+    previous_state = app.state._state.copy()
+    app.state.db_engine = db_engine
+    app.state.engine_config = EngineConfig(
+        base_model="test-model",
+        checkpoints_base=checkpoint_base,
+        database_url=f"sqlite:///{db_path}",
+    )
+    app.dependency_overrides[get_session] = override_get_session
+    client = TestClient(app)
+
+    try:
+        yield client, checkpoint_base, db_engine
+    finally:
+        client.close()
+        app.dependency_overrides.pop(get_session, None)
+        app.state._state.clear()
+        app.state._state.update(previous_state)
+        asyncio.run(db_engine.dispose())
+
+
+async def seed_model_and_checkpoints(
+    db_engine: AsyncEngine,
+    checkpoint_ids: list[str],
+    checkpoint_type: types.CheckpointType = types.CheckpointType.TRAINING,
+) -> None:
+    async with AsyncSession(db_engine) as session:
+        if await session.get(SessionDB, "session_fast_delete") is None:
+            session.add(SessionDB(session_id="session_fast_delete", tags=[], sdk_version="test"))
+        if await session.get(ModelDB, MODEL_ID) is None:
+            session.add(
+                ModelDB(
+                    model_id=MODEL_ID,
+                    base_model="test-model",
+                    lora_config={"rank": 1},
+                    status="created",
+                    request_id=1,
+                    session_id="session_fast_delete",
+                )
+            )
+        for checkpoint_id in checkpoint_ids:
+            session.add(
+                CheckpointDB(
+                    model_id=MODEL_ID,
+                    checkpoint_id=checkpoint_id,
+                    checkpoint_type=checkpoint_type,
+                    status=CheckpointStatus.COMPLETED,
+                    completed_at=datetime.now(timezone.utc),
+                )
+            )
+        await session.commit()
+
+
+def write_training_checkpoint(checkpoint_base: Path, checkpoint_id: str, directory: bool = False) -> Path:
+    checkpoint_path = checkpoint_base / MODEL_ID / f"{checkpoint_id}.tar.gz"
+    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+    if directory:
+        checkpoint_path.mkdir()
+        (checkpoint_path / "checkpoint_0").write_text("tiny")
+    else:
+        checkpoint_path.write_text("tiny")
+    return checkpoint_path
+
+
+def write_sampler_checkpoint(checkpoint_base: Path, checkpoint_id: str) -> Path:
+    checkpoint_path = checkpoint_base / MODEL_ID / "sampler_weights" / f"{checkpoint_id}.tar.gz"
+    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+    checkpoint_path.write_text("tiny")
+    return checkpoint_path
+
+
+def listed_checkpoint_ids(client: TestClient) -> set[str]:
+    response = client.get(f"/api/v1/training_runs/{MODEL_ID}/checkpoints")
+    assert response.status_code == 200
+    return {checkpoint["checkpoint_id"] for checkpoint in response.json()["checkpoints"]}
+
+
+def test_delete_checkpoint_removes_saved_artifact_and_list_entry(
+    checkpoint_api: tuple[TestClient, Path, AsyncEngine],
+) -> None:
+    client, checkpoint_base, db_engine = checkpoint_api
+    asyncio.run(seed_model_and_checkpoints(db_engine, ["delete_training"]))
+    training_checkpoint = write_training_checkpoint(checkpoint_base, "delete_training", directory=True)
+
+    response = client.delete(f"/api/v1/training_runs/{MODEL_ID}/checkpoints/weights/delete_training")
+
+    assert response.status_code == 204
+    assert not training_checkpoint.exists()
+    assert "delete_training" not in listed_checkpoint_ids(client)
+
+    asyncio.run(
+        seed_model_and_checkpoints(db_engine, ["delete_sampler"], checkpoint_type=types.CheckpointType.SAMPLER)
+    )
+    sampler_checkpoint = write_sampler_checkpoint(checkpoint_base, "delete_sampler")
+
+    response = client.delete(f"/api/v1/training_runs/{MODEL_ID}/checkpoints/delete_sampler")
+
+    assert response.status_code == 204
+    assert not sampler_checkpoint.exists()
+    assert "delete_sampler" not in listed_checkpoint_ids(client)
+
+
+def test_delete_even_checkpoints_leaves_odd_checkpoints_listed(
+    checkpoint_api: tuple[TestClient, Path, AsyncEngine],
+) -> None:
+    client, checkpoint_base, db_engine = checkpoint_api
+    checkpoint_ids = ["1", "2", "3", "4", "5"]
+    asyncio.run(seed_model_and_checkpoints(db_engine, checkpoint_ids))
+    checkpoint_files = {
+        checkpoint_id: write_training_checkpoint(checkpoint_base, checkpoint_id) for checkpoint_id in checkpoint_ids
+    }
+
+    assert listed_checkpoint_ids(client) == {"1", "2", "3", "4", "5"}
+
+    assert client.delete(f"/api/v1/training_runs/{MODEL_ID}/checkpoints/2").status_code == 204
+    assert client.delete(f"/api/v1/training_runs/{MODEL_ID}/checkpoints/4").status_code == 204
+
+    assert checkpoint_files["1"].exists()
+    assert not checkpoint_files["2"].exists()
+    assert checkpoint_files["3"].exists()
+    assert not checkpoint_files["4"].exists()
+    assert checkpoint_files["5"].exists()
+    assert listed_checkpoint_ids(client) == {"1", "3", "5"}