bge_m3_embedding_server/embedder/trt_cache/
paths.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//! Cache directory paths and startup snapshot type.
16
17use std::path::{Path, PathBuf};
18
19/// Subdirectory under `BGE_M3_CACHE_DIR` where ORT/TRT writes engine plan files.
20pub(crate) const TRT_ENGINE_SUBDIR: &str = "trt-engines";
21
22/// Subdirectory under `BGE_M3_CACHE_DIR` where TRT writes the timing cache.
23pub(crate) const TRT_TIMING_SUBDIR: &str = "trt-timing";
24
25/// Summary of the TRT engine cache state, produced at startup so the operator
26/// can see in `CloudWatch` whether the persistent volume is actually being
27/// reused between container restarts.
28#[derive(Debug, Clone)]
29pub(crate) struct TrtCacheInfo {
30    /// Absolute cache directory path (stable; no per-container ephemera).
31    pub path: PathBuf,
32    /// Number of `.engine` files found in the cache directory before this
33    /// container started doing any work. Zero means a cold cache.
34    pub engine_count: usize,
35    /// Number of `.profile` files (TRT input-shape profiles emitted alongside
36    /// each `.engine`). Reported for diagnostic completeness; the count is
37    /// expected to be `engine_count` once warmup completes.
38    pub profile_count: usize,
39}
40
41/// Returns the canonical TRT engine-cache directory for a given root cache.
42///
43/// Path is stable across container restarts as long as `cache_dir` is mounted
44/// at the same location — i.e. for ECS this means a persistent EFS access
45/// point or a host-bind mount. There is **no PID, hostname, or container
46/// identifier in the path** — that was already true before this change, but
47/// is now centralised here so future callers cannot reintroduce per-container
48/// ephemera.
49pub(crate) fn engine_cache_path(cache_dir: &Path) -> PathBuf {
50    cache_dir.join(TRT_ENGINE_SUBDIR)
51}
52
53/// Returns the canonical TRT timing-cache file path for a given root cache.
54///
55/// The timing cache is a single file, not a directory of per-shape files.
56pub(crate) fn timing_cache_path(cache_dir: &Path) -> PathBuf {
57    cache_dir.join(TRT_TIMING_SUBDIR)
58}
bge_m3_embedding_server/embedder/trt_cache/paths.rs

bge_m3_embedding_server/embedder/trt_cache/
paths.rs
