truss 0.10.0rc1__py3-none-any.whl → 0.60.0__py3-none-any.whl

truss/remote/baseten/service.py

@@ -1,12 +1,91 @@
-from typing import Dict, Optional
-
+import enum
+import time
+import urllib.parse
+import warnings
+from typing import Any, Dict, Iterator, NamedTuple, Optional
+
+import requests
+from tenacity import retry, stop_after_delay, wait_fixed
+from truss.base.errors import RemoteNetworkError
+from truss.remote.baseten.api import BasetenApi
 from truss.remote.baseten.auth import AuthService
 from truss.remote.truss_remote import TrussService
-from truss.truss_handle import TrussHandle
+from truss.truss_handle.truss_handle import TrussHandle
+
+# "classes created inside an enum will not become a member" -> intended here anyway.
+warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*enum.*")
 
 DEFAULT_STREAM_ENCODING = "utf-8"
 
 
+def _add_model_subdomain(rest_api_url: str, model_subdomain: str) -> str:
+    """E.g. `https://api.baseten.co` -> `https://{model_subdomain}.api.baseten.co`"""
+    parsed_url = urllib.parse.urlparse(rest_api_url)
+    new_netloc = f"{model_subdomain}.{parsed_url.netloc}"
+    model_url = parsed_url._replace(netloc=new_netloc)
+    return str(urllib.parse.urlunparse(model_url))
+
+
+class URLConfig(enum.Enum):
+    class Data(NamedTuple):
+        prefix: str
+        invoke_endpoint: str
+        app_endpoint: str
+
+    MODEL = Data("model", "predict", "models")
+    CHAIN = Data("chain", "run_remote", "chains")
+
+    @staticmethod
+    def invocation_url(
+        api_url: str,  # E.g. https://api.baseten.co
+        config: "URLConfig",
+        entity_id: str,
+        entity_version_id: str,
+        is_draft,
+    ) -> str:
+        """Get the URL for the predict/run_remote endpoint."""
+        # E.g. `https://api.baseten.co` -> `https://model-{model_id}.api.baseten.co`
+        url = _add_model_subdomain(api_url, f"{config.value.prefix}-{entity_id}")
+        if is_draft:
+            # "https://model-{model_id}.api.baseten.co/development".
+            url = f"{url}/development/{config.value.invoke_endpoint}"
+        else:
+            # "https://model-{model_id}.api.baseten.co/deployment/{deployment_id}".
+            url = f"{url}/deployment/{entity_version_id}/{config.value.invoke_endpoint}"
+        return url
+
+    @staticmethod
+    def status_page_url(
+        app_url: str,  # E.g. https://app.baseten.co/
+        config: "URLConfig",
+        entity_id: str,
+    ) -> str:
+        return f"{app_url}/{config.value.app_endpoint}/{entity_id}/overview"
+
+    @staticmethod
+    def model_logs_url(
+        app_url: str,  # E.g. https://app.baseten.co/
+        model_id: str,
+        model_version_id: str,
+    ) -> str:
+        return (
+            f"{app_url}/{URLConfig.MODEL.value.app_endpoint}/{model_id}/logs/"
+            f"{model_version_id}"
+        )
+
+    @staticmethod
+    def chainlet_logs_url(
+        app_url: str,  # E.g. https://app.baseten.co/
+        chain_id: str,
+        chain_deployment_id: str,
+        chainlet_id: str,
+    ) -> str:
+        return (
+            f"{app_url}/{URLConfig.CHAIN.value.app_endpoint}/{chain_id}/logs/"
+            f"{chain_deployment_id}/{chainlet_id}"
+        )
+
+
 class BasetenService(TrussService):
     def __init__(
         self,
@@ -15,12 +94,14 @@ class BasetenService(TrussService):
         is_draft: bool,
         api_key: str,
         service_url: str,
+        api: BasetenApi,
         truss_handle: Optional[TrussHandle] = None,
     ):
         super().__init__(is_draft=is_draft, service_url=service_url)
         self._model_id = model_id
         self._model_version_id = model_version_id
         self._auth_service = AuthService(api_key=api_key)
+        self._api = api
         self._truss_handle = truss_handle
 
     def is_live(self) -> bool:
@@ -41,12 +122,9 @@ class BasetenService(TrussService):
     def invocation_url(self) -> str:
         return f"{self._service_url}/predict"
 
