From 4722f965d116b36ea4658543e111d6fa17552eb1 Mon Sep 17 00:00:00 2001 From: Will Eaton Date: Tue, 31 Mar 2026 20:43:19 -0400 Subject: [PATCH 1/2] add LLMInferenceService support (serving.kserve.io/v1alpha1) - minimal hand-written CRD type with spec.replicas and serde flatten - label fast-path: app.kubernetes.io/part-of=llminferenceservice - owner-ref chain: Pod -> RS -> Deployment -> LLMInferenceService - scale patches both spec.replicas and spec.prefill.replicas to zero - new 'l' flag in enabled-resources (default "drsinl") - 10 unit tests covering bitflags, conversion, equality, hashing, meta --- README.md | 6 +- gpu-pruner/src/lib.rs | 186 ++++++++++++++++++++++++++- gpu-pruner/src/main.rs | 3 +- resources/src/lib.rs | 1 + resources/src/llminferenceservice.rs | 29 +++++ 5 files changed, 215 insertions(+), 10 deletions(-) create mode 100644 resources/src/llminferenceservice.rs diff --git a/README.md b/README.md index 30da3b9..88f849c 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # gpu-pruner -The `gpu-pruner` is a non-destructive idle culler that works with Red Hat OpenShift AI/Kubeflow provided APIs (`InferenceService` and `Notebook`), as well as generic `Deployment`, `ReplicaSet`, `StatefulSet` and `LeaderWorkerSet` (`leaderworkerset.x-k8s.io`). +The `gpu-pruner` is a non-destructive idle culler that works with Red Hat OpenShift AI/Kubeflow provided APIs (`InferenceService` and `Notebook`), as well as generic `Deployment`, `ReplicaSet`, `StatefulSet`, `LeaderWorkerSet` (`leaderworkerset.x-k8s.io`) and `LLMInferenceService` (`serving.kserve.io/v1alpha1`). The way it works is by querying cluster NVIDIA DCGM metrics and looking at a window of GPU utilization per pod. A scaling decision is made by looking up the pods metadata, and using owner-references to figure out the owning resource. @@ -33,9 +33,9 @@ Options: -e, --enabled-resources Specifcy enabled resources with a string of letters - - `d` for Deployment - `r` for ReplicaSet - `s` for StatefulSet - `i` for InferenceService - `n` for Notebook - `l` for LeaderWorkerSet + - `d` for Deployment - `r` for ReplicaSet - `s` for StatefulSet - `i` for InferenceService - `n` for Notebook - `l` for LeaderWorkerSet - `m` for LLMInferenceService - [default: drsinl] + [default: drsinlm] -c, --check-interval interval in seconds to check for idle pods, only used in daemon mode diff --git a/gpu-pruner/src/lib.rs b/gpu-pruner/src/lib.rs index 61b93ae..70d4f4b 100644 --- a/gpu-pruner/src/lib.rs +++ b/gpu-pruner/src/lib.rs @@ -13,7 +13,8 @@ use k8s_openapi::{ }; use kube::{Client, ResourceExt, api::PostParams}; use resources::{ - inferenceservice::InferenceService, leaderworkerset::LeaderWorkerSet, notebook::Notebook, + inferenceservice::InferenceService, leaderworkerset::LeaderWorkerSet, + llminferenceservice::LLMInferenceService, notebook::Notebook, }; use secrecy::ExposeSecret; use serde::Serialize; @@ -99,6 +100,7 @@ pub enum ScaleKind { ReplicaSet(ReplicaSet), StatefulSet(StatefulSet), InferenceService(Box), + LLMInferenceService(Box), Notebook(Notebook), LeaderWorkerSet(LeaderWorkerSet), } @@ -110,6 +112,9 @@ impl PartialEq for ScaleKind { (ScaleKind::ReplicaSet(a), ScaleKind::ReplicaSet(b)) => a == b, (ScaleKind::StatefulSet(a), ScaleKind::StatefulSet(b)) => a == b, (ScaleKind::InferenceService(a), ScaleKind::InferenceService(b)) => a.uid() == b.uid(), + (ScaleKind::LLMInferenceService(a), ScaleKind::LLMInferenceService(b)) => { + a.uid() == b.uid() + } (ScaleKind::Notebook(a), ScaleKind::Notebook(b)) => a.uid() == b.uid(), (ScaleKind::LeaderWorkerSet(a), ScaleKind::LeaderWorkerSet(b)) => a.uid() == b.uid(), // If they are different variants, they are not equal @@ -136,6 +141,9 @@ impl Hash for ScaleKind { ScaleKind::InferenceService(a) => { a.uid().hash(state); } + ScaleKind::LLMInferenceService(a) => { + a.uid().hash(state); + } ScaleKind::Notebook(a) => { a.uid().hash(state); } @@ -153,6 +161,7 @@ impl From for ResourceKind { ScaleKind::ReplicaSet(_) => ResourceKind::REPLICA_SET, ScaleKind::StatefulSet(_) => ResourceKind::STATEFUL_SET, ScaleKind::InferenceService(_) => ResourceKind::INFERENCE_SERVICE, + ScaleKind::LLMInferenceService(_) => ResourceKind::LLM_INFERENCE_SERVICE, ScaleKind::Notebook(_) => ResourceKind::NOTEBOOK, ScaleKind::LeaderWorkerSet(_) => ResourceKind::LEADER_WORKER_SET, } @@ -168,6 +177,7 @@ bitflags! { const INFERENCE_SERVICE = 0b01000; const NOTEBOOK = 0b10000; const LEADER_WORKER_SET = 0b100000; + const LLM_INFERENCE_SERVICE = 0b1000000; } } @@ -179,6 +189,7 @@ bitflags! { /// - `i` → InferenceService /// - `n` → Notebook /// - `l` → LeaderWorkerSet +/// - `m` → LLMInferenceService /// /// Unknown characters are silently ignored. pub fn get_enabled_resources(enabled_resources: &str) -> ResourceKind { @@ -191,6 +202,7 @@ pub fn get_enabled_resources(enabled_resources: &str) -> ResourceKind { 'i' => resource_kind |= ResourceKind::INFERENCE_SERVICE, 'n' => resource_kind |= ResourceKind::NOTEBOOK, 'l' => resource_kind |= ResourceKind::LEADER_WORKER_SET, + 'm' => resource_kind |= ResourceKind::LLM_INFERENCE_SERVICE, _ => {} } } @@ -381,6 +393,7 @@ macro_rules! delegate_resource_ext { ScaleKind::ReplicaSet(d) => d.$method(), ScaleKind::StatefulSet(d) => d.$method(), ScaleKind::InferenceService(d) => d.$method(), + ScaleKind::LLMInferenceService(d) => d.$method(), ScaleKind::Notebook(d) => d.$method(), ScaleKind::LeaderWorkerSet(d) => d.$method(), } @@ -404,6 +417,7 @@ impl Meta for ScaleKind { ScaleKind::Notebook(_) => "v1".to_string(), ScaleKind::InferenceService(_) => "v1beta1".to_string(), ScaleKind::LeaderWorkerSet(_) => "leaderworkerset.x-k8s.io/v1".to_string(), + ScaleKind::LLMInferenceService(_) => "serving.kserve.io/v1alpha1".to_string(), } } @@ -415,6 +429,7 @@ impl Meta for ScaleKind { ScaleKind::Notebook(_) => "Notebook".to_string(), ScaleKind::InferenceService(_) => "InferenceService".to_string(), ScaleKind::LeaderWorkerSet(_) => "LeaderWorkerSet".to_string(), + ScaleKind::LLMInferenceService(_) => "LLMInferenceService".to_string(), } } @@ -482,6 +497,14 @@ impl Scaler for ScaleKind { .await?; Ok(()) } + ScaleKind::LLMInferenceService(d) => { + let ns = d + .namespace() + .ok_or_else(|| anyhow::anyhow!("LLMInferenceService has no namespace"))?; + scale_llm_inference_service_to_zero(client.clone(), &d.name_unchecked(), &ns) + .await?; + Ok(()) + } }; if result.is_ok() @@ -551,6 +574,7 @@ fn workload_annotations( ScaleKind::Notebook(n) => n.metadata.annotations.as_ref(), ScaleKind::InferenceService(i) => i.metadata.annotations.as_ref(), ScaleKind::LeaderWorkerSet(l) => l.metadata.annotations.as_ref(), + ScaleKind::LLMInferenceService(l) => l.metadata.annotations.as_ref(), } } @@ -749,6 +773,12 @@ pub async fn fetch_workload( let api: Api = Api::namespaced(client, namespace); Ok(ScaleKind::InferenceService(Box::new(api.get(name).await?))) } + "LLMInferenceService" => { + let api: Api = Api::namespaced(client, namespace); + Ok(ScaleKind::LLMInferenceService(Box::new( + api.get(name).await?, + ))) + } _ => Err(anyhow::anyhow!("Unsupported resource kind: {}", kind)), } } @@ -787,6 +817,10 @@ pub async fn patch_workload( let api: Api = Api::namespaced(client, namespace); api.patch(name, ¶ms, &Patch::Merge(patch)).await?; } + "LLMInferenceService" => { + let api: Api = Api::namespaced(client, namespace); + api.patch(name, ¶ms, &Patch::Merge(patch)).await?; + } _ => { return Err(anyhow::anyhow!("Unsupported resource kind: {}", kind)); } @@ -809,9 +843,19 @@ pub async fn find_root_object( "Finding root object of {name:?} for scale-down.", name = &pod_meta.name ); - // first, check for the special kserve label - // if it exists, we can go directly to the InferenceService - // and scale it down + // fast-path: LLMInferenceService pods carry standard k8s app labels + if let Some(labels) = &pod_meta.labels + && labels.get("app.kubernetes.io/part-of").map(|v| v.as_str()) + == Some("llminferenceservice") + && let Some(llmis_name) = labels.get("app.kubernetes.io/name") + { + let namespace = pod_meta.namespace.clone().unwrap_or_default(); + let api: Api = Api::namespaced(client.clone(), &namespace); + let llmis = api.get(llmis_name).await?; + return Ok(ScaleKind::LLMInferenceService(Box::new(llmis))); + } + + // fast-path: InferenceService pods carry a kserve-specific label if let Some(labels) = &pod_meta.labels && let Some(ks_label) = labels.get("serving.kserve.io/inferenceservice") { @@ -838,6 +882,25 @@ pub async fn find_root_object( Api::namespaced(client.clone(), &namespace); let deployment = deployment_api.get(&rs_or.name).await?; + // check if this Deployment is owned by an LLMInferenceService + if let Some(dep_ors) = + deployment.metadata.owner_references.as_ref() + { + for dep_or in dep_ors { + if dep_or.kind == "LLMInferenceService" { + tracing::info!( + "Found LLMInferenceService owning Deployment!" + ); + let llmis_api: Api = + Api::namespaced(client.clone(), &namespace); + let llmis = llmis_api.get(&dep_or.name).await?; + return Ok(ScaleKind::LLMInferenceService( + Box::new(llmis), + )); + } + } + } + return Ok(ScaleKind::Deployment(deployment)); } } @@ -961,6 +1024,36 @@ async fn scale_inference_service_to_zero( Ok(res) } +/// Scale an LLMInferenceService to zero by patching spec.replicas. +/// +/// The LLMInferenceService CRD does not expose a /scale subresource, so we +/// patch spec.replicas directly. In disaggregated (prefill-decode) setups, +/// prefill has its own independent replica count at spec.prefill.replicas, +/// so we zero both to avoid leaving half the pipeline running. +#[tracing::instrument(skip(client))] +async fn scale_llm_inference_service_to_zero( + client: KubeClient, + name: &str, + namespace: &str, +) -> anyhow::Result { + let api: Api = Api::namespaced(client.clone(), namespace); + + let patch = serde_json::json!({ + "spec": { + "replicas": 0, + "prefill": { + "replicas": 0 + } + } + }); + + let res = api + .patch(name, &PatchParams::default(), &Patch::Merge(patch)) + .await?; + + Ok(res) +} + #[cfg(test)] mod tests { use std::collections::HashSet; @@ -973,6 +1066,8 @@ mod tests { notebook::NotebookSpec, }; + use resources::llminferenceservice::LLMInferenceService; + use crate::{ ACK_BY_ANNOTATION, ACK_UNTIL_ANNOTATION, Meta, NamespaceMentionMapper, Notebook, PENDING_SCALE_ANNOTATION, PendingScaleStatus, ResourceKind, SLACK_MENTIONS_ANNOTATION, @@ -1065,16 +1160,40 @@ mod tests { }) } + fn make_llm_inference_service(name: &str, ns: &str, uid: Option<&str>) -> ScaleKind { + let mut llmis: LLMInferenceService = serde_json::from_value(serde_json::json!({ + "metadata": { + "name": name, + "namespace": ns, + }, + "spec": {} + })) + .expect("valid LLMInferenceService JSON"); + llmis.metadata.uid = uid.map(Into::into); + ScaleKind::LLMInferenceService(Box::new(llmis)) + } + // ── get_enabled_resources ──────────────────────────────────────────── #[test] fn enabled_resources_all_flags() { - let rk = get_enabled_resources("drsin"); + let rk = get_enabled_resources("drsinlm"); assert!(rk.contains(ResourceKind::DEPLOYMENT)); assert!(rk.contains(ResourceKind::REPLICA_SET)); assert!(rk.contains(ResourceKind::STATEFUL_SET)); assert!(rk.contains(ResourceKind::INFERENCE_SERVICE)); assert!(rk.contains(ResourceKind::NOTEBOOK)); + assert!(rk.contains(ResourceKind::LEADER_WORKER_SET)); + assert!(rk.contains(ResourceKind::LLM_INFERENCE_SERVICE)); + } + + #[test] + fn enabled_resources_single_llm_inference_service() { + let rk = get_enabled_resources("m"); + assert!(rk.contains(ResourceKind::LLM_INFERENCE_SERVICE)); + assert!(!rk.contains(ResourceKind::DEPLOYMENT)); + assert!(!rk.contains(ResourceKind::INFERENCE_SERVICE)); + assert!(!rk.contains(ResourceKind::LEADER_WORKER_SET)); } #[test] @@ -1134,6 +1253,7 @@ mod tests { assert!(!empty.contains(ResourceKind::STATEFUL_SET)); assert!(!empty.contains(ResourceKind::INFERENCE_SERVICE)); assert!(!empty.contains(ResourceKind::NOTEBOOK)); + assert!(!empty.contains(ResourceKind::LLM_INFERENCE_SERVICE)); } // ── ScaleKind → ResourceKind conversion ────────────────────────────── @@ -1162,6 +1282,12 @@ mod tests { assert_eq!(rk, ResourceKind::INFERENCE_SERVICE); } + #[test] + fn scale_kind_to_resource_kind_llm_inference_service() { + let rk: ResourceKind = make_llm_inference_service("l", "ns", None).into(); + assert_eq!(rk, ResourceKind::LLM_INFERENCE_SERVICE); + } + #[test] fn scale_kind_to_resource_kind_notebook() { let rk: ResourceKind = make_notebook("n", "ns", None).into(); @@ -1205,6 +1331,27 @@ mod tests { assert_eq!(a, b); } + #[test] + fn llm_inference_service_equality_uses_uid() { + let a = make_llm_inference_service("llm-a", "ns", Some("uid-llm")); + let b = make_llm_inference_service("llm-b", "ns", Some("uid-llm")); + assert_eq!(a, b); + } + + #[test] + fn llm_inference_service_different_uid_not_equal() { + let a = make_llm_inference_service("llm", "ns", Some("uid-1")); + let b = make_llm_inference_service("llm", "ns", Some("uid-2")); + assert_ne!(a, b); + } + + #[test] + fn llm_inference_service_not_equal_to_inference_service() { + let llmis = make_llm_inference_service("x", "ns", Some("uid-1")); + let is = make_inference_service("x", "ns", Some("uid-1")); + assert_ne!(llmis, is); + } + // ── ScaleKind hashing / HashSet dedup ──────────────────────────────── #[test] @@ -1246,10 +1393,11 @@ mod tests { set.insert(make_replica_set("r1", "ns", Some("uid-r"))); set.insert(make_stateful_set("s1", "ns", Some("uid-s"))); set.insert(make_inference_service("i1", "ns", Some("uid-i"))); + set.insert(make_llm_inference_service("l1", "ns", Some("uid-l"))); set.insert(make_notebook("n1", "ns", Some("uid-n"))); // duplicate of first deployment set.insert(make_deployment("d1", "ns", Some("uid-d"))); - assert_eq!(set.len(), 5); + assert_eq!(set.len(), 6); } // ── Meta trait ─────────────────────────────────────────────────────── @@ -1304,6 +1452,16 @@ mod tests { assert_eq!(sk.api_version(), "v1beta1"); } + #[test] + fn meta_llm_inference_service() { + let sk = make_llm_inference_service("my-llmis", "genai", Some("llmis-uid")); + assert_eq!(sk.name(), "my-llmis"); + assert_eq!(sk.namespace(), Some("genai".into())); + assert_eq!(sk.kind(), "LLMInferenceService"); + assert_eq!(sk.uid(), Some("llmis-uid".into())); + assert_eq!(sk.api_version(), "serving.kserve.io/v1alpha1"); + } + // ── Event generation ───────────────────────────────────────────────── #[test] @@ -1373,6 +1531,22 @@ mod tests { assert_eq!(event.involved_object.api_version, Some("v1beta1".into())); } + #[test] + fn event_for_llm_inference_service() { + let sk = make_llm_inference_service("my-llmis", "genai", Some("llmis-uid")); + let event = sk.generate_scale_event().unwrap(); + + assert_eq!( + event.involved_object.kind, + Some("LLMInferenceService".into()) + ); + assert_eq!( + event.involved_object.api_version, + Some("serving.kserve.io/v1alpha1".into()) + ); + assert_eq!(event.involved_object.uid, Some("llmis-uid".into())); + } + #[test] fn event_names_are_unique() { let sk = make_notebook("nb", "ns", None); diff --git a/gpu-pruner/src/main.rs b/gpu-pruner/src/main.rs index c98d25a..6e2919f 100644 --- a/gpu-pruner/src/main.rs +++ b/gpu-pruner/src/main.rs @@ -66,7 +66,8 @@ struct Cli { /// - `i` for InferenceService /// - `n` for Notebook /// - `l` for LeaderWorkerSet - #[clap(short, long, default_value = "drsinl")] + /// - `m` for LLMInferenceService + #[clap(short, long, default_value = "drsinlm")] enabled_resources: String, /// interval in seconds to check for idle pods, only used in daemon mode diff --git a/resources/src/lib.rs b/resources/src/lib.rs index 36b797e..0ed9411 100644 --- a/resources/src/lib.rs +++ b/resources/src/lib.rs @@ -1,3 +1,4 @@ pub mod inferenceservice; pub mod leaderworkerset; +pub mod llminferenceservice; pub mod notebook; diff --git a/resources/src/llminferenceservice.rs b/resources/src/llminferenceservice.rs new file mode 100644 index 0000000..72e556e --- /dev/null +++ b/resources/src/llminferenceservice.rs @@ -0,0 +1,29 @@ +use kube::CustomResource; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +/// Minimal LLMInferenceService CRD definition. +/// +/// We only model `spec.replicas` because that is the sole field gpu-pruner +/// reads or patches. Everything else is captured by the `#[serde(flatten)]` +/// catch-all so the type round-trips through the API server without data loss. +/// +/// The upstream CRD is still v1alpha1 and evolving rapidly; keeping this +/// hand-written (rather than kopium-generated) avoids pulling in the entire +/// PodSpec tree and makes version bumps a one-line change. +#[derive(CustomResource, Serialize, Deserialize, Clone, Debug, JsonSchema)] +#[kube( + group = "serving.kserve.io", + version = "v1alpha1", + kind = "LLMInferenceService", + plural = "llminferenceservices" +)] +#[kube(namespaced)] +pub struct LLMInferenceServiceSpec { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub replicas: Option, + + #[serde(flatten)] + pub other: BTreeMap, +} From ed11ce7f2b86ff9d4751a74d14523bdc91110184 Mon Sep 17 00:00:00 2001 From: Will Eaton Date: Mon, 31 Aug 2026 08:43:49 -0400 Subject: [PATCH 2/2] resolve LLMISVC-owned LeaderWorkerSets, flatten owner-chain walk A LeaderWorkerSet created by an LLMInferenceService now resolves to the owning LLMInferenceService: scaling the LWS alone is undone by the KServe controller reconciling it back. find_owner/get_typed helpers replace the nested owner-ref loops, From impls give every workload type a ScaleKind conversion, and with_kind_api! collapses the seven-arm fetch/patch matches. CRD schemas for the opaque passthrough fields now carry an explicit object type with x-kubernetes-preserve-unknown-fields so LeaderWorkerSet and LLMInferenceService CRDs install cleanly; kind e2e tests cover the LWS and LLMISVC-owned-LWS owner chains. --- gpu-pruner/src/lib.rs | 281 +++++++++++++++------------ gpu-pruner/tests/e2e.rs | 189 ++++++++++++++++++ resources/src/leaderworkerset.rs | 10 + resources/src/llminferenceservice.rs | 2 + 4 files changed, 358 insertions(+), 124 deletions(-) diff --git a/gpu-pruner/src/lib.rs b/gpu-pruner/src/lib.rs index 70d4f4b..6f19d84 100644 --- a/gpu-pruner/src/lib.rs +++ b/gpu-pruner/src/lib.rs @@ -154,6 +154,42 @@ impl Hash for ScaleKind { } } +impl From for ScaleKind { + fn from(v: Deployment) -> Self { + ScaleKind::Deployment(v) + } +} +impl From for ScaleKind { + fn from(v: ReplicaSet) -> Self { + ScaleKind::ReplicaSet(v) + } +} +impl From for ScaleKind { + fn from(v: StatefulSet) -> Self { + ScaleKind::StatefulSet(v) + } +} +impl From for ScaleKind { + fn from(v: Notebook) -> Self { + ScaleKind::Notebook(v) + } +} +impl From for ScaleKind { + fn from(v: LeaderWorkerSet) -> Self { + ScaleKind::LeaderWorkerSet(v) + } +} +impl From for ScaleKind { + fn from(v: InferenceService) -> Self { + ScaleKind::InferenceService(Box::new(v)) + } +} +impl From for ScaleKind { + fn from(v: LLMInferenceService) -> Self { + ScaleKind::LLMInferenceService(Box::new(v)) + } +} + impl From for ResourceKind { fn from(kind: ScaleKind) -> Self { match kind { @@ -741,6 +777,44 @@ pub async fn clear_pending_scale_at( patch_workload(client, kind, name, namespace, &patch).await } +/// Dispatches a kind string to a typed `Api` bound as `$api`, then runs the +/// body. Unknown kinds return an error. +macro_rules! with_kind_api { + ($kind:expr, $client:expr, $namespace:expr, |$api:ident| $body:expr) => { + match $kind { + "Deployment" => { + let $api: Api = Api::namespaced($client, $namespace); + $body + } + "ReplicaSet" => { + let $api: Api = Api::namespaced($client, $namespace); + $body + } + "StatefulSet" => { + let $api: Api = Api::namespaced($client, $namespace); + $body + } + "LeaderWorkerSet" => { + let $api: Api = Api::namespaced($client, $namespace); + $body + } + "Notebook" => { + let $api: Api = Api::namespaced($client, $namespace); + $body + } + "InferenceService" => { + let $api: Api = Api::namespaced($client, $namespace); + $body + } + "LLMInferenceService" => { + let $api: Api = Api::namespaced($client, $namespace); + $body + } + other => Err(anyhow::anyhow!("Unsupported resource kind: {}", other)), + } + }; +} + #[tracing::instrument(skip(client))] pub async fn fetch_workload( client: KubeClient, @@ -748,39 +822,9 @@ pub async fn fetch_workload( name: &str, namespace: &str, ) -> anyhow::Result { - match kind { - "Deployment" => { - let api: Api = Api::namespaced(client, namespace); - Ok(ScaleKind::Deployment(api.get(name).await?)) - } - "ReplicaSet" => { - let api: Api = Api::namespaced(client, namespace); - Ok(ScaleKind::ReplicaSet(api.get(name).await?)) - } - "StatefulSet" => { - let api: Api = Api::namespaced(client, namespace); - Ok(ScaleKind::StatefulSet(api.get(name).await?)) - } - "LeaderWorkerSet" => { - let api: Api = Api::namespaced(client, namespace); - Ok(ScaleKind::LeaderWorkerSet(api.get(name).await?)) - } - "Notebook" => { - let api: Api = Api::namespaced(client, namespace); - Ok(ScaleKind::Notebook(api.get(name).await?)) - } - "InferenceService" => { - let api: Api = Api::namespaced(client, namespace); - Ok(ScaleKind::InferenceService(Box::new(api.get(name).await?))) - } - "LLMInferenceService" => { - let api: Api = Api::namespaced(client, namespace); - Ok(ScaleKind::LLMInferenceService(Box::new( - api.get(name).await?, - ))) - } - _ => Err(anyhow::anyhow!("Unsupported resource kind: {}", kind)), - } + with_kind_api!(kind, client, namespace, |api| Ok(ScaleKind::from( + api.get(name).await? + ))) } #[tracing::instrument(skip(client, patch))] @@ -791,41 +835,11 @@ pub async fn patch_workload( namespace: &str, patch: &serde_json::Value, ) -> anyhow::Result<()> { - let params = PatchParams::default(); - match kind { - "Deployment" => { - let api: Api = Api::namespaced(client, namespace); - api.patch(name, ¶ms, &Patch::Merge(patch)).await?; - } - "ReplicaSet" => { - let api: Api = Api::namespaced(client, namespace); - api.patch(name, ¶ms, &Patch::Merge(patch)).await?; - } - "StatefulSet" => { - let api: Api = Api::namespaced(client, namespace); - api.patch(name, ¶ms, &Patch::Merge(patch)).await?; - } - "LeaderWorkerSet" => { - let api: Api = Api::namespaced(client, namespace); - api.patch(name, ¶ms, &Patch::Merge(patch)).await?; - } - "Notebook" => { - let api: Api = Api::namespaced(client, namespace); - api.patch(name, ¶ms, &Patch::Merge(patch)).await?; - } - "InferenceService" => { - let api: Api = Api::namespaced(client, namespace); - api.patch(name, ¶ms, &Patch::Merge(patch)).await?; - } - "LLMInferenceService" => { - let api: Api = Api::namespaced(client, namespace); - api.patch(name, ¶ms, &Patch::Merge(patch)).await?; - } - _ => { - return Err(anyhow::anyhow!("Unsupported resource kind: {}", kind)); - } - } - Ok(()) + with_kind_api!(kind, client, namespace, |api| { + api.patch(name, &PatchParams::default(), &Patch::Merge(patch)) + .await?; + Ok(()) + }) } /// Crawl up the owner references to find the root Deployment or StatefulSet @@ -872,75 +886,57 @@ pub async fn find_root_object( match or.kind.as_str() { "ReplicaSet" => { tracing::info!("Found ReplicaSet!"); - let rs_api: Api = Api::namespaced(client.clone(), &namespace); - if let Ok(rs) = rs_api.get(&or.name).await { - if let Some(rs_meta) = rs.metadata.owner_references.as_ref() { - for rs_or in rs_meta { - if rs_or.kind == "Deployment" { - tracing::info!("Found Deployment owning ReplicaSet!"); - let deployment_api: Api = - Api::namespaced(client.clone(), &namespace); - let deployment = deployment_api.get(&rs_or.name).await?; - - // check if this Deployment is owned by an LLMInferenceService - if let Some(dep_ors) = - deployment.metadata.owner_references.as_ref() - { - for dep_or in dep_ors { - if dep_or.kind == "LLMInferenceService" { - tracing::info!( - "Found LLMInferenceService owning Deployment!" - ); - let llmis_api: Api = - Api::namespaced(client.clone(), &namespace); - let llmis = llmis_api.get(&dep_or.name).await?; - return Ok(ScaleKind::LLMInferenceService( - Box::new(llmis), - )); - } - } - } - - return Ok(ScaleKind::Deployment(deployment)); - } + if let Ok(rs) = get_typed::(&client, &namespace, &or.name).await { + if let Some(dep_or) = find_owner(&rs.metadata, "Deployment") { + tracing::info!("Found Deployment owning ReplicaSet!"); + let deployment = + get_typed::(&client, &namespace, &dep_or.name).await?; + + if let Some(llmis_or) = + find_owner(&deployment.metadata, "LLMInferenceService") + { + tracing::info!("Found LLMInferenceService owning Deployment!"); + let llmis = get_typed::( + &client, + &namespace, + &llmis_or.name, + ) + .await?; + return Ok(llmis.into()); } + + return Ok(deployment.into()); } // fallthrough, replica set with no owners - return Ok(ScaleKind::ReplicaSet(rs.clone())); + return Ok(rs.into()); } } "StatefulSet" => { tracing::info!("Found StatefulSet!"); - let ss_api: Api = Api::namespaced(client.clone(), &namespace); - if let Ok(ss) = ss_api.get(&or.name).await { - if let Some(ss_meta) = ss.metadata.owner_references.as_ref() { - for ss_or in ss_meta { - if ss_or.kind == "Notebook" { - tracing::info!("Found Notebook owning StatefulSet!"); - let nb_api: Api = - Api::namespaced(client.clone(), &namespace); - let nb = nb_api.get(&ss_or.name).await?; - - return Ok(ScaleKind::Notebook(nb)); - } else if ss_or.kind == "LeaderWorkerSet" { - tracing::info!("Found LeaderWorkerSet owning StatefulSet!"); - let lws_api: Api = - Api::namespaced(client.clone(), &namespace); - let lws = lws_api.get(&ss_or.name).await?; - - return Ok(ScaleKind::LeaderWorkerSet(lws)); - } - } + if let Ok(ss) = get_typed::(&client, &namespace, &or.name).await { + if let Some(nb_or) = find_owner(&ss.metadata, "Notebook") { + tracing::info!("Found Notebook owning StatefulSet!"); + let nb = + get_typed::(&client, &namespace, &nb_or.name).await?; + return Ok(nb.into()); + } + if let Some(lws_or) = find_owner(&ss.metadata, "LeaderWorkerSet") { + tracing::info!("Found LeaderWorkerSet owning StatefulSet!"); + let lws = + get_typed::(&client, &namespace, &lws_or.name) + .await?; + return resolve_lws(&client, &namespace, lws).await; } // fallthrough, statefulset with no owners - return Ok(ScaleKind::StatefulSet(ss)); + return Ok(ss.into()); } } "LeaderWorkerSet" => { tracing::info!("Found LeaderWorkerSet!"); - let lws_api: Api = Api::namespaced(client.clone(), &namespace); - if let Ok(lws) = lws_api.get(&or.name).await { - return Ok(ScaleKind::LeaderWorkerSet(lws)); + if let Ok(lws) = + get_typed::(&client, &namespace, &or.name).await + { + return resolve_lws(&client, &namespace, lws).await; } } "DaemonSet" | "Node" => { @@ -961,6 +957,43 @@ pub async fn find_root_object( )) } +fn find_owner<'a>( + meta: &'a ObjectMeta, + kind: &str, +) -> Option<&'a k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference> { + meta.owner_references + .as_ref()? + .iter() + .find(|or| or.kind == kind) +} + +async fn get_typed(client: &KubeClient, namespace: &str, name: &str) -> Result +where + K: kube::Resource + + Clone + + DeserializeOwned + + Debug, +{ + Api::::namespaced(client.clone(), namespace) + .get(name) + .await +} + +/// A LeaderWorkerSet created by an LLMInferenceService is re-reconciled by the +/// KServe controller, so scale the owning LLMInferenceService instead. +async fn resolve_lws( + client: &KubeClient, + namespace: &str, + lws: LeaderWorkerSet, +) -> Result { + if let Some(llmis_or) = find_owner(&lws.metadata, "LLMInferenceService") { + tracing::info!("Found LLMInferenceService owning LeaderWorkerSet!"); + let llmis = get_typed::(client, namespace, &llmis_or.name).await?; + return Ok(llmis.into()); + } + Ok(lws.into()) +} + /// Scale a resource to zero replicas via the /scale subresource endpoint #[tracing::instrument(skip(api))] async fn scale_to_zero(api: Api, name: &str) -> anyhow::Result<()> diff --git a/gpu-pruner/tests/e2e.rs b/gpu-pruner/tests/e2e.rs index 156c15e..610d75a 100644 --- a/gpu-pruner/tests/e2e.rs +++ b/gpu-pruner/tests/e2e.rs @@ -392,3 +392,192 @@ async fn hashset_dedup_with_real_uids() { delete_test_namespace(&client, &ns).await; } + +// ── CRD-backed owner chains (LeaderWorkerSet / LLMInferenceService) ────── + +async fn ensure_crd( + client: &Client, + crd: k8s_openapi::apiextensions_apiserver::pkg::apis::apiextensions::v1::CustomResourceDefinition, +) { + use k8s_openapi::apiextensions_apiserver::pkg::apis::apiextensions::v1::CustomResourceDefinition; + + let api: Api = Api::all(client.clone()); + let name = crd.metadata.name.clone().unwrap(); + match api.create(&PostParams::default(), &crd).await { + Ok(_) => {} + Err(kube::Error::Api(ae)) if ae.code == 409 => {} + Err(e) => panic!("failed to create CRD {name}: {e}"), + } + + for _ in 0..30 { + if let Ok(c) = api.get(&name).await + && let Some(status) = c.status + && let Some(conditions) = status.conditions + && conditions + .iter() + .any(|c| c.type_ == "Established" && c.status == "True") + { + return; + } + tokio::time::sleep(std::time::Duration::from_secs(1)).await; + } + panic!("CRD {name} never became established"); +} + +fn owner_reference( + api_version: &str, + kind: &str, + name: &str, + uid: &str, +) -> k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference { + k8s_openapi::apimachinery::pkg::apis::meta::v1::OwnerReference { + api_version: api_version.into(), + kind: kind.into(), + name: name.into(), + uid: uid.into(), + ..Default::default() + } +} + +/// Pod -> StatefulSet -> LeaderWorkerSet resolves to the LWS. +#[tokio::test] +#[ignore] +async fn find_root_object_lws_chain() { + use kube::CustomResourceExt; + use resources::leaderworkerset::{LeaderWorkerSet, LeaderWorkerSetSpec}; + + let client = Client::try_default().await.unwrap(); + ensure_crd(&client, LeaderWorkerSet::crd()).await; + let ns = create_test_namespace(&client, "gpu-pruner-e2e-lws").await; + + let svc_api: Api = Api::namespaced(client.clone(), &ns); + svc_api + .create( + &PostParams::default(), + &make_headless_service("e2e-lws-ss-svc", &ns), + ) + .await + .unwrap(); + + let lws_api: Api = Api::namespaced(client.clone(), &ns); + let lws = lws_api + .create( + &PostParams::default(), + &LeaderWorkerSet::new( + "e2e-lws", + LeaderWorkerSetSpec { + replicas: Some(1), + leader_worker_template: None, + }, + ), + ) + .await + .unwrap(); + + // no LWS controller in kind, so build the StatefulSet it would create + let mut ss = make_statefulset("e2e-lws-ss", &ns); + ss.metadata.owner_references = Some(vec![owner_reference( + "leaderworkerset.x-k8s.io/v1", + "LeaderWorkerSet", + "e2e-lws", + lws.metadata.uid.as_deref().unwrap(), + )]); + let ss_api: Api = Api::namespaced(client.clone(), &ns); + ss_api.create(&PostParams::default(), &ss).await.unwrap(); + wait_for_statefulset_ready(&ss_api, "e2e-lws-ss").await; + + let pod_api: Api = Api::namespaced(client.clone(), &ns); + let pods = pod_api.list(&Default::default()).await.unwrap(); + let pod = pods.items.first().expect("no pods found for statefulset"); + + let root = find_root_object(client.clone(), &pod.metadata) + .await + .expect("failed to find root object"); + match &root { + ScaleKind::LeaderWorkerSet(l) => assert_eq!(l.name_unchecked(), "e2e-lws"), + other => panic!("expected LeaderWorkerSet root, got {}", other.kind()), + } + + delete_test_namespace(&client, &ns).await; +} + +/// Pod -> StatefulSet -> LeaderWorkerSet -> LLMInferenceService resolves to +/// the LLMInferenceService: scaling the LWS alone is undone by the KServe +/// controller reconciling it back. +#[tokio::test] +#[ignore] +async fn find_root_object_llmis_owned_lws_chain() { + use kube::CustomResourceExt; + use resources::leaderworkerset::{LeaderWorkerSet, LeaderWorkerSetSpec}; + use resources::llminferenceservice::{LLMInferenceService, LLMInferenceServiceSpec}; + + let client = Client::try_default().await.unwrap(); + ensure_crd(&client, LeaderWorkerSet::crd()).await; + ensure_crd(&client, LLMInferenceService::crd()).await; + let ns = create_test_namespace(&client, "gpu-pruner-e2e-llmis-lws").await; + + let svc_api: Api = Api::namespaced(client.clone(), &ns); + svc_api + .create( + &PostParams::default(), + &make_headless_service("e2e-llmis-ss-svc", &ns), + ) + .await + .unwrap(); + + let llmis_api: Api = Api::namespaced(client.clone(), &ns); + let llmis = llmis_api + .create( + &PostParams::default(), + &LLMInferenceService::new( + "e2e-llmis", + LLMInferenceServiceSpec { + replicas: Some(1), + other: Default::default(), + }, + ), + ) + .await + .unwrap(); + + let lws_api: Api = Api::namespaced(client.clone(), &ns); + let mut lws = LeaderWorkerSet::new( + "e2e-llmis-lws", + LeaderWorkerSetSpec { + replicas: Some(1), + leader_worker_template: None, + }, + ); + lws.metadata.owner_references = Some(vec![owner_reference( + "serving.kserve.io/v1alpha1", + "LLMInferenceService", + "e2e-llmis", + llmis.metadata.uid.as_deref().unwrap(), + )]); + let lws = lws_api.create(&PostParams::default(), &lws).await.unwrap(); + + let mut ss = make_statefulset("e2e-llmis-ss", &ns); + ss.metadata.owner_references = Some(vec![owner_reference( + "leaderworkerset.x-k8s.io/v1", + "LeaderWorkerSet", + "e2e-llmis-lws", + lws.metadata.uid.as_deref().unwrap(), + )]); + let ss_api: Api = Api::namespaced(client.clone(), &ns); + ss_api.create(&PostParams::default(), &ss).await.unwrap(); + wait_for_statefulset_ready(&ss_api, "e2e-llmis-ss").await; + + let pod_api: Api = Api::namespaced(client.clone(), &ns); + let pods = pod_api.list(&Default::default()).await.unwrap(); + let pod = pods.items.first().expect("no pods found for statefulset"); + + let root = find_root_object(client.clone(), &pod.metadata) + .await + .expect("failed to find root object"); + match &root { + ScaleKind::LLMInferenceService(l) => assert_eq!(l.name_unchecked(), "e2e-llmis"), + other => panic!("expected LLMInferenceService root, got {}", other.kind()), + } + + delete_test_namespace(&client, &ns).await; +} diff --git a/resources/src/leaderworkerset.rs b/resources/src/leaderworkerset.rs index 3aef39d..ced5400 100644 --- a/resources/src/leaderworkerset.rs +++ b/resources/src/leaderworkerset.rs @@ -29,9 +29,19 @@ pub struct LeaderWorkerSetSpec { skip_serializing_if = "Option::is_none", rename = "leaderWorkerTemplate" )] + #[schemars(schema_with = "preserve_unknown_object")] pub leader_worker_template: Option, } +/// Schema for opaque pass-through fields: without an explicit type and the +/// preserve marker, the API server rejects the CRD as non-structural. +pub fn preserve_unknown_object(_gen: &mut schemars::SchemaGenerator) -> schemars::Schema { + schemars::json_schema!({ + "type": "object", + "x-kubernetes-preserve-unknown-fields": true + }) +} + #[derive(Serialize, Deserialize, Clone, Debug, JsonSchema)] pub struct LeaderWorkerSetStatus { #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/resources/src/llminferenceservice.rs b/resources/src/llminferenceservice.rs index 72e556e..287b113 100644 --- a/resources/src/llminferenceservice.rs +++ b/resources/src/llminferenceservice.rs @@ -20,10 +20,12 @@ use std::collections::BTreeMap; plural = "llminferenceservices" )] #[kube(namespaced)] +#[schemars(extend("x-kubernetes-preserve-unknown-fields" = true))] pub struct LLMInferenceServiceSpec { #[serde(default, skip_serializing_if = "Option::is_none")] pub replicas: Option, #[serde(flatten)] + #[schemars(skip)] pub other: BTreeMap, }