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

truss/remote/baseten/api.py

@@ -1,43 +1,103 @@
 import logging
 from enum import Enum
-from typing import Optional
+from typing import Any, List, Optional
 
 import requests
-from truss.remote.baseten.auth import AuthService
+import truss
+from truss.remote.baseten import custom_types as b10_types
+from truss.remote.baseten.auth import ApiKey, AuthService
 from truss.remote.baseten.error import ApiError
 from truss.remote.baseten.utils.transfer import base64_encoded_json_str
 
 logger = logging.getLogger(__name__)
 
+API_URL_MAPPING = {
+    "https://app.baseten.co": "https://api.baseten.co",
+    "https://app.staging.baseten.co": "https://api.staging.baseten.co",
+    "https://app.dev.baseten.co": "https://api.mc-dev.baseten.co",
+    # For local development, this is how we map URLs
+    "http://localhost:8000": "http://api.localhost:8000",
+}
+
+# If a non-standard domain is used with the baseten remote, default to
+# using the production api routes
+DEFAULT_API_DOMAIN = "https://api.baseten.co"
+
+
+def _oracle_data_to_graphql_mutation(oracle: b10_types.OracleData) -> str:
+    args = [
+        f'model_name: "{oracle.model_name}"',
+        f's3_key: "{oracle.s3_key}"',
+        f'encoded_config_str: "{oracle.encoded_config_str}"',
+        f"is_trusted: {str(oracle.is_trusted).lower()}",
+    ]
+
+    if oracle.semver_bump:
+        args.append(f'semver_bump: "{oracle.semver_bump}"')
+
+    if oracle.version_name:
+        args.append(f'version_name: "{oracle.version_name}"')
+
+    args_str = ",\n".join(args)
+
+    return f"""{{
+        {args_str}
+    }}"""
 
-class BasetenApi:
-    """
-    A client for the Baseten API.
 
-    Args:
-        api_url: The URL of the Baseten API.
-        auth_service: An AuthService instance.
-    """
+def _chainlet_data_atomic_to_graphql_mutation(
+    chainlet: b10_types.ChainletDataAtomic,
+) -> str:
+    oracle_data_string = _oracle_data_to_graphql_mutation(chainlet.oracle)
 
+    args = [f'name: "{chainlet.name}"', f"oracle: {oracle_data_string}"]
+
+    args_str = ",\n".join(args)
+
+    return f"""{{
+        {args_str}
+    }}"""
+
+
+class BasetenApi:
     class GraphQLErrorCodes(Enum):
         RESOURCE_NOT_FOUND = "RESOURCE_NOT_FOUND"
 
-    def __init__(self, api_url: str, auth_service: AuthService):
-        self._api_url = api_url
+    def __init__(self, remote_url: str, auth_service: AuthService) -> None:
+        graphql_api_url = f"{remote_url}/graphql/"
+        # Ensure we strip off trailing '/' to denormalize URLs.
+        rest_api_url = API_URL_MAPPING.get(remote_url.strip("/"), DEFAULT_API_DOMAIN)
+
+        self._remote_url = remote_url
+        self._graphql_api_url = graphql_api_url
+        self._rest_api_url = rest_api_url
         self._auth_service = auth_service
         self._auth_token = self._auth_service.authenticate()
 
+    @property
+    def app_url(self) -> str:
+        return self._remote_url
+
+    @property
+    def rest_api_url(self) -> str:
+        return self._rest_api_url
+
+    @property
+    def auth_token(self) -> ApiKey:
+        return self._auth_token
+
     def _post_graphql_query(self, query_string: str) -> dict:
         headers = self._auth_token.header()
+
         resp = requests.post(
-            self._api_url,
+            self._graphql_api_url,
             data={"query": query_string},
             headers=headers,
             timeout=120,
         )
 
         if not resp.ok:
