diff --git a/src/together/lib/cli/api/models/upload.py b/src/together/lib/cli/api/models/upload.py index a92cac105..0aaf64cd7 100644 --- a/src/together/lib/cli/api/models/upload.py +++ b/src/together/lib/cli/api/models/upload.py @@ -1,6 +1,5 @@ from __future__ import annotations -import sys from typing import Any, Literal, Optional, Annotated, cast from cyclopts import Parameter @@ -8,6 +7,7 @@ from together import omit from together._utils._json import openapi_dumps +from together.lib.cli.utils._exit import CliDiagnosticExit from together.lib.cli.utils.config import CLIConfigParameter from together.lib.cli.utils._prompt import PromptParameter from together.lib.cli.utils._console import console @@ -70,7 +70,7 @@ async def upload( if cast(Any, response).data is None: console.print(f"[red]X[/red] [bold]Error[/bold]") console.print(f" [white]{escape_rich_markup(response.message)}[/white]") - sys.exit(1) + raise CliDiagnosticExit("Model upload request was rejected") console.print("[bold green]Model upload job created successfully![/bold green]") table = ListTable("Upload Job") diff --git a/tests/cli/test_models.py b/tests/cli/test_models.py index c7a2834f3..2947fc594 100644 --- a/tests/cli/test_models.py +++ b/tests/cli/test_models.py @@ -2,7 +2,7 @@ import os import json -from typing import cast +from typing import Any, cast from textwrap import dedent import httpx @@ -10,6 +10,7 @@ from respx import MockRouter from tests.cli.utils import CliRunner +from together.lib.cli._track_cli import CliTrackingEvents base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") @@ -167,6 +168,36 @@ def test_upload_json(self, respx_mock: MockRouter, cli_runner: CliRunner) -> Non out = json.loads(result.output) assert out["message"] == _UPLOAD_BODY["message"] + @pytest.mark.respx(base_url=base_url) + def test_upload_rejection_preserves_telemetry_diagnostic( + self, + respx_mock: MockRouter, + cli_runner: CliRunner, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + events: list[tuple[CliTrackingEvents, dict[str, Any]]] = [] + + def capture(event: CliTrackingEvents, payload: dict[str, Any]) -> None: + events.append((event, payload)) + + monkeypatch.setattr("together.lib.cli.track_cli", capture) + respx_mock.post("/models").mock( + return_value=httpx.Response( + 200, + json={"data": None, "message": "Private upload rejection details"}, + ) + ) + + result = cli_runner.invoke( + ["models", "upload", "--model-name", "model-123", "--model-source", "s3://model-123"] + ) + + assert result.exit_code == 1 + assert "Private upload rejection details" in result.output + failure = next(payload for event, payload in events if event is CliTrackingEvents.CommandFailed) + assert failure["error"] == "Model upload request was rejected" + assert "Private upload rejection details" not in failure["error"] + class TestModelsListInvalid: def test_list_invalid_type_choice(self, cli_runner: CliRunner) -> None: