diff --git a/src/together/lib/cli/api/fine_tuning/list_checkpoints.py b/src/together/lib/cli/api/fine_tuning/list_checkpoints.py index fbca957e..e94555c6 100644 --- a/src/together/lib/cli/api/fine_tuning/list_checkpoints.py +++ b/src/together/lib/cli/api/fine_tuning/list_checkpoints.py @@ -7,6 +7,12 @@ from together.lib.cli.components.list import ListTable +def _format_registry_artifact(object_id: str | None, revision_id: str | None) -> str: + if object_id and revision_id: + return f"{object_id}@{revision_id}" + return object_id or revision_id or "" + + async def list_checkpoints( fine_tune_id: str, *, @@ -25,19 +31,34 @@ async def list_checkpoints( title="Checkpoints", empty_message=f"No checkpoints found for job {fine_tune_id}", ) - table.add_column("ID") + table.add_column("Download ID") table.add_column("Timestamp") + table.add_column("Registry Artifact", ratio=2) table.add_primary_column("Type") + registry_artifacts: list[tuple[str, str]] = [] for checkpoint in checkpoints.data: name = ( f"{fine_tune_id}:{checkpoint.step}" if "intermediate" in checkpoint.checkpoint_type.lower() else fine_tune_id ) - table.add_row(name, format_timestamp(checkpoint.created_at), checkpoint.checkpoint_type) + registry_artifact = _format_registry_artifact(checkpoint.object_id, checkpoint.object_revision_id) + if registry_artifact: + registry_artifacts.append((name, registry_artifact)) + table.add_row( + name, + format_timestamp(checkpoint.created_at), + registry_artifact, + checkpoint.checkpoint_type, + ) console.print(table) + if registry_artifacts: + console.print("\n[dim]Registry artifacts:[/dim]") + for name, registry_artifact in registry_artifacts: + console.print(f" [dim]{name}:[/dim] {registry_artifact}") + if checkpoints.data: console.print( "\n[bold dim]To download a checkpoint, use `together fine-tuning download \\[checkpoint-id]`[/bold dim]" diff --git a/tests/cli/test_fine_tuning.py b/tests/cli/test_fine_tuning.py index 5a17a6fc..6653118d 100644 --- a/tests/cli/test_fine_tuning.py +++ b/tests/cli/test_fine_tuning.py @@ -67,7 +67,10 @@ _FT_CHECKPOINT = { "checkpoint_type": "intermediate", + "checkpoint": "model", "created_at": "2024-01-01T00:00:00Z", + "object_id": "ml-checkpoint", + "object_revision_id": "rv-checkpoint", "path": "/p", "step": 5, } @@ -348,6 +351,8 @@ def test_list_checkpoints_table(self, respx_mock: MockRouter, cli_runner: CliRun result = cli_runner.invoke(["fine-tuning", "list-checkpoints", "ft-1"]) assert result.exit_code == 0 assert "ft-1:5" in result.output + assert "Registry artifacts" in result.output + assert "ml-checkpoint@rv-checkpoint" in result.output assert "intermediate" in result.output @pytest.mark.respx(base_url=base_url)