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

truss/remote/baseten/remote.py

@@ -1,49 +1,125 @@
+import enum
 import logging
 import re
 from pathlib import Path
-from typing import List, Optional, Tuple
+from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Type
 
-import click
-import rich
 import yaml
 from requests import ReadTimeout
+from truss.base.constants import PRODUCTION_ENVIRONMENT_NAME
+
+if TYPE_CHECKING:
+    from rich import console as rich_console
+    from rich import progress
+from truss.base import validation
+from truss.base.truss_config import ModelServer
 from truss.local.local_config_handler import LocalConfigHandler
+from truss.remote.baseten import custom_types
 from truss.remote.baseten.api import BasetenApi
 from truss.remote.baseten.auth import AuthService
 from truss.remote.baseten.core import (
+    ChainDeploymentHandleAtomic,
     ModelId,
     ModelIdentifier,
     ModelName,
     ModelVersionId,
     archive_truss,
+    create_chain_atomic,
     create_truss_service,
     exists_model,
     get_dev_version,
     get_dev_version_from_versions,
     get_model_versions,
     get_prod_version_from_versions,
+    get_truss_watch_state,
     upload_truss,
+    validate_truss_config,
 )
-from truss.remote.baseten.error import ApiError
-from truss.remote.baseten.service import BasetenService
+from truss.remote.baseten.error import ApiError, RemoteError
+from truss.remote.baseten.service import BasetenService, URLConfig
 from truss.remote.baseten.utils.transfer import base64_encoded_json_str
-from truss.remote.truss_remote import TrussRemote, TrussService
-from truss.truss_config import ModelServer
-from truss.truss_handle import TrussHandle
-from truss.util.path import is_ignored, load_trussignore_patterns
+from truss.remote.truss_remote import RemoteUser, TrussRemote
+from truss.truss_handle import build as truss_build
+from truss.truss_handle.truss_handle import TrussHandle
+from truss.util.path import is_ignored, load_trussignore_patterns_from_truss_dir
 from watchfiles import watch
 
 
+class PatchStatus(enum.Enum):
+    SUCCESS = enum.auto()
+    FAILED = enum.auto()
+    SKIPPED = enum.auto()
+
+
+class PatchResult(NamedTuple):
+    status: PatchStatus
+    message: str
+
+
+class FinalPushData(custom_types.OracleData):
+    class Config:
+        protected_namespaces = ()
+
+    is_draft: bool
+    model_id: Optional[str]
+    preserve_previous_prod_deployment: bool
+    origin: Optional[custom_types.ModelOrigin] = None
+    environment: Optional[str] = None
+    allow_truss_download: bool
+
+
 class BasetenRemote(TrussRemote):
-    def __init__(self, remote_url: str, api_key: str, **kwargs):
-        super().__init__(remote_url, **kwargs)
+    def __init__(self, remote_url: str, api_key: str):
+        super().__init__(remote_url)
         self._auth_service = AuthService(api_key=api_key)
-        self._api = BasetenApi(f"{self._remote_url}/graphql/", self._auth_service)
+        self._api = BasetenApi(remote_url, self._auth_service)
+
+    @property
+    def api(self) -> BasetenApi:
+        return self._api
+
+    def get_chainlets(
+        self, chain_deployment_id: str
+    ) -> List[custom_types.DeployedChainlet]:
+        return [
+            custom_types.DeployedChainlet(
+                name=chainlet["name"],
+                is_entrypoint=chainlet["is_entrypoint"],
+                is_draft=chainlet["oracle_version"]["is_draft"],
+                status=chainlet["oracle_version"]["current_model_deployment_status"][
+                    "status"
+                ],
+                logs_url=URLConfig.chainlet_logs_url(
+                    self.remote_url,
+                    chainlet["chain"]["id"],
+                    chain_deployment_id,
+                    chainlet["id"],
+                ),
+                oracle_predict_url=URLConfig.invocation_url(
+                    self._api.rest_api_url,
+                    URLConfig.MODEL,
+                    chainlet["oracle"]["id"],
+                    chainlet["oracle_version"]["id"],
+                    chainlet["oracle_version"]["is_draft"],
+                ),
+                oracle_name=chainlet["oracle"]["name"],
+            )
+            for chainlet in self._api.get_chainlets_by_deployment_id(
+                chain_deployment_id
+            )
+        ]
 
-    def authenticate(self):
-        return self._auth_service.validate()
+    def whoami(self) -> RemoteUser:
+        resp = self._api._post_graphql_query(
+            "query{organization{workspace_name}user{email}}"
+        )
+        workspace_name = resp["data"]["organization"]["workspace_name"]
+        user_email = resp["data"]["user"]["email"]
+        return RemoteUser(workspace_name, user_email)
 
-    def push(  # type: ignore
+    # Validate and finalize options.
+    # Upload Truss files to S3 and return S3 key.
+    def _prepare_push(
         self,
         truss_handle: TrussHandle,
         model_name: str,
@@ -51,20 +127,30 @@ class BasetenRemote(TrussRemote):
         trusted: bool = False,
         promote: bool = False,
         preserve_previous_prod_deployment: bool = False,
+        disable_truss_download: bool = False,
         deployment_name: Optional[str] = None,
-    ):
+        origin: Optional[custom_types.ModelOrigin] = None,
+        environment: Optional[str] = None,
+        progress_bar: Optional[Type["progress.Progress"]] = None,
+    ) -> FinalPushData:
         if model_name.isspace():
             raise ValueError("Model name cannot be empty")
 
-        model_id = exists_model(self._api, model_name)
+        if truss_handle.is_scattered():
+            truss_handle = TrussHandle(truss_handle.gather())
 
-        gathered_truss = TrussHandle(truss_handle.gather())
-        if gathered_truss.spec.model_server != ModelServer.TrussServer:
+        if truss_handle.spec.model_server != ModelServer.TrussServer:
             publish = True
 
         if promote:
-            # If we are promoting a model after deploy, it must be published.
-            # Draft models cannot be promoted.
+            environment = PRODUCTION_ENVIRONMENT_NAME
+
+        # If there is a target environment, it must be published.
+        # Draft models cannot be promoted.
+        if environment and not publish:
+            logging.info(
+                f"Automatically publishing model '{model_name}' based on environment setting."
+            )
             publish = True
 
         if not publish and deployment_name:
@@ -74,7 +160,8 @@ class BasetenRemote(TrussRemote):
 
         if not promote and preserve_previous_prod_deployment:
             raise ValueError(
-                "preserve-previous-production-deployment can only be used with the '--promote' option"
+                "preserve-previous-production-deployment can only be used "
+                "with the '--promote' option"
             )
 
         if deployment_name and not re.match(r"^[0-9a-zA-Z_\-\.]*$", deployment_name):
@@ -82,34 +169,146 @@ class BasetenRemote(TrussRemote):
                 "Deployment name must only contain alphanumeric, -, _ and . characters"
             )
 
+        model_id = exists_model(self._api, model_name)
+
+        if model_id is not None and disable_truss_download:
+            raise ValueError("disable-truss-download can only be used for new models")
+
+        temp_file = archive_truss(truss_handle, progress_bar)
+        s3_key = upload_truss(self._api, temp_file, progress_bar)
         encoded_config_str = base64_encoded_json_str(
-            gathered_truss._spec._config.to_dict()
+            truss_handle._spec._config.to_dict()
         )
 
-        temp_file = archive_truss(gathered_truss)
-        s3_key = upload_truss(self._api, temp_file)
+        validate_truss_config(self._api, encoded_config_str)
 
