diff --git a/src/together/lib/cli/api/beta/models/remote_uploads/_utils.py b/src/together/lib/cli/api/beta/models/remote_uploads/_utils.py index 233055c66..3f53f523c 100644 --- a/src/together/lib/cli/api/beta/models/remote_uploads/_utils.py +++ b/src/together/lib/cli/api/beta/models/remote_uploads/_utils.py @@ -1,14 +1,67 @@ from __future__ import annotations +import re from typing import TYPE_CHECKING +from urllib.parse import ParseResult, urlparse, urlunparse from together.lib.cli.utils._console import console +# Matches hermes ParseRemoteSource: org/model or org/model@revision. +_HF_REPO_ID = re.compile(r"^[a-zA-Z0-9_\-\.]+/[a-zA-Z0-9_\-\.]+(@[a-zA-Z0-9_\-\.]+)?$") +_HF_HOSTS = frozenset({"huggingface.co", "www.huggingface.co"}) + if TYPE_CHECKING: from together.types.beta.models.remote_upload_create_response import RemoteUploadCreateResponse from together.types.beta.models.remote_upload_retrieve_response import RemoteUploadRetrieveResponse +def apply_huggingface_revision(remote_url: str, revision: str | None) -> str: + """Pin a Hugging Face revision on a remote-upload source URL. + + CreateRemoteUpload has no revision field. Hermes reads the pin from an + ``@`` suffix on a Hugging Face repo id or URL. + """ + if not revision: + return remote_url + + if remote_url.startswith(("http://", "https://")): + parsed = urlparse(remote_url) + if parsed.hostname in _HF_HOSTS: + return _pin_hf_url_revision(remote_url, parsed, revision) + raise ValueError("--revision is only supported for Hugging Face sources") + + if _HF_REPO_ID.match(remote_url): + return _pin_repo_id_revision(remote_url, revision) + + raise ValueError("--revision is only supported for Hugging Face sources") + + +def _pin_repo_id_revision(remote_url: str, revision: str) -> str: + _, separator, embedded = remote_url.partition("@") + if separator: + if embedded != revision: + raise ValueError("conflicting revisions from --from (@…) and --revision") + return remote_url + return f"{remote_url}@{revision}" + + +def _pin_hf_url_revision(remote_url: str, parsed: ParseResult, revision: str) -> str: + path = parsed.path.strip("/") + parts = path.split("/", 2) + if len(parts) < 2 or not parts[0] or not parts[1]: + raise ValueError(f"invalid Hugging Face URL (expected https://huggingface.co/org/model): {remote_url}") + _, separator, embedded = f"{parts[0]}/{parts[1]}".partition("@") + if separator: + if embedded != revision: + raise ValueError("conflicting revisions from --from (@…) and --revision") + return remote_url + parts[1] = f"{parts[1]}@{revision}" + new_path = "/" + "/".join(parts) + if parsed.path.endswith("/") and not new_path.endswith("/"): + new_path += "/" + return urlunparse(parsed._replace(path=new_path)) + + def format_status(status: str | None) -> str: if not status: return "" diff --git a/src/together/lib/cli/api/beta/models/remote_uploads/create.py b/src/together/lib/cli/api/beta/models/remote_uploads/create.py index 9e713309e..b18c373d1 100644 --- a/src/together/lib/cli/api/beta/models/remote_uploads/create.py +++ b/src/together/lib/cli/api/beta/models/remote_uploads/create.py @@ -11,7 +11,10 @@ from together.lib.cli.utils._console import console from together.lib.cli.components.loader import show_loading_status from together.lib.cli.utils._assert_explicit_project_id import assert_explicit_project_id -from together.lib.cli.api.beta.models.remote_uploads._utils import print_remote_upload_detail +from together.lib.cli.api.beta.models.remote_uploads._utils import ( + apply_huggingface_revision, + print_remote_upload_detail, +) async def create( @@ -19,8 +22,16 @@ async def create( *, remote_url: Annotated[ str, - Parameter(name="--from", help="Hugging Face repository URL or presigned S3/GCS archive URL"), + Parameter( + name="--from", + help="Hugging Face repository URL, org/model, or presigned S3/GCS archive URL. " + "Pin an HF revision with @ or --revision", + ), ], + revision: Annotated[ + Optional[str], + Parameter(help="Hugging Face commit, tag, or branch to import"), + ] = None, token: Annotated[ Optional[str], Parameter(help="Source credential for a gated or private Hugging Face repository") ] = None, @@ -30,6 +41,8 @@ async def create( await assert_explicit_project_id(config) + remote_url = apply_huggingface_revision(remote_url, revision) + response = await show_loading_status( "Starting remote upload...", config.client.beta.models.remote_uploads.create( diff --git a/src/together/lib/cli/utils/_help_examples.py b/src/together/lib/cli/utils/_help_examples.py index b8709ab5c..73d03690b 100644 --- a/src/together/lib/cli/utils/_help_examples.py +++ b/src/together/lib/cli/utils/_help_examples.py @@ -553,6 +553,10 @@ [primary]tg beta models remote-uploads create ml_xxxxxxxxxxxx \\ --from https://huggingface.co/org/model[/primary] +[dim]-[/dim] Pin a Hugging Face commit, tag, or branch: + [primary]tg beta models remote-uploads create ml_xxxxxxxxxxxx \\ + --from https://huggingface.co/org/model --revision abc123[/primary] + [dim]-[/dim] Import a gated/private HF repo: [primary]tg beta models remote-uploads create ml_xxxxxxxxxxxx \\ --from https://huggingface.co/org/private-model --token "$HF_TOKEN"[/primary] diff --git a/tests/cli/test_beta_models.py b/tests/cli/test_beta_models.py index 3c37bc63f..93fd5c9ea 100644 --- a/tests/cli/test_beta_models.py +++ b/tests/cli/test_beta_models.py @@ -885,6 +885,134 @@ def test_create_remote_upload(self, respx_mock: MockRouter, cli_runner: CliRunne } assert json.loads(result.output)["id"] == "ru_1" + @pytest.mark.respx(base_url=base_url) + def test_create_remote_upload_with_revision(self, respx_mock: MockRouter, cli_runner: CliRunner) -> None: + route = respx_mock.post("/projects/proj/models/uploads").mock( + return_value=httpx.Response(200, json=_remote_upload_body()) + ) + + result = cli_runner.invoke( + [ + "beta", + "models", + "remote-uploads", + "create", + "ml_1", + "--project", + "proj", + "--from", + "https://huggingface.co/acme/model", + "--revision", + "abc123", + "--json", + ] + ) + + assert result.exit_code == 0, result.output + assert json.loads(cast(Call, route.calls[0]).request.content.decode()) == { + "modelId": "ml_1", + "remoteUrl": "https://huggingface.co/acme/model@abc123", + } + + @pytest.mark.respx(base_url=base_url) + def test_create_remote_upload_with_repo_id_revision(self, respx_mock: MockRouter, cli_runner: CliRunner) -> None: + route = respx_mock.post("/projects/proj/models/uploads").mock( + return_value=httpx.Response(200, json=_remote_upload_body()) + ) + + result = cli_runner.invoke( + [ + "beta", + "models", + "remote-uploads", + "create", + "ml_1", + "--project", + "proj", + "--from", + "acme/model", + "--revision", + "main", + "--json", + ] + ) + + assert result.exit_code == 0, result.output + assert json.loads(cast(Call, route.calls[0]).request.content.decode())["remoteUrl"] == "acme/model@main" + + @pytest.mark.respx(base_url=base_url) + def test_create_remote_upload_matching_embedded_revision( + self, respx_mock: MockRouter, cli_runner: CliRunner + ) -> None: + route = respx_mock.post("/projects/proj/models/uploads").mock( + return_value=httpx.Response(200, json=_remote_upload_body()) + ) + + result = cli_runner.invoke( + [ + "beta", + "models", + "remote-uploads", + "create", + "ml_1", + "--project", + "proj", + "--from", + "https://huggingface.co/acme/model@abc123", + "--revision", + "abc123", + "--json", + ] + ) + + assert result.exit_code == 0, result.output + assert ( + json.loads(cast(Call, route.calls[0]).request.content.decode())["remoteUrl"] + == "https://huggingface.co/acme/model@abc123" + ) + + def test_create_remote_upload_rejects_conflicting_revision(self, cli_runner: CliRunner) -> None: + result = cli_runner.invoke( + [ + "beta", + "models", + "remote-uploads", + "create", + "ml_1", + "--project", + "proj", + "--from", + "https://huggingface.co/acme/model@rev-a", + "--revision", + "rev-b", + "--json", + ] + ) + + assert result.exit_code != 0 + assert "conflicting revisions" in result.output + + def test_create_remote_upload_rejects_revision_on_presigned_url(self, cli_runner: CliRunner) -> None: + result = cli_runner.invoke( + [ + "beta", + "models", + "remote-uploads", + "create", + "ml_1", + "--project", + "proj", + "--from", + "https://bucket.s3.amazonaws.com/model.tar.gz", + "--revision", + "abc123", + "--json", + ] + ) + + assert result.exit_code != 0 + assert "only supported for Hugging Face sources" in result.output + @pytest.mark.respx(base_url=base_url) def test_retrieve_remote_upload(self, respx_mock: MockRouter, cli_runner: CliRunner) -> None: respx_mock.get("/projects/proj/models/uploads/ru_1").mock(