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

truss/base/trt_llm_config.py

@@ -68,6 +68,7 @@ class TrussTRTLLMQuantizationType(str, Enum):
     FP8_KV = "fp8_kv"
     FP4 = "fp4"
     FP4_KV = "fp4_kv"
+    FP4_MLP_ONLY = "fp4_mlp_only"
 
 
 class TrussTRTLLMPluginConfiguration(PydanticTrTBaseModel):
@@ -713,7 +714,9 @@ def trt_llm_common_validation(config: "TrussConfig"):
             "accelerators or newer (CUDA_COMPUTE>=89)"
         )
     elif trt_llm_config.build.quantization_type in [
-        TrussTRTLLMQuantizationType.FP4
+        TrussTRTLLMQuantizationType.FP4,
+        TrussTRTLLMQuantizationType.FP4_KV,
+        TrussTRTLLMQuantizationType.FP4_MLP_ONLY,
     ] and config.resources.accelerator.accelerator in [
         truss_config.Accelerator.H100,
         truss_config.Accelerator.L4,

truss/templates/server.Dockerfile.jinja

@@ -69,7 +69,7 @@ COPY --chown={{ default_owner }} ./{{ config.data_dir }} ${APP_HOME}/data
 
 {%- if model_cache_v2 %}
 {# v0.0.9, keep synced with server_requirements.txt #}
-RUN curl -sSL --fail --retry 5 --retry-delay 2 -o /usr/local/bin/truss-transfer-cli https://github.com/basetenlabs/truss/releases/download/v0.11.12rc4/truss-transfer-cli-v0.11.12rc4-linux-x86_64-unknown-linux-musl
+RUN curl -sSL --fail --retry 5 --retry-delay 2 -o /usr/local/bin/truss-transfer-cli https://github.com/basetenlabs/truss/releases/download/v0.11.13rc3/truss-transfer-cli-v0.11.13rc3-linux-x86_64-unknown-linux-musl
 RUN chmod +x /usr/local/bin/truss-transfer-cli
 RUN mkdir /static-bptr
 RUN echo "hash {{model_cache_hash}}"

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,128 @@ 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 METRICS_REGISTERED:
+            logging.info(
+                "Model cache metrics already registered, skipping."
+            )  # this should never happen
+            return
+        else:
+            # Ensure metrics are only registered once
+            METRICS_REGISTERED = True
+
+        # Define metrics with model_cache prefix
+        manifest_size_gauge = Gauge(
+            "model_cache_manifest_size_bytes", "Total manifest size in bytes"
+        )
+        # histograms have intentially wide buckets to capture a variety of download times
+        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 +316,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 +330,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.13rc4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.13rc2
+Version: 0.11.13rc4
 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.40,>=0.0.36
+Requires-Dist: truss-transfer<0.0.40,>=0.0.37
 Requires-Dist: watchfiles<0.20,>=0.19.0
 Description-Content-Type: text/markdown
 

{truss-0.11.13rc2.dist-info → truss-0.11.13rc4.dist-info}/RECORD

@@ -5,7 +5,7 @@ truss/base/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss/base/constants.py,sha256=sExArdnuGg83z83XMgaQ4b8SS3V_j_bJEpOATDGJzpE,3600
 truss/base/custom_types.py,sha256=FUSIT2lPOQb6gfg6IzT63YBV8r8L6NIZ0D74Fp3e_jQ,2835
 truss/base/errors.py,sha256=zDVLEvseTChdPP0oNhBBQCtQUtZJUaof5zeWMIjqz6o,691
-truss/base/trt_llm_config.py,sha256=81ZZxRQF3o29HLCX6nlXtPwALejcdns6c4mbrExwASk,32958
+truss/base/trt_llm_config.py,sha256=rEtBVFg2QnNMxnaz11s5Z69dJB1w7Bpt48Wf6jSsVZI,33087
 truss/base/truss_config.py,sha256=7CtiJIwMHtDU8Wzn8UTJUVVunD0pWFl4QUVycK2aIpY,28055
 truss/base/truss_spec.py,sha256=jFVF79CXoEEspl2kXBAPyi-rwISReIGTdobGpaIhwJw,5979
 truss/cli/chains_commands.py,sha256=Kpa5mCg6URAJQE2ZmZfVQFhjBHEitKT28tKiW0H6XAI,17406
@@ -71,7 +71,7 @@ truss/templates/cache.Dockerfile.jinja,sha256=1qZqDo1phrcqi-Vwol-VafYJkADsBbQWU6
 truss/templates/cache_requirements.txt,sha256=xoPoJ-OVnf1z6oq_RVM3vCr3ionByyqMLj7wGs61nUs,87
 truss/templates/copy_cache_files.Dockerfile.jinja,sha256=Os5zFdYLZ_AfCRGq4RcpVTObOTwL7zvmwYcvOzd_Zqo,126
 truss/templates/docker_server_requirements.txt,sha256=PyhOPKAmKW1N2vLvTfLMwsEtuGpoRrbWuNo7tT6v2Mc,18
-truss/templates/server.Dockerfile.jinja,sha256=BQpo2Mt_fBrdin1qD8HBKBo2N3Yr2lXrvV_a7J5WSzE,7071
+truss/templates/server.Dockerfile.jinja,sha256=FdxCIXBjBMXVQy0fFxEY2acL-MAZGOD8JKKWAhu3M24,7071
 truss/templates/control/requirements.txt,sha256=tJGr83WoE0CZm2FrloZ9VScK84q-_FTuVXjDYrexhW0,250
 truss/templates/control/control/application.py,sha256=5Kam6M-XtfKGaXQz8cc3d0bwDkB80o2MskABWROx1gk,5321
 truss/templates/control/control/endpoints.py,sha256=KzqsLVNJE6r6TCPW8D5FMCtsfHadTwR15A3z_viGxmM,11782
@@ -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=2QS_0Qac5MMJYhzW-DGgs9_Wh7BtIGLfBtnm29I7X8o,13818
 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.13rc4.dist-info/METADATA,sha256=xxM0NsWyJwukE2KXOuzkZ--EuVl7gRNiJw2wBFwpPT0,6681
+truss-0.11.13rc4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.13rc4.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.13rc4.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.13rc4.dist-info/RECORD,,