fix(cli): preserve model download diagnostics

src/together/lib/cli/api/beta/models/download.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import os
-import sys
 import shutil
 import asyncio
 import tempfile
@@ -15,6 +14,7 @@
 
 from together import omit
 from together._utils import path_template
+from together.lib.cli.utils._exit import CliDiagnosticExit
 from together.lib.cli.utils.config import CLIConfig, CLIConfigParameter
 from together.lib.cli.utils._console import console
 from together.lib.cli.api.beta.models.upload import (
@@ -58,6 +58,21 @@
     return model_id, revision
 
 
+def _download_failure_diagnostic(exc: ValueError) -> str:
+    message = str(exc)
+    if message.startswith("HuggingFace layout requires project/name"):
+        return "Hugging Face model download requires a project-qualified model name"
+    if message == "download response missing file path":
+        return "Model download response is missing a file path"
+    if message.startswith("insufficient disk space"):
+        return "Insufficient disk space for model download"
+    if message == "could not resolve latest revision":
+        return "Could not resolve latest model revision"
+    if message.startswith("download after retries"):
+        return "Model download failed after retries"
+    return "Model download failed"
+
+
 def _normalize_files(files: list[str] | None) -> list[str]:
     if not files:
         return []
@@ -135,7 +150,7 @@
 
 
 def _write_at(path: Path, offset: int, data: bytes) -> None:
     with path.open("r+b") as file:
         file.seek(offset)
         file.write(data)
 
@@ -338,7 +353,7 @@
         object_id, resolved_revision = _parse_object_and_revision(model_id, revision)
     except ValueError as exc:
         console.print(f"[red]Error:[/red] {exc}")
-        sys.exit(1)
+        raise CliDiagnosticExit("Invalid model download request") from exc
 
     # Ensure we use the v2 apis for this
     config.client.base_url = "https://api.together.ai/v2"
@@ -355,7 +370,7 @@
             )
         except ValueError as exc:
             console.print(f"[red]Error:[/red] {exc}")
-            sys.exit(1)
+            raise CliDiagnosticExit(_download_failure_diagnostic(exc)) from exc
         console.print_json(data=result)
         return
 
@@ -372,5 +387,5 @@
         )
     except ValueError as exc:
         console.print(f"[red]Error:[/red] {exc}")
-        sys.exit(1)
+        raise CliDiagnosticExit(_download_failure_diagnostic(exc)) from exc
     console.print("Download complete")

tests/cli/test_beta_models.py

@@ -12,6 +12,7 @@
 from respx.models import Call
 
 from tests.cli.utils import CliRunner
+from together.lib.cli.api.beta.models.download import _download_failure_diagnostic
 
 base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
 
@@ -55,7 +56,7 @@
             {
                 "path": "weights.bin",
                 "sizeBytes": "11",
                 "hash": hashlib.md5(b"hello world").hexdigest(),
             }
         ],
         "next_cursor": None,
@@ -704,7 +705,7 @@
             "files": [
                 {
                     "path": "weights.bin",
                     "hash": hashlib.md5(b"hello world").hexdigest(),
                     "numParts": 1,
                 }
             ],
@@ -717,7 +718,7 @@
             "files": [
                 {
                     "path": "weights.bin",
                     "hash": hashlib.md5(b"hello world").hexdigest(),
                     "uploadDetails": {
                         "uploadId": "upload-1",
                         "parts": [{"partNumber": 1, "hash": "etag-1"}],
@@ -766,7 +767,7 @@
             "files": [
                 {
                     "path": "weights.bin",
                     "hash": hashlib.md5(b"cached").hexdigest(),
                     "skipUpload": True,
                 }
             ],
@@ -774,6 +775,23 @@
 
 
 class TestBetaModelsDownload:
+    @pytest.mark.parametrize(
+        ("message", "diagnostic"),
+        [
+            (
+                'HuggingFace layout requires project/name, got "ml_private"',
+                "Hugging Face model download requires a project-qualified model name",
+            ),
+            ("download response missing file path", "Model download response is missing a file path"),
+            ("insufficient disk space: need 10 GB, have 1 GB", "Insufficient disk space for model download"),
+            ("could not resolve latest revision", "Could not resolve latest model revision"),
+            ("download after retries: private response", "Model download failed after retries"),
+            ("customer-specific failure", "Model download failed"),
+        ],
+    )
+    def test_download_failure_diagnostic_is_stable(self, message: str, diagnostic: str) -> None:
+        assert _download_failure_diagnostic(ValueError(message)) == diagnostic
+
     @pytest.mark.respx(base_url=base_url)
     def test_download_writes_files(
         self,
@@ -794,7 +812,7 @@
                     "files": [
                         {
                             "path": "weights.bin",
                             "hash": hashlib.md5(content).hexdigest(),
                             "sizeBytes": len(content),
                             "parts": [
                                 {

tests/cli/test_command_telemetry.py

@@ -125,6 +125,42 @@ async def test_interactive_missing_required_argument_preserves_diagnostic(
     assert failed["error"] == "Missing required argument: --model"
 
 
+@pytest.mark.usefixtures("isolated_cli_config")
+@pytest.mark.asyncio
+async def test_model_download_validation_preserves_diagnostic(
+    track_cli_capture: list[tuple[CliTrackingEvents, dict[str, Any]]],
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    from together.lib.cli import launcher
+
+    monkeypatch.setenv("TOGETHER_DISABLE_VERSION_CHECK", "1")
+
+    with pytest.raises(SystemExit) as exc_info:
+        await launcher(
+            "beta",
+            "models",
+            "download",
+            "ml_example@rev-a",
+            str(tmp_path),
+            "--revision",
+            "rev-b",
+            "--json",
+            api_key="0000000000000000000000000000000000000000",
+            project_id="project",
+        )
+
+    assert exc_info.value.code == 1
+    assert _event_kinds(track_cli_capture) == [
+        CliTrackingEvents.CommandStarted.value,
+        CliTrackingEvents.CommandFailed.value,
+    ]
+    failed = track_cli_capture[1][1]
+    assert failed["command"] == "models download"
+    assert failed["is_beta_command"] is True
+    assert failed["error"] == "Invalid model download request"
+
+
 @pytest.mark.usefixtures("isolated_cli_config")
 def test_command_system_exit_zero_emits_started_then_completed(
     track_cli_capture: list[tuple[CliTrackingEvents, dict[str, Any]]],