RFC: move trajectory logging from the backend to the frontend

[corbt]

Currently, we save logged trajectories in the backend. However, this doesn't really work with TinkerBackend. I propose we instead provide the user with an explicit function to log their trajectories to the local filesystem in their frontend, and don't keep logging them automatically in the backend.

[corbt]

Proposed Implementation Plan

This RFC proposes moving trajectory logging from the backend to the frontend (Model class). This simplifies backends and keeps the frontend responsible for most state management, while the backend focuses purely on GPU operations. The only state that remains in the backend is checkpoint/optimizer state, since it's impractical to ship those back and forth.

Design Philosophy

Concern	Frontend (Model)	Backend
Write trajectory parquet	✅	❌
Write history.jsonl	✅	❌
Log to wandb	✅	❌
Read history.jsonl	✅	❌
Calculate LR multiplier	✅	❌
Find best checkpoint step	✅	❌
Pack tensors & train	❌	✅
Delete checkpoint files	❌	✅
Manage vLLM server	❌	✅
Checkpoint/optimizer state	❌	✅

---

API Backward Compatibility

No changes required to existing user code. The following calls work exactly as before:

await model.log(trajectories, split="val")
await model.train(trajectory_groups, config=art.TrainConfig(learning_rate=5e-5))

New optional attributes allow customization:

model = art.TrainableModel(
    ...,
    base_path=".art",              # Custom output path (default: ".art")
    report_metrics=["wandb"],      # Explicit control (default: None = wandb if key present)
)

---

Key Changes

1. Add New Attributes to Model Class

class Model:
    base_path: str = ".art"  # Same default as LocalBackend for backward compat
    report_metrics: list[str] | None = None  # None = default (wandb if key present)
    _wandb_run: Run | None = None  # Private, for lazy wandb initialization

2. Rewrite model.log() Implementation

Move all logic from LocalBackend._log() and LocalBackend._log_metrics() directly into Model.log():

async def log(self, trajectories, split="val"):
    trajectory_groups = # ... convert to list[TrajectoryGroup] ...
    step = await self.get_step() if self.trainable else 0
    
    # 1. Write parquet
    path = f"{self.base_path}/{self.project}/models/{self.name}/trajectories/{split}/{step:04d}.parquet"
    os.makedirs(os.path.dirname(path), exist_ok=True)
    write_trajectory_groups_parquet(trajectory_groups, path)
    
    # 2. Calculate aggregate metrics (reward, exception_rate, reward_std_dev, custom metrics)
    
    # 3. Write to history.jsonl
    prefixed = {f"{split}/{k}": v for k, v in metrics.items()}
    history_path = f"{self.base_path}/{self.project}/models/{self.name}/history.jsonl"
    with open(history_path, "a") as f:
        f.write(json.dumps({**prefixed, "step": step, "recorded_at": ...}) + "\n")
    
    # 4. Log to wandb (if enabled based on report_metrics)
    if should_log_wandb:
        if self._wandb_run is None:
            self._wandb_run = wandb.init(project=self.project, name=self.name, ...)
        self._wandb_run.log({"training_step": step, **prefixed}, step=step)

3. Modify model.train() to Handle Logging and LR Scaling

async def train(self, trajectory_groups, config, _config=None, ...):
    groups_list = list(trajectory_groups)
    
    # 1. Log trajectories first (frontend handles this now)
    await self.log(groups_list, split="train")
    
    # 2. Apply LR scaling if enabled (moved from backend)
    if _config and _config.get("scale_learning_rate_by_reward_std_dev"):
        multiplier = self._calculate_lr_multiplier()  # Reads history.jsonl
        config = config.model_copy(update={"learning_rate": config.learning_rate * multiplier})
    
    # 3. Train (backend no longer logs internally)
    training_metrics = []
    async for metrics in self.backend()._train_model(self, groups_list, config, ...):
        training_metrics.append(metrics)
    
    # 4. Log training metrics (loss, gradient norms, etc.) to history.jsonl and wandb

4. Modify model.delete_checkpoints() to Read History

Move history.jsonl reading to frontend, backend just deletes files:

async def delete_checkpoints(self, best_checkpoint_metric="val/reward"):
    output_dir = f"{self.base_path}/{self.project}/models/{self.name}"
    steps_to_keep = [await self.get_step()]  # Keep latest
    
    # Read history.jsonl to find best step (moved from backend)
    try:
        best_step = pl.read_ndjson(f"{output_dir}/history.jsonl")...
        steps_to_keep.append(best_step)
    except (FileNotFoundError, pl.exceptions.ColumnNotFoundError):
        pass
    
    # Backend only does file deletion
    await self.backend()._delete_checkpoint_files(self, steps_to_keep)

5. Simplify Backend

In LocalBackend._train_model():

• Remove internal self._log() call
• Remove self._log_metrics() call
• Remove scale_learning_rate_by_reward_std_dev handling (now in frontend)
• Keep yielding metrics for frontend to collect

Remove these backend methods entirely:

• Backend._log()
• LocalBackend._log()
• LocalBackend._log_metrics()
• LocalBackend._get_reward_std_dev_learning_rate_multiplier()
• ServerlessBackend._log()

Simplify _delete_checkpoints() -> _delete_checkpoint_files():

• Remove history.jsonl reading logic
• Just take steps_to_keep list and delete checkpoint directories

---

Path Structure (Unchanged)

{base_path}/{project}/models/{name}/trajectories/{split}/{step:04d}.parquet
{base_path}/{project}/models/{name}/history.jsonl

Default base_path=".art" matches current LocalBackend behavior exactly.

---

Files to Modify

• src/art/model.py - Add attributes, rewrite log(), modify train(), modify delete_checkpoints()
• src/art/local/backend.py - Remove logging methods, simplify _train_model() and _delete_checkpoints()
• src/art/serverless/backend.py - Remove _log()
• src/art/backend.py - Remove _log()

---

Breaking Changes

1. ServerlessBackend: Trajectories now saved locally instead of sent to API
2. Backend interface: _log() method removed, _delete_checkpoints() signature changes (internal API, not user-facing)

What Stays the Same for LocalBackend Users

• Same API calls (model.log(), model.train())
• Same file locations (.art/{project}/models/{name}/...)
• Same wandb behavior (auto-logs if WANDB_API_KEY present)
• Same metrics logged
• Same history.jsonl format