-            logger.error(f"GraphQL endpoint failed with error: {resp.content}")  # type: ignore
+            logger.error(f"GraphQL endpoint failed with error: {str(resp.content)}")
             resp.raise_for_status()
 
         resp_dict = resp.json()
@@ -73,7 +133,9 @@ class BasetenApi:
         semver_bump: str,
         client_version: str,
         is_trusted: bool,
+        allow_truss_download: bool = True,
         deployment_name: Optional[str] = None,
+        origin: Optional[b10_types.ModelOrigin] = None,
     ):
         query_string = f"""
         mutation {{
@@ -83,8 +145,10 @@ class BasetenApi:
                 config: "{config}",
                 semver_bump: "{semver_bump}",
                 client_version: "{client_version}",
-                is_trusted: {'true' if is_trusted else 'false'},
+                is_trusted: {"true" if is_trusted else "false"},
+                allow_truss_download: {"true" if allow_truss_download else "false"},
                 {f'version_name: "{deployment_name}"' if deployment_name else ""}
+                {f"model_origin: {origin.value}" if origin else ""}
             ) {{
                 id,
                 name,
@@ -92,6 +156,7 @@ class BasetenApi:
             }}
         }}
         """
+
         resp = self._post_graphql_query(query_string)
         return resp["data"]["create_model_from_truss"]
 
@@ -103,9 +168,9 @@ class BasetenApi:
         semver_bump: str,
         client_version: str,
         is_trusted: bool,
-        promote: bool = False,
         preserve_previous_prod_deployment: bool = False,
         deployment_name: Optional[str] = None,
+        environment: Optional[str] = None,
     ):
         query_string = f"""
         mutation {{
@@ -115,15 +180,16 @@ class BasetenApi:
                 config: "{config}",
                 semver_bump: "{semver_bump}",
                 client_version: "{client_version}",
-                is_trusted: {'true' if is_trusted else 'false'},
-                promote_after_deploy: {'true' if promote else 'false'},
-                scale_down_old_production: {'false' if preserve_previous_prod_deployment else 'true'},
+                is_trusted: {"true" if is_trusted else "false"},
+                scale_down_old_production: {"false" if preserve_previous_prod_deployment else "true"},
                 {f'name: "{deployment_name}"' if deployment_name else ""}
+                {f'environment_name: "{environment}"' if environment else ""}
             ) {{
                 id
             }}
         }}
         """
+
         resp = self._post_graphql_query(query_string)
         return resp["data"]["create_model_version_from_truss"]
 
@@ -134,6 +200,8 @@ class BasetenApi:
         config,
         client_version,
         is_trusted=False,
+        allow_truss_download=True,
+        origin: Optional[b10_types.ModelOrigin] = None,
     ):
         query_string = f"""
         mutation {{
@@ -141,7 +209,9 @@ class BasetenApi:
                     s3_key: "{s3_key}",
                     config: "{config}",
                     client_version: "{client_version}",
-                    is_trusted: {'true' if is_trusted else 'false'}
+                    is_trusted: {"true" if is_trusted else "false"},
+                    allow_truss_download: {"true" if allow_truss_download else "false"},
+                    {f"model_origin: {origin.value}" if origin else ""}
     ) {{
             id,
             name,
@@ -152,6 +222,127 @@ class BasetenApi:
         resp = self._post_graphql_query(query_string)
         return resp["data"]["deploy_draft_truss"]
 
+    def deploy_chain_atomic(
+        self,
+        entrypoint: b10_types.ChainletDataAtomic,
+        dependencies: List[b10_types.ChainletDataAtomic],
+        chain_id: Optional[str] = None,
+        chain_name: Optional[str] = None,
+        environment: Optional[str] = None,
+        is_draft: bool = False,
+    ):
+        entrypoint_str = _chainlet_data_atomic_to_graphql_mutation(entrypoint)
+
+        dependencies_str = ", ".join(
+            [
+                _chainlet_data_atomic_to_graphql_mutation(dependency)
+                for dependency in dependencies
+            ]
+        )
+
+        query_string = f"""
+            mutation {{
+                deploy_chain_atomic(
+                    {f'chain_id: "{chain_id}"' if chain_id else ""}
+                    {f'chain_name: "{chain_name}"' if chain_name else ""}
+                    {f'environment: "{environment}"' if environment else ""}
+                    is_draft: {str(is_draft).lower()}
+                    entrypoint: {entrypoint_str}
+                    dependencies: [{dependencies_str}]
+                    client_version: "{truss.version()}"
+                ) {{
+                    chain_id
+                    chain_deployment_id
+                    entrypoint_model_id
+                    entrypoint_model_version_id
+                }}
+            }}
+        """
+
+        resp = self._post_graphql_query(query_string)
+
+        return resp["data"]["deploy_chain_atomic"]
+
+    def get_chains(self):
+        query_string = """
+        {
+            chains {
+                id
+                name
+            }
+        }
+        """
+
+        resp = self._post_graphql_query(query_string)
+        return resp["data"]["chains"]
+
+    def get_chain_deployments(self, chain_id: str):
+        query_string = f"""
+        {{
+          chain(id: "{chain_id}") {{
+            deployments {{
+              id
+              created
+              is_draft
+            }}
+          }}
+        }}
+        """
+        resp = self._post_graphql_query(query_string)
+        return resp["data"]["chain"]["deployments"]
+
+    def get_chainlets_by_deployment_id(self, chain_deployment_id: str):
+        query_string = f"""
+        {{
+            chain_deployment(id:"{chain_deployment_id}") {{
+                chainlets {{
+                    name
+                    id
+                    is_entrypoint
+                    oracle {{
+                        id
+                        name
+                    }}
+                    oracle_version {{
+                        id
+                        is_draft
+                        current_model_deployment_status {{
+                            status
+                        }}
+                    }}
+                }}
+                chain {{
+                  id
+                }}
+            }}
+        }}
+        """
+        resp = self._post_graphql_query(query_string)
+        chainlets = resp["data"]["chain_deployment"]["chainlets"]
+        for chainlet in chainlets:
+            chainlet["chain"] = {"id": resp["data"]["chain_deployment"]["chain"]["id"]}
+        return chainlets
+
+    def delete_chain(self, chain_id: str) -> Any:
+        url = f"{self._rest_api_url}/v1/chains/{chain_id}"
+        headers = self._auth_token.header()
+        resp = requests.delete(url, headers=headers)
+        if not resp.ok:
+            resp.raise_for_status()
+
+        deployment = resp.json()
+        return deployment
+
+    def delete_chain_deployment(self, chain_id: str, chain_deployment_id: str) -> Any:
+        url = f"{self._rest_api_url}/v1/chains/{chain_id}/deployments/{chain_deployment_id}"
+        headers = self._auth_token.header()
+        resp = requests.delete(url, headers=headers)
+        if not resp.ok:
+            resp.raise_for_status()
+
+        deployment = resp.json()
+        return deployment
+
     def models(self):
         query_string = """
         {
@@ -171,6 +362,25 @@ class BasetenApi:
         resp = self._post_graphql_query(query_string)
         return resp["data"]
 
+    def get_truss_watch_state(self, model_name: str):
+        query_string = f"""
+        {{
+            truss_watch_state(name: "{model_name}") {{
+                is_container_built_from_push
+                django_patch_state {{
+                    current_hash
+                    current_signature
+                }}
+                container_patch_state {{
+                    current_hash
+                    current_signature
+                }}
+            }}
+        }}
+        """
+        resp = self._post_graphql_query(query_string)
+        return resp["data"]
+
     def get_model(self, model_name):
         query_string = f"""
         {{
@@ -221,8 +431,12 @@ class BasetenApi:
         {{
             model_version(id: "{model_version_id}") {{
                 id
+                is_draft
+                truss_hash
+                truss_signature
                 oracle{{
                     id
+                    name
                 }}
             }}
           }}
@@ -230,12 +444,12 @@ class BasetenApi:
         resp = self._post_graphql_query(query_string)
         return resp["data"]
 
-    def patch_draft_truss(self, model_name, patch_request):
+    def patch_draft_truss_two_step(self, model_name, patch_request):
         patch = base64_encoded_json_str(patch_request.to_dict())
         query_string = f"""
         mutation {{
-        patch_draft_truss(name: "{model_name}",
-                    client_version: "TRUSS",
+        stage_patch_for_draft_truss(name: "{model_name}",
+                    client_version: "{truss.version()}",
                     patch: "{patch}",
     ) {{
             id,
@@ -248,4 +462,75 @@ class BasetenApi:
         }}
         """
         resp = self._post_graphql_query(query_string)
-        return resp["data"]["patch_draft_truss"]
+        result = resp["data"]["stage_patch_for_draft_truss"]
+        if not result["succeeded"]:
+            logging.debug(f"Failed to stage patch: {result}")
+            return result
+        logging.debug("Succesfully staged patch. Syncing patch to truss...")
+
+        return self.sync_draft_truss(model_name)
+
+    def sync_draft_truss(self, model_name):
+        query_string = f"""
+        mutation {{
+        sync_draft_truss(name: "{model_name}",
+                    client_version: "{truss.version()}",
+    ) {{
+            id,
+            name,
+            version_id
+            succeeded
+            needs_full_deploy
+            error
+        }}
+        }}
+        """
+        resp = self._post_graphql_query(query_string)
+        result = resp["data"]["sync_draft_truss"]
+        if not result["succeeded"]:
+            logging.debug(f"Failed to sync patch: {result}")
+        return result
+
+    def validate_truss(self, client_version: str, config: str):
+        query_string = f"""{{
+            truss_validation(client_version: "{client_version}", config: "{config}") {{
+                success
+                details
+            }}
+        }}
+        """
+        resp = self._post_graphql_query(query_string)
+        return resp["data"]["truss_validation"]
+
+    def get_deployment(self, model_id: str, deployment_id: str) -> Any:
+        headers = self._auth_token.header()
+        resp = requests.get(
+            f"{self._rest_api_url}/v1/models/{model_id}/deployments/{deployment_id}",
+            headers=headers,
+        )
+        if not resp.ok:
+            resp.raise_for_status()
+
+        deployment = resp.json()
+        return deployment
+
+    def upsert_secret(self, name: str, value: str) -> Any:
+        headers = self._auth_token.header()
+        data = {"name": name, "value": value}
+        resp = requests.post(
+            f"{self._rest_api_url}/v1/secrets", headers=headers, json=data
+        )
+        if not resp.ok:
+            resp.raise_for_status()
+
+        secret_info = resp.json()
+        return secret_info
+
+    def get_all_secrets(self) -> Any:
+        headers = self._auth_token.header()
+        resp = requests.get(f"{self._rest_api_url}/v1/secrets", headers=headers)
+        if not resp.ok:
+            resp.raise_for_status()
+
+        secrets_info = resp.json()
+        return secrets_info

truss/remote/baseten/auth.py

@@ -10,7 +10,7 @@ logger = logging.getLogger(__name__)
 class ApiKey:
     value: str
 
-    def __init__(self, value: str):
+    def __init__(self, value: str) -> None:
         self.value = value
 
     def header(self):
@@ -18,12 +18,12 @@ class ApiKey:
 
 
 class AuthService:
-    def __init__(self, api_key: Optional[str] = None):
+    def __init__(self, api_key: Optional[str] = None) -> None:
         if not api_key:
             api_key = os.environ.get("BASETEN_API_KEY", None)
         self._api_key = api_key
 
-    def validate(self):
+    def validate(self) -> None:
         if not self._api_key:
             raise AuthorizationError("No API key provided.")