together 1.5.17__py3-none-any.whl → 2.0.0a8__py3-none-any.whl

together/{cli → lib/cli}/api/endpoints.py

@@ -1,34 +1,30 @@
 from __future__ import annotations
 
-import json
 import sys
+import json
+from typing import Any, Dict, Literal, TypeVar, Callable, cast
 from functools import wraps
-from typing import Any, Callable, Dict, List, Literal, TypeVar, Union
 
 import click
 
-from together import Together
-from together.error import InvalidRequestError
-from together.types import DedicatedEndpoint, ListEndpoint
+from together import Together, omit
+from together.types import DedicatedEndpoint
+from together._exceptions import APIError
+from together.lib.utils.serializer import datetime_serializer
+from together.types.endpoint_list_response import Data as DedicatedEndpointListItem
 
 
-def print_endpoint(
-    endpoint: Union[DedicatedEndpoint, ListEndpoint],
-) -> None:
+def print_endpoint(endpoint: DedicatedEndpoint | DedicatedEndpointListItem) -> None:
     """Print endpoint details in a Docker-like format or JSON."""
 
     # Print header info
     click.echo(f"ID:\t\t{endpoint.id}")
     click.echo(f"Name:\t\t{endpoint.name}")
 
-    # Print type-specific fields
     if isinstance(endpoint, DedicatedEndpoint):
         click.echo(f"Display Name:\t{endpoint.display_name}")
         click.echo(f"Hardware:\t{endpoint.hardware}")
-        click.echo(
-            f"Autoscaling:\tMin={endpoint.autoscaling.min_replicas}, "
-            f"Max={endpoint.autoscaling.max_replicas}"
-        )
+        click.echo(f"Autoscaling:\tMin={endpoint.autoscaling.min_replicas}, Max={endpoint.autoscaling.max_replicas}")
 
     click.echo(f"Model:\t\t{endpoint.model}")
     click.echo(f"Type:\t\t{endpoint.type}")
@@ -40,15 +36,10 @@ def print_endpoint(
 F = TypeVar("F", bound=Callable[..., Any])
 
 
-def print_api_error(
-    e: InvalidRequestError,
-) -> None:
-    error_details = e.api_response.message
+def print_api_error(e: APIError) -> None:
+    error_details = cast(Dict[str, Any], e.body)["error"]["message"]
 
-    if error_details and (
-        "credentials" in error_details.lower()
-        or "authentication" in error_details.lower()
-    ):
+    if error_details and ("credentials" in error_details.lower() or "authentication" in error_details.lower()):
         click.echo("Error: Invalid API key or authentication failed", err=True)
     else:
         click.echo(f"Error: {error_details}", err=True)
@@ -61,7 +52,7 @@ def handle_api_errors(f: F) -> F:
     def wrapper(*args: Any, **kwargs: Any) -> Any:
         try:
             return f(*args, **kwargs)
-        except InvalidRequestError as e:
+        except APIError as e:
             print_api_error(e)
             sys.exit(1)
         except Exception as e:
@@ -133,8 +124,11 @@ def endpoints(ctx: click.Context) -> None:
     help="Number of minutes of inactivity after which the endpoint will be automatically stopped. Set to 0 to disable.",
 )
 @click.option(
-    "--wait",
-    is_flag=True,
+    "--availability-zone",
+    help="Start endpoint in specified availability zone (e.g., us-central-4b)",
+)
+@click.option(
+    "--wait/--no-wait",
     default=True,
     help="Wait for the endpoint to be ready after creation",
 )
@@ -148,10 +142,11 @@ def create(
     gpu: str,
     gpu_count: int,
     display_name: str | None,
-    no_prompt_cache: bool,
-    no_speculative_decoding: bool,
+    no_prompt_cache: bool | None,
+    no_speculative_decoding: bool | None,
     no_auto_start: bool,
     inactive_timeout: int | None,
+    availability_zone: str | None,
     wait: bool,
 ) -> None:
     """Create a new dedicated inference endpoint."""
@@ -170,20 +165,21 @@ def create(
         response = client.endpoints.create(
             model=model,
             hardware=hardware_id,
-            min_replicas=min_replicas,
-            max_replicas=max_replicas,
-            display_name=display_name,
-            disable_prompt_cache=no_prompt_cache,
-            disable_speculative_decoding=no_speculative_decoding,
+            autoscaling={
+                "min_replicas": min_replicas,
+                "max_replicas": max_replicas,
+            },
+            display_name=display_name or omit,
+            disable_prompt_cache=no_prompt_cache or omit,
+            disable_speculative_decoding=no_speculative_decoding or omit,
             state="STOPPED" if no_auto_start else "STARTED",
             inactive_timeout=inactive_timeout,
+            extra_query={"availability_zone": availability_zone or omit},
         )
-    except InvalidRequestError as e:
+    except APIError as e:
         print_api_error(e)
-        if "check the hardware api" in str(e).lower():
-            fetch_and_print_hardware_options(
-                client=client, model=model, print_json=False, available=True
-            )
+        if "check the hardware api" in str(e.args[0]).lower() or "invalid hardware provided" in str(e.args[0]).lower():
+            fetch_and_print_hardware_options(client=client, model=model, print_json=False, available=True)
 
         sys.exit(1)
 
@@ -203,6 +199,8 @@ def create(
         click.echo("  Auto-start: disabled", err=True)
     if inactive_timeout is not None:
         click.echo(f"  Inactive timeout: {inactive_timeout} minutes", err=True)
+    if availability_zone:
+        click.echo(f"  Availability zone: {availability_zone}", err=True)
 
     click.echo(f"Endpoint created successfully, id: {response.id}", err=True)
 
@@ -210,7 +208,7 @@ def create(
         import time
 
         click.echo("Waiting for endpoint to be ready...", err=True)
-        while client.endpoints.get(response.id).state != "STARTED":
+        while client.endpoints.retrieve(response.id).state != "STARTED":
             time.sleep(1)
         click.echo("Endpoint ready", err=True)
 
@@ -225,11 +223,11 @@ def create(
 @handle_api_errors
 def get(client: Together, endpoint_id: str, json: bool) -> None:
     """Get a dedicated inference endpoint."""
-    endpoint = client.endpoints.get(endpoint_id)
+    endpoint = client.endpoints.retrieve(endpoint_id)
     if json:
         import json as json_lib
 
-        click.echo(json_lib.dumps(endpoint.model_dump(), indent=2))
+        click.echo(json_lib.dumps(endpoint.model_dump(), indent=2, default=datetime_serializer))
     else:
         print_endpoint(endpoint)
 
@@ -249,35 +247,31 @@ def hardware(client: Together, model: str | None, json: bool, available: bool) -
     fetch_and_print_hardware_options(client, model, json, available)
 
 
-def fetch_and_print_hardware_options(
-    client: Together, model: str | None, print_json: bool, available: bool
-) -> None:
+def fetch_and_print_hardware_options(client: Together, model: str | None, print_json: bool, available: bool) -> None:
     """Print hardware options for a model."""
 
     message = "Available hardware options:" if available else "All hardware options:"
     click.echo(message, err=True)
-    hardware_options = client.endpoints.list_hardware(model)
+    hardware_options = client.hardware.list(model=model or omit)
+    # hardware_options = client.endpoints.list_hardware(model)
     if available:
-        hardware_options = [
+        hardware_options.data = [
             hardware
-            for hardware in hardware_options
-            if hardware.availability is not None
-            and hardware.availability.status == "available"
+            for hardware in hardware_options.data
+            if hardware.availability is not None and hardware.availability.status == "available"
         ]
 
     if print_json:
-        json_output = [hardware.model_dump() for hardware in hardware_options]
-        click.echo(json.dumps(json_output, indent=2))
+        json_output = [hardware.model_dump() for hardware in hardware_options.data]
+        click.echo(json.dumps(json_output, default=datetime_serializer, indent=2))
     else:
-        for hardware in hardware_options:
+        for hardware in hardware_options.data:
             click.echo(f"  {hardware.id}", err=True)
 
 
 @endpoints.command()
 @click.argument("endpoint-id", required=True)
-@click.option(
-    "--wait", is_flag=True, default=True, help="Wait for the endpoint to stop"
-)
+@click.option("--wait/--no-wait", default=True, help="Wait for the endpoint to stop")
 @click.pass_obj
 @handle_api_errors
 def stop(client: Together, endpoint_id: str, wait: bool) -> None:
@@ -289,7 +283,7 @@ def stop(client: Together, endpoint_id: str, wait: bool) -> None:
         import time
 
         click.echo("Waiting for endpoint to stop...", err=True)
-        while client.endpoints.get(endpoint_id).state != "STOPPED":
+        while client.endpoints.retrieve(endpoint_id).state != "STOPPED":
             time.sleep(1)
         click.echo("Endpoint stopped", err=True)
 
@@ -298,9 +292,7 @@ def stop(client: Together, endpoint_id: str, wait: bool) -> None:
 
 @endpoints.command()
 @click.argument("endpoint-id", required=True)
-@click.option(
-    "--wait", is_flag=True, default=True, help="Wait for the endpoint to start"
-)
+@click.option("--wait/--no-wait", default=True, help="Wait for the endpoint to start")
 @click.pass_obj
 @handle_api_errors
 def start(client: Together, endpoint_id: str, wait: bool) -> None:
@@ -312,7 +304,7 @@ def start(client: Together, endpoint_id: str, wait: bool) -> None:
         import time
 
         click.echo("Waiting for endpoint to start...", err=True)
-        while client.endpoints.get(endpoint_id).state != "STARTED":
+        while client.endpoints.retrieve(endpoint_id).state != "STARTED":
             time.sleep(1)
         click.echo("Endpoint started", err=True)
 
@@ -337,13 +329,32 @@ def delete(client: Together, endpoint_id: str) -> None:
     type=click.Choice(["dedicated", "serverless"]),
     help="Filter by endpoint type",
 )
+@click.option(
+    "--mine",
+    type=click.BOOL,
+    default=None,
+    help="true (only mine), default=all",
+)
+@click.option(
+    "--usage-type",
+    type=click.Choice(["on-demand", "reserved"]),
+    help="Filter by endpoint usage type",
+)
 @click.pass_obj
 @handle_api_errors
 def list(
-    client: Together, json: bool, type: Literal["dedicated", "serverless"] | None
+    client: Together,
+    json: bool,
+    type: Literal["dedicated", "serverless"] | None,
+    usage_type: Literal["on-demand", "reserved"] | None,
+    mine: bool | None,
 ) -> None:
     """List all inference endpoints (includes both dedicated and serverless endpoints)."""
-    endpoints: List[ListEndpoint] = client.endpoints.list(type=type)
+    endpoints = client.endpoints.list(
+        type=type or omit,
+        usage_type=usage_type or omit,
+        mine=mine if mine is not None else omit,
+    )
 
     if not endpoints:
         click.echo("No dedicated endpoints found", err=True)
@@ -354,10 +365,12 @@ def list(
         import json as json_lib
 
         click.echo(
-            json_lib.dumps([endpoint.model_dump() for endpoint in endpoints], indent=2)
+            json_lib.dumps(
+                [endpoint.model_dump() for endpoint in endpoints.data], default=datetime_serializer, indent=2
+            )
         )
     else:
-        for endpoint in endpoints:
+        for endpoint in endpoints.data:
             print_endpoint(
                 endpoint,
             )
@@ -400,35 +413,55 @@ def update(
         click.echo("Error: At least one update option must be specified", err=True)
         sys.exit(1)
 
-    # If only one of min/max replicas is specified, we need both for the update
-    if (min_replicas is None) != (max_replicas is None):
-        click.echo(
-            "Error: Both --min-replicas and --max-replicas must be specified together",
-            err=True,
-        )
-        sys.exit(1)
-
     # Build kwargs for the update
     kwargs: Dict[str, Any] = {}
     if display_name is not None:
         kwargs["display_name"] = display_name
-    if min_replicas is not None and max_replicas is not None:
-        kwargs["min_replicas"] = min_replicas
-        kwargs["max_replicas"] = max_replicas
+
+    if min_replicas is not None or max_replicas is not None:
+        kwargs["autoscaling"] = {}
+        if min_replicas is not None:
+            kwargs["autoscaling"]["min_replicas"] = min_replicas
+        if max_replicas is not None:
+            kwargs["autoscaling"]["max_replicas"] = max_replicas
+
     if inactive_timeout is not None:
         kwargs["inactive_timeout"] = inactive_timeout
 
-    _response = client.endpoints.update(endpoint_id, **kwargs)
+    client.endpoints.update(endpoint_id, **kwargs)
 
     # Print what was updated
     click.echo("Updated endpoint configuration:", err=True)
     if display_name:
         click.echo(f"  Display name: {display_name}", err=True)
-    if min_replicas is not None and max_replicas is not None:
+    if min_replicas:
         click.echo(f"  Min replicas: {min_replicas}", err=True)
+    if max_replicas:
         click.echo(f"  Max replicas: {max_replicas}", err=True)
     if inactive_timeout is not None:
         click.echo(f"  Inactive timeout: {inactive_timeout} minutes", err=True)
 
     click.echo("Successfully updated endpoint", err=True)
     click.echo(endpoint_id)
+
+
+@endpoints.command()
+@click.option("--json", is_flag=True, help="Print output in JSON format")
+@click.pass_obj
+@handle_api_errors
+def availability_zones(client: Together, json: bool) -> None:
+    """List all availability zones."""
+    avzones = client.endpoints.list_avzones()
+
+    if not avzones:
+        click.echo("No availability zones found", err=True)
+        return
+
+    if json:
+        import json as json_lib
+
+        click.echo(json_lib.dumps(avzones.model_dump(), indent=2))
+    else:
+        click.echo("Available zones:", err=True)
+        for availability_zone in sorted(avzones.avzones):
+            click.echo(f"  {availability_zone}")