-    def predict(
-        self,
-        model_request_body: Dict,
-    ):
+    def predict(self, model_request_body: Dict) -> Any:
         response = self._send_request(
-            self.invocation_url, "POST", data=model_request_body, stream=True
+            self.predict_url, "POST", data=model_request_body, stream=True
         )
 
         if response.headers.get("transfer-encoding") == "chunked":
@@ -73,10 +151,39 @@ class BasetenService(TrussService):
             # will be a json with an `error` key.
             return parsed_response
 
-        return response.json()["model_output"]
+        return response.json()
 
     def authenticate(self) -> dict:
         return self._auth_service.authenticate().header()
 
-    def logs_url(self, base_url: str) -> str:
-        return f"{base_url}/models/{self._model_id}/logs/{self._model_version_id}"
+    @property
+    def logs_url(self) -> str:
+        return URLConfig.model_logs_url(
+            self._api.app_url, self.model_id, self.model_version_id
+        )
+
+    @property
+    def predict_url(self) -> str:
+        return URLConfig.invocation_url(
+            self._api.rest_api_url,
+            URLConfig.MODEL,
+            self.model_id,
+            self._model_version_id,
+            self.is_draft,
+        )
+
+    @retry(stop=stop_after_delay(60), wait=wait_fixed(1), reraise=True)
+    def _fetch_deployment(self) -> Any:
+        return self._api.get_deployment(self._model_id, self._model_version_id)
+
+    def poll_deployment_status(self, sleep_secs: int = 1) -> Iterator[str]:
+        """
+        Wait for the service to be deployed.
+        """
+        while True:
+            time.sleep(sleep_secs)
+            try:
+                deployment = self._fetch_deployment()
+                yield deployment["status"]
+            except requests.exceptions.RequestException:
+                raise RemoteNetworkError("Could not reach backend.")

truss/remote/baseten/utils/status.py

@@ -0,0 +1,29 @@
+STATUS_TO_DISPLAYABLE = {
+    "BUILDING_MODEL": "BUILDING",
+    "DEPLOYING_MODEL": "DEPLOYING",
+    "MODEL_DEPLOY_FAILED": "DEPLOY_FAILED",
+    "MODEL_LOADING": "LOADING_MODEL",
+    "MODEL_READY": "ACTIVE",
+    "MODEL_UNHEALTHY": "UNHEALTHY",
+    "BUILDING_MODEL_FAILED": "BUILD_FAILED",
+    "BUILDING_MODEL_STOPPED": "BUILD_STOPPED",
+    "DEACTIVATING_MODEL": "DEACTIVATING",
+    "DEACTIVATED_MODEL": "INACTIVE",
+    "MODEL_DNE_ERROR": "FAILED",
+    "UPDATING": "UPDATING",
+    "MIGRATING_WORKLOAD_PLANES": "UPDATING",
+    "SCALED_TO_ZERO": "SCALED_TO_ZERO",
+    "SCALING_FROM_ZERO": "WAKING_UP",
+}
+
+
+def get_displayable_status(status: str) -> str:
+    """
+    TODO: Remove this method once Chains is supported in the REST API
+
+    This is used by the `truss chains deploy` command right now to
+    print the right status. Once Chains are supported by the REST API, the
+    Baseten REST API will return status strings matching the ones here, so we don't
+    need to do any mapping.
+    """
+    return STATUS_TO_DISPLAYABLE.get(status, "UNKNOWN")

truss/remote/baseten/utils/tar.py

@@ -1,9 +1,12 @@
+import contextlib
 import tarfile
 import tempfile
 from pathlib import Path
-from typing import IO, Any, Callable, List
+from typing import IO, TYPE_CHECKING, Any, Callable, List, Optional, Type
+
+if TYPE_CHECKING:
+    from rich import progress
 
-from rich.progress import Progress
 from truss.util.path import is_ignored
 
 
@@ -24,37 +27,46 @@ class ReadProgressIndicatorFileHandle:
 
 
 def create_tar_with_progress_bar(
-    source_dir: Path, ignore_patterns: List[str] = [], delete=True
+    source_dir: Path,
+    ignore_patterns: Optional[List[str]] = None,
+    delete=True,
+    progress_bar: Optional[Type["progress.Progress"]] = None,
 ):
-    # Exclude files that match the ignore_patterns
     files_to_include = [
         f
         for f in source_dir.rglob("*")
-        if f.is_file() and not is_ignored(f, ignore_patterns, source_dir)
+        if f.is_file() and not is_ignored(f, ignore_patterns or [], source_dir)
     ]
 
     total_size = sum(f.stat().st_size for f in files_to_include)
     temp_file = tempfile.NamedTemporaryFile(suffix=".tgz", delete=delete)
 
-    with tarfile.open(temp_file.name, "w:") as tar:
-
-        progress = Progress()
-
-        task_id = progress.add_task("[cyan]Compressing...", total=total_size)
+    progress_context = (
+        progress_bar(transient=True) if progress_bar else contextlib.nullcontext()
+    )
+    # Trailing spaces are to align with `multipart_upload_boto3` message.
+    task_id = (
+        progress_context.add_task("[cyan]Packing Truss  ", total=total_size)
+        if not isinstance(progress_context, contextlib.nullcontext)
+        else None
+    )
 
-        with progress:
+    def file_read_progress_callback(bytes_read: int):
+        if not isinstance(progress_context, contextlib.nullcontext):
+            assert task_id is not None
+            progress_context.update(task_id, advance=bytes_read)
 
-            def file_read_progress_callback(bytes_read: int):
-                progress.update(task_id, advance=bytes_read)
-
-            for file_path in files_to_include:
-                arcname = str(file_path.relative_to(source_dir))
-                with file_path.open("rb") as file_obj:
-                    file_obj_with_progress = ReadProgressIndicatorFileHandle(
+    with tarfile.open(temp_file.name, "w:") as tar, progress_context:
+        for file_path in files_to_include:
+            arcname = str(file_path.relative_to(source_dir))
+            with file_path.open("rb") as file_obj:
+                file_obj_with_progress = (
+                    ReadProgressIndicatorFileHandle(
                         file_obj, file_read_progress_callback
                     )
-                    tarinfo = tar.gettarinfo(name=str(file_path), arcname=arcname)
-                    tar.addfile(
-                        tarinfo=tarinfo, fileobj=file_obj_with_progress  # type: ignore[arg-type]
-                    )
+                    if progress_bar
+                    else file_obj
+                )
+                tarinfo = tar.gettarinfo(name=str(file_path), arcname=arcname)
+                tar.addfile(tarinfo=tarinfo, fileobj=file_obj_with_progress)  # type: ignore[arg-type]  # `ReadProgressIndicatorFileHandle` implements `IO[bytes]`.
     return temp_file

truss/remote/baseten/utils/transfer.py

@@ -1,37 +1,50 @@
 import base64
+import contextlib
 import json
 import os
+from typing import TYPE_CHECKING, Optional, Type
 
 import boto3
 from boto3.s3.transfer import TransferConfig
-from rich.progress import Progress
+from truss.util.env_vars import override_env_vars
+
+if TYPE_CHECKING:
+    from rich import progress
 
 
 def base64_encoded_json_str(obj):
     return base64.b64encode(str.encode(json.dumps(obj))).decode("utf-8")
 
 
-def multipart_upload_boto3(file_path, bucket_name, key, credentials):
-    s3_resource = boto3.resource("s3", **credentials)
-    filesize = os.stat(file_path).st_size
-
-    # Create a new progress bar
-    progress = Progress()
-
-    # Add a new task to the progress bar
-    task_id = progress.add_task("[cyan]Uploading...", total=filesize)
-
-    with progress:
+def multipart_upload_boto3(
+    file_path,
+    bucket_name: str,
+    key: str,
+    credentials: dict,
+    progress_bar: Optional[Type["progress.Progress"]],
+) -> None:
+    # In the CLI flow, ignore any local ~/.aws/config files,
+    # which can interfere with uploading the Truss to S3.
+    with override_env_vars({"AWS_CONFIG_FILE": ""}):
+        s3_resource = boto3.resource("s3", **credentials)
+        filesize = os.stat(file_path).st_size
+
+        progress_context = (
+            progress_bar(transient=True) if progress_bar else contextlib.nullcontext()
+        )
+        task_id = (
+            progress_context.add_task("[cyan]Uploading Truss", total=filesize)
+            if not isinstance(progress_context, contextlib.nullcontext)
+            else None
+        )
 
         def callback(bytes_transferred):
-            # Update the progress bar
-            progress.update(task_id, advance=bytes_transferred)
-
-        s3_resource.Object(bucket_name, key).upload_file(
-            file_path,
-            Config=TransferConfig(
-                max_concurrency=10,
-                use_threads=True,
-            ),
-            Callback=callback,
-        )
+            if progress_bar:
+                progress_context.update(task_id, advance=bytes_transferred)
+
+        with progress_context:
+            s3_resource.Object(bucket_name, key).upload_file(
+                file_path,
+                Config=TransferConfig(max_concurrency=10, use_threads=True),
+                Callback=callback,
+            )

truss/remote/remote_factory.py

@@ -1,5 +1,14 @@
 import inspect
-from configparser import DEFAULTSECT, SafeConfigParser
+import os
+
+try:
+    from configparser import DEFAULTSECT, ConfigParser  # type: ignore
+except ImportError:
+    # We need to do this for old python.
+    from configparser import DEFAULTSECT
+    from configparser import SafeConfigParser as ConfigParser
+
+
 from functools import partial
 from operator import is_not
 from pathlib import Path
@@ -8,16 +17,16 @@ from typing import Dict, List, Type
 from truss.remote.baseten import BasetenRemote
 from truss.remote.truss_remote import RemoteConfig, TrussRemote
 
-USER_TRUSSRC_PATH = Path("~/.trussrc").expanduser()
+USER_TRUSSRC_PATH = Path(os.environ.get("USER_TRUSSRC_PATH", "~/.trussrc")).expanduser()
 
 
-def load_config() -> SafeConfigParser:
-    config = SafeConfigParser()
+def load_config() -> ConfigParser:
+    config = ConfigParser()
     config.read(USER_TRUSSRC_PATH)
     return config
 
 
-def update_config(config: SafeConfigParser):
+def update_config(config: ConfigParser):
     with open(USER_TRUSSRC_PATH, "w") as configfile:
         config.write(configfile)
 

truss/remote/truss_remote.py

@@ -1,16 +1,31 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional
 
 import requests
-from truss.truss_handle import TrussHandle
+
+if TYPE_CHECKING:
+    from rich import console as rich_console
+from truss.truss_handle.truss_handle import TrussHandle
+
+
+class RemoteUser:
+    """Class to hold information about the remote user"""
+
+    workspace_name: str
+    user_email: str
+
+    def __init__(self, workspace_name: str, user_email: str):
+        self.workspace_name = workspace_name
+        self.user_email = user_email
 
 
 class TrussService(ABC):
     """
     Define the abstract base class for a TrussService.
 
-    A TrussService interacts with a service at a given URL and can either be in draft or non-dreaf mode.
+    A TrussService interacts with a service at a given URL and can either be in
+    draft or non-draft mode.
 
     Attributes:
         _service_url: The URL of the service.
@@ -23,7 +38,7 @@ class TrussService(ABC):
 
     """
 
-    def __init__(self, service_url: str, is_draft: bool, **kwargs):
+    def __init__(self, service_url: str, is_draft: bool, **kwargs) -> None:
         self._service_url = service_url
         self._is_draft = is_draft
 
@@ -34,7 +49,7 @@ class TrussService(ABC):
         headers: Optional[Dict] = None,
         data: Optional[Dict] = None,
         stream: Optional[bool] = False,
-    ):
+    ) -> Any:
         """
         Send a HTTP request.
 
@@ -75,7 +90,7 @@ class TrussService(ABC):
         return response
 
     @property
-    def is_draft(self):
+    def is_draft(self) -> bool:
         """
         Check if the service is in draft mode.
 
@@ -89,7 +104,8 @@ class TrussService(ABC):
         """
         Check if the service is live.
 
-        Sends a GET request to the root of the service and returns whether it is successful.
+        Sends a GET request to the root of the service and returns whether it
+        is successful.
 
         Returns:
             A boolean indicating if the service is live.
@@ -103,7 +119,8 @@ class TrussService(ABC):
         """
         Check if the service is ready.
 
-        Sends a GET request to the model path of the service and returns whether it is successful.
+        Sends a GET request to the model path of the service and returns whether it
+        is successful.
 
         Returns:
             A boolean indicating if the service is ready.
@@ -112,21 +129,20 @@ class TrussService(ABC):
         response = self._send_request(readiness_url, "GET", {})
         return response.status_code == 200
 
-    def predict(self, model_request_body: Dict):
+    def predict(self, model_request_body: Dict) -> Any:
         """
         Send a prediction request to the service.
 
         Args:
-            model_request_body: A dictionary representing the body of the prediction request.
+            model_request_body: A dictionary representing the body of the
+            prediction request.
 
         Returns:
             A Response object resulting from the prediction request.
         """
-        invocation_url = f"{self._service_url}/v1/models/model:predict"
-        response = self._send_request(invocation_url, "POST", data=model_request_body)
-        return response
+        return self._send_request(self.predict_url, "POST", data=model_request_body)
 
-    def patch(self):
+    def patch(self) -> None:
         """
         Patch the service. TrussServices in draft mode can be patched.
         """
@@ -137,25 +153,42 @@ class TrussService(ABC):
         """
         Authenticate to the service.
 
-        This method should be implemented in subclasses and return a dictionary of headers
-        to include in requests to the service with authentication information.
+        This method should be implemented in subclasses and return a dictionary
+        of headers to include in requests to the service with authentication
+        information.
         """
         return {}
 
+    @property
     @abstractmethod
-    def logs_url(self, base_url: str) -> str:
+    def logs_url(self) -> str:
         """
         Get the URL for the service logs.
         """
         pass
 
+    @property
+    @abstractmethod
+    def predict_url(self) -> str:
+        """
+        Get the URL for the prediction endpoint.
+        """
+        pass
+
+    @abstractmethod
+    def poll_deployment_status(self, sleep_secs: int = 1) -> Iterator[str]:
+        """
+        Poll for a deployment status.
+        """
+        pass
+
 
 class TrussRemote(ABC):
     """
     Define the abstract base class for a remote Truss service.
 
-    A remote Truss service is a service that can push a TrussHandle to a remote location.
-    The `push` and `authenticate` methods should be implemented in subclasses.
+    A remote Truss service is a service that can push a TrussHandle to a remote
+    location. The `push` and `authenticate` methods should be implemented in subclasses.
 
     Attributes:
         _remote_url: The URL of the remote service.
@@ -166,11 +199,15 @@ class TrussRemote(ABC):
 
     """
 
-    def __init__(self, remote_url: str, **kwargs):
+    def __init__(self, remote_url: str) -> None:
         self._remote_url = remote_url
 
+    @property
+    def remote_url(self) -> str:
+        return self._remote_url
+
     @abstractmethod
-    def push(self, truss_handle: TrussHandle, **kwargs):
+    def push(self, truss_handle: TrussHandle, **kwargs) -> TrussService:
         """
         Push a TrussHandle to the remote service.
 
@@ -182,24 +219,19 @@ class TrussRemote(ABC):
             **kwargs: Additional keyword arguments for the push operation.
 
         """
-        pass
 
     @abstractmethod
-    def authenticate(self, **kwargs):
+    def whoami(self) -> RemoteUser:
         """
-        Authenticate the user to push to the remote service.
+        Returns account information for the current user.
+
+        This method should be implemented in subclasses and return a RemoteUser.
 
-        This method should be implemented in subclasses. It should check whether
-        the user has valid authentication credentials to push to the remote service.
-        If not, it should raise an exception.
 
-        Args:
-            **kwargs: Additional keyword arguments for the authentication operation.
         """
-        pass
 
     @abstractmethod
-    def get_service(self, **kwargs):
+    def get_service(self, **kwargs) -> TrussService:
         """
         Get a TrussService object for interacting with the remote service.
 
@@ -211,20 +243,15 @@ class TrussRemote(ABC):
             **kwargs: Keyword arguments for the get_service operation.
 
         """
-        pass
 
     @abstractmethod
-    def get_remote_logs_url(self, service: TrussService) -> str:
-        """
-        Get the URL for the remote service logs.
-
-        Args:
-            service: The TrussService object for interacting with the remote service.
-        """
-        pass
-
-    @abstractmethod
-    def sync_truss_to_dev_version_by_name(self, model_name: str, target_directory: str):
+    def sync_truss_to_dev_version_by_name(
+        self,
+        model_name: str,
+        target_directory: str,
+        console: "rich_console.Console",
+        error_console: "rich_console.Console",
+    ) -> None:
         """
         This method watches for changes to files in the `target_directory`,
         and syncs them to the development version of the model, identified
@@ -233,10 +260,10 @@ class TrussRemote(ABC):
         Args:
             model_name: The name of the model to sync to the dev version
             target_directory: The directory to sync the model to
+            console: For printing informative output.
+            error_console: For printing errors.
         """
 
-        pass
-
 
 @dataclass
 class RemoteConfig: