truss 0.10.9rc535__tar.gz → 0.10.10rc0__tar.gz

{truss-0.10.9rc535 → truss-0.10.10rc0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.10.9rc535
+Version: 0.10.10rc0
 Summary: A seamless bridge from model development to model delivery
 Project-URL: Repository, https://github.com/basetenlabs/truss
 Project-URL: Homepage, https://truss.baseten.co
@@ -37,7 +37,7 @@ Requires-Dist: rich<14,>=13.4.2
 Requires-Dist: ruff>=0.4.8
 Requires-Dist: tenacity>=8.0.1
 Requires-Dist: tomlkit>=0.13.2
-Requires-Dist: truss-transfer==0.0.27
+Requires-Dist: truss-transfer==0.0.28
 Requires-Dist: watchfiles<0.20,>=0.19.0
 Description-Content-Type: text/markdown
 

{truss-0.10.9rc535 → truss-0.10.10rc0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truss"
-version = "0.10.9rc535"
+version = "0.10.10rc0"
 description = "A seamless bridge from model development to model delivery"
 authors = [
     { name = "Pankaj Gupta", email = "no-reply@baseten.co" },
@@ -43,7 +43,7 @@ dependencies = [
     "ruff>=0.4.8",
     "tenacity>=8.0.1",
     "watchfiles>=0.19.0,<0.20",
-    "truss-transfer==0.0.27",
+    "truss-transfer==0.0.28",
 ]
 
 [project.urls]
@@ -78,6 +78,7 @@ dev = [
     "types-aiofiles>=24.1.0",
     "types-requests==2.31.0.2",
     "types-setuptools>=69.0.0.0,<70",
+    "types-python-dateutil>=2.9.0.20250822",
 ]
 dev-server = [
     "aiohttp>3.11.13",

{truss-0.10.9rc535 → truss-0.10.10rc0}/truss/cli/logs/base_watcher.py

@@ -9,7 +9,7 @@ from truss.remote.baseten.api import BasetenApi
 
 POLL_INTERVAL_SEC = 2
 # NB(nikhil): This helps account for (1) log processing delays (2) clock skews
-CLOCK_SKEW_BUFFER_MS = 10000
+CLOCK_SKEW_BUFFER_MS = 60000
 
 
 class LogWatcher(ABC):

{truss-0.10.9rc535 → truss-0.10.10rc0}/truss/cli/train/deploy_checkpoints/deploy_checkpoints.py

@@ -33,6 +33,10 @@ from .deploy_lora_checkpoints import (
     hydrate_lora_checkpoint,
     render_vllm_lora_truss_config,
 )
+from .deploy_whisper_checkpoints import (
+    hydrate_whisper_checkpoint,
+    render_vllm_whisper_truss_config,
+)
 
 HF_TOKEN_ENVVAR_NAME = "HF_TOKEN"
 # If we change this, make sure to update the logic in backend codebase
@@ -178,6 +182,8 @@ def hydrate_checkpoint(
         return hydrate_lora_checkpoint(job_id, checkpoint_id, checkpoint)
     elif checkpoint_type.lower() == ModelWeightsFormat.FULL.value:
         return hydrate_full_checkpoint(job_id, checkpoint_id, checkpoint)
+    elif checkpoint_type.lower() == ModelWeightsFormat.WHISPER.value:
+        return hydrate_whisper_checkpoint(job_id, checkpoint_id, checkpoint)
     else:
         raise ValueError(
             f"Unsupported checkpoint type: {checkpoint_type}. Contact Baseten for support with other checkpoint types."
@@ -196,6 +202,8 @@ def _render_truss_config_for_checkpoint_deployment(
         return render_vllm_lora_truss_config(checkpoint_deploy)
     elif checkpoint_deploy.model_weight_format == ModelWeightsFormat.FULL:
         return render_vllm_full_truss_config(checkpoint_deploy)
+    elif checkpoint_deploy.model_weight_format == ModelWeightsFormat.WHISPER:
+        return render_vllm_whisper_truss_config(checkpoint_deploy)
     else:
         raise ValueError(
             f"Unsupported model weight format: {checkpoint_deploy.model_weight_format}. Please upgrade to the latest Truss version to access the latest supported formats. Contact Baseten if you would like us to support additional formats."
@@ -288,18 +296,6 @@ def _get_checkpoint_ids_to_deploy(
     return checkpoint_ids
 
 
-def _select_single_checkpoint(checkpoint_id_options: List[str]) -> List[str]:
-    """Select a single checkpoint using interactive prompt."""
-    checkpoint_id = inquirer.select(
-        message="Select the checkpoint to deploy:", choices=checkpoint_id_options
-    ).execute()
-
-    if not checkpoint_id:
-        raise click.UsageError("A checkpoint must be selected.")
-
-    return [checkpoint_id]
-
-
 def _select_multiple_checkpoints(checkpoint_id_options: List[str]) -> List[str]:
     """Select multiple checkpoints using interactive checkbox."""
     checkpoint_ids = inquirer.checkbox(
@@ -351,6 +347,8 @@ def _get_base_model_id(user_input: Optional[str], checkpoint: dict) -> Optional[
         )
     elif checkpoint.get("checkpoint_type") == ModelWeightsFormat.FULL.value.lower():
         return None
+    elif checkpoint.get("checkpoint_type") == ModelWeightsFormat.WHISPER.value.lower():
+        return None
     else:
         base_model_id = inquirer.text(message="Enter the base model id.").execute()
     if not base_model_id:
@@ -416,18 +414,28 @@ def _validate_selected_checkpoints(
             "Unable to infer model weight format. Reach out to Baseten for support."
         )
 
-    has_full_checkpoint = any(
-        response_checkpoints[checkpoint_id].get("checkpoint_type")
-        == ModelWeightsFormat.FULL.value
-        for checkpoint_id in checkpoint_ids
-    )
-
-    if has_full_checkpoint and len(checkpoint_ids) > 1:
-        # vLLM does not support multiple checkpoints when any checkpoint is full model weights.
-        raise ValueError(
-            "Full checkpoints are not supported for multiple checkpoints. Please select a single checkpoint."
+    validation_rules = {
+        ModelWeightsFormat.FULL.value: {
+            "error_message": "Full checkpoints are not supported for multiple checkpoints. Please select a single checkpoint.",
+            "reason": "vLLM does not support multiple checkpoints when any checkpoint is full model weights.",
+        },
+        ModelWeightsFormat.WHISPER.value: {
+            "error_message": "Whisper checkpoints are not supported for multiple checkpoints. Please select a single checkpoint.",
+            "reason": "vLLM does not support multiple checkpoints when any checkpoint is whisper model weights.",
+        },
+    }
+
+    # Check each checkpoint type that has restrictions
+    for checkpoint_type, rule in validation_rules.items():
+        has_restricted_checkpoint = any(
+            response_checkpoints[checkpoint_id].get("checkpoint_type")
+            == checkpoint_type
+            for checkpoint_id in checkpoint_ids
         )
 
+        if has_restricted_checkpoint and len(checkpoint_ids) > 1:
+            raise ValueError(rule["error_message"])
+
 
 def get_hf_secret_name(user_input: Union[str, SecretReference, None]) -> str:
     """Get HuggingFace secret name from user input or prompt for it."""

{truss-0.10.9rc535 → truss-0.10.10rc0}/truss/cli/train/deploy_checkpoints/deploy_checkpoints_helpers.py

@@ -3,7 +3,7 @@ from pathlib import Path
 
 from truss.base import truss_config
 from truss.cli.train.types import DeployCheckpointsConfigComplete
-from truss_train.definitions import SecretReference
+from truss_train.definitions import ModelWeightsFormat, SecretReference
 
 START_COMMAND_ENVVAR_NAME = "BT_DOCKER_SERVER_START_CMD"
 
@@ -12,8 +12,14 @@ def setup_base_truss_config(
     checkpoint_deploy: DeployCheckpointsConfigComplete,
 ) -> truss_config.TrussConfig:
     """Set up the base truss config with common properties."""
+    truss_deploy_config = None
+    truss_base_file = (
+        "deploy_from_checkpoint_config_whisper.yml"
+        if checkpoint_deploy.model_weight_format == ModelWeightsFormat.WHISPER
+        else "deploy_from_checkpoint_config.yml"
+    )
     truss_deploy_config = truss_config.TrussConfig.from_yaml(
-        Path(os.path.dirname(__file__), "..", "deploy_from_checkpoint_config.yml")
+        Path(os.path.dirname(__file__), "..", truss_base_file)
     )
     if not truss_deploy_config.docker_server:
         raise ValueError(

{truss-0.10.9rc535 → truss-0.10.10rc0}/truss/cli/train/deploy_checkpoints/deploy_full_checkpoints.py

@@ -14,12 +14,19 @@ from .deploy_checkpoints_helpers import (
     setup_environment_variables_and_secrets,
 )
 
+# NB(aghilan): Transformers was recently changed to save a chat_template.jinja file instead of inside the tokenizer_config.json file.
+# Old Models will not have this file, so we check for it and use it if it exists.
+# vLLM will not automatically resolve the chat_template.jinja file, so we need to pass it to the start command.
+# This logic is needed for any models trained using Transformers v4.51.3 or later
 VLLM_FULL_START_COMMAND = Template(
-    'sh -c "{%if envvars %}{{ envvars }} {% endif %}vllm serve {{ model_path }}'
-    + " --port 8000"
-    + " --tensor-parallel-size {{ specify_tensor_parallelism }}"
-    + " --dtype bfloat16"
-    + '"'
+    "sh -c '{% if envvars %}{{ envvars }} {% endif %}"
+    'HF_TOKEN="$$(cat /secrets/hf_access_token)" && export HF_TOKEN && '
+    "if [ -f {{ model_path }}/chat_template.jinja ]; then "
+    "  vllm serve {{ model_path }} --chat-template {{ model_path }}/chat_template.jinja "
+    "  --port 8000 --tensor-parallel-size {{ specify_tensor_parallelism }} --dtype bfloat16; "
+    "else "
+    "  vllm serve {{ model_path }} --port 8000 --tensor-parallel-size {{ specify_tensor_parallelism }} --dtype bfloat16; "
+    "fi'"
 )
 
 
@@ -33,7 +40,7 @@ def render_vllm_full_truss_config(
         truss_deploy_config, checkpoint_deploy
     )
 
-    checkpoint_str = _build_full_checkpoint_string(truss_deploy_config)
+    checkpoint_str = build_full_checkpoint_string(truss_deploy_config)
 
     accelerator = checkpoint_deploy.compute.accelerator
 
@@ -64,7 +71,7 @@ def hydrate_full_checkpoint(
     return FullCheckpoint(training_job_id=job_id, paths=paths)
 
 
-def _build_full_checkpoint_string(truss_deploy_config) -> str:
+def build_full_checkpoint_string(truss_deploy_config) -> str:
     """Build checkpoint string from artifact references for full checkpoints.
 
     Args:

truss-0.10.10rc0/truss/cli/train/deploy_checkpoints/deploy_whisper_checkpoints.py

@@ -0,0 +1,63 @@
+from jinja2 import Template
+
+from truss.base import truss_config
+from truss.cli.train.deploy_checkpoints.deploy_checkpoints_helpers import (
+    START_COMMAND_ENVVAR_NAME,
+)
+from truss.cli.train.deploy_checkpoints.deploy_full_checkpoints import (
+    build_full_checkpoint_string,
+)
+from truss.cli.train.types import DeployCheckpointsConfigComplete
+from truss_train.definitions import WhisperCheckpoint
+
+from .deploy_checkpoints_helpers import (
+    setup_base_truss_config,
+    setup_environment_variables_and_secrets,
+)
+
+VLLM_WHISPER_START_COMMAND = Template(
+    "sh -c '{% if envvars %}{{ envvars }} {% endif %}"
+    'HF_TOKEN="$$(cat /secrets/hf_access_token)" && export HF_TOKEN && '
+    "vllm serve {{ model_path }} --port 8000 --tensor-parallel-size {{ specify_tensor_parallelism }}'"
+)
+
+
+def render_vllm_whisper_truss_config(
+    checkpoint_deploy: DeployCheckpointsConfigComplete,
+) -> truss_config.TrussConfig:
+    """Render truss config specifically for whisper checkpoints using vLLM."""
+    truss_deploy_config = setup_base_truss_config(checkpoint_deploy)
+
+    start_command_envvars = setup_environment_variables_and_secrets(
+        truss_deploy_config, checkpoint_deploy
+    )
+
+    checkpoint_str = build_full_checkpoint_string(truss_deploy_config)
+
+    accelerator = checkpoint_deploy.compute.accelerator
+
+    start_command_args = {
+        "model_path": checkpoint_str,
+        "envvars": start_command_envvars,
+        "specify_tensor_parallelism": accelerator.count if accelerator else 1,
+    }
+    # Note: we set the start command as an environment variable in supervisord config.
+    # This is so that we don't have to change the supervisord config when the start command changes.
+    # Our goal is to reduce the number of times we need to rebuild the image, and allow us to deploy faster.
+    start_command = VLLM_WHISPER_START_COMMAND.render(**start_command_args)
+    truss_deploy_config.environment_variables[START_COMMAND_ENVVAR_NAME] = start_command
+    # Note: supervisord uses the convention %(ENV_VAR_NAME)s to access environment variable VAR_NAME
+    truss_deploy_config.docker_server.start_command = (  # type: ignore[union-attr]
+        f"%(ENV_{START_COMMAND_ENVVAR_NAME})s"
+    )
+
+    return truss_deploy_config
+
+
+def hydrate_whisper_checkpoint(
+    job_id: str, checkpoint_id: str, checkpoint: dict
+) -> WhisperCheckpoint:
+    """Create a Checkpoint object for whisper model weights."""
+    # NOTE: Slash at the end is important since it means the checkpoint is a directory
+    paths = [f"rank-0/{checkpoint_id}/"]
+    return WhisperCheckpoint(training_job_id=job_id, paths=paths)

truss-0.10.10rc0/truss/cli/train/deploy_from_checkpoint_config_whisper.yml

@@ -0,0 +1,17 @@
+base_image:
+  image:  vllm/vllm-openai:latest
+
+docker_server:
+  start_command: sh -c "" # replaced when deploying
+  readiness_endpoint: /health
+  liveness_endpoint: /health
+  predict_endpoint: /v1/audio/transcriptions
+  server_port: 8000
+runtime:
+  predict_concurrency : 256
+environment_variables:
+  VLLM_LOGGING_LEVEL: WARNING
+  VLLM_USE_V1: 0
+  HF_HUB_ENABLE_HF_TRANSFER: 1
+requirements:
+  - vllm[audio]

truss-0.10.10rc0/truss/cli/train/metrics_watcher.py

@@ -0,0 +1,331 @@
+import signal
+import time
+import traceback
+from typing import Any, Dict, List, Optional, Tuple, cast
+
+from rich.columns import Columns
+from rich.layout import Layout
+from rich.live import Live
+from rich.table import Table
+from rich.text import Text
+
+from truss.cli.train.poller import TrainingPollerMixin
+from truss.cli.utils import common
+from truss.cli.utils.output import console
+from truss.remote.baseten.api import BasetenApi
+
+METRICS_POLL_INTERVAL_SEC = 30
+
+
+class MetricsWatcher(TrainingPollerMixin):
+    live: Optional[Live]
+
+    def __init__(self, api: BasetenApi, project_id: str, job_id: str):
+        super().__init__(api, project_id, job_id)
+
+        self.live = None
+        signal.signal(signal.SIGINT, self._handle_sigint)
+
+    def _handle_sigint(self, signum: int, frame: Any) -> None:
+        if self.live:
+            self.live.stop()
+        msg = f"\n\nExiting training job metrics. To stop the job, run `truss train stop --job-id {self.job_id}`"
+        console.print(msg, style="yellow")
+        raise KeyboardInterrupt()
+
+    def _format_bytes(self, bytes_val: float) -> Tuple[str, str]:
+        """Convert bytes to human readable format"""
+        color_map = {"MB": "green", "GB": "cyan", "TB": "magenta"}
+        unit = "MB"
+        if bytes_val > 1024 * 1024 * 1024 * 1024:
+            unit = "TB"
+        elif bytes_val > 1024 * 1024 * 1024:
+            unit = "GB"
+
+        if unit == "MB":
+            return f"{bytes_val / (1024 * 1024):.2f} MB", color_map[unit]
+        elif unit == "GB":
+            return f"{bytes_val / (1024 * 1024 * 1024):.2f} GB", color_map[unit]
+        return f"{bytes_val:.2f} bytes", color_map[unit]
+
+    def _format_storage_utilization(self, utilization: float) -> Tuple[str, str]:
+        percent = round(utilization * 100, 4)
+        if percent > 90:
+            return f"{percent}%", "red"
+        elif percent > 70:
+            return f"{percent}%", "yellow"
+        return f"{percent}%", "green"
+
+    def _get_latest_metric(self, metrics: List[Dict]) -> Optional[float]:
+        """Get the most recent metric value"""
+        if not metrics:
+            return None
+        return metrics[-1].get("value")
+
+    def _get_latest_storage_metrics(
+        self, storage_data: Optional[Dict[str, List[Dict]]]
+    ) -> Optional[Tuple[int, float]]:
+        if not storage_data:
+            return None
+        usage_data = storage_data.get("usage_bytes")
+        utilization_data = storage_data.get("utilization")
+        if not usage_data or not utilization_data:
+            return None
+        usage_value = usage_data[-1].get("value", None)
+        utilization_value = utilization_data[-1].get("value", None)
+        if not usage_value or not utilization_value:
+            return None
+        return cast(int, usage_value), cast(float, utilization_value)
+
+    def _maybe_format_storage_table_row(
+        self, table: Table, label: str, storage_data: Optional[Dict[str, List[Dict]]]
+    ) -> bool:
+        if not storage_data:
+            return False
+        maybe_values = self._get_latest_storage_metrics(storage_data)
+        if not maybe_values:
+            return False
+        raw_usage, raw_utilization = maybe_values
+        usage_value, usage_color = self._format_bytes(raw_usage)
+        utilization_value, utilization_color = self._format_storage_utilization(
+            raw_utilization
+        )
+        table.add_row(
+            label,
+            Text(usage_value, style=usage_color),
+            Text(utilization_value, style=utilization_color),
+        )
+        return True
+
+    def create_metrics_table(self, metrics_data: Dict) -> Layout:
+        """Create a Rich table with the metrics"""
+        tables = []
+
+        timestamp = self._get_timestamp_from_metrics(metrics_data)
+
+        node_tables = self._create_unified_node_metrics_tables(metrics_data)
+        tables.extend(node_tables)
+
+        storage_tables = self._create_storage_tables(metrics_data)
+        tables.extend(storage_tables)
+
+        columns = Columns(tables, title="Training Job Metrics")
+
+        layout = Layout()
+
+        if timestamp:
+            from rich.panel import Panel
+
+            layout.split_column(
+                Layout(
+                    Panel(
+                        f"🕐 Last Updated: {timestamp}\n💡 Press Ctrl+C to exit",
+                        style="bold cyan",
+                    ),
+                    size=4,
+                ),
+                Layout(columns),
+            )
+        else:
+            layout.split_column(Layout(columns))
+
+        return layout
+
+    def _get_timestamp_from_metrics(self, metrics_data: Dict) -> Optional[str]:
+        """Extract timestamp from metrics data for display"""
+        # Try to get timestamp from per_node_metrics first. Fall back to main metrics if not there.
+        per_node_metrics = metrics_data.get("per_node_metrics", [])
+        if per_node_metrics and len(per_node_metrics) > 0:
+            first_node_metrics = per_node_metrics[0].get("metrics", {})
+            cpu_usage_data = first_node_metrics.get("cpu_usage", [])
+            if cpu_usage_data and len(cpu_usage_data) > 0:
+                timestamp = cpu_usage_data[-1].get("timestamp")
+                if timestamp:
+                    return common.format_localized_time(timestamp)
+
+        cpu_usage_data = metrics_data.get("cpu_usage", [])
+        if cpu_usage_data and len(cpu_usage_data) > 0:
+            timestamp = cpu_usage_data[-1].get("timestamp")
+            if timestamp:
+                return common.format_localized_time(timestamp)
+
+        return None
+
+    def _create_unified_node_metrics_tables(self, metrics_data: Dict) -> List[Table]:
+        """Create tables for node metrics, handling both single and multi-node scenarios"""
+        tables = []
+
+        per_node_metrics = metrics_data.get("per_node_metrics", [])
+
+        if not per_node_metrics:
+            # Job is likely just starting up - it takes some type for the
+            # the metrics to become available after the job starts running.
+            from rich.text import Text
+
+            waiting_table = Table(title="Training Job Status")
+            waiting_table.add_column("Status")
+            waiting_table.add_column("Message")
+
+            waiting_table.add_row(
+                "Status",
+                Text("⏳ Waiting for metrics to become available...", style="yellow"),
+            )
+            waiting_table.add_row(
+                "Note",
+                Text(
+                    "Metrics will appear once the training job starts running.",
+                    style="dim",
+                ),
+            )
+
+            tables.append(waiting_table)
+            return tables
+
+        for node_metrics in per_node_metrics:
+            node_id = node_metrics.get("node_id", "Unknown")
+            metrics = node_metrics.get("metrics", {})
+
+            if not metrics:
+                continue
+
+            table = self._create_node_table(node_id, metrics)
+            tables.append(table)
+
+        return tables
+
+    def _create_node_table(self, node_id: str, metrics: Dict) -> Table:
+        """Create a table for a single node's metrics"""
+        table = Table(title=f"Node: {node_id}")
+        table.add_column("Metric")
+        table.add_column("Value")
+
+        cpu_usage = self._get_latest_metric(metrics.get("cpu_usage", []))
+        if cpu_usage is not None:
+            table.add_row("CPU usage", f"{cpu_usage:.2f} cores")
+
+        cpu_memory = self._get_latest_metric(metrics.get("cpu_memory_usage_bytes", []))
+        if cpu_memory is not None:
+            formatted_value, color = self._format_bytes(cpu_memory)
+            table.add_row("CPU memory", Text(formatted_value, style=color))
+
+        if cpu_usage is not None or cpu_memory is not None:
+            table.add_section()
+
+        gpu_utilization = metrics.get("gpu_utilization", {})
+        gpu_memory = metrics.get("gpu_memory_usage_bytes", {})
+
+        # API should return same GPU IDs for utilization and memory
+        keys = gpu_utilization.keys()
+        for idx, gpu_id in enumerate(keys):
+            latest_util = self._get_latest_metric(gpu_utilization.get(gpu_id, []))
+            if latest_util is not None:
+                table.add_row(f"GPU {gpu_id} utilization", f"{latest_util * 100:.1f}%")
+
+            latest_memory = self._get_latest_metric(gpu_memory.get(gpu_id, []))
+            if latest_memory is not None:
+                formatted_value, color = self._format_bytes(latest_memory)
+                table.add_row(
+                    f"GPU {gpu_id} memory", Text(formatted_value, style=color)
+                )
+
+            if idx != len(keys) - 1:
+                table.add_section()
+
+        ephemeral_storage = metrics.get("ephemeral_storage")
+        if ephemeral_storage:
+            if gpu_utilization or gpu_memory:
+                table.add_section()
+
+            usage_bytes = self._get_latest_metric(
+                ephemeral_storage.get("usage_bytes", [])
+            )
+            utilization = self._get_latest_metric(
+                ephemeral_storage.get("utilization", [])
+            )
+
+            if usage_bytes is not None:
+                formatted_value, color = self._format_bytes(usage_bytes)
+                table.add_row("Eph. storage usage", Text(formatted_value, style=color))
+
+            if utilization is not None:
+                utilization_percent = utilization * 100
+                if utilization_percent > 90:
+                    color = "red"
+                elif utilization_percent > 70:
+                    color = "yellow"
+                else:
+                    color = "green"
+                table.add_row(
+                    "Eph. storage utilization",
+                    Text(f"{utilization_percent:.1f}%", style=color),
+                )
+
+        return table
+
+    def _create_storage_tables(self, metrics_data: Dict) -> List[Table]:
+        """Create storage tables - only cache per job (ephemeral is now in node tables)"""
+        tables = []
+
+        # Create cache storage table (job-level, shown once)
+        cache_storage = metrics_data.get("cache")
+        if cache_storage:
+            table = self._create_cache_storage_table(cache_storage)
+            if table:
+                tables.append(table)
+
+        return tables
+
+    def _create_cache_storage_table(self, cache_storage: Dict) -> Optional[Table]:
+        """Create table for cache storage metrics (job-level)"""
+        usage_bytes = self._get_latest_metric(cache_storage.get("usage_bytes", []))
+        utilization = self._get_latest_metric(cache_storage.get("utilization", []))
+
+        if usage_bytes is None and utilization is None:
+            return None
+
+        table = Table(title="Cache storage")
+        table.add_column("Storage Type")
+        table.add_column("Usage")
+        table.add_column("Utilization")
+
+        self._maybe_format_storage_table_row(table, "Cache storage", cache_storage)
+
+        return table
+
+    def watch(self, refresh_rate: int = METRICS_POLL_INTERVAL_SEC):
+        """Display continuously updating metrics"""
+        self.before_polling()
+        with Live(auto_refresh=False) as live:
+            self.live = live
+            while True:
+                # our first instance of fetching metrics passes no explicit time range. We do this so that we can fetch metrics
+                # for inactive jobs, using the job's completion time to set the time range.
+                # Subsequent queries will fetch only the most recent data to avoid unnecessary load on VM
+                metrics = self.api.get_training_job_metrics(
+                    self.project_id, self.job_id
+                )
+                try:
+                    # range of one minute since we only want the last recording
+                    table = self.create_metrics_table(metrics)
+                    live.update(table, refresh=True)
+                    if not self.should_poll_again():
+                        live.stop()
+                        break
+                    time.sleep(refresh_rate)
+                    end_epoch_millis = int(time.time() * 1000)
+                    start_epoch_millis = end_epoch_millis - 60 * 1000
+                    metrics = self.api.get_training_job_metrics(
+                        self.project_id,
+                        self.job_id,
+                        end_epoch_millis=end_epoch_millis,
+                        start_epoch_millis=start_epoch_millis,
+                    )
+                    self.post_poll()
+                except Exception as e:
+                    live.stop()
+                    console.print(
+                        f"Error fetching metrics: {e}: {traceback.format_exc()}",
+                        style="red",
+                    )
+                    break
+        self.after_polling()

{truss-0.10.9rc535 → truss-0.10.10rc0}/truss/cli/train_commands.py

@@ -13,6 +13,7 @@ from truss.cli.train import common as train_common
 from truss.cli.train import core
 from truss.cli.utils import common
 from truss.cli.utils.output import console, error_console
+from truss.remote.baseten.core import get_training_job_logs_with_pagination
 from truss.remote.baseten.remote import BasetenRemote
 from truss.remote.remote_factory import RemoteFactory
 
@@ -72,8 +73,11 @@ def _prepare_click_context(f: click.Command, params: dict) -> click.Context:
 @click.argument("config", type=Path, required=True)
 @click.option("--remote", type=str, required=False, help="Remote to use")
 @click.option("--tail", is_flag=True, help="Tail for status + logs after push.")
+@click.option("--job-name", type=str, required=False, help="Name of the training job.")
 @common.common_options()
-def push_training_job(config: Path, remote: Optional[str], tail: bool):
+def push_training_job(
+    config: Path, remote: Optional[str], tail: bool, job_name: Optional[str]
+):
     """Run a training job"""
     from truss_train import deployment
 
@@ -84,7 +88,9 @@ def push_training_job(config: Path, remote: Optional[str], tail: bool):
         remote_provider: BasetenRemote = cast(
             BasetenRemote, RemoteFactory.create(remote=remote)
         )
-        job_resp = deployment.create_training_job_from_file(remote_provider, config)
+        job_resp = deployment.create_training_job_from_file(
+            remote_provider, config, job_name
+        )
 
     # Note: This post create logic needs to happen outside the context
     # of the above context manager, as only one console session can be active
@@ -138,7 +144,9 @@ def get_job_logs(
     )
 
     if not tail:
-        logs = remote_provider.api.get_training_job_logs(project_id, job_id)
+        logs = get_training_job_logs_with_pagination(
+            remote_provider.api, project_id, job_id
+        )
         for log in cli_log_utils.parse_logs(logs):
             cli_log_utils.output_log(log)
     else: