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

truss/api/__init__.py

@@ -65,6 +65,7 @@ def push(
     progress_bar: Optional[Type["progress.Progress"]] = None,
     include_git_info: bool = False,
     preserve_env_instance_type: bool = True,
+    deploy_timeout_minutes: Optional[int] = None,
 ) -> definitions.ModelDeployment:
     """
     Pushes a Truss to Baseten.
@@ -77,13 +78,13 @@ def push(
             promote the truss to production after deploy completes.
         promote: Push the truss as a published deployment. Even if a production deployment exists,
             promote the truss to production after deploy completes.
-        preserve_previous_production_deployment: Preserve the previous production deployment’s autoscaling
+        preserve_previous_production_deployment: Preserve the previous production deployment's autoscaling
             setting. When not specified, the previous production deployment will be updated to allow it to
             scale to zero. Can only be use in combination with `promote` option.
         trusted: [DEPRECATED]
         deployment_name: Name of the deployment created by the push. Can only be
             used in combination with `publish` or `promote`. Deployment name must
-            only contain alphanumeric, ’.’, ’-’ or ’_’ characters.
+            only contain alphanumeric, '.', '-' or '_' characters.
         environment: Name of stable environment on baseten.
         progress_bar: Optional `rich.progress.Progress` if output is desired.
         include_git_info: Whether to attach git versioning info (sha, branch, tag) to
@@ -92,6 +93,7 @@ def push(
         preserve_env_instance_type: When pushing a truss to an environment, whether to use the resources
           specified in the truss config to resolve the instance type or preserve the instance type
           configured in the specified environment.
+        deploy_timeout_minutes: Optional timeout in minutes for the deployment operation.
 
     Returns:
         The newly created ModelDeployment.
@@ -135,6 +137,7 @@ def push(
         progress_bar=progress_bar,
         include_git_info=include_git_info,
         preserve_env_instance_type=preserve_env_instance_type,
+        deploy_timeout_minutes=deploy_timeout_minutes,
     )  # type: ignore
 
     return definitions.ModelDeployment(cast(BasetenService, service))

truss/base/truss_config.py

@@ -147,7 +147,7 @@ class ModelRepo(custom_types.ConfigModel):
     volume_folder: Optional[
         Annotated[str, pydantic.StringConstraints(min_length=1)]
     ] = None
-    use_volume: bool = False
+    use_volume: bool
     kind: ModelRepoSourceKind = ModelRepoSourceKind.HF
     runtime_secret_name: str = "hf_access_token"
 
@@ -163,7 +163,7 @@ class ModelRepo(custom_types.ConfigModel):
             return v
         if v.get("kind") == ModelRepoSourceKind.HF.value and v.get("revision") is None:
             logger.warning(
-                "the key `revision: str` is required for use_volume=True huggingface repos."
+                "the key `revision: str` is required for use_volume=True huggingface repos. For S3/GCS/Azure repos, set it to any non-empty string."
             )
             raise_insufficent_revision(v.get("repo_id"), v.get("revision"))
         if v.get("volume_folder") is None or len(v["volume_folder"]) == 0:
@@ -202,7 +202,14 @@ class ModelCache(pydantic.RootModel[list[ModelRepo]]):
             )
 
 
-class CacheInternal(ModelCache): ...
+class ModelRepoCacheInternal(ModelRepo):
+    use_volume: bool = False  # override
+
+
+class CacheInternal(pydantic.RootModel[list[ModelRepoCacheInternal]]):
+    @property
+    def models(self) -> list[ModelRepoCacheInternal]:
+        return self.root
 
 
 class HealthChecks(custom_types.ConfigModel):

truss/cli/chains_commands.py

@@ -1,6 +1,6 @@
 import time
 from pathlib import Path
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, cast
 
 import rich
 import rich.live
@@ -14,10 +14,13 @@ from rich import progress
 
 from truss.cli import remote_cli
 from truss.cli.cli import truss_cli
+from truss.cli.resolvers.chain_team_resolver import resolve_chain_team_name
 from truss.cli.utils import common, output
 from truss.cli.utils.output import console
 from truss.remote.baseten.core import ACTIVE_STATUS, DEPLOYING_STATUSES
+from truss.remote.baseten.remote import BasetenRemote
 from truss.remote.baseten.utils.status import get_displayable_status
+from truss.remote.remote_factory import RemoteFactory
 from truss.util import user_config
 from truss.util.log_utils import LogInterceptor
 
@@ -219,6 +222,22 @@ def _create_chains_table(service) -> Tuple[rich.table.Table, List[str]]:
     default=False,
     help="Disable downloading of pushed chain source code from the UI.",
 )
+@click.option(
+    "--deployment-name",
+    type=str,
+    required=False,
+    help=(
+        "Name of the deployment created by the publish. Can only be used "
+        "in combination with '--publish' or '--promote'."
+    ),
+)
+@click.option(
+    "--team",
+    "provided_team_name",
+    type=str,
+    required=False,
+    help="Team name for the chain deployment",
+)
 @click.pass_context
 @common.common_options()
 def push_chain(
@@ -236,6 +255,8 @@ def push_chain(
     experimental_watch_chainlet_names: Optional[str],
     include_git_info: bool = False,
     disable_chain_download: bool = False,
+    deployment_name: Optional[str] = None,
+    provided_team_name: Optional[str] = None,
 ) -> None:
     """
     Deploys a chain remotely.
@@ -280,10 +301,24 @@ def push_chain(
     if not include_git_info:
         include_git_info = user_config.settings.include_git_info
 
+    # Resolve team if not in dryrun mode
+    team_id = None
     with framework.ChainletImporter.import_target(source, entrypoint) as entrypoint_cls:
         chain_name = (
             name or entrypoint_cls.meta_data.chain_name or entrypoint_cls.display_name
         )
+
+        remote_provider = None
+        if not dryrun and remote:
+            remote_provider = cast(BasetenRemote, RemoteFactory.create(remote=remote))
+            existing_teams = remote_provider.api.get_teams()
+            _, team_id = resolve_chain_team_name(
+                remote_provider,
+                provided_team_name,
+                existing_chain_name=chain_name,
+                existing_teams=existing_teams,
+            )
+
         options = chains_def.PushOptionsBaseten.create(
             chain_name=chain_name,
             promote=promote,
@@ -294,6 +329,9 @@ def push_chain(
             include_git_info=include_git_info,
             working_dir=source.parent if source.is_file() else source.resolve(),
             disable_chain_download=disable_chain_download,
+            deployment_name=deployment_name,
+            team_id=team_id,
+            remote_provider=remote_provider,
         )
         service = deployment_client.push(
             entrypoint_cls, options, progress_bar=progress.Progress

truss/cli/cli.py

@@ -19,6 +19,7 @@ from truss.base.truss_config import Build, ModelServer, TransportKind
 from truss.cli import remote_cli
 from truss.cli.logs import utils as cli_log_utils
 from truss.cli.logs.model_log_watcher import ModelDeploymentLogWatcher
+from truss.cli.resolvers.model_team_resolver import resolve_model_team_name
 from truss.cli.utils import common
 from truss.cli.utils.output import console, error_console
 from truss.remote.baseten.core import (
@@ -462,7 +463,7 @@ def run_python(script, target_directory):
     required=False,
     help=(
         "Name of the deployment created by the push. Can only be "
-        "used in combination with '--publish' or '--promote'."
+        "used in combination with --publish or --promote."
     ),
 )
 @click.option(
@@ -501,6 +502,19 @@ def run_python(script, target_directory):
         "Default is --preserve-env-instance-type."
     ),
 )
+@click.option(
+    "--deploy-timeout-minutes",
+    type=int,
+    required=False,
+    help="Timeout in minutes for the deploy operation.",
+)
+@click.option(
+    "--team",
+    "provided_team_name",
+    type=str,
+    required=False,
+    help="Team name for the model",
+)
 @common.common_options()
 def push(
     target_directory: str,
@@ -518,6 +532,8 @@ def push(
     include_git_info: bool = False,
     tail: bool = False,
     preserve_env_instance_type: bool = True,
+    deploy_timeout_minutes: Optional[int] = None,
+    provided_team_name: Optional[str] = None,
 ) -> None:
     """
     Pushes a truss to a TrussRemote.
@@ -547,6 +563,17 @@ def push(
     if not model_name:
         model_name = remote_cli.inquire_model_name()
 
+    # Resolve team_id if BasetenRemote
+    team_id = None
+    if isinstance(remote_provider, BasetenRemote):
+        existing_teams = remote_provider.api.get_teams()
+        _, team_id = resolve_model_team_name(
+            remote_provider=remote_provider,
+            provided_team_name=provided_team_name,
+            existing_model_name=model_name,
+            existing_teams=existing_teams,
+        )
+
     if promote and environment:
         raise click.UsageError(
             "'promote' flag and 'environment' flag cannot both be specified."
@@ -612,11 +639,12 @@ def push(
             console.print(fp8_and_num_builder_gpus_text, style="yellow")
 
     source = Path(target_directory)
-    # TODO(Abu): This needs to be refactored to be more generic
+    working_dir = source.parent if source.is_file() else source.resolve()
+
     service = remote_provider.push(
-        tr,
+        truss_handle=tr,
         model_name=model_name,
-        working_dir=source.parent if source.is_file() else source.resolve(),
+        working_dir=working_dir,
         publish=publish,
         promote=promote,
         preserve_previous_prod_deployment=preserve_previous_production_deployment,
@@ -626,7 +654,9 @@ def push(
         progress_bar=progress.Progress,
         include_git_info=include_git_info,
         preserve_env_instance_type=preserve_env_instance_type,
-    )  # type: ignore
+        deploy_timeout_minutes=deploy_timeout_minutes,
+        team_id=team_id,
+    )
 
     click.echo(f"✨ Model {model_name} was successfully pushed ✨")
 

truss/cli/remote_cli.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from InquirerPy import inquirer
 from InquirerPy.validator import ValidationError, Validator
 
@@ -56,3 +58,30 @@ def inquire_remote_name() -> str:
 
 def inquire_model_name() -> str:
     return inquirer.text("📦 Name this model:", qmark="").execute()
+
+
+def get_team_id_from_name(
+    teams: dict[str, dict[str, str]], team_name: str
+) -> Optional[str]:
+    team = teams.get(team_name)
+    return team["id"] if team else None
+
+
+def format_available_teams(teams: dict[str, dict[str, str]]) -> str:
+    team_names = list(teams.keys())
+    return ", ".join(team_names) if team_names else "none"
+
+
+def inquire_team(
+    existing_teams: Optional[dict[str, dict[str, str]]] = None,
+) -> Optional[str]:
+    if existing_teams is not None:
+        selected_team_name = inquirer.select(
+            "👥 Which team do you want to push to?",
+            qmark="",
+            choices=list[str](existing_teams.keys()),
+        ).execute()
+        return selected_team_name
+
+    # If no existing teams, return None (don't propagate team param)
+    return None

truss/cli/resolvers/chain_team_resolver.py

@@ -0,0 +1,82 @@
+"""Team resolution logic for chains."""
+
+from typing import Optional
+
+import click
+
+from truss.cli import remote_cli
+from truss.remote.baseten.remote import BasetenRemote
+
+
+def resolve_chain_team_name(
+    remote_provider: BasetenRemote,
+    provided_team_name: Optional[str],
+    existing_chain_name: Optional[str] = None,
+    existing_teams: Optional[dict[str, dict[str, str]]] = None,
+) -> tuple[Optional[str], Optional[str]]:
+    """Resolve team name and team_id from provided team name or by prompting the user.
+    Returns a tuple of (team_name, team_id).
+    This function handles 8 distinct scenarios organized into 3 high-level categories:
+
+    HIGH-LEVEL SCENARIO 1: --team PROVIDED
+        SCENARIO 1: Valid team name, user has access
+            → Returns (team_name, team_id) for that team (no prompt, no error)
+        SCENARIO 2: Invalid team name (does not exist)
+            → Raises ClickException with error message listing available teams
+
+    HIGH-LEVEL SCENARIO 2: --team NOT PROVIDED, Chain does not exist
+        SCENARIO 3: User has multiple teams, no existing chain
+            → Prompts user to select a team via inquire_team()
+        SCENARIO 6: User has exactly one team, no existing chain
+            → Returns (team_name, team_id) for the single team automatically (no prompt)
+
+    HIGH-LEVEL SCENARIO 3: --team NOT PROVIDED, Chain exists
+        SCENARIO 4: User has multiple teams, existing chain in exactly one team
+            → Auto-detects and returns (team_name, team_id) for that team (no prompt)
+        SCENARIO 5: User has multiple teams, existing chain exists in multiple teams
+            → Prompts user to select a team via inquire_team()
+        SCENARIO 7: User has exactly one team, existing chain matches the team
+            → Auto-detects and returns (team_name, team_id) for the single team (no prompt)
+        SCENARIO 8: User has exactly one team, existing chain exists in different team
+            → Returns (team_name, team_id) for the single team automatically (no prompt, uses user's only team)
+    """
+    if existing_teams is None:
+        existing_teams = remote_provider.api.get_teams()
+
+    def _get_team_id(team_name: Optional[str]) -> Optional[str]:
+        if team_name and existing_teams:
+            team_data = existing_teams.get(team_name)
+            return team_data["id"] if team_data else None
+        return None
+
+    if provided_team_name is not None:
+        if provided_team_name not in existing_teams:
+            available_teams_str = remote_cli.format_available_teams(existing_teams)
+            raise click.ClickException(
+                f"Team '{provided_team_name}' does not exist. Available teams: {available_teams_str}"
+            )
+        return (provided_team_name, _get_team_id(provided_team_name))
+
+    existing_chains = None
+    if existing_chain_name is not None:
+        existing_chains = remote_provider.api.get_chains()
+        matching_chains = [
+            c for c in existing_chains if c.get("name") == existing_chain_name
+        ]
+
+        if len(matching_chains) > 1:
+            selected_team_name = remote_cli.inquire_team(existing_teams=existing_teams)
+            return (selected_team_name, _get_team_id(selected_team_name))
+
+        if len(matching_chains) == 1:
+            chain_team = matching_chains[0].get("team")
+            chain_team_name = chain_team.get("name") if chain_team else None
+            if chain_team_name and chain_team_name in existing_teams:
+                return (chain_team_name, _get_team_id(chain_team_name))
+
+    if len(existing_teams) == 1:
+        single_team_name = list(existing_teams.keys())[0]
+        return (single_team_name, _get_team_id(single_team_name))
+
+    selected_team_name = remote_cli.inquire_team(existing_teams=existing_teams)
+    return (selected_team_name, _get_team_id(selected_team_name))

truss/cli/resolvers/model_team_resolver.py

@@ -0,0 +1,90 @@
+"""Team resolution logic for models."""
+
+from typing import Optional
+
+import click
+
+from truss.cli import remote_cli
+from truss.remote.baseten.remote import BasetenRemote
+
+
+def resolve_model_team_name(
+    remote_provider: BasetenRemote,
+    provided_team_name: Optional[str],
+    existing_model_name: Optional[str] = None,
+    existing_teams: Optional[dict[str, dict[str, str]]] = None,
+) -> tuple[Optional[str], Optional[str]]:
+    """Resolve team name and team_id from provided team name or by prompting the user.
+    Returns a tuple of (team_name, team_id).
+    This function handles 8 distinct scenarios organized into 3 high-level categories:
+
+    HIGH-LEVEL SCENARIO 1: --team PROVIDED
+        SCENARIO 1: Valid team name, user has access
+            → Returns (team_name, team_id) for that team (no prompt, no error)
+        SCENARIO 2: Invalid team name (does not exist)
+            → Raises ClickException with error message listing available teams
+
+    HIGH-LEVEL SCENARIO 2: --team NOT PROVIDED, Model does not exist
+        SCENARIO 3: User has multiple teams, no existing model
+            → Prompts user to select a team via inquire_team()
+        SCENARIO 6: User has exactly one team, no existing model
+            → Returns (team_name, team_id) for the single team automatically (no prompt)
+
+    HIGH-LEVEL SCENARIO 3: --team NOT PROVIDED, Model exists
+        SCENARIO 4: User has multiple teams, existing model in exactly one team
+            → Auto-detects and returns (team_name, team_id) for that team (no prompt)
+        SCENARIO 5: User has multiple teams, existing model exists in multiple teams
+            → Prompts user to select a team via inquire_team()
+        SCENARIO 7: User has exactly one team, existing model matches the team
+            → Auto-detects and returns (team_name, team_id) for the single team (no prompt)
+        SCENARIO 8: User has exactly one team, existing model exists in different team
+            → Returns (team_name, team_id) for the single team automatically (no prompt, uses user's only team)
+    """
+    if existing_teams is None:
+        existing_teams = remote_provider.api.get_teams()
+
+    def _get_team_id(team_name: Optional[str]) -> Optional[str]:
+        if team_name and existing_teams:
+            team_data = existing_teams.get(team_name)
+            return team_data["id"] if team_data else None
+        return None
+
+    def _get_matching_models_in_accessible_teams(model_name: str) -> list[dict]:
+        """Get models matching the name that are in teams the user has access to."""
+        all_models_data = remote_provider.api.models()
+        accessible_team_ids = {team_data["id"] for team_data in existing_teams.values()}
+
+        return [
+            m
+            for m in all_models_data.get("models", [])
+            if m.get("name") == model_name
+            and m.get("team", {}).get("id") in accessible_team_ids
+        ]
+
+    if provided_team_name is not None:
+        if provided_team_name not in existing_teams:
+            available_teams_str = remote_cli.format_available_teams(existing_teams)
+            raise click.ClickException(
+                f"Team '{provided_team_name}' does not exist. Available teams: {available_teams_str}"
+            )
+        return (provided_team_name, _get_team_id(provided_team_name))
+
+    if existing_model_name is not None:
+        matching_models = _get_matching_models_in_accessible_teams(existing_model_name)
+
+        if len(matching_models) == 1:
+            # Exactly one model in an accessible team - auto-detect
+            team = matching_models[0].get("team", {})
+            model_team_name = team.get("name")
+            model_team_id = team.get("id")
+            if model_team_name and model_team_name in existing_teams:
+                return (model_team_name, model_team_id)
+        # If len > 1, multiple models exist - fall through to prompt logic
+        # If len == 0, no models exist - fall through to prompt logic
+
+    if len(existing_teams) == 1:
+        single_team_name = list(existing_teams.keys())[0]
+        return (single_team_name, _get_team_id(single_team_name))
+
+    selected_team_name = remote_cli.inquire_team(existing_teams=existing_teams)
+    return (selected_team_name, _get_team_id(selected_team_name))

truss/cli/resolvers/training_project_team_resolver.py

@@ -0,0 +1,81 @@
+"""Team resolution logic for training projects."""
+
+from typing import Optional
+
+import click
+
+from truss.cli import remote_cli
+from truss.remote.baseten.remote import BasetenRemote
+
+
+def resolve_training_project_team_name(
+    remote_provider: BasetenRemote,
+    provided_team_name: Optional[str],
+    existing_project_name: Optional[str] = None,
+    existing_teams: Optional[dict[str, dict[str, str]]] = None,
+) -> tuple[Optional[str], Optional[str]]:
+    """Resolve team name and team_id from provided team name or by prompting the user.
+    Returns a tuple of (team_name, team_id).
+    This function handles 8 distinct scenarios organized into 3 high-level categories:
+
+    HIGH-LEVEL SCENARIO 1: --team PROVIDED
+        SCENARIO 1: Valid team name, user has access
+            → Returns (team_name, team_id) for that team (no prompt, no error)
+        SCENARIO 2: Invalid team name (does not exist)
+            → Raises ClickException with error message listing available teams
+
+    HIGH-LEVEL SCENARIO 2: --team NOT PROVIDED, Training project does not exist
+        SCENARIO 3: User has multiple teams, no existing project
+            → Prompts user to select a team via inquire_team()
+        SCENARIO 6: User has exactly one team, no existing project
+            → Returns (team_name, team_id) for the single team automatically (no prompt)
+
+    HIGH-LEVEL SCENARIO 3: --team NOT PROVIDED, Training project exists
+        SCENARIO 4: User has multiple teams, existing project in exactly one team
+            → Auto-detects and returns (team_name, team_id) for that team (no prompt)
+        SCENARIO 5: User has multiple teams, existing project exists in multiple teams
+            → Prompts user to select a team via inquire_team()
+        SCENARIO 7: User has exactly one team, existing project matches the team
+            → Auto-detects and returns (team_name, team_id) for the single team (no prompt)
+        SCENARIO 8: User has exactly one team, existing project exists in different team
+            → Returns (team_name, team_id) for the single team automatically (no prompt, uses user's only team)
+    """
+    if existing_teams is None:
+        existing_teams = remote_provider.api.get_teams()
+
+    def _get_team_id(team_name: Optional[str]) -> Optional[str]:
+        if team_name and existing_teams:
+            team_data = existing_teams.get(team_name)
+            return team_data["id"] if team_data else None
+        return None
+
+    if provided_team_name is not None:
+        if provided_team_name not in existing_teams:
+            available_teams_str = remote_cli.format_available_teams(existing_teams)
+            raise click.ClickException(
+                f"Team '{provided_team_name}' does not exist. Available teams: {available_teams_str}"
+            )
+        return (provided_team_name, _get_team_id(provided_team_name))
+
+    existing_projects = None
+    if existing_project_name is not None:
+        existing_projects = remote_provider.api.list_training_projects()
+        matching_projects = [
+            p for p in existing_projects if p.get("name") == existing_project_name
+        ]
+
+        if len(matching_projects) > 1:
+            selected_team_name = remote_cli.inquire_team(existing_teams=existing_teams)
+            return (selected_team_name, _get_team_id(selected_team_name))
+
+        if len(matching_projects) == 1:
+            project_team_name = matching_projects[0].get("team_name")
+            if project_team_name in existing_teams:
+                return (project_team_name, _get_team_id(project_team_name))
+
+    if len(existing_teams) == 1:
+        single_team_name = list(existing_teams.keys())[0]
+        return (single_team_name, _get_team_id(single_team_name))
+
+    selected_team_name = remote_cli.inquire_team(existing_teams=existing_teams)
+    return (selected_team_name, _get_team_id(selected_team_name))