Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 4 additions & 51 deletions modelexpress_client/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

use modelexpress_common::{
Result as CommonResult,
cache::{CacheConfig, CacheStats, resolve_model_path},
cache::{CacheConfig, resolve_model_path},
client_config::ClientConfig as Config,
constants, download,
grpc::{
Expand Down Expand Up @@ -231,57 +231,10 @@ impl Client {
Ok(client)
}

/// Get cache configuration
pub fn get_cache_config(&self) -> Option<&CacheConfig> {
self.cache_config.as_ref()
}

/// Set cache configuration
pub fn set_cache_config(&mut self, cache_config: CacheConfig) {
self.cache_config = Some(cache_config);
}

/// List cached models
pub fn list_cached_models(&self) -> CommonResult<CacheStats> {
let cache_config = self.cache_config.as_ref().ok_or_else(|| {
modelexpress_common::Error::Server("Cache not configured".to_string())
})?;

cache_config.get_cache_stats().map_err(|e| {
modelexpress_common::Error::Server(format!("Failed to get cache stats: {e}")).into()
})
}

/// Clear specific model from cache for a given provider.
pub fn clear_cached_model(
&self,
model_name: &str,
provider: ModelProvider,
) -> CommonResult<()> {
let cache_config = self.cache_config.as_ref().ok_or_else(|| {
modelexpress_common::Error::Server("Cache not configured".to_string())
})?;

cache_config.clear_model(model_name, provider).map_err(|e| {
modelexpress_common::Error::Server(format!("Failed to clear model: {e}")).into()
})
}

/// Clear entire cache
pub fn clear_all_cached_models(&self) -> CommonResult<()> {
let cache_config = self.cache_config.as_ref().ok_or_else(|| {
modelexpress_common::Error::Server("Cache not configured".to_string())
})?;

cache_config.clear_all().map_err(|e| {
modelexpress_common::Error::Server(format!("Failed to clear cache: {e}")).into()
})
}

/// Delete a model's record from the server-side registry. This complements
/// `clear_cached_model` (which only removes local files) so a cleared model does
/// not leave behind a stale `DOWNLOADED` record that would satisfy a later download
/// request without re-fetching the files.
/// removing the model's local files, so a cleared model does not leave behind a
/// stale `DOWNLOADED` record that would satisfy a later download request without
/// re-fetching the files.
pub async fn delete_model_on_server(
&mut self,
model_name: &str,
Expand Down
127 changes: 0 additions & 127 deletions modelexpress_common/src/artifact_manifest.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
use crate::grpc::p2p::{
ArtifactManifest as ProtoArtifactManifest, ArtifactManifestChunk as ProtoArtifactManifestChunk,
ArtifactManifestFile as ProtoArtifactManifestFile, ArtifactSourceMetadata,
GetArtifactManifestChunksResponse, GetArtifactManifestHeaderResponse,
};
use anyhow::{Context, Result, anyhow, bail};
use crc32c::{crc32c, crc32c_append};
Expand All @@ -20,10 +19,6 @@ use std::{

pub const ARTIFACT_MANIFEST_VERSION: u32 = 1;
pub const MAX_ARTIFACT_TRANSFER_CHUNK_SIZE: u64 = 4 * 1024 * 1024 * 1024;
// Number of chunk metadata records per GetArtifactManifestChunks response.
// This is not the artifact byte chunk size; 1024 keeps metadata responses
// bounded while avoiding one RPC per transfer chunk.
const ARTIFACT_CHUNK_METADATA_PAGE_SIZE: u32 = 1024;

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ArtifactManifest {
Expand Down Expand Up @@ -173,75 +168,6 @@ impl SealedArtifactManifest {
node_rank: 0,
})
}

pub fn to_header_response(
&self,
mx_source_id: impl Into<String>,
metadata_endpoint: impl Into<String>,
agent_name: impl Into<String>,
worker_rank: u32,
) -> Result<GetArtifactManifestHeaderResponse> {
Ok(GetArtifactManifestHeaderResponse {
mx_source_id: mx_source_id.into(),
artifact_id: self.artifact_id.clone(),
manifest_version: self.manifest.manifest_version,
mx_source_type: self.manifest.mx_source_type,
total_size: self.manifest.total_size()?,
file_count: u32::try_from(self.manifest.files.len())
.context("artifact manifest file count exceeds u32")?,
chunk_count: self.manifest.chunk_count()?,
chunk_size: self.manifest.chunk_size,
metadata_endpoint: metadata_endpoint.into(),
agent_name: agent_name.into(),
worker_rank,
files: self
.manifest
.files
.iter()
.map(ArtifactManifestFile::to_proto)
.collect(),
})
}

pub fn to_chunks_response(
&self,
mx_source_id: impl Into<String>,
start_chunk_index: u32,
max_chunks: u32,
) -> Result<GetArtifactManifestChunksResponse> {
let start = usize::try_from(start_chunk_index)
.context("artifact manifest start chunk index exceeds usize")?;
if start > self.manifest.chunks.len() {
bail!(
"artifact manifest start_chunk_index {} exceeds chunk_count {}",
start_chunk_index,
self.manifest.chunks.len()
);
}
let max_chunks = if max_chunks == 0 {
ARTIFACT_CHUNK_METADATA_PAGE_SIZE
} else {
max_chunks.min(ARTIFACT_CHUNK_METADATA_PAGE_SIZE)
};
let max =
usize::try_from(max_chunks).context("artifact manifest page size exceeds usize")?;
let end = start.saturating_add(max).min(self.manifest.chunks.len());
let next_page_token = if end < self.manifest.chunks.len() {
end.to_string()
} else {
String::new()
};
Ok(GetArtifactManifestChunksResponse {
mx_source_id: mx_source_id.into(),
artifact_id: self.artifact_id.clone(),
start_chunk_index,
chunks: self.manifest.chunks[start..end]
.iter()
.map(ArtifactManifestChunk::to_proto)
.collect(),
next_page_token,
})
}
}

impl ArtifactManifestFile {
Expand Down Expand Up @@ -590,36 +516,6 @@ mod tests {
assert_eq!(manifest.chunk_count().expect("chunk count"), 0);
}

#[test]
fn chunks_response_uses_default_page_size_for_zero_max_chunks() {
let sealed = SealedArtifactManifest {
artifact_id: "artifact".to_string(),
manifest: manifest_with_1025_chunks(),
};

let response = sealed
.to_chunks_response("source-123", 0, 0)
.expect("chunks response");

assert_eq!(response.chunks.len(), 1024);
assert_eq!(response.next_page_token, "1024");
}

#[test]
fn chunks_response_caps_requested_page_size() {
let sealed = SealedArtifactManifest {
artifact_id: "artifact".to_string(),
manifest: manifest_with_1025_chunks(),
};

let response = sealed
.to_chunks_response("source-123", 0, 2048)
.expect("chunks response");

assert_eq!(response.chunks.len(), 1024);
assert_eq!(response.next_page_token, "1024");
}

#[test]
fn pinned_artifact_manifest_id_cross_checked_with_python() {
let manifest = ArtifactManifest {
Expand Down Expand Up @@ -700,27 +596,4 @@ mod tests {
checksum: checksum.to_string(),
}
}

fn manifest_with_1025_chunks() -> ArtifactManifest {
ArtifactManifest {
manifest_version: ARTIFACT_MANIFEST_VERSION,
mx_source_type: MxSourceType::TorchCompileCache as i32,
chunk_size: 1,
files: vec![ArtifactManifestFile {
file_index: 0,
path: "/tmp/artifact.bin".to_string(),
size: 1025,
checksum: "file".to_string(),
}],
chunks: (0..1025)
.map(|index| ArtifactManifestChunk {
chunk_index: index,
file_index: 0,
file_offset: u64::from(index),
length: 1,
checksum: format!("chunk-{index}"),
})
.collect(),
}
}
}
7 changes: 0 additions & 7 deletions modelexpress_common/src/cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -191,13 +191,6 @@ impl CacheConfig {
))
}

/// Query server for cache information
pub fn from_server() -> Result<Self> {
// This would typically make an HTTP request to the server
// For now, we'll return an error to indicate server is not available
Err(anyhow::anyhow!("Server not available for cache discovery"))
}

/// Get cache path from command line arguments
fn get_cache_path_from_args() -> Option<String> {
let args: Vec<String> = env::args().collect();
Expand Down
8 changes: 0 additions & 8 deletions modelexpress_common/src/client_config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -237,14 +237,6 @@ impl ClientConfig {
}
}

/// Apply cache path override if provided
pub fn with_cache_path(mut self, cache_path: Option<PathBuf>) -> Self {
if let Some(path) = cache_path {
self.cache.local_path = path;
}
self
}

/// Set timeout for the connection
pub fn with_timeout(mut self, timeout_secs: u64) -> Self {
self.connection.timeout_secs = Some(timeout_secs);
Expand Down
16 changes: 0 additions & 16 deletions modelexpress_common/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -197,22 +197,6 @@ impl FromStr for LogFormat {
}
}

/// Base trait for configuration loading with layered approach
pub trait ConfigLoader<T> {
/// Load configuration from multiple sources in order of precedence:
/// 1. Command line arguments (highest priority)
/// 2. Environment variables
/// 3. Configuration file
/// 4. Default values (lowest priority)
fn load_layered(
config_file: Option<PathBuf>,
env_prefix: &str,
defaults: T,
) -> Result<T, ConfigError>
where
T: serde::de::DeserializeOwned + Default;
}

/// Load configuration file strictly without any fallbacks to defaults.
/// This function will return an error if the file doesn't exist, has invalid syntax,
/// or contains invalid values. Use this for validation purposes.
Expand Down
6 changes: 0 additions & 6 deletions modelexpress_common/src/envs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -107,10 +107,6 @@ pub const REDIS_PORT: &str = "REDIS_PORT";
pub const MX_METADATA_NAMESPACE: &str = "MX_METADATA_NAMESPACE";
/// Kubernetes namespace injected via the downward API for in-cluster pods.
pub const POD_NAMESPACE: &str = "POD_NAMESPACE";
/// Kubernetes pod name injected via the downward API (used by clients).
pub const POD_NAME: &str = "POD_NAME";
/// Kubernetes pod UID injected via the downward API (used by clients).
pub const POD_UID: &str = "POD_UID";
Comment on lines -110 to -113

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are in use by ownerReference

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fair, and the doc comment i deleted said "used by clients" which should have stopped me.

but as far as i can tell the ownerReference path runs entirely through python: the downward api sets these on the worker, envs.py declares them, client.py sends pod_name and pod_uid on the request, and pod_owner_references takes them off the request rather than off the server's own environment. i can't find a rust reader for either constant, so the removal should be behaviourally inert.

which leaves the question of what envs.rs is for. POD_NAMESPACE is read in rust at backend_config.rs, so those three sit together today with two of them declaration-only. is envs.rs meant to be the canonical registry for python-only var names as well, or should those live only in envs.py? if it's the registry, i'll put both back and leave the trio alone. if it's rust readers only, POD_NAMESPACE is the one that's correctly there and the other two were always redundant with envs.py.

holding the removal as-is until you say which.


// ── Reaper (server) ─────────────────────────────────────────────────────────
/// Interval (seconds) between reaper scans for stale/GC worker sweeps.
Expand Down Expand Up @@ -360,8 +356,6 @@ mod tests {
assert_eq!(HOME, "HOME");
assert_eq!(USERPROFILE, "USERPROFILE");
assert_eq!(KUBECONFIG, "KUBECONFIG");
assert_eq!(POD_NAME, "POD_NAME");
assert_eq!(POD_UID, "POD_UID");
}

#[test]
Expand Down
6 changes: 0 additions & 6 deletions modelexpress_common/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,9 +56,6 @@ pub struct Response<T> {
/// Common error types that both client and server can use
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("Network error: {0}")]
Network(String),

#[error("Server returned error: {0}")]
Server(String),

Expand Down Expand Up @@ -370,9 +367,6 @@ mod tests {

#[test]
fn test_error_types() {
let network_error = Error::Network("Connection failed".to_string());
assert!(network_error.to_string().contains("Network error"));

let server_error = Error::Server("Internal error".to_string());
assert!(server_error.to_string().contains("Server returned error"));

Expand Down
38 changes: 0 additions & 38 deletions modelexpress_server/src/cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -409,44 +409,6 @@ impl CacheEvictionService {
Ok(())
}

/// Manually trigger eviction for specific models
pub async fn manual_evict(
&self,
model_names: &[String],
) -> Result<EvictionResult, Box<dyn std::error::Error + Send + Sync>> {
info!(
"Manual eviction requested for models: {models:?}",
models = model_names
);

let mut successfully_evicted = Vec::new();
for model_name in model_names {
match self.evict_model(model_name).await {
Ok(()) => {
successfully_evicted.push(model_name.clone());
info!(
"Successfully evicted model: {model_name}",
model_name = model_name
);
}
Err(e) => {
warn!(
"Failed to evict model '{model_name}': {e}",
model_name = model_name,
e = e
);
}
}
}

Ok(EvictionResult {
evicted_count: successfully_evicted.len() as u32,
evicted_models: successfully_evicted,
bytes_freed: None,
reason: EvictionReason::Manual,
})
}

/// Get statistics about the current cache state
pub async fn get_cache_stats(
&self,
Expand Down
Loading
Loading