-        model_id, model_version_id = create_truss_service(
-            api=self._api,
+        return FinalPushData(
             model_name=model_name,
             s3_key=s3_key,
-            config=encoded_config_str,
+            encoded_config_str=encoded_config_str,
             is_draft=not publish,
             model_id=model_id,
             is_trusted=trusted,
+            preserve_previous_prod_deployment=preserve_previous_prod_deployment,
+            version_name=deployment_name,
+            origin=origin,
+            environment=environment,
+            allow_truss_download=not disable_truss_download,
+        )
+
+    def push(  # type: ignore
+        self,
+        truss_handle: TrussHandle,
+        model_name: str,
+        publish: bool = True,
+        trusted: bool = False,
+        promote: bool = False,
+        preserve_previous_prod_deployment: bool = False,
+        disable_truss_download: bool = False,
+        deployment_name: Optional[str] = None,
+        origin: Optional[custom_types.ModelOrigin] = None,
+        environment: Optional[str] = None,
+        progress_bar: Optional[Type["progress.Progress"]] = None,
+    ) -> BasetenService:
+        push_data = self._prepare_push(
+            truss_handle=truss_handle,
+            model_name=model_name,
+            publish=publish,
+            trusted=trusted,
             promote=promote,
             preserve_previous_prod_deployment=preserve_previous_prod_deployment,
+            disable_truss_download=disable_truss_download,
             deployment_name=deployment_name,
+            origin=origin,
+            environment=environment,
+            progress_bar=progress_bar,
+        )
+
+        # TODO(Tyron): This set of args is duplicated across
+        # many functions. We should consolidate them into a
+        # data class with standardized default values so
+        # we're not drilling these arguments everywhere.
+        model_id, model_version_id = create_truss_service(
+            api=self._api,
+            model_name=push_data.model_name,
+            s3_key=push_data.s3_key,
+            config=push_data.encoded_config_str,
+            is_draft=push_data.is_draft,
+            model_id=push_data.model_id,
+            is_trusted=push_data.is_trusted,
+            preserve_previous_prod_deployment=push_data.preserve_previous_prod_deployment,
+            allow_truss_download=push_data.allow_truss_download,
+            deployment_name=push_data.version_name,
+            origin=push_data.origin,
+            environment=push_data.environment,
         )
 
         return BasetenService(
             model_id=model_id,
             model_version_id=model_version_id,
-            is_draft=not publish,
+            is_draft=push_data.is_draft,
             api_key=self._auth_service.authenticate().value,
             service_url=f"{self._remote_url}/model_versions/{model_version_id}",
             truss_handle=truss_handle,
+            api=self._api,
+        )
+
+    def push_chain_atomic(
+        self,
+        chain_name: str,
+        entrypoint_artifact: custom_types.ChainletArtifact,
+        dependency_artifacts: List[custom_types.ChainletArtifact],
+        publish: bool = False,
+        environment: Optional[str] = None,
+        progress_bar: Optional[Type["progress.Progress"]] = None,
+    ) -> ChainDeploymentHandleAtomic:
+        # If we are promoting a model to an environment after deploy, it must be published.
+        # Draft models cannot be promoted.
+        if environment and not publish:
+            publish = True
+
+        chainlet_data: List[custom_types.ChainletDataAtomic] = []
+
+        for artifact in [entrypoint_artifact, *dependency_artifacts]:
+            truss_handle = truss_build.load(str(artifact.truss_dir))
+            model_name = truss_handle.spec.config.model_name
+
+            assert model_name and validation.is_valid_model_name(model_name)
+
+            push_data = self._prepare_push(
+                truss_handle=truss_handle,
+                model_name=model_name,
+                # Models must be trusted to use the API KEY secret.
+                trusted=True,
+                publish=publish,
+                origin=custom_types.ModelOrigin.CHAINS,
+                progress_bar=progress_bar,
+            )
+            oracle_data = custom_types.OracleData(
+                model_name=push_data.model_name,
+                s3_key=push_data.s3_key,
+                encoded_config_str=push_data.encoded_config_str,
+                is_draft=push_data.is_draft,
+                model_id=push_data.model_id,
+                is_trusted=push_data.is_trusted,
+                version_name=push_data.version_name,
+            )
+            chainlet_data.append(
+                custom_types.ChainletDataAtomic(
+                    name=artifact.display_name, oracle=oracle_data
+                )
+            )
+
+        chain_deployment_handle = create_chain_atomic(
+            api=self._api,
+            chain_name=chain_name,
+            entrypoint=chainlet_data[0],
+            dependencies=chainlet_data[1:],
+            is_draft=not publish,
+            environment=environment,
         )
+        logging.info("Successfully pushed to baseten. Chain is building and deploying.")
+        return chain_deployment_handle
 
     @staticmethod
     def _get_matching_version(model_versions: List[dict], published: bool) -> dict:
@@ -117,7 +316,7 @@ class BasetenRemote(TrussRemote):
             # Return the development model version.
             dev_version = get_dev_version_from_versions(model_versions)
             if not dev_version:
-                raise click.UsageError(
+                raise RemoteError(
                     "No development model found. Run `truss push` then try again."
                 )
             return dev_version
@@ -125,7 +324,7 @@ class BasetenRemote(TrussRemote):
         # Return the production deployment version.
         prod_version = get_prod_version_from_versions(model_versions)
         if not prod_version:
-            raise click.UsageError(
+            raise RemoteError(
                 "No production model found. Run `truss push --publish` then try again."
             )
         return prod_version
@@ -138,9 +337,7 @@ class BasetenRemote(TrussRemote):
             try:
                 model_version = api.get_model_version_by_id(model_identifier.value)
             except ApiError:
-                raise click.UsageError(
-                    f"Model version {model_identifier.value} not found."
-                )
+                raise RemoteError(f"Model version {model_identifier.value} not found.")
             model_version_id = model_version["model_version"]["id"]
             model_id = model_version["model_version"]["oracle"]["id"]
             service_url_path = f"/model_versions/{model_version_id}"
@@ -159,14 +356,15 @@ class BasetenRemote(TrussRemote):
             try:
                 model = api.get_model_by_id(model_identifier.value)
             except ApiError:
-                raise click.UsageError(f"Model {model_identifier.value} not found.")
+                raise RemoteError(f"Model {model_identifier.value} not found.")
             model_id = model["model"]["id"]
             model_version_id = model["model"]["primary_version"]["id"]
             service_url_path = f"/models/{model_id}"
         else:
             # Model identifier is of invalid type.
-            raise click.UsageError(
-                "You must either be inside of a Truss directory, or provide --model-deployment or --model options."
+            raise RemoteError(
+                "You must either be inside of a Truss directory, or provide "
+                "--model-deployment or --model options."
             )
 
         return service_url_path, model_id, model_version_id
@@ -178,12 +376,10 @@ class BasetenRemote(TrussRemote):
             raise ValueError("Baseten Service requires a model_identifier")
 
         published = kwargs.get("published", False)
-        (
-            service_url_path,
-            model_id,
-            model_version_id,
-        ) = self._get_service_url_path_and_model_ids(
-            self._api, model_identifier, published
+        (service_url_path, model_id, model_version_id) = (
+            self._get_service_url_path_and_model_ids(
+                self._api, model_identifier, published
+            )
         )
 
         return BasetenService(
@@ -192,119 +388,174 @@ class BasetenRemote(TrussRemote):
             is_draft=not published,
             api_key=self._auth_service.authenticate().value,
             service_url=f"{self._remote_url}{service_url_path}",
+            api=self._api,
         )
 
-    def get_remote_logs_url(
-        self,
-        service: TrussService,
-    ) -> str:
-        return service.logs_url(self._remote_url)
-
     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:
         # verify that development deployment exists for given model name
-        dev_version = get_dev_version(
-            self._api, model_name
-        )  # pylint: disable=protected-access
+        dev_version = get_dev_version(self._api, model_name)  # pylint: disable=protected-access
         if not dev_version:
-            raise click.UsageError(
+            raise RemoteError(
                 "No development model found. Run `truss push` then try again."
             )
 
         watch_path = Path(target_directory)
-        truss_ignore_patterns = load_trussignore_patterns()
+        truss_ignore_patterns = load_trussignore_patterns_from_truss_dir(watch_path)
 
         def watch_filter(_, path):
-            return not is_ignored(
-                Path(path),
-                truss_ignore_patterns,
-            )
+            return not is_ignored(Path(path), truss_ignore_patterns)
 
         # disable watchfiles logger
         logging.getLogger("watchfiles.main").disabled = True
 
-        rich.print(f"🚰 Attempting to sync truss at '{watch_path}' with remote")
-        self.patch(watch_path, truss_ignore_patterns)
+        console.print(f"🚰 Attempting to sync truss at '{watch_path}' with remote")
+        self.patch(watch_path, truss_ignore_patterns, console, error_console)
 
-        rich.print(f"👀 Watching for changes to truss at '{watch_path}' ...")
+        console.print(f"👀 Watching for changes to truss at '{watch_path}' ...")
         for _ in watch(watch_path, watch_filter=watch_filter, raise_interrupt=False):
-            self.patch(watch_path, truss_ignore_patterns)
+            self.patch(watch_path, truss_ignore_patterns, console, error_console)
 
-    def patch(
+    def _patch(
         self,
         watch_path: Path,
         truss_ignore_patterns: List[str],
-    ):
-        from truss.cli.console import console, error_console
-
+        console: Optional["rich_console.Console"] = None,
+    ) -> PatchResult:
         try:
             truss_handle = TrussHandle(watch_path)
-        except yaml.parser.ParserError:
-            error_console.print("Unable to parse config file")
-            return
-        except ValueError:
-            error_console.print(f"Error when reading truss from directory {watch_path}")
-            return
+        except yaml.parser.ParserError as e:
+            return PatchResult(PatchStatus.FAILED, f"Unable to parse config file. {e}")
+        except ValueError as e:
+            return PatchResult(
+                PatchStatus.FAILED,
+                f"Error when reading truss from directory {watch_path}. {e}",
+            )
+
         model_name = truss_handle.spec.config.model_name
         dev_version = get_dev_version(self._api, model_name)  # type: ignore
         if not dev_version:
-            error_console.print(
-                f"No development deployment found with model name: {model_name}"
+            return PatchResult(
+                PatchStatus.FAILED,
+                f"No development deployment found for model: {model_name}.",
             )
-            return
+
         truss_hash = dev_version.get("truss_hash", None)
         truss_signature = dev_version.get("truss_signature", None)
         if not (truss_hash and truss_signature):
-            error_console.print(
-                """Failed to inspect a running remote deployment to watch for changes.
-Ensure that there exists a running remote deployment before attempting to watch for changes
-            """
+            return PatchResult(
+                PatchStatus.FAILED,
+                (
+                    "Failed to inspect a running remote deployment to watch for "
+                    "changes.  Ensure that there exists a running remote deployment"
+                    " before attempting to watch for changes."
+                ),
+            )
+
+        truss_watch_state = get_truss_watch_state(self._api, model_name)  # type: ignore
+        # Make sure the patches are calculated against the current django patch state, if it exists.
+        # This is important to ensure that the sequence of patches for a given sesion forms a
+        # valid patch sequence (via a linked list)
+        if truss_watch_state.patches:
+            truss_hash = truss_watch_state.patches.django_patch_state.current_hash
+            truss_signature = (
+                truss_watch_state.patches.django_patch_state.current_signature
+            )
+            logging.debug(f"db patch hash: {truss_hash}")
+            logging.debug(
+                f"container_patch_hash: {truss_watch_state.patches.container_patch_state.current_hash}"
             )
-            return
         LocalConfigHandler.add_signature(truss_hash, truss_signature)
         try:
             patch_request = truss_handle.calc_patch(truss_hash, truss_ignore_patterns)
-        except Exception:
-            error_console.print("Failed to calculate patch, bailing on patching")
-            return
-        if patch_request:
-            if (
-                patch_request.prev_hash == patch_request.next_hash
-                or len(patch_request.patch_ops) == 0
-            ):
-                console.print("No changes observed, skipping patching")
-                return
-            try:
+        except Exception as e:
+            return PatchResult(PatchStatus.FAILED, f"Failed to calculate patch. {e}")
+        if not patch_request:
+            return PatchResult(
+                PatchStatus.FAILED,
+                "Failed to calculate patch. Change type might not be supported.",
+            )
+
+        django_has_unapplied_patches = (
+            not truss_watch_state.is_container_built_from_push
+            and truss_watch_state.patches
+            and (
+                truss_watch_state.patches.django_patch_state.current_hash
+                != truss_watch_state.patches.container_patch_state.current_hash
+            )
+        )
+        should_create_patch = (
+            patch_request.prev_hash != patch_request.next_hash
+            and len(patch_request.patch_ops) > 0
+        )
+        is_synced = not django_has_unapplied_patches and not should_create_patch
+        if is_synced:
+            return PatchResult(
+                PatchStatus.SKIPPED, "No changes observed, skipping patching."
+            )
+
+        def do_patch():
+            if should_create_patch:
+                resp = self._api.patch_draft_truss_two_step(model_name, patch_request)
+            else:
+                resp = self._api.sync_draft_truss(model_name)
+            return resp
+
+        try:
+            if console:
                 with console.status("Applying patch..."):
-                    resp = self._api.patch_draft_truss(model_name, patch_request)
-            except ReadTimeout:
-                error_console.print(
-                    "Read Timeout when attempting to connect to remote. Bailing on patching"
-                )
-                return
-            except Exception:
-                error_console.print(
-                    "Failed to patch draft deployment, bailing on patching"
+                    resp = do_patch()
+            else:
+                resp = do_patch()
+
+        except ReadTimeout:
+            return PatchResult(
+                PatchStatus.FAILED, "Read Timeout when attempting to patch remote."
+            )
+        except Exception as e:
+            return PatchResult(
+                PatchStatus.FAILED, f"Failed to patch draft deployment. {e}"
+            )
+        if not resp["succeeded"]:
+            needs_full_deploy = resp.get("needs_full_deploy", None)
+            if needs_full_deploy:
+                message = (
+                    f"Model {model_name} is not able to be patched: `{resp['error']}`. "
+                    f"Use `truss push` to deploy."
                 )
-                return
-            if not resp["succeeded"]:
-                needs_full_deploy = resp.get("needs_full_deploy", None)
-                if needs_full_deploy:
-                    error_console.print(
-                        f"Model {model_name} is not able to be patched, use `truss push` to deploy"
-                    )
-                else:
-                    error_console.print(
-                        f"Failed to patch: `{resp['error']}`. Model left in original state"
-                    )
             else:
-                console.print(
-                    resp.get(
-                        "success_message",
-                        f"Model {model_name} patched successfully",
-                    ),
-                    style="green",
+                message = (
+                    f"Failed to patch. Server error: `{resp['error']}`. "
+                    "Model left in original state."
                 )
+            return PatchResult(PatchStatus.FAILED, message)
+        else:
+            return PatchResult(
+                PatchStatus.SUCCESS,
+                resp.get(
+                    "success_message", f"Model {model_name} patched successfully."
+                ),
+            )
+
+    def patch(
+        self,
+        watch_path: Path,
+        truss_ignore_patterns: List[str],
+        console: "rich_console.Console",
+        error_console: "rich_console.Console",
+    ):
+        result = self._patch(watch_path, truss_ignore_patterns)
+        if result.status in (PatchStatus.SUCCESS, PatchStatus.SKIPPED):
+            console.print(result.message, style="green")
+        else:
+            error_console.print(result.message)
+
+    def patch_for_chainlet(
+        self, watch_path: Path, truss_ignore_patterns: List[str]
+    ) -> PatchResult:
+        return self._patch(watch_path, truss_ignore_patterns, console=None)