feat(truss): add training checkpoint viewer to truss cli by William-Gao1 · Pull Request #2271 · basetenlabs/truss

William-Gao1 · 2026-03-06T00:42:28Z

🚀 What

Adds the truss train checkpoints list <job_id> CLI command, which lists checkpoints for a training job with sorting, multiple output formats, and an interactive file explorer.

Features:

Non-interactive mode: Table, CSV, or JSON output with sort/order options
Interactive mode: Fuzzy-searchable file explorer with directory navigation, colored directory entries, checkpoint type annotations, safetensor header inspection, syntax-highlighted file viewing, and presigned URL display

💻 How

truss/cli/train/checkpoint.py — Viewer classes (CLI table, CSV, JSON) behind an ABC, interactive file explorer using InquirerPy's FuzzyPrompt with a custom _ColoredFuzzyControl for per-choice styling, safetensor header parsing via range requests, and pygments syntax highlighting
truss/cli/train_commands.py — truss train checkpoints list click command wiring

🔬 Testing

16 unit tests covering sort ordering, all output formats, empty states, API errors, directory listing with/without checkpoint annotations, nested directories, file fetching
Manual testing of interactive explorer against live training jobs

…navigation Replace the flat file list view with a directory-aware file explorer for checkpoint drill-down. Users can navigate nested directories (e.g. rank-0/, rank-1/), view small files in a pager, and get download URLs for large files. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…irs, and safetensor viewing Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

rcano-baseten

If i wanted to use this with the cache summary (which wouldn't support file previews), how hard would that be to make things work?

rcano-baseten · 2026-03-06T21:40:16Z

truss/cli/train/common.py

        return f"{bytes} B"
+
+
+def _normalize_iso_timestamp(iso_timestamp: str) -> str:


Can you check if we do this in the codebase already?

Good catch, i'll use that instead

rcano-baseten · 2026-03-06T21:41:10Z

truss/cli/train/checkpoint.py

+
+def _show_url(url: str) -> None:
+    """Display a download URL and wait for left-arrow to dismiss."""
+    from prompt_toolkit import Application


Is there a reason we prefer inline imports here?

rcano-baseten · 2026-03-06T21:42:41Z

truss/cli/train/checkpoint.py

+        choices: list[dict] = []
+        if path_stack:
+            choices.append({"name": "..", "value": ("back", None)})
+        for d in sorted(dirs, key=lambda x: x["name"]):
+            size_str = cli_common.format_bytes_to_human_readable(d["total_size"])
+            label = f"{d['name']}/ ({size_str}, {d['file_count']} files)"
+            if d.get("checkpoint_type"):
+                ckpt_type = d["checkpoint_type"]
+                base_model = d.get("base_model", "")
+                annotation_parts = [ckpt_type]
+                if base_model:
+                    annotation_parts.append(base_model)
+                label += f" [{' \u00b7 '.join(annotation_parts)}]"
+            choices.append({"name": label, "value": ("dir", d["name"])})
+        for f in sorted(dir_files, key=lambda x: x["_rel_path"]):
+            name = f["_rel_path"].split("/")[-1]
+            size_str = cli_common.format_bytes_to_human_readable(f.get("size_bytes", 0))
+            choices.append({"name": f"{name} ({size_str})", "value": ("file", f)})
+        choices.append({"name": EXIT_OPTION, "value": ("exit", None)})


it might read more cleanly if you you put this in a subroutine

rcano-baseten · 2026-03-06T21:43:23Z

truss/cli/train/checkpoint.py

+            path_stack.append(payload)
+        elif action == "file":
+            file_name = payload.get("relative_file_name", "")
+            if file_name.endswith(".safetensors"):


there are other weights files - do they provide the same headers as safetensors?

No only safetensors provides this metadata

rcano-baseten · 2026-03-06T21:43:48Z

truss/cli/train/checkpoint.py

+        size_resp.raise_for_status()
+        header_size = struct.unpack("<Q", size_resp.content)[0]
+
+        if header_size > 10_000_000:


what is this constant? Why does it need to exist? How frequently would we expect to see something bigger?

In case the safetensors file is corrupt. It should really never be above a few kb so 10mb is more than enough

rcano-baseten · 2026-03-06T21:50:20Z

truss/cli/train/checkpoint.py

+    if metadata:
+        lines.append("Metadata:")
+        for k, v in sorted(metadata.items()):
+            lines.append(f"  {k}: {v}")
+        lines.append("")
+
+    lines.append(
+        f"Tensors: {len(tensors)}  |  Parameters: {total_params:,}  |  Size: {cli_common.format_bytes_to_human_readable(total_bytes)}"
+    )
+    lines.append("")


it might be better to build an object here and then define a serialization method on the class (e.g. str)

rcano-baseten · 2026-03-06T21:54:11Z

truss/cli/train/checkpoint.py

+
+def _highlight_content(content: str, file_name: str) -> str:
+    """Apply syntax highlighting with ANSI escape codes if a lexer is available."""
+    try:


why do we need try/except here?

get_lexer_for_filename can raise class not found error

rcano-baseten · 2026-03-06T21:54:56Z

truss/cli/train/checkpoint.py

+        num_params = 1
+        for dim in shape:
+            num_params *= dim
+        size_bytes = num_params * DTYPE_SIZES.get(dtype, 0)


if 0 a sane default here? Would it be better for us to communicate that that we weren't able to calc the size?

changed to None so unknown dtypes surface as "?" in the size column

rcano-baseten · 2026-03-06T21:55:10Z

truss/cli/train/checkpoint.py

+        for dim in shape:
+            num_params *= dim
+        size_bytes = num_params * DTYPE_SIZES.get(dtype, 0)
+        tensors.append((name, dtype, shape, num_params, size_bytes))


it would be good to have a class, e.g. TensorSummary

rcano-baseten · 2026-03-06T21:56:42Z

truss/cli/train/checkpoint.py

+                pass
+
+        _open_in_pager(content, file_name)
+    except requests.RequestException as e:


if we are in a cli viewer and the fails, what does the console print look like? Does it hijack the cli viewer? Does it emit a log at the bottom? Would be good to see what this looks like (you can just console.print from a happy path to see)

consol will clear first, error will print and then prompt the user to press entrer to go back to file viewer

rcano-baseten

Some comments on scoping of the command

rcano-baseten · 2026-03-09T21:27:52Z

truss/cli/train/checkpoint.py

might be worth scoping this to checkpoint_viewer.py instead of checkpoint.py

rcano-baseten · 2026-03-09T21:30:23Z

truss/cli/train_commands.py

    )


