Skip to content

Add loops commands - #50

Open
cretz wants to merge 1 commit into
mainfrom
cretz/loops-commands
Open

Add loops commands#50
cretz wants to merge 1 commit into
mainfrom
cretz/loops-commands

Conversation

@cretz

@cretz cretz commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

🚀 What

Adds the baseten loops command surface, ported from truss loops. Nouns follow the backend (run, sampler, checkpoint); truss's deprecated spellings and its per-command output-format flag are dropped, the latter because the CLI has a global --output.

  • loops run create (truss loops push): creates a session, a run, and its paired sampler. Renamed because nothing is uploaded. --base-model required; --name, --replicas, --team optional. Returns once provisioned rather than blocking on readiness, which is the SDK's job.
  • loops run list (truss loops view, plus the deprecated loops runs view): runs newest-first. --org adds an OWNER column, --all includes inactive runs, --base-model filters, --direction sorts.
  • loops run describe (no truss equivalent): one run, with its paired sampler's URL and model/deployment ids.
  • loops run deactivate (truss loops deactivate): tears down trainer and sampler. Prompts unless --yes; --yes required when stdin is not a TTY.
  • loops run logs (truss loops logs): trainer logs by default, --sampler for the sampler's separate stream. Uses the shared logs flow, so --tail, --since/--start/--end, and --limit all work. Only --min-level is offered as a filter, since the trainer endpoint has no message or replica filters.
  • loops usage (truss loops usage): GPU summary plus one row per trainer and per standalone sampler. --org, --user, --all to include allocations holding no live GPUs. Also covers the deprecated loops samplers view.
  • loops checkpoint list (truss loops checkpoints view): top-level rather than nested, matching REST, since checkpoints are queryable across runs. --run-id or --base-model, one required.
  • loops checkpoint files (no truss equivalent, SDK-only): presigned download URLs for a checkpoint's files.
  • No loops checkpoint deploy yet: it is the only GraphQL-backed command in the surface and takes an evaluated Python file as config, so it likely needs to delegate to truss. Deferred until the train port, which hits the same question with train deploy_checkpoints.

💻 How

  • All eight commands go through the generated baseten-go management client. No GraphQL.
  • --org/--user map to the REST scope=org param; owner filtering for --user is client-side.
  • loops usage joins trainer deployments with standalone samplers client-side, matching paired samplers by id before the owner filter so a filtered-out trainer's sampler is not counted as standalone. The summary totals every row, including ones the table hides.
  • loops checkpoint files walks every page and streams records via NewJSONArrayWriter, since the URLs are short-lived. Dropping the response envelope also drops total_count from its JSON.
  • Log-flag validation is split: validateLogFlags holds the standalone checks and every logs command calls it before resolving the entity it fetches logs for, so a flag mistake no longer costs an API call or hides behind a lookup error. runLogsCommand keeps only the checks needing the window resolved against now. This also fixes model environment logs --model-name, which had the same ordering.
  • go.mod pins a baseten-go branch pseudo-version pending Update API specs, regenerate clients, and fix nullability baseten-go#19; re-pointed once that merges.

🔬 Testing

  • 53 httptest-backed unit tests in internal/cmd/command.loops_test.go covering all eight commands: team-scoped creation, scope/filter/direction handling, the trainer-versus-sampler logs split, tail termination on trainer status, the usage pairing/summary/hidden-row logic, and the checkpoint-files page walk in table, JSON, and JSONL form.
  • Flag-validation tests assert the mock server received no requests at all, so an ordering regression fails.
  • Full internal/cmd package passes.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds a new baseten loops command surface (run, usage, checkpoint) backed by the generated baseten-go management client, and refactors the shared logs flow to validate flags before performing entity lookups/network calls. This fits into the CLI by expanding the command tree and reusing existing output/log streaming infrastructure (tables/JSON/JSONL, shared logs tail/paging).

Changes:

  • Add baseten loops commands: run create/list/describe/deactivate/logs, usage, checkpoint list/files (+ extensive httptest-backed unit tests).
  • Refactor shared logs handling: introduce validateLogFlags, generalize tail gating via tailStatus, and update model/environment logs + model push to use it.
  • Update environment deployment nil-handling and bump github.com/basetenlabs/baseten-go pseudo-version.

Reviewed changes

Copilot reviewed 11 out of 12 changed files in this pull request and generated no comments.

Show a summary per file
File Description
internal/ssh/connect.go Fix nil-handling when resolving an environment’s current deployment.
internal/cmd/command.model_push.go Update tail result handling to use the new tailStatus accessor.
internal/cmd/command.model_environment.go Make environment list/describe resilient to CurrentDeployment == nil.
internal/cmd/command.model_environment_logs.go Validate log flags before lookups; adapt tail status fetcher to tailStatus.
internal/cmd/command.model_environment_logs_test.go Add regression test ensuring flag validation happens before model lookup.
internal/cmd/command.model_deployment_logs.go Introduce validateLogFlags, tailStatus, and refactor tail gating + final status reporting.
internal/cmd/command.loops.go Implement Loops command handlers (runs/usage/checkpoints) using management REST endpoints.
internal/cmd/command.loops_test.go Add comprehensive unit tests for all Loops commands and log/usage behaviors.
cmd/command.loops.go Add loops command tree + flag structs + JSON/text output descriptions.
cmd/command.go Register loops at the root command level.
go.mod Bump baseten-go dependency to a newer pseudo-version.
go.sum Update sums for the bumped baseten-go version.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants