diff --git a/src/together/lib/cli/api/beta/models/download.py b/src/together/lib/cli/api/beta/models/download.py index 28a6ac44b..49e690843 100644 --- a/src/together/lib/cli/api/beta/models/download.py +++ b/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 @@ def _parse_object_and_revision(model_id: str, revision: str | None) -> tuple[str 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 [] @@ -338,7 +353,7 @@ async def download( 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 @@ async def download( ) 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 @@ async def download( ) 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") diff --git a/tests/cli/test_beta_models.py b/tests/cli/test_beta_models.py index 3c37bc63f..0d7659edd 100644 --- a/tests/cli/test_beta_models.py +++ b/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") @@ -774,6 +775,23 @@ def test_upload_skips_existing_files( 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, diff --git a/tests/cli/test_command_telemetry.py b/tests/cli/test_command_telemetry.py index 83e990112..cb0bb04eb 100644 --- a/tests/cli/test_command_telemetry.py +++ b/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]]],