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

truss/remote/baseten/core.py

@@ -1,17 +1,34 @@
+import datetime
+import json
 import logging
-from typing import IO, List, Optional, Tuple
+import textwrap
+from typing import IO, TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Type
+
+from truss.base.errors import ValidationError
+
+if TYPE_CHECKING:
+    from rich import progress
 
 import truss
+from truss.base.constants import PRODUCTION_ENVIRONMENT_NAME
+from truss.remote.baseten import custom_types as b10_types
 from truss.remote.baseten.api import BasetenApi
 from truss.remote.baseten.error import ApiError
 from truss.remote.baseten.utils.tar import create_tar_with_progress_bar
 from truss.remote.baseten.utils.transfer import multipart_upload_boto3
-from truss.truss_handle import TrussHandle
-from truss.util.path import load_trussignore_patterns
+from truss.truss_handle.truss_handle import TrussHandle
+from truss.util.path import load_trussignore_patterns_from_truss_dir
 
 logger = logging.getLogger(__name__)
 
 
+DEPLOYING_STATUSES = ["BUILDING", "DEPLOYING", "LOADING_MODEL", "UPDATING"]
+ACTIVE_STATUS = "ACTIVE"
+NO_ENVIRONMENTS_EXIST_ERROR_MESSAGING = (
+    "Model hasn't been deployed yet. No evironments exist."
+)
+
+
 class ModelIdentifier:
     value: str
 
@@ -31,6 +48,126 @@ class ModelVersionId(ModelIdentifier):
         self.value = model_version_id
 
 
+class PatchState(NamedTuple):
+    current_hash: str
+    current_signature: str
+
+
+class TrussPatches(NamedTuple):
+    django_patch_state: PatchState
+    container_patch_state: PatchState
+
+
+class TrussWatchState(NamedTuple):
+    is_container_built_from_push: bool
+    patches: Optional[TrussPatches]
+
+
+class ChainDeploymentHandleAtomic(NamedTuple):
+    chain_id: str
+    chain_deployment_id: str
+    is_draft: bool
+    entrypoint_model_id: str
+    entrypoint_model_version_id: str
+
+
+def get_chain_id_by_name(api: BasetenApi, chain_name: str) -> Optional[str]:
+    """
+    Check if a chain with the given name exists in the Baseten remote.
+
+    Args:
+        api: BasetenApi instance
+        chain_name: Name of the chain to check for existence
+
+    Returns:
+        chain_id if present, otherwise None
+    """
+    chains = api.get_chains()
+
+    chain_name_id_mapping = {chain["name"]: chain["id"] for chain in chains}
+    return chain_name_id_mapping.get(chain_name)
+
+
+def get_dev_chain_deployment(api: BasetenApi, chain_id: str):
+    chain_deployments = api.get_chain_deployments(chain_id)
+    dev_deployments = [
+        deployment for deployment in chain_deployments if deployment["is_draft"]
+    ]
+    if not dev_deployments:
+        return None
+    newest_draft_deployment = max(
+        dev_deployments, key=lambda d: datetime.datetime.fromisoformat(d["created"])
+    )
+    return newest_draft_deployment
+
+
+def create_chain_atomic(
+    api: BasetenApi,
+    chain_name: str,
+    entrypoint: b10_types.ChainletDataAtomic,
+    dependencies: List[b10_types.ChainletDataAtomic],
+    is_draft: bool,
+    environment: Optional[str],
+) -> ChainDeploymentHandleAtomic:
+    if environment and is_draft:
+        logging.info(
+            f"Automatically publishing Chain `{chain_name}` based on "
+            "environment setting."
+        )
+        is_draft = False
+
+    chain_id = get_chain_id_by_name(api, chain_name)
+
+    # TODO(Tyron): Refactor for better readability:
+    # 1. Prepare all arguments for `deploy_chain_atomic`.
+    # 2. Validate argument combinations.
+    # 3. Make a single invocation to `deploy_chain_atomic`.
+    if is_draft:
+        res = api.deploy_chain_atomic(
+            chain_name=chain_name,
+            is_draft=True,
+            entrypoint=entrypoint,
+            dependencies=dependencies,
+        )
+    elif chain_id:
+        # This is the only case where promote has relevance, since
+        # if there is no chain already, the first deployment will
+        # already be production, and only published deployments can
+        # be promoted.
+        try:
+            res = api.deploy_chain_atomic(
+                chain_id=chain_id,
+                environment=environment,
+                entrypoint=entrypoint,
+                dependencies=dependencies,
+            )
+        except ApiError as e:
+            if (
+                e.graphql_error_code
+                == BasetenApi.GraphQLErrorCodes.RESOURCE_NOT_FOUND.value
+            ):
+                raise ValueError(
+                    f"Environment `{environment}` does not exist. You can "
+                    f"create environments in the Chains UI."
+                ) from e
+
+            raise e
+    elif environment and environment != PRODUCTION_ENVIRONMENT_NAME:
+        raise ValueError(NO_ENVIRONMENTS_EXIST_ERROR_MESSAGING)
+    else:
+        res = api.deploy_chain_atomic(
+            chain_name=chain_name, entrypoint=entrypoint, dependencies=dependencies
+        )
+
+    return ChainDeploymentHandleAtomic(
+        chain_id=res["chain_id"],
+        chain_deployment_id=res["chain_deployment_id"],
+        entrypoint_model_id=res["entrypoint_model_id"],
+        entrypoint_model_version_id=res["entrypoint_model_version_id"],
+        is_draft=is_draft,
+    )
+
+
 def exists_model(api: BasetenApi, model_name: str) -> Optional[str]:
     """
     Check if a model with the given name exists in the Baseten remote.
@@ -58,7 +195,7 @@ def exists_model(api: BasetenApi, model_name: str) -> Optional[str]:
 
 def get_model_versions(api: BasetenApi, model_name: ModelName) -> Tuple[str, List]:
     query_result = api.get_model(model_name.value)["model"]
-    return (query_result["id"], query_result["versions"])
+    return query_result["id"], query_result["versions"]
 
 
 def get_dev_version_from_versions(versions: List[dict]) -> Optional[dict]:
@@ -93,6 +230,36 @@ def get_dev_version(api: BasetenApi, model_name: str) -> Optional[dict]:
     return get_dev_version_from_versions(versions)
 
 
+def get_truss_watch_state(api: BasetenApi, model_name: str) -> TrussWatchState:
+    response = api.get_truss_watch_state(model_name)["truss_watch_state"]
+    django_patch_state = (
+        None
+        if response["django_patch_state"] is None
+        else PatchState(
+            current_hash=response["django_patch_state"]["current_hash"],
+            current_signature=response["django_patch_state"]["current_signature"],
+        )
+    )
+    container_patch_state = (
+        None
+        if response["container_patch_state"] is None
+        else PatchState(
+            current_hash=response["container_patch_state"]["current_hash"],
+            current_signature=response["container_patch_state"]["current_signature"],
+        )
+    )
+    patches = None
+    if django_patch_state and container_patch_state:
+        patches = TrussPatches(
+            django_patch_state=django_patch_state,
+            container_patch_state=container_patch_state,
+        )
+    return TrussWatchState(
+        is_container_built_from_push=response["is_container_built_from_push"],
+        patches=patches,
+    )
+
+
 def get_prod_version_from_versions(versions: List[dict]) -> Optional[dict]:
     # Loop over versions instead of using the primary_version field because
     # primary_version is set to the development version ID if no published
@@ -103,38 +270,37 @@ def get_prod_version_from_versions(versions: List[dict]) -> Optional[dict]:
     return None
 
 
-def archive_truss(truss_handle: TrussHandle) -> IO:
-    """
-    Archive a TrussHandle into a tar file.
-
-    Args:
-        b10_truss: TrussHandle to archive
+def archive_truss(
+    truss_handle: TrussHandle, progress_bar: Optional[Type["progress.Progress"]]
+) -> IO:
+    """Archive a TrussHandle into a tar file.
 
     Returns:
         A file-like object containing the tar file
     """
