truss 0.11.18rc500__py3-none-any.whl → 0.11.24rc2__py3-none-any.whl

truss/remote/baseten/api.py

@@ -200,6 +200,8 @@ class BasetenApi:
         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) {{
@@ -213,6 +215,8 @@ class BasetenApi:
                     {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
@@ -244,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) {{
@@ -257,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
@@ -286,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) {{
@@ -295,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
@@ -327,6 +337,8 @@ class BasetenApi:
         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
@@ -350,10 +362,14 @@ class BasetenApi:
             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)
 
@@ -382,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"""
@@ -456,6 +478,10 @@ class BasetenApi:
             models {
                 id,
                 name
+                team {
+                    id
+                    name
+                }
                 versions{
                     id,
                     semver,
@@ -495,6 +521,10 @@ class BasetenApi:
                 id
                 name
                 hostname
+                team {{
+                    id
+                    name
+                }}
                 versions {{
                     id
                     semver
@@ -647,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"]
 
@@ -903,3 +937,22 @@ class BasetenApi:
         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

@@ -92,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)
 
@@ -132,6 +134,8 @@ def create_chain_atomic(
     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(
@@ -140,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`.
@@ -156,6 +160,8 @@ def create_chain_atomic(
             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
@@ -171,6 +177,8 @@ def create_chain_atomic(
                 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 (
@@ -193,6 +201,8 @@ def create_chain_atomic(
             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(
@@ -397,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.
@@ -412,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.
@@ -424,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(
@@ -448,6 +463,8 @@ def create_truss_service(
             deployment_name=deployment_name,
             origin=origin,
             environment=environment,
+            deploy_timeout_minutes=deploy_timeout_minutes,
+            team_id=team_id,
         )
 
         return ModelVersionHandle(
@@ -472,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/remote.py

@@ -69,6 +69,7 @@ class FinalPushData(custom_types.OracleData):
     origin: Optional[custom_types.ModelOrigin] = None
     environment: Optional[str] = None
     allow_truss_download: bool
+    team_id: Optional[str] = None
 
 
 class BasetenRemote(TrussRemote):
@@ -127,6 +128,8 @@ class BasetenRemote(TrussRemote):
         origin: Optional[custom_types.ModelOrigin] = None,
         environment: Optional[str] = None,
         progress_bar: Optional[Type["progress.Progress"]] = None,
+        deploy_timeout_minutes: Optional[int] = None,
+        team_id: Optional[str] = None,
     ) -> FinalPushData:
         if model_name.isspace():
             raise ValueError("Model name cannot be empty")
@@ -164,6 +167,13 @@ class BasetenRemote(TrussRemote):
                 "Deployment name must only contain alphanumeric, -, _ and . characters"
             )
 
+        if deploy_timeout_minutes is not None and (
+            deploy_timeout_minutes < 10 or deploy_timeout_minutes > 1440
+        ):
+            raise ValueError(
+                "deploy-timeout-minutes must be between 10 minutes and 1440 minutes (24 hours)"
+            )
+
         model_id = exists_model(self._api, model_name)
 
         if model_id is not None and disable_truss_download:
@@ -188,6 +198,7 @@ class BasetenRemote(TrussRemote):
             origin=origin,
             environment=environment,
             allow_truss_download=not disable_truss_download,
+            team_id=team_id,
         )
 
     def push(  # type: ignore
@@ -205,6 +216,8 @@ class BasetenRemote(TrussRemote):
         progress_bar: Optional[Type["progress.Progress"]] = None,
         include_git_info: bool = False,
         preserve_env_instance_type: bool = True,
+        deploy_timeout_minutes: Optional[int] = None,
+        team_id: Optional[str] = None,
     ) -> BasetenService:
         push_data = self._prepare_push(
             truss_handle=truss_handle,
@@ -217,6 +230,8 @@ class BasetenRemote(TrussRemote):
             origin=origin,
             environment=environment,
             progress_bar=progress_bar,
+            deploy_timeout_minutes=deploy_timeout_minutes,
+            team_id=team_id,
         )
 
         if include_git_info:
@@ -242,6 +257,8 @@ class BasetenRemote(TrussRemote):
             environment=push_data.environment,
             truss_user_env=truss_user_env,
             preserve_env_instance_type=preserve_env_instance_type,
+            deploy_timeout_minutes=deploy_timeout_minutes,
+            team_id=push_data.team_id,
         )
 
         if model_version_handle.instance_type_name:
@@ -269,6 +286,8 @@ class BasetenRemote(TrussRemote):
         environment: Optional[str] = None,
         progress_bar: Optional[Type["progress.Progress"]] = None,
         disable_chain_download: bool = False,
+        deployment_name: Optional[str] = None,
+        team_id: Optional[str] = None,
     ) -> ChainDeploymentHandleAtomic:
         # If we are promoting a model to an environment after deploy, it must be published.
         # Draft models cannot be promoted.
@@ -289,6 +308,7 @@ class BasetenRemote(TrussRemote):
                 origin=custom_types.ModelOrigin.CHAINS,
                 progress_bar=progress_bar,
                 disable_truss_download=disable_chain_download,
+                deployment_name=deployment_name,
             )
             oracle_data = custom_types.OracleData(
                 model_name=push_data.model_name,
@@ -326,6 +346,8 @@ class BasetenRemote(TrussRemote):
             environment=environment,
             original_source_artifact_s3_key=raw_chain_s3_key,
             allow_truss_download=not disable_chain_download,
+            deployment_name=deployment_name,
+            team_id=team_id,
         )
         logging.info("Successfully pushed to baseten. Chain is building and deploying.")
         return chain_deployment_handle
@@ -589,5 +611,5 @@ class BasetenRemote(TrussRemote):
     ) -> PatchResult:
         return self._patch(watch_path, truss_ignore_patterns, console=None)
 
-    def upsert_training_project(self, training_project):
-        return self._api.upsert_training_project(training_project)
+    def upsert_training_project(self, training_project, team_id=None):
+        return self._api.upsert_training_project(training_project, team_id=team_id)

truss/templates/control/control/helpers/inference_server_process_controller.py

@@ -49,7 +49,9 @@ class InferenceServerProcessController:
 
     def _terminate_children_and_process(self):
         """Kill child processes first, then parent. Prevents port binding conflicts."""
-        kill_child_processes(self._inference_server_process.pid)
+        # Use a shorter timeout than the truss patch read timeout (=120s):
+        # see remote/baseten/api.py:_post_graphql_query()
+        kill_child_processes(self._inference_server_process.pid, timeout_seconds=30)
         self._inference_server_process.terminate()
 
     def stop(self):

truss/templates/server/requirements.txt

@@ -18,7 +18,7 @@ psutil>=5.9.4
 python-json-logger>=2.0.2
 pyyaml>=6.0.0
 requests>=2.31.0
-truss-transfer==0.0.37
+truss-transfer==0.0.38
 uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0

truss/templates/server.Dockerfile.jinja

@@ -56,12 +56,6 @@ RUN mkdir -p {{ dst.parent }}; curl -L "{{ url }}" -o {{ dst }}
 {% endfor %} {#- endfor external_data_files #}
 {%- endif %} {#- endif external_data_files #}
 
-{%- if build_commands %}
-{% for command in build_commands %}
-RUN {% for secret,path in config.build.secret_to_path_mapping.items() %} --mount=type=secret,id={{ secret }},target={{ path }}{%- endfor %} {{ command }}
-{% endfor %} {#- endfor build_commands #}
-{%- endif %} {#- endif build_commands #}
-
 {# Copy data before code for better caching #}
 {%- if data_dir_exists %}
 COPY --chown={{ default_owner }} ./{{ config.data_dir }} ${APP_HOME}/data
@@ -109,7 +103,13 @@ USER {{ app_username }}
 {%- endif %} {#- endif non_root_user #}
 {%- endmacro -%}
 
-    {%- if config.docker_server %}
+{%- if build_commands %}
+{% for command in build_commands %}
+RUN {% for secret,path in config.build.secret_to_path_mapping.items() %} --mount=type=secret,id={{ secret }},target={{ path }}{%- endfor %} {{ command }}
+{% endfor %} {#- endfor build_commands #}
+{%- endif %} {#- endif build_commands #}
+
+{%- if config.docker_server %}
 RUN apt-get update -y && apt-get install -y --no-install-recommends \
         curl nginx && rm -rf /var/lib/apt/lists/*
 COPY --chown={{ default_owner }} ./docker_server_requirements.txt ${APP_HOME}/docker_server_requirements.txt
@@ -131,7 +131,7 @@ RUN rm -f /etc/nginx/sites-enabled/default
 {{ chown_and_switch_to_regular_user_if_enabled(["/var/lib/nginx", "/var/log/nginx", "/run"]) }}
 ENTRYPOINT ["/docker_server/.venv/bin/supervisord", "-c", "{{ supervisor_config_path }}"]
 
-    {%- elif requires_live_reload %} {#- elif requires_live_reload #}
+{%- elif requires_live_reload %} {#- elif requires_live_reload #}
 ENV HASH_TRUSS="{{ truss_hash }}"
 ENV CONTROL_SERVER_PORT="8080"
 ENV INFERENCE_SERVER_PORT="8090"
@@ -139,11 +139,11 @@ ENV SERVER_START_CMD="/control/.env/bin/python /control/control/server.py"
 {{ chown_and_switch_to_regular_user_if_enabled() }}
 ENTRYPOINT ["/control/.env/bin/python", "/control/control/server.py"]
 
-    {%- else %} {#- else (default inference server) #}
+{%- else %} {#- else (default inference server) #}
 ENV INFERENCE_SERVER_PORT="8080"
 ENV SERVER_START_CMD="{{ python_executable }} /app/main.py"
 {{ chown_and_switch_to_regular_user_if_enabled() }}
 ENTRYPOINT ["{{ python_executable }}", "/app/main.py"]
-    {%- endif %} {#- endif config.docker_server / live_reload #}
+{%- endif %} {#- endif config.docker_server / live_reload #}
 
 {% endblock %} {#- endblock run #}

truss/templates/shared/util.py

@@ -1,7 +1,7 @@
 import multiprocessing
 import os
 import sys
-from typing import List
+from typing import List, Optional
 
 import psutil
 
@@ -62,7 +62,10 @@ def all_processes_dead(procs: List[multiprocessing.Process]) -> bool:
     return True
 
 
-def kill_child_processes(parent_pid: int):
+def kill_child_processes(
+    parent_pid: int,
+    timeout_seconds: Optional[float] = CHILD_PROCESS_WAIT_TIMEOUT_SECONDS,
+):
     try:
         parent = psutil.Process(parent_pid)
     except psutil.NoSuchProcess:
@@ -70,8 +73,6 @@ def kill_child_processes(parent_pid: int):
     children = parent.children(recursive=True)
     for process in children:
         process.terminate()
-    gone, alive = psutil.wait_procs(
-        children, timeout=CHILD_PROCESS_WAIT_TIMEOUT_SECONDS
-    )
+    gone, alive = psutil.wait_procs(children, timeout=timeout_seconds)
     for process in alive:
         process.kill()