truss 0.10.11rc2__tar.gz → 0.10.12__tar.gz

{truss-0.10.11rc2 → truss-0.10.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.10.11rc2
+Version: 0.10.12
 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

{truss-0.10.11rc2 → truss-0.10.12}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truss"
-version = "0.10.11rc2"
+version = "0.10.12"
 description = "A seamless bridge from model development to model delivery"
 authors = [
     { name = "Pankaj Gupta", email = "no-reply@baseten.co" },

{truss-0.10.11rc2 → truss-0.10.12}/truss/cli/train/common.py

@@ -5,6 +5,12 @@ import rich_click as click
 from truss.cli.utils.output import console
 from truss.remote.baseten import BasetenRemote
 
+# Byte size constants
+BYTES_PER_KB = 1000
+BYTES_PER_MB = BYTES_PER_KB * 1000
+BYTES_PER_GB = BYTES_PER_MB * 1000
+BYTES_PER_TB = BYTES_PER_GB * 1000
+
 
 def get_most_recent_job(
     remote_provider: BasetenRemote, project_id: Optional[str], job_id: Optional[str]
@@ -27,3 +33,16 @@ def get_most_recent_job(
         project_id = cast(str, job["training_project"]["id"])
         job_id = cast(str, job["id"])
     return project_id, job_id
+
+
+def format_bytes_to_human_readable(bytes: int) -> str:
+    if bytes > BYTES_PER_TB:
+        return f"{bytes / BYTES_PER_TB:.2f} TB"
+    if bytes > BYTES_PER_GB:
+        return f"{bytes / BYTES_PER_GB:.2f} GB"
+    elif bytes > BYTES_PER_MB:
+        return f"{bytes / BYTES_PER_MB:.2f} MB"
+    elif bytes > BYTES_PER_KB:
+        return f"{bytes / BYTES_PER_KB:.2f} KB"
+    else:
+        return f"{bytes} B"

{truss-0.10.11rc2 → truss-0.10.12}/truss/cli/train/core.py

@@ -20,6 +20,14 @@ from truss.remote.baseten.remote import BasetenRemote
 from truss_train import loader
 from truss_train.definitions import DeployCheckpointsConfig
 
+SORT_BY_FILEPATH = "filepath"
+SORT_BY_SIZE = "size"
+SORT_BY_MODIFIED = "modified"
+SORT_BY_TYPE = "type"
+SORT_BY_PERMISSIONS = "permissions"
+SORT_ORDER_ASC = "asc"
+SORT_ORDER_DESC = "desc"
+
 ACTIVE_JOB_STATUSES = [
     "TRAINING_JOB_RUNNING",
     "TRAINING_JOB_CREATED",
@@ -401,3 +409,131 @@ def download_checkpoint_artifacts(
 
 def status_page_url(remote_url: str, training_job_id: str) -> str:
     return f"{remote_url}/training/jobs/{training_job_id}"
+
+
+def fetch_project_by_name_or_id(
+    remote_provider: BasetenRemote, project_identifier: str
+) -> dict:
+    """Fetch a training project by name or ID.
+
+    Args:
+        remote_provider: The remote provider instance
+        project_identifier: Either a project ID or project name
+
+    Returns:
+        The project object as a dictionary
+
+    Raises:
+        click.ClickException: If the project is not found
+    """
+    try:
+        projects = remote_provider.api.list_training_projects()
+        projects_by_name = {project.get("name"): project for project in projects}
+        projects_by_id = {project.get("id"): project for project in projects}
+        if project_identifier in projects_by_id:
+            return projects_by_id[project_identifier]
+        if project_identifier in projects_by_name:
+            return projects_by_name[project_identifier]
+        valid_project_ids_and_names = ", ".join(
+            [f"{project.get('id')} ({project.get('name')})" for project in projects]
+        )
+        raise click.ClickException(
+            f"Project '{project_identifier}' not found. Valid project IDs and names: {valid_project_ids_and_names}"
+        )
+    except click.ClickException:
+        raise
+    except Exception as e:
+        raise click.ClickException(f"Error fetching project: {str(e)}")
+
+
+def view_cache_summary(
+    remote_provider: BasetenRemote,
+    project_id: str,
+    sort_by: str = SORT_BY_FILEPATH,
+    order: str = SORT_ORDER_ASC,
+):
+    """View cache summary for a training project."""
+    try:
+        cache_data = remote_provider.api.get_cache_summary(project_id)
+
+        if not cache_data:
+            console.print("No cache summary found for this project.", style="yellow")
+            return
+
+        table = rich.table.Table(title=f"Cache summary for project: {project_id}")
+        table.add_column("File Path", style="cyan")
+        table.add_column("Size", style="green")
+        table.add_column("Modified", style="yellow")
+        table.add_column("Type")
+        table.add_column("Permissions", style="magenta")
+
+        files = cache_data.get("file_summaries", [])
+        if not files:
+            console.print("No files found in cache.", style="yellow")
+            return
+
+        reverse = order == SORT_ORDER_DESC
+
+        if sort_by == SORT_BY_FILEPATH:
+            files.sort(key=lambda x: x.get("path", ""), reverse=reverse)
+        elif sort_by == SORT_BY_SIZE:
+            files.sort(key=lambda x: x.get("size_bytes", 0), reverse=reverse)
+        elif sort_by == SORT_BY_MODIFIED:
+            files.sort(key=lambda x: x.get("modified", ""), reverse=reverse)
+        elif sort_by == SORT_BY_TYPE:
+            files.sort(key=lambda x: x.get("file_type", ""), reverse=reverse)
+        elif sort_by == SORT_BY_PERMISSIONS:
+            files.sort(key=lambda x: x.get("permissions", ""), reverse=reverse)
+
+        total_size = 0
+        for file_info in files:
+            total_size += file_info.get("size_bytes", 0)
+
+        total_size_str = common.format_bytes_to_human_readable(total_size)
+
+        console.print(
+            f"📅 Cache captured at: {cache_data.get('timestamp', 'Unknown')}",
+            style="bold blue",
+        )
+        console.print(
+            f"📁 Project ID: {cache_data.get('project_id', 'Unknown')}",
+            style="bold blue",
+        )
+        console.print()
+        console.print(f"📊 Total files: {len(files)}", style="bold green")
+        console.print(f"💾 Total size: {total_size_str}", style="bold green")
+        console.print()
+
+        for file_info in files:
+            size_bytes = file_info.get("size_bytes", 0)
+
+            size_str = cli_common.format_bytes_to_human_readable(int(size_bytes))
+
+            modified_str = cli_common.format_localized_time(
+                file_info.get("modified", "Unknown")
+            )
+
+            table.add_row(
+                file_info.get("path", "Unknown"),
+                size_str,
+                modified_str,
+                file_info.get("file_type", "Unknown"),
+                file_info.get("permissions", "Unknown"),
+            )
+
+        console.print(table)
+
+    except Exception as e:
+        console.print(f"Error fetching cache summary: {str(e)}", style="red")
+        raise
+
+
+def view_cache_summary_by_project(
+    remote_provider: BasetenRemote,
+    project_identifier: str,
+    sort_by: str = SORT_BY_FILEPATH,
+    order: str = SORT_ORDER_ASC,
+):
+    """View cache summary for a training project by ID or name."""
+    project = fetch_project_by_name_or_id(remote_provider, project_identifier)
+    view_cache_summary(remote_provider, project["id"], sort_by, order)

{truss-0.10.11rc2 → truss-0.10.12}/truss/cli/train/metrics_watcher.py

@@ -35,18 +35,17 @@ class MetricsWatcher(TrainingPollerMixin):
 
     def _format_bytes(self, bytes_val: float) -> Tuple[str, str]:
         """Convert bytes to human readable format"""
+        default_color = "green"
         color_map = {"MB": "green", "GB": "cyan", "TB": "magenta"}
-        unit = "MB"
-        if bytes_val > 1024 * 1024 * 1024 * 1024:
+        unit = "B"
+        if bytes_val > 1000 * 1000 * 1000 * 1000:
             unit = "TB"
-        elif bytes_val > 1024 * 1024 * 1024:
+        elif bytes_val > 1000 * 1000 * 1000:
             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]
+        elif bytes_val > 1000 * 1000:
+            unit = "MB"
+        color = color_map.get(unit, default_color)
+        return (common.format_bytes_to_human_readable(int(bytes_val)), color)
 
     def _format_storage_utilization(self, utilization: float) -> Tuple[str, str]:
         percent = round(utilization * 100, 4)

{truss-0.10.11rc2 → truss-0.10.12}/truss/cli/train_commands.py

@@ -11,11 +11,21 @@ from truss.cli.logs import utils as cli_log_utils
 from truss.cli.logs.training_log_watcher import TrainingLogWatcher
 from truss.cli.train import common as train_common
 from truss.cli.train import core
+from truss.cli.train.core import (
+    SORT_BY_FILEPATH,
+    SORT_BY_MODIFIED,
+    SORT_BY_PERMISSIONS,
+    SORT_BY_SIZE,
+    SORT_BY_TYPE,
+    SORT_ORDER_ASC,
+    SORT_ORDER_DESC,
+)
 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
+from truss_train import TrainingJob
 
 
 @click.group()
@@ -27,15 +37,30 @@ truss_cli.add_command(train)
 
 
 def _print_training_job_success_message(
-    job_id: str, remote_provider: BasetenRemote
+    job_id: str,
+    project_name: str,
+    job_object: TrainingJob,
+    remote_provider: BasetenRemote,
 ) -> None:
     """Print success message and helpful commands for a training job."""
     console.print("✨ Training job successfully created!", style="green")
+    should_print_cache_summary = (
+        job_object.runtime.enable_cache
+        or job_object.runtime.cache_config
+        and job_object.runtime.cache_config.enabled
+    )
+    cache_summary_snippet = ""
+    if should_print_cache_summary:
+        cache_summary_snippet = (
+            f"📁 View cache summary via "
+            f"[cyan]'truss train cache summarize --project {project_name}'[/cyan]\n"
+        )
     console.print(
         f"🪵 View logs for your job via "
         f"[cyan]'truss train logs --job-id {job_id} --tail'[/cyan]\n"
         f"🔍 View metrics for your job via "
         f"[cyan]'truss train metrics --job-id {job_id}'[/cyan]\n"
+        f"{cache_summary_snippet}"
         f"🌐 Status page: {common.format_link(core.status_page_url(remote_provider.remote_url, job_id))}"
     )
 
@@ -44,6 +69,7 @@ def _handle_post_create_logic(
     job_resp: dict, remote_provider: BasetenRemote, tail: bool
 ) -> None:
     project_id, job_id = job_resp["training_project"]["id"], job_resp["id"]
+    project_name = job_resp["training_project"]["name"]
 
     if job_resp.get("current_status", None) == "TRAINING_JOB_QUEUED":
         console.print(
@@ -51,7 +77,9 @@ def _handle_post_create_logic(
             style="green",
         )
     else:
-        _print_training_job_success_message(job_id, remote_provider)
+        _print_training_job_success_message(
+            job_id, project_name, job_resp["job_object"], remote_provider
+        )
 
     if tail:
         watcher = TrainingLogWatcher(remote_provider.api, project_id, job_id)
@@ -351,3 +379,46 @@ def download_checkpoint_artifacts(job_id: Optional[str], remote: Optional[str])
     except Exception as e:
         error_console.print(f"Failed to download checkpoint artifacts data: {str(e)}")
         sys.exit(1)
+
+
+@train.group(name="cache")
+def cache():
+    """Cache-related subcommands for truss train"""
+
+
+@cache.command(name="summarize")
+@click.option(
+    "--project", type=str, required=True, help="Project ID or name to view cache for."
+)
+@click.option("--remote", type=str, required=False, help="Remote to use")
+@click.option(
+    "--sort",
+    type=click.Choice(
+        [
+            SORT_BY_FILEPATH,
+            SORT_BY_SIZE,
+            SORT_BY_MODIFIED,
+            SORT_BY_TYPE,
+            SORT_BY_PERMISSIONS,
+        ]
+    ),
+    default=SORT_BY_FILEPATH,
+    help="Sort files by filepath, size, modified date, file type, or permissions.",
+)
+@click.option(
+    "--order",
+    type=click.Choice([SORT_ORDER_ASC, SORT_ORDER_DESC]),
+    default=SORT_ORDER_ASC,
+    help="Sort order: ascending or descending.",
+)
+@common.common_options()
+def view_cache_summary(project: str, remote: Optional[str], sort: str, order: str):
+    """View cache summary for a training project"""
+    if not remote:
+        remote = remote_cli.inquire_remote_name()
+
+    remote_provider: BasetenRemote = cast(
+        BasetenRemote, RemoteFactory.create(remote=remote)
+    )
+
+    train_cli.view_cache_summary_by_project(remote_provider, project, sort, order)

{truss-0.10.11rc2 → truss-0.10.12}/truss/cli/utils/common.py

@@ -186,4 +186,17 @@ def format_localized_time(iso_timestamp: str) -> str:
         iso_timestamp = iso_timestamp.replace("Z", "+00:00")
     utc_time = datetime.datetime.fromisoformat(iso_timestamp)
     local_time = utc_time.astimezone()
-    return local_time.strftime("%Y-%m-%d %H:%M")
+    return local_time.strftime("%Y-%m-%d %H:%M:%S")
+
+
+def format_bytes_to_human_readable(bytes: int) -> str:
+    if bytes > 1000 * 1000 * 1000 * 1000:
+        return f"{bytes / (1000 * 1000 * 1000 * 1000):.2f} TB"
+    if bytes > 1000 * 1000 * 1000:
+        return f"{bytes / (1000 * 1000 * 1000):.2f} GB"
+    elif bytes > 1000 * 1000:
+        return f"{bytes / (1000 * 1000):.2f} MB"
+    elif bytes > 1000:
+        return f"{bytes / 1000:.2f} KB"
+    else:
+        return f"{bytes} B"

{truss-0.10.11rc2 → truss-0.10.12}/truss/remote/baseten/api.py

@@ -669,6 +669,13 @@ class BasetenApi:
         # NB(nikhil): reverse order so latest logs are at the end
         return resp_json["logs"][::-1]
 
+    def get_cache_summary(self, project_id: str):
+        """Get cache summary for a training project."""
+        resp_json = self._rest_api_client.get(
+            f"v1/training_projects/{project_id}/cache/summary"
+        )
+        return resp_json
+
     def _fetch_log_batch(
         self, project_id: str, job_id: str, query_params: Dict[str, Any]
     ) -> List[Any]:

{truss-0.10.11rc2 → truss-0.10.12}/truss/remote/baseten/custom_types.py

@@ -120,3 +120,31 @@ class TrussUserEnv(pydantic.BaseModel):
 class BlobType(Enum):
     MODEL = "model"
     TRAIN = "train"
+
+
+class FileSummary(pydantic.BaseModel):
+    """Information about a file in the cache."""
+
+    path: str = pydantic.Field(description="Relative path of the file in the cache")
+    size_bytes: int = pydantic.Field(description="Size of the file in bytes")
+    modified: str = pydantic.Field(description="Last modification time of the file")
+    file_type: Optional[str] = pydantic.Field(
+        default=None,
+        description="Type of the file (e.g., 'file', 'directory', 'symlink')",
+    )
+    permissions: Optional[str] = pydantic.Field(
+        default=None,
+        description="File permissions in Unix symbolic format (e.g., 'drwxr-xr-x', '-rw-r--r--')",
+    )
+
+
+class GetCacheSummaryResponseV1(pydantic.BaseModel):
+    """Response for getting cache summary."""
+
+    timestamp: str = pydantic.Field(
+        description="Timestamp when the cache summary was captured"
+    )
+    project_id: str = pydantic.Field(description="Project ID associated with the cache")
+    file_summaries: list[FileSummary] = pydantic.Field(
+        description="List of files in the cache"
+    )