-    truss_dir = truss_handle._spec.truss_dir
-    ignore_patterns = []
+    truss_dir = truss_handle._truss_dir
 
     # check for a truss_ignore file and read the ignore patterns if it exists
-    truss_ignore_file = truss_dir / ".truss_ignore"
-    if truss_ignore_file.exists():
-        ignore_patterns = load_trussignore_patterns(truss_ignore_file=truss_ignore_file)
-    else:
-        ignore_patterns = load_trussignore_patterns()
+    ignore_patterns = load_trussignore_patterns_from_truss_dir(truss_dir)
 
     try:
-        temp_file = create_tar_with_progress_bar(truss_dir, ignore_patterns)
+        temp_file = create_tar_with_progress_bar(
+            truss_dir, ignore_patterns, progress_bar=progress_bar
+        )
     except PermissionError:
         # workaround for Windows bug with Tempfile that causes PermissionErrors
         temp_file = create_tar_with_progress_bar(
-            truss_dir, ignore_patterns, delete=False
+            truss_dir, ignore_patterns, delete=False, progress_bar=progress_bar
         )
     temp_file.file.seek(0)
     return temp_file
 
 
-def upload_truss(api: BasetenApi, serialize_file: IO) -> str:
+def upload_truss(
+    api: BasetenApi,
+    serialize_file: IO,
+    progress_bar: Optional[Type["progress.Progress"]],
+) -> str:
     """
     Upload a TrussHandle to the Baseten remote.
 
@@ -149,7 +315,7 @@ def upload_truss(api: BasetenApi, serialize_file: IO) -> str:
     s3_key = temp_credentials_s3_upload.pop("s3_key")
     s3_bucket = temp_credentials_s3_upload.pop("s3_bucket")
     multipart_upload_boto3(
-        serialize_file.name, s3_bucket, s3_key, temp_credentials_s3_upload
+        serialize_file.name, s3_bucket, s3_key, temp_credentials_s3_upload, progress_bar
     )
     return s3_key
 
@@ -161,25 +327,29 @@ def create_truss_service(
     config: str,
     semver_bump: str = "MINOR",
     is_trusted: bool = False,
-    promote: bool = False,
     preserve_previous_prod_deployment: bool = False,
+    allow_truss_download: bool = False,
     is_draft: Optional[bool] = False,
     model_id: Optional[str] = None,
     deployment_name: Optional[str] = None,
+    origin: Optional[b10_types.ModelOrigin] = None,
+    environment: Optional[str] = None,
 ) -> Tuple[str, str]:
     """
     Create a model in the Baseten remote.
 
     Args:
-        api: BasetenApi instance
-        model_name: Name of the model to create
-        s3_key: S3 key of the uploaded TrussHandle
-        config: Base64 encoded JSON string of the Truss config
-        semver_bump: Semver bump type, defaults to "MINOR"
-        is_trusted: Whether the model is trusted, defaults to False
-        promote: Whether to promote the model after deploy, defaults to False
-        preserve_previous_prod_deployment: Wheter to scale old production deployment to zero
-        deployment_name: Name to apply to the created deployment. Not applied to development model
+        api: BasetenApi instance.
+        model_name: Name of the model to create.
+        s3_key: S3 key of the uploaded TrussHandle.
+        config: Base64 encoded JSON string of the Truss config.
+        semver_bump: Semver bump type, defaults to "MINOR".
+        is_trusted: Whether the model is trusted, defaults to False.
+        promote: Whether to promote the model after deploy, defaults to False.
+        preserve_previous_prod_deployment: Whether to scale old production deployment
+            to zero.
+        deployment_name: Name to apply to the created deployment. Not applied to
+            development model.
 
     Returns:
         A tuple of the model ID and version ID
@@ -189,35 +359,72 @@ def create_truss_service(
             model_name,
             s3_key,
             config,
-            f"truss=={truss.version()}",
-            is_trusted,
+            truss.version(),
+            is_trusted=is_trusted,
+            allow_truss_download=allow_truss_download,
+            origin=origin,
         )
 
-        return (model_version_json["id"], model_version_json["version_id"])
+        return model_version_json["id"], model_version_json["version_id"]
 
     if model_id is None:
+        if environment and environment != PRODUCTION_ENVIRONMENT_NAME:
+            raise ValueError(NO_ENVIRONMENTS_EXIST_ERROR_MESSAGING)
         model_version_json = api.create_model_from_truss(
             model_name=model_name,
             s3_key=s3_key,
             config=config,
             semver_bump=semver_bump,
-            client_version=f"truss=={truss.version()}",
+            client_version=truss.version(),
             is_trusted=is_trusted,
+            allow_truss_download=allow_truss_download,
             deployment_name=deployment_name,
+            origin=origin,
         )
-        return (model_version_json["id"], model_version_json["version_id"])
-
-    # Case where there is a model id already, create another version
-    model_version_json = api.create_model_version_from_truss(
-        model_id=model_id,
-        s3_key=s3_key,
-        config=config,
-        semver_bump=semver_bump,
-        client_version=f"truss=={truss.version()}",
-        is_trusted=is_trusted,
-        promote=promote,
-        preserve_previous_prod_deployment=preserve_previous_prod_deployment,
-        deployment_name=deployment_name,
-    )
+        return model_version_json["id"], model_version_json["version_id"]
+
+    try:
+        model_version_json = api.create_model_version_from_truss(
+            model_id=model_id,
+            s3_key=s3_key,
+            config=config,
+            semver_bump=semver_bump,
+            client_version=truss.version(),
+            is_trusted=is_trusted,
+            preserve_previous_prod_deployment=preserve_previous_prod_deployment,
+            deployment_name=deployment_name,
+            environment=environment,
+        )
+    except ApiError as e:
+        if (
+            e.graphql_error_code
+            == BasetenApi.GraphQLErrorCodes.RESOURCE_NOT_FOUND.value
+        ):
+            raise ValueError(
+                f'Environment "{environment}" does not exist. You can create environments in the Baseten UI.'
+            ) from e
+        raise e
     model_version_id = model_version_json["id"]
-    return (model_id, model_version_id)
+    return model_id, model_version_id
+
+
+def validate_truss_config(api: BasetenApi, config: str):
+    """
+    Validate a truss config as well as the truss version.
+
+    Args:
+        api: BasetenApi instance
+        config: Base64 encoded JSON string of the Truss config
+
+    Returns:
+        None if the config is valid, otherwise raises an error message
+    """
+    valid_config = api.validate_truss(truss.version(), config)
+    if not valid_config.get("success"):
+        details = json.loads(valid_config.get("details"))
+        errors = details.get("errors", [])
+        if errors:
+            error_messages = "\n".join(textwrap.indent(error, "  ") for error in errors)
+            raise ValidationError(
+                f"Validation failed with the following errors:\n{error_messages}"
+            )

truss/remote/baseten/custom_types.py

@@ -0,0 +1,44 @@
+import pathlib
+from enum import Enum
+from typing import Optional
+
+import pydantic
+
+
+class DeployedChainlet(pydantic.BaseModel):
+    name: str
+    is_entrypoint: bool
+    is_draft: bool
+    status: str
+    logs_url: str
+    oracle_predict_url: str
+    oracle_name: str
+
+
+class ChainletArtifact(pydantic.BaseModel):
+    truss_dir: pathlib.Path
+    display_name: str
+    name: str
+
+
+class ModelOrigin(Enum):
+    BASETEN = "BASETEN"
+    CHAINS = "CHAINS"
+
+
+class OracleData(pydantic.BaseModel):
+    class Config:
+        protected_namespaces = ()
+
+    model_name: str
+    s3_key: str
+    encoded_config_str: str
+    semver_bump: Optional[str] = "MINOR"
+    is_trusted: bool
+    version_name: Optional[str] = None
+
+
+# This corresponds to `ChainletInputAtomicGraphene` in the backend.
+class ChainletDataAtomic(pydantic.BaseModel):
+    name: str
+    oracle: OracleData

truss/remote/baseten/error.py

@@ -21,3 +21,7 @@ class AuthorizationError(Error):
     """Raised in places where the user needs to be logged in and is not."""
 
     pass
+
+
+class RemoteError(Error):
+    pass