bge_m3_embedding_server/embedder/trt_cache/
prewarm_log.rs

1// Copyright (c) 2026 J. Patrick Fulton
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Operator-visible prewarm basename logging.
16
17use std::path::Path;
18
19use super::enumerate::engine_basenames_for_sm;
20
21pub(crate) fn log_engine_basenames_before_prewarm_for_sm(engine_dir: &Path, sm: Option<&str>) {
22    const MAX_LIST: usize = 64;
23    let basenames = match engine_basenames_for_sm(engine_dir, sm) {
24        Ok(v) => v,
25        Err(e) => {
26            tracing::info!(
27                cache_path = %engine_dir.display(),
28                detected_sm = sm.unwrap_or("unfiltered"),
29                error = %e,
30                "trt prewarm: could not read engine cache directory for basename listing"
31            );
32            return;
33        }
34    };
35
36    let matching_total = basenames.len();
37    let mut listed = basenames;
38    let truncated = matching_total > MAX_LIST;
39    if truncated {
40        listed.truncate(MAX_LIST);
41    }
42    // Also surface the unfiltered total so heterogeneous-SM situations are
43    // obvious at a glance. When `sm` is `None` this equals `matching_total`.
44    let unfiltered_total = if sm.is_some() {
45        engine_basenames_for_sm(engine_dir, None).map_or(matching_total, |v| v.len())
46    } else {
47        matching_total
48    };
49
50    if matching_total == 0 {
51        tracing::info!(
52            cache_path = %engine_dir.display(),
53            detected_sm = sm.unwrap_or("unfiltered"),
54            engine_basename_count = 0,
55            engine_basename_total_count = unfiltered_total,
56            "trt prewarm: cache engine basenames (none for this SM)"
57        );
58    } else {
59        tracing::info!(
60            cache_path = %engine_dir.display(),
61            detected_sm = sm.unwrap_or("unfiltered"),
62            engine_basename_count = matching_total,
63            engine_basename_total_count = unfiltered_total,
64            truncated,
65            engine_basenames = ?listed,
66            "trt prewarm: cache engine basenames (for operator correlation; not shape-specific)"
67        );
68    }
69}
70
71/// Backwards-compatible wrapper for the unfiltered prewarm basename log.
72///
73/// Delegates to [`log_engine_basenames_before_prewarm_for_sm`] with
74/// `sm = None` — equivalent to today's behaviour for callers that have not
75/// yet been updated to pass the worker's detected SM.
76pub(crate) fn log_engine_basenames_before_prewarm(engine_dir: &Path) {
77    log_engine_basenames_before_prewarm_for_sm(engine_dir, None);
78}
bge_m3_embedding_server/embedder/trt_cache/prewarm_log.rs

bge_m3_embedding_server/embedder/trt_cache/
prewarm_log.rs