+@train.group(name="checkpoints")


we're putting this at the highest level e.g. truss checkpoints - what's the reason for not putting this in truss train? I think without the train , job_id becomes unintuitive

rcano-baseten · 2026-03-09T21:30:36Z

truss/cli/train_commands.py

+
+
+@checkpoints.command(name="list")
+@click.argument("job_id", type=str)


is job_id or job-id more canonical?

why force a job id? Why not let them list all checkpoints?

rcano-baseten · 2026-03-09T21:33:23Z

truss/cli/train_commands.py

+    default=checkpoint_mod.SORT_BY_CREATED,
+    help="Sort checkpoints by checkpoint-id, size, created date, or type.",
+)
+@click.option(


do we want them to be able to look at a specific checkpoint? e.g. --checkpoint-name

William-Gao1 requested review from brollb, ervinwang-jpg and rcano-baseten March 6, 2026 00:42

William-Gao1 and others added 2 commits March 6, 2026 11:46

feat(truss): Enhance checkpoint explorer with fuzzy search, colored d…

fc31db1

…irs, and safetensor viewing Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

William-Gao1 force-pushed the will/train-checkpoint-cli branch from 5bf0cef to 5dea218 Compare March 6, 2026 19:46

rcano-baseten reviewed Mar 9, 2026

View reviewed changes

refactor(truss): remove circular import and rename checkpoints command

26dee61

William-Gao1 force-pushed the will/train-checkpoint-cli branch from 5dea218 to 26dee61 Compare March 9, 2026 21:14

rcano-baseten reviewed Mar 9, 2026

View reviewed changes

		return f"{bytes} B"


		def _normalize_iso_timestamp(iso_timestamp: str) -> str:



		@checkpoints.command(name="list")
		@click.argument("job_id", type=str)

		)


		@train.group(name="checkpoints")

Conversation

William-Gao1 commented Mar 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🚀 What

💻 How

🔬 Testing

Uh oh!

rcano-baseten left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rcano-baseten left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

William-Gao1 commented Mar 6, 2026 •

edited

Loading