truss 0.11.6rc102__py3-none-any.whl → 0.11.24rc2__py3-none-any.whl

truss/contexts/image_builder/serving_image_builder.py

@@ -32,6 +32,7 @@ from truss.base.constants import (
     FILENAME_CONSTANTS_MAP,
     MODEL_CACHE_PATH,
     MODEL_DOCKERFILE_NAME,
+    NO_BUILD_DOCKERFILE_TEMPLATE_NAME,
     REQUIREMENTS_TXT_FILENAME,
     SERVER_CODE_DIR,
     SERVER_DOCKERFILE_TEMPLATE_NAME,
@@ -577,7 +578,10 @@ class ServingImageBuilder(ImageBuilder):
                 else:
                     self.prepare_trtllm_decoder_build_dir(build_dir=build_dir)
 
-        if config.docker_server is not None:
+        if (
+            config.docker_server is not None
+            and config.docker_server.no_build is not True
+        ):
             self._copy_into_build_dir(
                 TEMPLATES_DIR / "docker_server_requirements.txt",
                 build_dir,
@@ -750,12 +754,21 @@ class ServingImageBuilder(ImageBuilder):
         build_commands: List[str],
     ):
         config = self._spec.config
+
         data_dir = build_dir / config.data_dir
         model_dir = build_dir / config.model_module_dir
         bundled_packages_dir = build_dir / config.bundled_packages_dir
-        dockerfile_template = read_template_from_fs(
-            TEMPLATES_DIR, SERVER_DOCKERFILE_TEMPLATE_NAME
-        )
+
+        # Note: no-build deployment template doesn't use most of the template variables,
+        # because it tries to run the base image as-is to the extent possible.
+        if config.docker_server and config.docker_server.no_build:
+            dockerfile_template = read_template_from_fs(
+                TEMPLATES_DIR, NO_BUILD_DOCKERFILE_TEMPLATE_NAME
+            )
+        else:
+            dockerfile_template = read_template_from_fs(
+                TEMPLATES_DIR, SERVER_DOCKERFILE_TEMPLATE_NAME
+            )
         python_version = truss_config.to_dotted_python_version(config.python_version)
         if config.base_image:
             base_image_name_and_tag = config.base_image.image

truss/remote/baseten/api.py

@@ -3,7 +3,9 @@ from enum import Enum
 from typing import Any, Dict, List, Mapping, Optional
 
 import requests
+from pydantic import BaseModel, Field
 
+from truss.base.custom_types import SafeModel
 from truss.remote.baseten import custom_types as b10_types
 from truss.remote.baseten.auth import ApiKey, AuthService
 from truss.remote.baseten.custom_types import APIKeyCategory
@@ -12,6 +14,62 @@ from truss.remote.baseten.rest_client import RestAPIClient
 from truss.remote.baseten.utils.transfer import base64_encoded_json_str
 
 logger = logging.getLogger(__name__)
+PARAMS_INDENT = "\n                    "
+
+
+class ChainAWSCredential(SafeModel):
+    aws_access_key_id: str
+    aws_secret_access_key: str
+    aws_session_token: str
+
+
+class ChainUploadCredentials(SafeModel):
+    s3_bucket: str
+    s3_key: str
+    aws_access_key_id: str
+    aws_secret_access_key: str
+    aws_session_token: str
+
+    @property
+    def aws_credentials(self) -> ChainAWSCredential:
+        return ChainAWSCredential(
+            aws_access_key_id=self.aws_access_key_id,
+            aws_secret_access_key=self.aws_secret_access_key,
+            aws_session_token=self.aws_session_token,
+        )
+
+
+class InstanceTypeV1(BaseModel):
+    """An instance type."""
+
+    id: str = Field(description="Identifier string for the instance type")
+    name: str = Field(description="Name of the instance type")
+    display_name: str = Field(
+        alias="displayName", description="Display name of the instance type"
+    )
+    gpu_count: int = Field(
+        alias="gpuCount", description="Number of GPUs on the instance type"
+    )
+    default: bool = Field(description="Whether this is the default instance type")
+    gpu_memory: Optional[int] = Field(alias="gpuMemory", description="GPU memory in MB")
+    node_count: int = Field(alias="nodeCount", description="Number of nodes")
+    gpu_type: Optional[str] = Field(
+        alias="gpuType", description="Type of GPU on the instance type"
+    )
+    millicpu_limit: int = Field(
+        alias="millicpuLimit", description="CPU limit of the instance type in millicpu"
+    )
+    memory_limit: int = Field(
+        alias="memoryLimit", description="Memory limit of the instance type in MB"
+    )
+    price: Optional[float] = Field(description="Price of the instance type")
+    limited_capacity: Optional[bool] = Field(
+        alias="limitedCapacity",
+        description="Whether this instance type has limited capacity",
+    )
+
+    class Config:
+        populate_by_name = True
 
 
 API_URL_MAPPING = {
@@ -141,6 +199,9 @@ class BasetenApi:
         allow_truss_download: bool = True,
         deployment_name: Optional[str] = None,
         origin: Optional[b10_types.ModelOrigin] = None,
+        environment: Optional[str] = None,
+        deploy_timeout_minutes: Optional[int] = None,
+        team_id: Optional[str] = None,
     ):
         query_string = f"""
             mutation ($trussUserEnv: String) {{
@@ -153,6 +214,9 @@ class BasetenApi:
                     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 ""}
+                    {f'environment_name: "{environment}"' if environment else ""}
+                    {f"deploy_timeout_minutes: {deploy_timeout_minutes}" if deploy_timeout_minutes is not None else ""}
+                    {f'team_id: "{team_id}"' if team_id else ""}
                 ) {{
                     model_version {{
                         id
@@ -184,6 +248,7 @@ class BasetenApi:
         deployment_name: Optional[str] = None,
         environment: Optional[str] = None,
         preserve_env_instance_type: bool = True,
+        deploy_timeout_minutes: Optional[int] = None,
     ):
         query_string = f"""
             mutation ($trussUserEnv: String) {{
@@ -197,6 +262,7 @@ class BasetenApi:
                     preserve_env_instance_type: {"true" if preserve_env_instance_type else "false"}
                     {f'name: "{deployment_name}"' if deployment_name else ""}
                     {f'environment_name: "{environment}"' if environment else ""}
+                    {f"deploy_timeout_minutes: {deploy_timeout_minutes}" if deploy_timeout_minutes is not None else ""}
                 ) {{
                     model_version {{
                         id
@@ -226,6 +292,8 @@ class BasetenApi:
         truss_user_env: b10_types.TrussUserEnv,
         allow_truss_download=True,
         origin: Optional[b10_types.ModelOrigin] = None,
+        deploy_timeout_minutes: Optional[int] = None,
+        team_id: Optional[str] = None,
     ):
         query_string = f"""
             mutation ($trussUserEnv: String) {{
@@ -235,6 +303,8 @@ class BasetenApi:
                     truss_user_env: $trussUserEnv
                     allow_truss_download: {"true" if allow_truss_download else "false"}
                     {f"model_origin: {origin.value}" if origin else ""}
+                    {f"deploy_timeout_minutes: {deploy_timeout_minutes}" if deploy_timeout_minutes is not None else ""}
+                    {f'team_id: "{team_id}"' if team_id else ""}
                 ) {{
                     model_version {{
                         id
@@ -265,7 +335,13 @@ class BasetenApi:
         chain_name: Optional[str] = None,
         environment: Optional[str] = None,
         is_draft: bool = False,
+        original_source_artifact_s3_key: Optional[str] = None,
+        allow_truss_download: Optional[bool] = True,
+        deployment_name: Optional[str] = None,
+        team_id: Optional[str] = None,
     ):
+        if allow_truss_download is None:
+            allow_truss_download = True
         entrypoint_str = _chainlet_data_atomic_to_graphql_mutation(entrypoint)
 
         dependencies_str = ", ".join(
@@ -275,13 +351,32 @@ class BasetenApi:
             ]
         )
 
+        params = []
+        if chain_id:
+            params.append(f'chain_id: "{chain_id}"')
+        if chain_name:
+            params.append(f'chain_name: "{chain_name}"')
+        if environment:
+            params.append(f'environment: "{environment}"')
+        if original_source_artifact_s3_key:
+            params.append(
+                f'original_source_artifact_s3_key: "{original_source_artifact_s3_key}"'
+            )
+        if team_id:
+            params.append(f'team_id: "{team_id}"')
+
+        params.append(f"is_draft: {str(is_draft).lower()}")
+        if allow_truss_download is False:
+            params.append("allow_truss_download: false")
+        if deployment_name:
+            params.append(f'deployment_name: "{deployment_name}"')
+
+        params_str = PARAMS_INDENT.join(params)
+
         query_string = f"""
             mutation ($trussUserEnv: String) {{
                 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()}
+                    {params_str}
                     entrypoint: {entrypoint_str}
                     dependencies: [{dependencies_str}]
                     truss_user_env: $trussUserEnv
@@ -303,18 +398,24 @@ class BasetenApi:
 
         return resp["data"]["deploy_chain_atomic"]
 
-    def get_chains(self):
+    def get_chains(self, team_id: Optional[str] = None):
         query_string = """
         {
             chains {
                 id
                 name
+                team {
+                    name
+                }
             }
         }
         """
 
         resp = self._post_graphql_query(query_string)
-        return resp["data"]["chains"]
+        chains = resp["data"]["chains"]
+
+        # TODO(COR-492): Filter by team_id in the backend
+        return chains
 
     def get_chain_deployments(self, chain_id: str):
         query_string = f"""
@@ -377,6 +478,10 @@ class BasetenApi:
             models {
                 id,
                 name
+                team {
+                    id
+                    name
+                }
                 versions{
                     id,
                     semver,
@@ -416,6 +521,10 @@ class BasetenApi:
                 id
                 name
                 hostname
+                team {{
+                    id
+                    name
+                }}
                 versions {{
                     id
                     semver
@@ -568,10 +677,14 @@ class BasetenApi:
             "v1/api_keys", body={"type": api_key_type.value, "name": name}
         )
 
-    def upsert_training_project(self, training_project):
+    def upsert_training_project(self, training_project, team_id: Optional[str] = None):
+        if team_id:
+            endpoint = f"v1/teams/{team_id}/training_projects"
+        else:
+            endpoint = "v1/training_projects"
         resp_json = self._rest_api_client.post(
-            "v1/training_projects",
-            body={"training_project": training_project.model_dump()},
+            endpoint,
+            body={"training_project": training_project.model_dump(exclude_none=True)},
         )
         return resp_json["training_project"]
 
@@ -623,8 +736,29 @@ class BasetenApi:
         return resp_json["training_projects"]
 
     def get_blob_credentials(self, blob_type: b10_types.BlobType):
+        if blob_type == b10_types.BlobType.CHAIN:
+            return self.get_chain_s3_upload_credentials()
         return self._rest_api_client.get(f"v1/blobs/credentials/{blob_type.value}")
 
+    def get_chain_s3_upload_credentials(self) -> ChainUploadCredentials:
+        """Get chain artifact credentials using GraphQL query."""
+        query = """
+        query {
+            chain_s3_upload_credentials {
+                s3_bucket
+                s3_key
+                aws_access_key_id
+                aws_secret_access_key
+                aws_session_token
+            }
+        }
+        """
+        response = self._post_graphql_query(query)
+
+        return ChainUploadCredentials.model_validate(
+            response["data"]["chain_s3_upload_credentials"]
+        )
+
     def get_training_job_metrics(
         self,
         project_id: str,
@@ -750,3 +884,75 @@ class BasetenApi:
 
         # NB(nikhil): reverse order so latest logs are at the end
         return resp_json["logs"][::-1]
+
+    def create_model_version_from_inference_template(self, request_data: dict):
+        """
+        Create a model version from an inference template using GraphQL mutation.
+
+        Args:
+            request_data: Dictionary containing the request structure with metadata,
+                         weights_sources, inference_stack, and instance_type_id
+        """
+        query_string = """
+        mutation ($request: CreateModelVersionFromInferenceTemplateRequest!) {
+            create_model_version_from_inference_template(request: $request) {
+                model_version {
+                    id
+                    name
+                }
+                truss_config
+            }
+        }
+        """
+        resp = self._post_graphql_query(
+            query_string, variables={"request": request_data}
+        )
+        return resp["data"]["create_model_version_from_inference_template"]
+
+    def get_instance_types(self) -> List[InstanceTypeV1]:
+        """
+        Get all available instance types via GraphQL API.
+        """
+        query_string = """
+        query Instances {
+            listedInstances: listed_instances {
+                id
+                name
+                millicpuLimit: millicpu_limit
+                memoryLimit: memory_limit
+                gpuCount: gpu_count
+                gpuType: gpu_type
+                gpuMemory: gpu_memory
+                default
+                displayName: display_name
+                nodeCount: node_count
+                price
+                limitedCapacity: limited_capacity
+            }
+        }
+        """
+
+        resp = self._post_graphql_query(query_string)
+        instance_types_data = resp["data"]["listedInstances"]
+        return [
+            InstanceTypeV1(**instance_type) for instance_type in instance_types_data
+        ]
+
+    def get_teams(self) -> Dict[str, Dict[str, str]]:
+        """
+        Get all available teams via GraphQL API.
+        Returns a dictionary mapping team name to team data (with 'id' and 'name' keys).
+        """
+        query_string = """
+        query Teams {
+            teams {
+                id
+                name
+            }
+        }
+        """
+
+        resp = self._post_graphql_query(query_string)
+        teams_data = resp["data"]["teams"]
+        # Convert list to dict mapping team_name -> team
+        return {team["name"]: team for team in teams_data}

truss/remote/baseten/core.py

@@ -8,6 +8,7 @@ from typing import IO, TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Tup
 import requests
 
 from truss.base.errors import ValidationError
+from truss.util.error_utils import handle_client_error
 
 if TYPE_CHECKING:
     from rich import progress
@@ -80,6 +81,8 @@ class ChainDeploymentHandleAtomic(NamedTuple):
     chain_id: str
     chain_deployment_id: str
     is_draft: bool
+    original_source_artifact_s3_key: Optional[str] = None
+    allow_truss_download: Optional[bool] = True
 
 
 class ModelVersionHandle(NamedTuple):
@@ -89,19 +92,21 @@ class ModelVersionHandle(NamedTuple):
     instance_type_name: Optional[str] = None
 
 
-def get_chain_id_by_name(api: BasetenApi, chain_name: str) -> Optional[str]:
+def get_chain_id_by_name(
+    api: BasetenApi, chain_name: str, team_id: Optional[str] = None
+) -> 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
+        team_id: Optional team_id to filter chains by team
 
     Returns:
         chain_id if present, otherwise None
     """
-    chains = api.get_chains()
-
+    chains = api.get_chains(team_id=team_id)
     chain_name_id_mapping = {chain["name"]: chain["id"] for chain in chains}
     return chain_name_id_mapping.get(chain_name)
 
@@ -127,6 +132,10 @@ def create_chain_atomic(
     is_draft: bool,
     truss_user_env: b10_types.TrussUserEnv,
     environment: Optional[str],
+    original_source_artifact_s3_key: Optional[str] = None,
+    allow_truss_download: bool = True,
+    deployment_name: Optional[str] = None,
+    team_id: Optional[str] = None,
 ) -> ChainDeploymentHandleAtomic:
     if environment and is_draft:
         logging.info(
@@ -135,7 +144,7 @@ def create_chain_atomic(
         )
         is_draft = False
 
-    chain_id = get_chain_id_by_name(api, chain_name)
+    chain_id = get_chain_id_by_name(api, chain_name, team_id=team_id)
 
     # TODO(Tyron): Refactor for better readability:
     # 1. Prepare all arguments for `deploy_chain_atomic`.
@@ -149,6 +158,10 @@ def create_chain_atomic(
             chain_name=chain_name,
             is_draft=True,
             truss_user_env=truss_user_env,
+            original_source_artifact_s3_key=original_source_artifact_s3_key,
+            allow_truss_download=allow_truss_download,
+            deployment_name=deployment_name,
+            team_id=team_id,
         )
     elif chain_id:
         # This is the only case where promote has relevance, since
@@ -162,6 +175,10 @@ def create_chain_atomic(
                 chain_id=chain_id,
                 environment=environment,
                 truss_user_env=truss_user_env,
+                original_source_artifact_s3_key=original_source_artifact_s3_key,
+                allow_truss_download=allow_truss_download,
+                deployment_name=deployment_name,
+                team_id=team_id,
             )
         except ApiError as e:
             if (
@@ -182,6 +199,10 @@ def create_chain_atomic(
             dependencies=dependencies,
             chain_name=chain_name,
             truss_user_env=truss_user_env,
+            original_source_artifact_s3_key=original_source_artifact_s3_key,
+            allow_truss_download=allow_truss_download,
+            deployment_name=deployment_name,
+            team_id=team_id,
         )
 
     return ChainDeploymentHandleAtomic(
@@ -189,6 +210,8 @@ def create_chain_atomic(
         chain_id=res["chain_deployment"]["chain"]["id"],
         hostname=res["chain_deployment"]["chain"]["hostname"],
         is_draft=is_draft,
+        original_source_artifact_s3_key=original_source_artifact_s3_key,
+        allow_truss_download=allow_truss_download,
     )
 
 
@@ -342,6 +365,33 @@ def upload_truss(
     return s3_key
 
 
+def upload_chain_artifact(
+    api: BasetenApi,
+    serialize_file: IO,
+    progress_bar: Optional[Type["progress.Progress"]],
+) -> str:
+    """
+    Upload a chain artifact to the Baseten remote.
+
+    Args:
+        api: BasetenApi instance
+        serialize_file: File-like object containing the serialized chain artifact
+
+    Returns:
+        The S3 key of the uploaded file
+    """
+    credentials = api.get_chain_s3_upload_credentials()
+    with handle_client_error("Uploading chain source"):
+        multipart_upload_boto3(
+            serialize_file.name,
+            credentials.s3_bucket,
+            credentials.s3_key,
+            credentials.aws_credentials.model_dump(),
+            progress_bar,
+        )
+    return credentials.s3_key
+
+
 def create_truss_service(
     api: BasetenApi,
     model_name: str,
@@ -357,6 +407,8 @@ def create_truss_service(
     origin: Optional[b10_types.ModelOrigin] = None,
     environment: Optional[str] = None,
     preserve_env_instance_type: bool = True,
+    deploy_timeout_minutes: Optional[int] = None,
+    team_id: Optional[str] = None,
 ) -> ModelVersionHandle:
     """
     Create a model in the Baseten remote.
@@ -372,6 +424,7 @@ def create_truss_service(
             to zero.
         deployment_name: Name to apply to the created deployment. Not applied to
             development model.
+        team_id: ID of the team to create the model in.
 
     Returns:
         A Model Version handle.
@@ -384,6 +437,8 @@ def create_truss_service(
             truss_user_env,
             allow_truss_download=allow_truss_download,
             origin=origin,
+            deploy_timeout_minutes=deploy_timeout_minutes,
+            team_id=team_id,
         )
 
         return ModelVersionHandle(
@@ -398,9 +453,6 @@ def create_truss_service(
         )
 
     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,
             s3_key,
@@ -410,6 +462,9 @@ def create_truss_service(
             allow_truss_download=allow_truss_download,
             deployment_name=deployment_name,
             origin=origin,
+            environment=environment,
+            deploy_timeout_minutes=deploy_timeout_minutes,
+            team_id=team_id,
         )
 
         return ModelVersionHandle(
@@ -434,6 +489,7 @@ def create_truss_service(
             deployment_name=deployment_name,
             environment=environment,
             preserve_env_instance_type=preserve_env_instance_type,
+            deploy_timeout_minutes=deploy_timeout_minutes,
         )
     except ApiError as e:
         if (

truss/remote/baseten/custom_types.py

@@ -120,6 +120,7 @@ class TrussUserEnv(pydantic.BaseModel):
 class BlobType(Enum):
     MODEL = "model"
     TRAIN = "train"
+    CHAIN = "chain"
 
 
 class FileSummary(pydantic.BaseModel):