truss 0.11.13rc2__py3-none-any.whl → 0.11.13rc3__py3-none-any.whl

truss/templates/shared/lazy_data_resolver.py

@@ -1,6 +1,7 @@
 import atexit
 import json
 import logging
+import os
 import time
 from dataclasses import dataclass
 from functools import lru_cache
@@ -8,6 +9,14 @@ from pathlib import Path
 from threading import Lock, Thread
 from typing import Optional, Union
 
+try:
+    from prometheus_client import Counter, Gauge, Histogram
+
+    PROMETHEUS_AVAILABLE = True
+except ImportError:
+    PROMETHEUS_AVAILABLE = False
+METRICS_REGISTERED = False
+
 
 @dataclass(frozen=True)
 class FileDownloadMetric:
@@ -61,6 +70,121 @@ class TrussTransferStats:
         except Exception:
             return None
 
+    def publish_to_prometheus(self):
+        """Publish transfer stats to Prometheus metrics. Only runs once."""
+        if not PROMETHEUS_AVAILABLE:
+            return
+        global METRICS_REGISTERED
+        if not METRICS_REGISTERED:
+            # Ensure metrics are only registered once
+            METRICS_REGISTERED = True
+
+        # Define metrics with model_cache label
+        manifest_size_gauge = Gauge(
+            "model_cache_manifest_size_bytes", "Total manifest size in bytes"
+        )
+        download_time_histogram = Histogram(
+            "model_cache_download_time_seconds",
+            "Total download time in seconds",
+            buckets=[
+                2**i
+                for i in range(-3, 11)  # = [0.125, .. 2048] seconds
+            ],
+        )
+        download_speed_gauge = Gauge(
+            "model_cache_download_speed_mbps", "Aggregated download speed in MB/s"
+        )
+
+        # File download metrics (aggregated)
+        files_downloaded_counter = Counter(
+            "model_cache_files_downloaded_total", "Total number of files downloaded"
+        )
+        total_file_size_counter = Counter(
+            "model_cache_file_size_bytes_total",
+            "Total size of downloaded files in bytes",
+        )
+        file_download_time_histogram = Histogram(
+            "model_cache_file_download_time_seconds",
+            "File download time distribution",
+            buckets=[
+                2**i
+                for i in range(-3, 11)  # = [0.125, .. 2048] seconds
+            ],
+        )
+        file_download_speed_histogram = Histogram(
+            "model_cache_file_download_speed_mbps",
+            "File download speed distribution",
+            buckets=[
+                2**i
+                for i in range(-1, 12)  # = [0.5, .. 4096] MB/s
+            ],
+        )
+
+        # B10FS specific metrics
+        b10fs_enabled_gauge = Gauge(
+            "model_cache_b10fs_enabled", "Whether B10FS is enabled"
+        )
+        b10fs_decision_gauge = Gauge(
+            "model_cache_b10fs_decision_to_use", "Whether B10FS was chosen for use"
+        )
+        b10fs_read_speed_gauge = Gauge(
+            "model_cache_b10fs_read_speed_mbps", "B10FS read speed in Mbps"
+        )
+        b10fs_hot_files_gauge = Gauge(
+            "model_cache_b10fs_hot_starts_files", "Number of hot start files"
+        )
+        b10fs_hot_bytes_gauge = Gauge(
+            "model_cache_b10fs_hot_starts_bytes", "Number of hot start bytes"
+        )
+        b10fs_cold_files_gauge = Gauge(
+            "model_cache_b10fs_cold_starts_files", "Number of cold start files"
+        )
+        b10fs_cold_bytes_gauge = Gauge(
+            "model_cache_b10fs_cold_starts_bytes", "Number of cold start bytes"
+        )
+
+        # Transfer success metric
+        transfer_success_counter = Counter(
+            "model_cache_transfer_success_total",
+            "Total successful transfers",
+            ["success"],
+        )
+
+        # Set main transfer metrics
+        manifest_size_gauge.set(self.total_manifest_size_bytes)
+        download_time_histogram.observe(self.total_download_time_secs)
+
+        if self.total_aggregated_mb_s is not None:
+            download_speed_gauge.set(self.total_aggregated_mb_s)
+
+        # Aggregate file download metrics
+        total_files = len(self.file_downloads)
+        total_file_bytes = sum(fd.file_size_bytes for fd in self.file_downloads)
+
+        files_downloaded_counter.inc(total_files)
+        total_file_size_counter.inc(total_file_bytes)
+
+        # Record individual file metrics for distribution
+        for fd in self.file_downloads:
+            if fd.file_size_bytes > 1 * 1024 * 1024:  # Only log files larger than 1MB
+                file_download_time_histogram.observe(fd.download_time_secs)
+                file_download_speed_histogram.observe(fd.download_speed_mb_s)
+
+        # B10FS metrics
+        b10fs_enabled_gauge.set(1 if self.b10fs_enabled else 0)
+        b10fs_decision_gauge.set(1 if self.b10fs_decision_to_use else 0)
+
+        if self.b10fs_read_speed_mbps is not None:
+            b10fs_read_speed_gauge.set(self.b10fs_read_speed_mbps)
+
+        b10fs_hot_files_gauge.set(self.b10fs_hot_starts_files)
+        b10fs_hot_bytes_gauge.set(self.b10fs_hot_starts_bytes)
+        b10fs_cold_files_gauge.set(self.b10fs_cold_starts_files)
+        b10fs_cold_bytes_gauge.set(self.b10fs_cold_starts_bytes)
+
+        # Success metric
+        transfer_success_counter.labels(success=str(self.success)).inc()
+
 
 LAZY_DATA_RESOLVER_PATH = [
     # synced with pub static LAZY_DATA_RESOLVER_PATHS: &[&str]
@@ -185,6 +309,9 @@ class LazyDataResolverV2:
 
         """
         start_lock = time.time()
+        publish_stats = (
+            log_stats and not self._is_collected_by_user
+        )  # only publish results once per resolver
         self._is_collected_by_user = issue_collect or self._is_collected_by_user
         with self._lock:
             result = self._fetch()
@@ -196,13 +323,20 @@ class LazyDataResolverV2:
                 # TODO: instument the stats, which are written to /tmp/truss_transfer_stats.json
                 # also add fetch time, and blocking time
                 # TrussTransferStats
+                fetch_t = time.time() - self._start_time
+                start_lock_t = time.time() - start_lock
                 stats = TrussTransferStats.from_json_file(
                     Path("/tmp/truss_transfer_stats.json")
                 )
-                if stats is None:
+                if stats and publish_stats:
                     self.logger.info(f"model_cache: {stats}")
+                    # Publish stats to Prometheus
+                    if (
+                        os.getenv("TRUSS_MODEL_CACHE_PROMETHEUS", "0") == "1"
+                    ):  # Hide behind feature flag for core-product to enabled.
+                        stats.publish_to_prometheus()
                 self.logger.info(
-                    f"model_cache: Fetch took {time.time() - self._start_time:.2f} seconds, of which {time.time() - start_lock:.2f} seconds were spent blocking."
+                    f"model_cache: Fetch took {fetch_t:.2f} seconds, of which {start_lock_t:.2f} seconds were spent blocking."
                 )
             return result
 

{truss-0.11.13rc2.dist-info → truss-0.11.13rc3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.13rc2
+Version: 0.11.13rc3
 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.11.13rc2.dist-info → truss-0.11.13rc3.dist-info}/RECORD

@@ -107,7 +107,7 @@ truss/templates/server/common/tracing.py,sha256=XSTXNoRtV8vXwveJoX3H32go0JKnLmzn
 truss/templates/server/common/patches/whisper/patch.py,sha256=kDECQ-wmEpeAZFhUTQP457ofueeMsm7DgNy9tqinhJQ,2383
 truss/templates/shared/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss/templates/shared/dynamic_config_resolver.py,sha256=75s42NFhQI5jL7BqlJH_UkuQS7ptbtFh13f2nh6X5Wo,920
-truss/templates/shared/lazy_data_resolver.py,sha256=eOq7Fgr9QkAWpsxyDWkeZic1Z2S4Mt-drB1A7zNEkYE,8368
+truss/templates/shared/lazy_data_resolver.py,sha256=czfggu9DZ_qDnE2MxOdE2R8aZyJe2G1Cd-PL0AUGx-I,13561
 truss/templates/shared/log_config.py,sha256=l9udyu4VKHZePlfK9LQEd5TOUUodPuehypsXRSUL4Ac,5411
 truss/templates/shared/secrets_resolver.py,sha256=3prDe3Q06NTmUEe7KCW-W4TD1CzGck9lpDG789209z4,2110
 truss/templates/shared/serialization.py,sha256=_WC_2PPkRi-MdTwxwjG8LKQptnHi4sANfpOlKWevqWc,3736
@@ -370,8 +370,8 @@ truss_train/deployment.py,sha256=lWWANSuzBWu2M4oK4qD7n-oVR1JKdmw2Pn5BJQHg-Ck,307
 truss_train/loader.py,sha256=0o66EjBaHc2YY4syxxHVR4ordJWs13lNXnKjKq2wq0U,1630
 truss_train/public_api.py,sha256=9N_NstiUlmBuLUwH_fNG_1x7OhGCytZLNvqKXBlStrM,1220
 truss_train/restore_from_checkpoint.py,sha256=8hdPm-WSgkt74HDPjvCjZMBpvA9MwtoYsxVjOoa7BaM,1176
-truss-0.11.13rc2.dist-info/METADATA,sha256=GCadCR-s-rOXbH__XXvO5wFJEp-sPcFjpRC1tlcRBwo,6681
-truss-0.11.13rc2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.13rc2.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.13rc2.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.13rc2.dist-info/RECORD,,
+truss-0.11.13rc3.dist-info/METADATA,sha256=cPpD-bEoveXxM_dTFQAJPui0DFO940QYrTTWlO7aivc,6681
+truss-0.11.13rc3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.13rc3.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.13rc3.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.13rc3.dist-info/RECORD,,