Skip to content
Merged
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
14 changes: 12 additions & 2 deletions crates/mesh-llm-host-runtime/src/runtime/local.rs
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,7 @@ pub(super) struct LocalRuntimeModelStartSpec<'a> {
pub(super) mmproj_override: Option<&'a Path>,
pub(super) ctx_size_override: Option<u32>,
pub(super) pinned_gpu: Option<&'a crate::runtime::StartupPinnedGpuTarget>,
pub(super) device_override: Option<String>,
pub(super) capacity_budget_bytes: Option<u64>,
pub(super) cache_type_k_override: Option<&'a str>,
pub(super) cache_type_v_override: Option<&'a str>,
Expand All @@ -250,6 +251,7 @@ pub(super) struct LocalOpenAiModelStartSpec<'a> {
pub(super) mmproj_override: Option<&'a Path>,
pub(super) ctx_size_override: Option<u32>,
pub(super) pinned_gpu: Option<&'a crate::runtime::StartupPinnedGpuTarget>,
pub(super) device_override: Option<String>,
pub(super) capacity_budget_bytes: u64,
pub(super) cache_type_k_override: Option<&'a str>,
pub(super) cache_type_v_override: Option<&'a str>,
Expand Down Expand Up @@ -350,6 +352,9 @@ pub(super) fn resolve_local_openai_skippy_config(
if let Some(gpu) = spec.pinned_gpu {
resolved.hardware.device = Some(gpu.backend_device.clone());
}
if let Some(device) = &spec.device_override {
resolved.hardware.device = Some(device.clone());
}
Ok(resolved)
}

Expand Down Expand Up @@ -564,6 +569,7 @@ pub(super) async fn start_runtime_local_model(
mmproj_override: spec.mmproj_override,
ctx_size_override: spec.ctx_size_override,
pinned_gpu: spec.pinned_gpu,
device_override: spec.device_override,
capacity_budget_bytes: local_capacity_bytes,
cache_type_k_override: spec.cache_type_k_override,
cache_type_v_override: spec.cache_type_v_override,
Expand Down Expand Up @@ -738,7 +744,9 @@ async fn start_local_skippy_model(
.with_openai_guardrails(skippy::skippy_openai_guardrails_for_policy_handle(
spec.openai_guardrail_policy.clone(),
));
if let Some(gpu) = spec.pinned_gpu {
if spec.device_override.is_none()
&& let Some(gpu) = spec.pinned_gpu
{
options = options.with_selected_device(pinned_skippy_device(gpu));
}
let _ = emit_event(OutputEvent::ModelLoading {
Expand Down Expand Up @@ -852,7 +860,9 @@ async fn start_local_layer_package_model(
runtime_options.config.ctx_size = context_length;
runtime_options.config.lane_count = plan.slots as u32;
runtime_options.config.filter_tensors_on_load = true;
if let Some(gpu) = spec.pinned_gpu {
if spec.device_override.is_none()
&& let Some(gpu) = spec.pinned_gpu
{
runtime_options.config.selected_device = Some(pinned_stage_device(gpu));
}
runtime_options.config.load_mode = LoadMode::LayerPackage;
Expand Down
8 changes: 5 additions & 3 deletions crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,10 @@ use super::{
SkippyNativeLogForwardingGuard, acquire_instance_runtime,
apply_runtime_cli_speculative_overrides, apply_runtime_config_options,
build_startup_model_specs, cleanup_run_auto_runtime_dir, configure_run_auto_process_state,
emit_shutdown, openai_guardrail_policy_handle, preflight_config_owned_startup_models,
emit_shutdown, openai_guardrail_policy_handle, preflight_pinned_startup_models,
resolve_local_model_only_startup_models, runtime_model_required_bytes,
skippy_telemetry_options, start_local_openai_model, wait_shutdown_signal,
skippy_telemetry_options, start_local_openai_model, startup_device_override,
wait_shutdown_signal,
};
use crate::inference::election;
use crate::plugin;
Expand Down Expand Up @@ -119,7 +120,7 @@ pub(super) async fn run_local_model_only(mut options: RuntimeOptions) -> Result<
"--local-model-only requires exactly one startup model"
);
let mut startup_models = resolve_local_model_only_startup_models(&startup_specs).await?;
preflight_config_owned_startup_models(
preflight_pinned_startup_models(
&config,
&startup_specs,
&mut startup_models,
Expand Down Expand Up @@ -172,6 +173,7 @@ pub(super) async fn run_local_model_only(mut options: RuntimeOptions) -> Result<
mmproj_override: model.mmproj_path.as_deref(),
ctx_size_override: model.ctx_size,
pinned_gpu: model.pinned_gpu.as_ref(),
device_override: startup_device_override(model.gpu_id.as_deref()),
capacity_budget_bytes: local_capacity_bytes,
cache_type_k_override: model.cache_type_k.as_deref(),
cache_type_v_override: model.cache_type_v.as_deref(),
Expand Down
2 changes: 2 additions & 0 deletions crates/mesh-llm-host-runtime/src/runtime/local_split.rs
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,7 @@ pub(super) async fn start_runtime_split_model(
flash_attention_override: spec.flash_attention_override,
openai_guardrail_policy: spec.openai_guardrail_policy.clone(),
pinned_gpu: spec.pinned_gpu,
device_override: spec.device_override.as_deref(),
slots,
skippy_telemetry: spec.skippy_telemetry.clone(),
survey_telemetry: spec.survey_telemetry.clone(),
Expand Down Expand Up @@ -275,6 +276,7 @@ pub(super) async fn start_runtime_split_model(
flash_attention_override: spec.flash_attention_override,
openai_guardrail_policy: spec.openai_guardrail_policy.clone(),
pinned_gpu: spec.pinned_gpu.cloned(),
device_override: spec.device_override.clone(),
slots,
skippy_telemetry: spec.skippy_telemetry.clone(),
survey_telemetry: spec.survey_telemetry.clone(),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,7 @@ pub(super) struct SplitTopologyCoordinator {
pub(super) flash_attention_override: FlashAttentionType,
pub(super) openai_guardrail_policy: OpenAiGuardrailPolicyHandle,
pub(super) pinned_gpu: Option<crate::runtime::StartupPinnedGpuTarget>,
pub(super) device_override: Option<String>,
pub(super) slots: usize,
pub(super) skippy_telemetry: skippy::SkippyTelemetryOptions,
pub(super) survey_telemetry: survey::SurveyTelemetry,
Expand Down Expand Up @@ -642,6 +643,7 @@ impl SplitTopologyCoordinator {
flash_attention_override: self.flash_attention_override,
openai_guardrail_policy: self.openai_guardrail_policy.clone(),
pinned_gpu: self.pinned_gpu.as_ref(),
device_override: self.device_override.as_deref(),
slots: self.slots,
skippy_telemetry: self.skippy_telemetry.clone(),
survey_telemetry: self.survey_telemetry.clone(),
Expand Down
24 changes: 21 additions & 3 deletions crates/mesh-llm-host-runtime/src/runtime/local_split/loading.rs
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ pub(super) struct SplitGenerationLoadSpec<'a> {
pub(super) ctx_size: u32,
pub(super) compact_meta: &'a models::gguf::GgufCompactMeta,
pub(super) pinned_gpu: Option<&'a crate::runtime::StartupPinnedGpuTarget>,
pub(super) device_override: Option<&'a str>,
pub(super) slots: usize,
pub(super) cache_type_k_override: Option<&'a str>,
pub(super) cache_type_v_override: Option<&'a str>,
Expand Down Expand Up @@ -245,9 +246,11 @@ pub(super) async fn load_split_runtime_generation_inner(
runtime_options.config.ctx_size = spec.ctx_size;
runtime_options.config.lane_count = spec.slots as u32;
runtime_options.config.filter_tensors_on_load = true;
if let Some(gpu) = spec.pinned_gpu {
runtime_options.config.selected_device = Some(pinned_stage_device(gpu));
}
apply_split_generation_pinned_device(
&mut runtime_options.config,
spec.pinned_gpu,
spec.device_override,
);
runtime_options.config.load_mode = settings.load_mode.clone();
runtime_options.config.bind_addr = stage0_return_endpoint;
runtime_options.config.upstream = None;
Expand Down Expand Up @@ -600,6 +603,9 @@ pub(super) async fn split_generation_load_settings<'a>(
if let Some(gpu) = spec.pinned_gpu {
resolved.hardware.device = Some(gpu.backend_device.clone());
}
if let Some(device) = spec.device_override {
resolved.hardware.device = Some(device.to_string());
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
let embedded_openai = resolved.to_embedded_openai_args(activation_width, true)?;
let lifecycle = configured_stage_lifecycle_intervals(spec.mesh_config, spec.config_model_id);
let runtime_options = resolved.to_embedded_runtime_options(
Expand All @@ -624,6 +630,18 @@ pub(super) async fn split_generation_load_settings<'a>(
})
}

pub(super) fn apply_split_generation_pinned_device(
config: &mut skippy_protocol::StageConfig,
pinned_gpu: Option<&crate::runtime::StartupPinnedGpuTarget>,
device_override: Option<&str>,
) {
if device_override.is_none()
&& let Some(gpu) = pinned_gpu
{
config.selected_device = Some(pinned_stage_device(gpu));
}
}

pub(super) fn split_generation_load_mode(package: &skippy::SkippyPackageIdentity) -> LoadMode {
if skippy::is_layer_package_ref(&package.package_ref) {
LoadMode::LayerPackage
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,36 @@ pub(super) fn stage_load_request(load_mode: LoadMode) -> skippy::StageLoadReques
}
}

#[test]
fn split_generation_cli_device_override_survives_pinned_stage_selection() {
let mut config = skippy_protocol::StageConfig {
selected_device: Some(skippy_protocol::StageDevice {
backend_device: "CPU".to_string(),
stable_id: None,
index: None,
vram_bytes: None,
}),
..Default::default()
};
let pinned_gpu = crate::runtime::StartupPinnedGpuTarget {
index: 0,
stable_id: "pci:0000:65:00.0".to_string(),
backend_device: "CUDA0".to_string(),
vram_bytes: 24_000_000_000,
reserved_bytes: None,
};

apply_split_generation_pinned_device(&mut config, Some(&pinned_gpu), Some("CPU"));

assert_eq!(
config
.selected_device
.as_ref()
.map(|device| device.backend_device.as_str()),
Some("CPU")
);
}

pub(super) fn split_test_peer(
seed: u8,
model_name: &str,
Expand Down Expand Up @@ -427,6 +457,7 @@ stop = ["END"]
ctx_size: 8192,
compact_meta: &compact_meta,
pinned_gpu: None,
device_override: None,
slots: 4,
cache_type_k_override: None,
cache_type_v_override: None,
Expand Down Expand Up @@ -554,6 +585,7 @@ async fn split_stage_load_guards_family_kv_default_with_planned_metadata() {
ctx_size: 4096,
compact_meta: &incompatible_meta,
pinned_gpu: None,
device_override: None,
slots: 1,
cache_type_k_override: None,
cache_type_v_override: None,
Expand Down Expand Up @@ -601,7 +633,7 @@ async fn split_stage_load_guards_family_kv_default_with_planned_metadata() {
}

#[tokio::test]
async fn runtime_resolver_uses_config_model_id_but_preserves_served_model_id() {
async fn runtime_resolver_uses_config_identity_and_honors_device_override() {
let node = mesh::Node::new_for_tests(NodeRole::Host { http_port: 9337 })
.await
.unwrap();
Expand All @@ -615,6 +647,7 @@ model = "other/model-ref"

[models.hardware]
model_path = "{model_path}"
device = "CUDA1"

[models.throughput]
threads = 17
Expand Down Expand Up @@ -645,6 +678,7 @@ max_tokens = 222
mmproj_override: None,
ctx_size_override: None,
pinned_gpu: None,
device_override: Some("CPU".to_string()),
capacity_budget_bytes: node.vram_bytes(),
cache_type_k_override: None,
cache_type_v_override: None,
Expand Down Expand Up @@ -680,6 +714,55 @@ max_tokens = 222
assert_eq!(resolved.request_defaults.max_tokens, 222);
assert_eq!(resolved.model_fit.ctx_size, 4096);
assert_eq!(resolved.throughput.parallel, 3);
assert_eq!(resolved.hardware.device.as_deref(), Some("CPU"));
// An explicit CLI artifact may use the same served name as a configured
// model, but must not inherit that entry's path or runtime tuning.
let cli_model_path = temp_dir.path().join("cli-selected.gguf");
write_fake_gguf_model(&cli_model_path);
let cli_model_bytes = fs::metadata(&cli_model_path).unwrap().len();
let cli_spec = LocalOpenAiModelStartSpec {
mesh_config: &mesh_config,
config_model_id: None,
model_path: &cli_model_path,
model_bytes: cli_model_bytes,
mmproj_override: None,
ctx_size_override: None,
pinned_gpu: None,
device_override: None,
capacity_budget_bytes: node.vram_bytes(),
cache_type_k_override: None,
cache_type_v_override: None,
n_batch_override: None,
n_ubatch_override: None,
flash_attention_override: FlashAttentionType::Auto,
parallel_override: None,
planning_profile: RuntimeResourcePlanningProfile::DedicatedLocal,
openai_guardrail_policy: openai_guardrail_policy_handle(
openai_frontend::GuardrailMode::Disabled,
),
skippy_telemetry: skippy::SkippyTelemetryOptions::off(),
survey_telemetry: survey::SurveyTelemetry::disabled(),
hook_policy: None,
serving_hooks_factory: None,
http_bind_addr: "127.0.0.1:0".parse().expect("valid loopback address"),
};
let cli_resolved = resolve_local_openai_skippy_config(
&cli_spec,
"configured/model-ref",
cli_model_bytes,
4096,
3,
None,
None,
)
.expect("explicit CLI runtime config should not consult model entries");
assert_eq!(cli_resolved.hardware.resolved_model_path, cli_model_path);
assert_eq!(cli_resolved.throughput.threads, None);
assert_eq!(cli_resolved.throughput.threads_batch, None);
assert_eq!(
cli_resolved.request_defaults.max_tokens,
skippy_server::CONTEXT_BUDGET_MAX_TOKENS
);
}

#[test]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1311,6 +1311,7 @@ async fn load_split_runtime_generation_stops_candidate_stages_after_partial_load
ctx_size: 4096,
compact_meta: &compact_meta,
pinned_gpu: None,
device_override: None,
slots: 1,
cache_type_k_override: None,
cache_type_v_override: None,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@ pub(crate) async fn run_auto_load_runtime_model(
mmproj_override: None,
ctx_size_override,
pinned_gpu: None,
device_override: None,
capacity_budget_bytes: Some(capacity_budget_bytes),
cache_type_k_override: model_overrides.and_then(|m| m.cache_type_k.as_deref()),
cache_type_v_override: model_overrides.and_then(|m| m.cache_type_v.as_deref()),
Expand Down
6 changes: 5 additions & 1 deletion crates/mesh-llm-host-runtime/src/runtime/run_auto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,8 @@ use super::{
runtime_data_producer_for_console, runtime_startup_requirements, setup_run_auto_console_state,
setup_run_auto_serving_surface, spawn_embedded_runtime_control_forwarder,
spawn_run_auto_additional_model_tasks, spawn_run_auto_discovery_publisher,
start_run_auto_bootstrap_proxy, startup_local_model_loop, swarm_capture_observer_requested,
start_run_auto_bootstrap_proxy, startup_device_override, startup_local_model_loop,
swarm_capture_observer_requested,
};
use crate::api;
use crate::inference::{election, skippy};
Expand Down Expand Up @@ -1125,6 +1126,8 @@ pub(super) async fn spawn_run_auto_startup_model_tasks(ctx: RunAutoStartupTasksC
let primary_mmproj = primary_startup_model.and_then(|model| model.mmproj_path.clone());
let primary_ctx_size = primary_startup_model.and_then(|model| model.ctx_size);
let primary_pinned_gpu = primary_startup_model.and_then(|model| model.pinned_gpu.clone());
let primary_device_override =
primary_startup_model.and_then(|model| startup_device_override(model.gpu_id.as_deref()));
let primary_cache_type_k = primary_startup_model.and_then(|model| model.cache_type_k.clone());
let primary_cache_type_v = primary_startup_model.and_then(|model| model.cache_type_v.clone());
let primary_n_batch = primary_startup_model.and_then(|model| model.n_batch);
Expand Down Expand Up @@ -1162,6 +1165,7 @@ pub(super) async fn spawn_run_auto_startup_model_tasks(ctx: RunAutoStartupTasksC
mmproj_path: primary_mmproj,
ctx_size: primary_ctx_size,
pinned_gpu: primary_pinned_gpu,
device_override: primary_device_override,
runtime_capacity_ledger: runtime_capacity_ledger.clone(),
cache_type_k: primary_cache_type_k,
cache_type_v: primary_cache_type_v,
Expand Down
3 changes: 3 additions & 0 deletions crates/mesh-llm-host-runtime/src/runtime/serving_surface.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1156,6 +1156,9 @@ pub(super) async fn spawn_run_auto_additional_model_tasks(ctx: RunAutoAdditional
mmproj_path: extra_model.mmproj_path.clone(),
ctx_size: extra_model.ctx_size,
pinned_gpu: extra_model.pinned_gpu.clone(),
device_override: super::startup_models::startup_device_override(
extra_model.gpu_id.as_deref(),
),
runtime_capacity_ledger: ctx.runtime_capacity_ledger.clone(),
cache_type_k: extra_model.cache_type_k.clone(),
cache_type_v: extra_model.cache_type_v.clone(),
Expand Down
3 changes: 3 additions & 0 deletions crates/mesh-llm-host-runtime/src/runtime/startup_handles.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ pub(super) struct StartupLocalModelTask {
pub(super) mmproj_path: Option<PathBuf>,
pub(super) ctx_size: Option<u32>,
pub(super) pinned_gpu: Option<StartupPinnedGpuTarget>,
pub(super) device_override: Option<String>,
pub(super) runtime_capacity_ledger: RuntimeCapacityLedger,
pub(super) cache_type_k: Option<String>,
pub(super) cache_type_v: Option<String>,
Expand Down Expand Up @@ -576,6 +577,7 @@ pub(super) async fn startup_launch_runtime(
mmproj_path,
ctx_size,
pinned_gpu,
device_override,
runtime_capacity_ledger,
cache_type_k,
cache_type_v,
Expand Down Expand Up @@ -605,6 +607,7 @@ pub(super) async fn startup_launch_runtime(
mmproj_override: mmproj_path.map(PathBuf::as_path),
ctx_size_override: ctx_size,
pinned_gpu,
device_override: device_override.clone(),
capacity_budget_bytes: None,
cache_type_k_override: cache_type_k,
cache_type_v_override: cache_type_v,
Expand Down
Loading
Loading