mirror of
https://github.com/kata-containers/kata-containers.git
synced 2026-05-04 04:00:07 +00:00
pod-resources-rs: Add kubelet Pod Resources API client
Add a gRPC client crate that speaks the kubelet PodResourcesLister service (v1). The runtime-rs VFIO cold-plug path needs this to discover which GPU devices the kubelet has assigned to a pod so they can be passed through to the guest before the VM boots. The crate is intentionally kept minimal: it wraps the upstream pod_resources.proto, exposes a Unix-domain-socket client, and re-exports the generated types. Signed-off-by: Alex Lyn <alex.lyn@antgroup.com> Signed-off-by: Fabiano Fidêncio <ffidencio@nvidia.com>
This commit is contained in:
Alex Lyn
committed by
Fabiano Fidêncio
Fabiano Fidêncio
parent
86e5975ad6
commit
c2e7785321
145
Cargo.lock
generated
145
Cargo.lock
generated
@@ -507,7 +507,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3b829e4e32b91e643de6eafe82b1d90675f5874230191a4ffbc1b336dec4d6bf"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum-core",
|
||||
"axum-core 0.3.4",
|
||||
"bitflags 1.3.2",
|
||||
"bytes 1.11.1",
|
||||
"futures-util",
|
||||
@@ -515,7 +515,7 @@ dependencies = [
|
||||
"http-body 0.4.6",
|
||||
"hyper 0.14.32",
|
||||
"itoa",
|
||||
"matchit",
|
||||
"matchit 0.7.3",
|
||||
"memchr",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
@@ -528,6 +528,31 @@ dependencies = [
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum"
|
||||
version = "0.8.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8"
|
||||
dependencies = [
|
||||
"axum-core 0.5.6",
|
||||
"bytes 1.11.1",
|
||||
"futures-util",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"itoa",
|
||||
"matchit 0.8.4",
|
||||
"memchr",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"serde_core",
|
||||
"sync_wrapper 1.0.2",
|
||||
"tower 0.5.3",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-core"
|
||||
version = "0.3.4"
|
||||
@@ -545,6 +570,24 @@ dependencies = [
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-core"
|
||||
version = "0.5.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
|
||||
dependencies = [
|
||||
"bytes 1.11.1",
|
||||
"futures-core",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"mime",
|
||||
"pin-project-lite",
|
||||
"sync_wrapper 1.0.2",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "backon"
|
||||
version = "1.6.0"
|
||||
@@ -1140,7 +1183,7 @@ dependencies = [
|
||||
"prost 0.11.9",
|
||||
"prost-types 0.11.9",
|
||||
"tokio",
|
||||
"tonic",
|
||||
"tonic 0.9.2",
|
||||
"tonic-build 0.9.2",
|
||||
"tower 0.4.13",
|
||||
]
|
||||
@@ -2458,7 +2501,7 @@ dependencies = [
|
||||
"tar",
|
||||
"tempfile",
|
||||
"tokio",
|
||||
"tonic",
|
||||
"tonic 0.9.2",
|
||||
"tower 0.4.13",
|
||||
]
|
||||
|
||||
@@ -2840,6 +2883,7 @@ dependencies = [
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"httparse",
|
||||
"httpdate",
|
||||
"itoa",
|
||||
"pin-project-lite",
|
||||
"pin-utils",
|
||||
@@ -3456,7 +3500,7 @@ checksum = "9f1ac03a0ee89d53fc350989682a56915a4f93fe7b51801a1066cb3caeb2a23f"
|
||||
dependencies = [
|
||||
"prost 0.11.9",
|
||||
"serde",
|
||||
"tonic",
|
||||
"tonic 0.9.2",
|
||||
"tonic-build 0.8.4",
|
||||
]
|
||||
|
||||
@@ -4057,6 +4101,12 @@ version = "0.7.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94"
|
||||
|
||||
[[package]]
|
||||
name = "matchit"
|
||||
version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
|
||||
|
||||
[[package]]
|
||||
name = "mem-agent"
|
||||
version = "0.2.0"
|
||||
@@ -5329,6 +5379,24 @@ version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6"
|
||||
|
||||
[[package]]
|
||||
name = "pod-resources-rs"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"container-device-interface",
|
||||
"hyper-util",
|
||||
"oci-spec 0.8.4",
|
||||
"prost 0.14.3",
|
||||
"slog",
|
||||
"slog-scope",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"tonic 0.14.5",
|
||||
"tonic-prost",
|
||||
"tower 0.5.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "polling"
|
||||
version = "2.8.0"
|
||||
@@ -5574,6 +5642,16 @@ dependencies = [
|
||||
"prost-derive 0.11.9",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "prost"
|
||||
version = "0.14.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568"
|
||||
dependencies = [
|
||||
"bytes 1.11.1",
|
||||
"prost-derive 0.14.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "prost-build"
|
||||
version = "0.8.0"
|
||||
@@ -5640,6 +5718,19 @@ dependencies = [
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "prost-derive"
|
||||
version = "0.14.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"itertools",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "prost-types"
|
||||
version = "0.8.0"
|
||||
@@ -8023,7 +8114,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3082666a3a6433f7f511c7192923fa1fe07c69332d3c6a2e6bb040b569199d5a"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum",
|
||||
"axum 0.6.20",
|
||||
"base64 0.21.7",
|
||||
"bytes 1.11.1",
|
||||
"futures-core",
|
||||
@@ -8044,6 +8135,35 @@ dependencies = [
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic"
|
||||
version = "0.14.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum 0.8.8",
|
||||
"base64 0.22.1",
|
||||
"bytes 1.11.1",
|
||||
"h2 0.4.13",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"hyper 1.8.1",
|
||||
"hyper-timeout 0.5.2",
|
||||
"hyper-util",
|
||||
"percent-encoding",
|
||||
"pin-project",
|
||||
"socket2 0.6.3",
|
||||
"sync_wrapper 1.0.2",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tower 0.5.3",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic-build"
|
||||
version = "0.8.4"
|
||||
@@ -8070,6 +8190,17 @@ dependencies = [
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost"
|
||||
version = "0.14.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"
|
||||
dependencies = [
|
||||
"bytes 1.11.1",
|
||||
"prost 0.14.3",
|
||||
"tonic 0.14.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower"
|
||||
version = "0.4.13"
|
||||
@@ -8098,7 +8229,9 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"indexmap 2.13.0",
|
||||
"pin-project-lite",
|
||||
"slab",
|
||||
"sync_wrapper 1.0.2",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
|
||||
@@ -11,6 +11,7 @@ members = [
|
||||
"src/libs/kata-types",
|
||||
"src/libs/logging",
|
||||
"src/libs/mem-agent",
|
||||
"src/libs/pod-resources-rs",
|
||||
"src/libs/protocols",
|
||||
"src/libs/runtime-spec",
|
||||
"src/libs/safe-path",
|
||||
@@ -123,6 +124,7 @@ wasm_container = { path = "src/runtime-rs/crates/runtimes/wasm_container" }
|
||||
|
||||
# Local dependencies from `src/lib`
|
||||
kata-sys-util = { path = "src/libs/kata-sys-util" }
|
||||
pod-resources-rs = { path = "src/libs/pod-resources-rs" }
|
||||
kata-types = { path = "src/libs/kata-types", features = ["safe-path"] }
|
||||
logging = { path = "src/libs/logging" }
|
||||
mem-agent = { path = "src/libs/mem-agent" }
|
||||
|
||||
20
src/libs/pod-resources-rs/Cargo.toml
Normal file
20
src/libs/pod-resources-rs/Cargo.toml
Normal file
@@ -0,0 +1,20 @@
|
||||
[package]
|
||||
name = "pod-resources-rs"
|
||||
version = "0.1.0"
|
||||
edition = { workspace = true }
|
||||
license = { workspace = true }
|
||||
|
||||
[dependencies]
|
||||
anyhow = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tokio-util = "0.7.17"
|
||||
tower = "0.5"
|
||||
hyper-util = { version = "0.1", features = ["tokio"] }
|
||||
# gRPC dependencies for kubelet pod-resources API
|
||||
tonic = "0.14"
|
||||
prost = "0.14"
|
||||
tonic-prost = "0.14"
|
||||
oci-spec = { workspace = true }
|
||||
container-device-interface = "0.1.2"
|
||||
slog = { workspace = true }
|
||||
slog-scope = { workspace = true }
|
||||
120
src/libs/pod-resources-rs/proto/pod_resources.proto
Normal file
120
src/libs/pod-resources-rs/proto/pod_resources.proto
Normal file
@@ -0,0 +1,120 @@
|
||||
// Copyright (c) 2026 The Kubernetes Authors
|
||||
//
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
//
|
||||
syntax = "proto3";
|
||||
|
||||
package v1;
|
||||
option go_package = "k8s.io/kubelet/pkg/apis/podresources/v1";
|
||||
|
||||
|
||||
// PodResourcesLister is a service provided by the kubelet that provides information about the
|
||||
// node resources consumed by pods and containers on the node
|
||||
service PodResourcesLister {
|
||||
/// List returns the node resources assigned to pods and containers.
|
||||
rpc List(ListPodResourcesRequest) returns (ListPodResourcesResponse) {}
|
||||
/// GetAllocatableResources returns the node resources that are available for assignment to pods and containers.
|
||||
rpc GetAllocatableResources(AllocatableResourcesRequest) returns (AllocatableResourcesResponse) {}
|
||||
/// Get returns the node resources assigned to a specific pod.
|
||||
rpc Get(GetPodResourcesRequest) returns (GetPodResourcesResponse) {}
|
||||
}
|
||||
|
||||
// AllocatableResourcesRequest is the request made to the GetAllocatableResources service
|
||||
message AllocatableResourcesRequest {}
|
||||
|
||||
// AllocatableResourcesResponses contains informations about all the devices known by the kubelet
|
||||
message AllocatableResourcesResponse {
|
||||
repeated ContainerDevices devices = 1;
|
||||
repeated int64 cpu_ids = 2;
|
||||
repeated ContainerMemory memory = 3;
|
||||
}
|
||||
|
||||
// ListPodResourcesRequest is the request made to the PodResources service
|
||||
message ListPodResourcesRequest {}
|
||||
|
||||
// ListPodResourcesResponse is the response returned by List function
|
||||
message ListPodResourcesResponse {
|
||||
repeated PodResources pod_resources = 1;
|
||||
}
|
||||
|
||||
// GetPodResourcesRequest is the request made to the Get service
|
||||
message GetPodResourcesRequest {
|
||||
string pod_name = 1;
|
||||
string pod_namespace = 2;
|
||||
}
|
||||
|
||||
// GetPodResourcesResponse is the response returned by Get function
|
||||
message GetPodResourcesResponse {
|
||||
PodResources pod_resources = 1;
|
||||
}
|
||||
|
||||
// PodResources contains information about the node resources assigned to a pod
|
||||
message PodResources {
|
||||
string name = 1;
|
||||
string namespace = 2;
|
||||
repeated ContainerResources containers = 3;
|
||||
}
|
||||
|
||||
// ContainerResources contains information about the resources assigned to a container
|
||||
message ContainerResources {
|
||||
string name = 1;
|
||||
repeated ContainerDevices devices = 2;
|
||||
repeated int64 cpu_ids = 3;
|
||||
repeated ContainerMemory memory = 4;
|
||||
repeated DynamicResource dynamic_resources = 5;
|
||||
}
|
||||
|
||||
// ContainerDevices contains information about the devices assigned to a container
|
||||
message ContainerDevices {
|
||||
string resource_name = 1;
|
||||
repeated string device_ids = 2;
|
||||
TopologyInfo topology = 3;
|
||||
}
|
||||
|
||||
// ContainerMemory contains information about memory and hugepages assigned to a container
|
||||
message ContainerMemory {
|
||||
string memory_type = 1;
|
||||
uint64 size = 2;
|
||||
TopologyInfo topology = 3;
|
||||
}
|
||||
|
||||
// DynamicResource contains information about the devices assigned to a container by DRA
|
||||
message DynamicResource {
|
||||
// tombstone: removed in 1.31 because claims are no longer associated with one class
|
||||
// string class_name = 1;
|
||||
string claim_name = 2;
|
||||
string claim_namespace = 3;
|
||||
repeated ClaimResource claim_resources = 4;
|
||||
}
|
||||
|
||||
// ClaimResource contains resource information. The driver name/pool name/device name
|
||||
// triplet uniquely identifies the device. Should DRA get extended to other kinds
|
||||
// of resources, then device_name will be empty and other fields will get added.
|
||||
// Each device at the DRA API level may map to zero or more CDI devices.
|
||||
message ClaimResource {
|
||||
repeated CDIDevice cdi_devices = 1;
|
||||
string driver_name = 2;
|
||||
string pool_name = 3;
|
||||
string device_name = 4;
|
||||
}
|
||||
|
||||
// Topology describes hardware topology of the resource
|
||||
message TopologyInfo {
|
||||
repeated NUMANode nodes = 1;
|
||||
}
|
||||
|
||||
// NUMA representation of NUMA node
|
||||
message NUMANode {
|
||||
int64 ID = 1;
|
||||
}
|
||||
|
||||
// CDIDevice specifies a CDI device information
|
||||
message CDIDevice {
|
||||
// Fully qualified CDI device name
|
||||
// for example: vendor.com/gpu=gpudevice1
|
||||
// see more details in the CDI specification:
|
||||
// https://github.com/container-orchestrated-devices/container-device-interface/blob/main/SPEC.md
|
||||
string name = 1;
|
||||
}
|
||||
|
||||
|
||||
80
src/libs/pod-resources-rs/src/lib.rs
Normal file
80
src/libs/pod-resources-rs/src/lib.rs
Normal file
@@ -0,0 +1,80 @@
|
||||
// Copyright (c) 2026 Ant Group
|
||||
//
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
//
|
||||
|
||||
pub mod pod_resources;
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use cdi::specs::config::DeviceNode;
|
||||
// use cdi::container_edits::DeviceNode;
|
||||
use cdi::cache::{CdiOption, new_cache, with_auto_refresh};
|
||||
use cdi::spec_dirs::with_spec_dirs;
|
||||
use container_device_interface as cdi;
|
||||
|
||||
use slog::info;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// DEFAULT_DYNAMIC_CDI_SPEC_PATH is the default directory for dynamic CDI Specs,
|
||||
/// which can be overridden by specifying a different path when creating the cache.
|
||||
const DEFAULT_DYNAMIC_CDI_SPEC_PATH: &str = "/var/run/cdi";
|
||||
/// DEFAULT_STATIC_CDI_SPEC_PATH is the default directory for static CDI Specs,
|
||||
/// which can be overridden by specifying a different path when creating the cache.
|
||||
const DEFAULT_STATIC_CDI_SPEC_PATH: &str = "/etc/cdi";
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! sl {
|
||||
() => {
|
||||
slog_scope::logger()
|
||||
};
|
||||
}
|
||||
|
||||
pub async fn handle_cdi_devices(devices: &[String]) -> Result<Vec<DeviceNode>> {
|
||||
if devices.is_empty() {
|
||||
info!(sl!(), "no pod CDI devices requested.");
|
||||
return Ok(vec![]);
|
||||
}
|
||||
// Explicitly set the cache options to disable auto-refresh and
|
||||
// to use the default spec dirs for dynamic and static CDI Specs
|
||||
let options: Vec<CdiOption> = vec![
|
||||
with_auto_refresh(false),
|
||||
with_spec_dirs(&[DEFAULT_DYNAMIC_CDI_SPEC_PATH, DEFAULT_STATIC_CDI_SPEC_PATH]),
|
||||
];
|
||||
let cache: Arc<std::sync::Mutex<cdi::cache::Cache>> = new_cache(options);
|
||||
|
||||
let target_devices = {
|
||||
let mut target_devices = vec![];
|
||||
// Lock cache within this scope, std::sync::Mutex has no Send
|
||||
// and await will not work with time::sleep
|
||||
let mut cache = cache.lock().unwrap();
|
||||
match cache.refresh() {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
return Err(anyhow!("Refreshing cache failed: {:?}", e));
|
||||
}
|
||||
}
|
||||
|
||||
for dev in devices.iter() {
|
||||
info!(sl!(), "Requested CDI device with FQN: {}", dev);
|
||||
match cache.get_device(dev) {
|
||||
Some(device) => {
|
||||
info!(sl!(), "Target CDI device: {}", device.get_qualified_name());
|
||||
if let Some(devnodes) = device.edits().container_edits.device_nodes {
|
||||
target_devices.extend(devnodes.iter().cloned());
|
||||
}
|
||||
}
|
||||
None => {
|
||||
return Err(anyhow!(
|
||||
"Failed to get device node for CDI device: {} in cache",
|
||||
dev
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
target_devices
|
||||
};
|
||||
info!(sl!(), "target CDI devices to inject: {:?}", target_devices);
|
||||
|
||||
Ok(target_devices)
|
||||
}
|
||||
136
src/libs/pod-resources-rs/src/pod_resources/mod.rs
Normal file
136
src/libs/pod-resources-rs/src/pod_resources/mod.rs
Normal file
@@ -0,0 +1,136 @@
|
||||
// Copyright (c) 2026 Ant Group
|
||||
//
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
//
|
||||
|
||||
pub mod v1;
|
||||
|
||||
use v1::pod_resources_lister_client::PodResourcesListerClient;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::convert::TryFrom;
|
||||
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use hyper_util::rt::TokioIo;
|
||||
use tokio::net::UnixStream;
|
||||
use tokio::time::{Duration, timeout};
|
||||
use tonic::transport::{Channel, Endpoint, Uri};
|
||||
use tower::service_fn;
|
||||
|
||||
use crate::pod_resources::v1::GetPodResourcesRequest;
|
||||
|
||||
// containerd CRI annotations
|
||||
const SANDBOX_NAME_ANNOTATION: &str = "io.kubernetes.cri.sandbox-name";
|
||||
const SANDBOX_NAMESPACE_ANNOTATION: &str = "io.kubernetes.cri.sandbox-namespace";
|
||||
|
||||
// CRI-O annotations (fallback)
|
||||
const CRIO_NAME_ANNOTATION: &str = "io.kubernetes.cri-o.KubeName";
|
||||
const CRIO_NAMESPACE_ANNOTATION: &str = "io.kubernetes.cri-o.Namespace";
|
||||
pub const DEFAULT_POD_RESOURCES_PATH: &str = "/var/lib/kubelet/pod-resources";
|
||||
pub const DEFAULT_POD_RESOURCES_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
pub const CDI_K8S_PREFIX: &str = "cdi.k8s.io/";
|
||||
const MAX_RECV_MSG_SIZE: usize = 16 * 1024 * 1024; // 16MB
|
||||
|
||||
// Create a gRPC channel to the specified Unix socket
|
||||
async fn create_grpc_channel(socket_path: &str) -> Result<Channel> {
|
||||
let socket_path = socket_path.trim_start_matches("unix://");
|
||||
let socket_path_owned = socket_path.to_string();
|
||||
|
||||
// Create a gRPC endpoint with a timeout
|
||||
let endpoint = Endpoint::try_from("http://[::]:50051")
|
||||
.context("failed to create endpoint")?
|
||||
.timeout(DEFAULT_POD_RESOURCES_TIMEOUT);
|
||||
|
||||
// Connect to the Unix socket using a custom connector
|
||||
let channel = endpoint
|
||||
.connect_with_connector(service_fn(move |_: Uri| {
|
||||
let socket_path = socket_path_owned.clone();
|
||||
async move {
|
||||
let stream = UnixStream::connect(&socket_path).await.map_err(|e| {
|
||||
std::io::Error::new(
|
||||
e.kind(),
|
||||
format!("failed to connect to {}: {}", socket_path, e),
|
||||
)
|
||||
})?;
|
||||
Ok::<_, std::io::Error>(TokioIo::new(stream))
|
||||
}
|
||||
}))
|
||||
.await
|
||||
.context("failed to connect to unix socket")?;
|
||||
|
||||
Ok(channel)
|
||||
}
|
||||
|
||||
pub async fn get_pod_cdi_devices(
|
||||
socket: &str,
|
||||
annotations: &HashMap<String, String>,
|
||||
) -> Result<Vec<String>> {
|
||||
let pod_name = annotations
|
||||
.get(SANDBOX_NAME_ANNOTATION)
|
||||
.or_else(|| annotations.get(CRIO_NAME_ANNOTATION))
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"cold plug: missing annotation {} or {}",
|
||||
SANDBOX_NAME_ANNOTATION,
|
||||
CRIO_NAME_ANNOTATION
|
||||
)
|
||||
})?;
|
||||
|
||||
let pod_namespace = annotations
|
||||
.get(SANDBOX_NAMESPACE_ANNOTATION)
|
||||
.or_else(|| annotations.get(CRIO_NAMESPACE_ANNOTATION))
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"cold plug: missing annotation {} or {}",
|
||||
SANDBOX_NAMESPACE_ANNOTATION,
|
||||
CRIO_NAMESPACE_ANNOTATION
|
||||
)
|
||||
})?;
|
||||
|
||||
// Create gRPC channel to kubelet pod-resources socket
|
||||
let channel = create_grpc_channel(socket)
|
||||
.await
|
||||
.context("cold plug: failed to connect to kubelet")?;
|
||||
|
||||
// Create PodResourcesLister client
|
||||
let mut client = PodResourcesListerClient::new(channel)
|
||||
.max_decoding_message_size(MAX_RECV_MSG_SIZE)
|
||||
.max_encoding_message_size(MAX_RECV_MSG_SIZE);
|
||||
|
||||
// Prepare and send GetPodResources request
|
||||
let request = tonic::Request::new(GetPodResourcesRequest {
|
||||
pod_name: pod_name.to_string(),
|
||||
pod_namespace: pod_namespace.to_string(),
|
||||
});
|
||||
|
||||
// Await response with timeout
|
||||
let response = timeout(DEFAULT_POD_RESOURCES_TIMEOUT, client.get(request))
|
||||
.await
|
||||
.context("cold plug: GetPodResources timeout")?
|
||||
.context("cold plug: GetPodResources RPC failed")?;
|
||||
|
||||
// Extract PodResources from response
|
||||
let pod_resources = response
|
||||
.into_inner()
|
||||
.pod_resources
|
||||
.ok_or_else(|| anyhow!("cold plug: PodResources is nil"))?;
|
||||
|
||||
// Format device specifications
|
||||
let format_cdi_device_ids = |resource_name: &str, device_ids: &[String]| -> Vec<String> {
|
||||
device_ids
|
||||
.iter()
|
||||
.map(|id| format!("{}={}", resource_name, id))
|
||||
.collect()
|
||||
};
|
||||
|
||||
// Collect all device specifications from all containers
|
||||
let mut devices = Vec::new();
|
||||
for container in &pod_resources.containers {
|
||||
for device in &container.devices {
|
||||
let cdi_devices = format_cdi_device_ids(&device.resource_name, &device.device_ids);
|
||||
devices.extend(cdi_devices);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(devices)
|
||||
}
|
||||
306
src/libs/pod-resources-rs/src/pod_resources/v1.rs
Normal file
306
src/libs/pod-resources-rs/src/pod_resources/v1.rs
Normal file
@@ -0,0 +1,306 @@
|
||||
// Copyright (c) 2026 Ant Group
|
||||
//
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
//
|
||||
// This file is @generated by prost-build.
|
||||
/// AllocatableResourcesRequest is the request made to the GetAllocatableResources service
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct AllocatableResourcesRequest {}
|
||||
/// AllocatableResourcesResponses contains informations about all the devices known by the kubelet
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct AllocatableResourcesResponse {
|
||||
#[prost(message, repeated, tag = "1")]
|
||||
pub devices: ::prost::alloc::vec::Vec<ContainerDevices>,
|
||||
#[prost(int64, repeated, tag = "2")]
|
||||
pub cpu_ids: ::prost::alloc::vec::Vec<i64>,
|
||||
#[prost(message, repeated, tag = "3")]
|
||||
pub memory: ::prost::alloc::vec::Vec<ContainerMemory>,
|
||||
}
|
||||
/// ListPodResourcesRequest is the request made to the PodResources service
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct ListPodResourcesRequest {}
|
||||
/// ListPodResourcesResponse is the response returned by List function
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct ListPodResourcesResponse {
|
||||
#[prost(message, repeated, tag = "1")]
|
||||
pub pod_resources: ::prost::alloc::vec::Vec<PodResources>,
|
||||
}
|
||||
/// GetPodResourcesRequest is the request made to the Get service
|
||||
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct GetPodResourcesRequest {
|
||||
#[prost(string, tag = "1")]
|
||||
pub pod_name: ::prost::alloc::string::String,
|
||||
#[prost(string, tag = "2")]
|
||||
pub pod_namespace: ::prost::alloc::string::String,
|
||||
}
|
||||
/// GetPodResourcesResponse is the response returned by Get function
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct GetPodResourcesResponse {
|
||||
#[prost(message, optional, tag = "1")]
|
||||
pub pod_resources: ::core::option::Option<PodResources>,
|
||||
}
|
||||
/// PodResources contains information about the node resources assigned to a pod
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct PodResources {
|
||||
#[prost(string, tag = "1")]
|
||||
pub name: ::prost::alloc::string::String,
|
||||
#[prost(string, tag = "2")]
|
||||
pub namespace: ::prost::alloc::string::String,
|
||||
#[prost(message, repeated, tag = "3")]
|
||||
pub containers: ::prost::alloc::vec::Vec<ContainerResources>,
|
||||
}
|
||||
/// ContainerResources contains information about the resources assigned to a container
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct ContainerResources {
|
||||
#[prost(string, tag = "1")]
|
||||
pub name: ::prost::alloc::string::String,
|
||||
#[prost(message, repeated, tag = "2")]
|
||||
pub devices: ::prost::alloc::vec::Vec<ContainerDevices>,
|
||||
#[prost(int64, repeated, tag = "3")]
|
||||
pub cpu_ids: ::prost::alloc::vec::Vec<i64>,
|
||||
#[prost(message, repeated, tag = "4")]
|
||||
pub memory: ::prost::alloc::vec::Vec<ContainerMemory>,
|
||||
#[prost(message, repeated, tag = "5")]
|
||||
pub dynamic_resources: ::prost::alloc::vec::Vec<DynamicResource>,
|
||||
}
|
||||
/// ContainerDevices contains information about the devices assigned to a container
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct ContainerDevices {
|
||||
#[prost(string, tag = "1")]
|
||||
pub resource_name: ::prost::alloc::string::String,
|
||||
#[prost(string, repeated, tag = "2")]
|
||||
pub device_ids: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
|
||||
#[prost(message, optional, tag = "3")]
|
||||
pub topology: ::core::option::Option<TopologyInfo>,
|
||||
}
|
||||
/// ContainerMemory contains information about memory and hugepages assigned to a container
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct ContainerMemory {
|
||||
#[prost(string, tag = "1")]
|
||||
pub memory_type: ::prost::alloc::string::String,
|
||||
#[prost(uint64, tag = "2")]
|
||||
pub size: u64,
|
||||
#[prost(message, optional, tag = "3")]
|
||||
pub topology: ::core::option::Option<TopologyInfo>,
|
||||
}
|
||||
/// DynamicResource contains information about the devices assigned to a container by DRA
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct DynamicResource {
|
||||
/// tombstone: removed in 1.31 because claims are no longer associated with one class
|
||||
/// string class_name = 1;
|
||||
#[prost(string, tag = "2")]
|
||||
pub claim_name: ::prost::alloc::string::String,
|
||||
#[prost(string, tag = "3")]
|
||||
pub claim_namespace: ::prost::alloc::string::String,
|
||||
#[prost(message, repeated, tag = "4")]
|
||||
pub claim_resources: ::prost::alloc::vec::Vec<ClaimResource>,
|
||||
}
|
||||
/// ClaimResource contains resource information. The driver name/pool name/device name
|
||||
/// triplet uniquely identifies the device. Should DRA get extended to other kinds
|
||||
/// of resources, then device_name will be empty and other fields will get added.
|
||||
/// Each device at the DRA API level may map to zero or more CDI devices.
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct ClaimResource {
|
||||
#[prost(message, repeated, tag = "1")]
|
||||
pub cdi_devices: ::prost::alloc::vec::Vec<CdiDevice>,
|
||||
#[prost(string, tag = "2")]
|
||||
pub driver_name: ::prost::alloc::string::String,
|
||||
#[prost(string, tag = "3")]
|
||||
pub pool_name: ::prost::alloc::string::String,
|
||||
#[prost(string, tag = "4")]
|
||||
pub device_name: ::prost::alloc::string::String,
|
||||
}
|
||||
/// Topology describes hardware topology of the resource
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
pub struct TopologyInfo {
|
||||
#[prost(message, repeated, tag = "1")]
|
||||
pub nodes: ::prost::alloc::vec::Vec<NumaNode>,
|
||||
}
|
||||
/// NUMA representation of NUMA node
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct NumaNode {
|
||||
#[prost(int64, tag = "1")]
|
||||
pub id: i64,
|
||||
}
|
||||
/// CDIDevice specifies a CDI device information
|
||||
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct CdiDevice {
|
||||
/// Fully qualified CDI device name
|
||||
/// for example: vendor.com/gpu=gpudevice1
|
||||
/// see more details in the CDI specification:
|
||||
/// <https://github.com/container-orchestrated-devices/container-device-interface/blob/main/SPEC.md>
|
||||
#[prost(string, tag = "1")]
|
||||
pub name: ::prost::alloc::string::String,
|
||||
}
|
||||
/// Generated client implementations.
|
||||
pub mod pod_resources_lister_client {
|
||||
#![allow(
|
||||
unused_variables,
|
||||
dead_code,
|
||||
missing_docs,
|
||||
clippy::wildcard_imports,
|
||||
clippy::let_unit_value,
|
||||
)]
|
||||
use std::convert::TryInto;
|
||||
use tonic::codegen::*;
|
||||
use tonic::codegen::http::Uri;
|
||||
/// PodResourcesLister is a service provided by the kubelet that provides information about the
|
||||
/// node resources consumed by pods and containers on the node
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PodResourcesListerClient<T> {
|
||||
inner: tonic::client::Grpc<T>,
|
||||
}
|
||||
impl PodResourcesListerClient<tonic::transport::Channel> {
|
||||
/// Attempt to create a new client by connecting to a given endpoint.
|
||||
pub async fn connect<D>(dst: D) -> Result<Self, tonic::transport::Error>
|
||||
where
|
||||
D: TryInto<tonic::transport::Endpoint>,
|
||||
D::Error: Into<StdError>,
|
||||
{
|
||||
let conn = tonic::transport::Endpoint::new(dst)?.connect().await?;
|
||||
Ok(Self::new(conn))
|
||||
}
|
||||
}
|
||||
impl<T> PodResourcesListerClient<T>
|
||||
where
|
||||
T: tonic::client::GrpcService<tonic::body::Body>,
|
||||
T::Error: Into<StdError>,
|
||||
T::ResponseBody: Body<Data = Bytes> + std::marker::Send + 'static,
|
||||
<T::ResponseBody as Body>::Error: Into<StdError> + std::marker::Send,
|
||||
{
|
||||
pub fn new(inner: T) -> Self {
|
||||
let inner = tonic::client::Grpc::new(inner);
|
||||
Self { inner }
|
||||
}
|
||||
pub fn with_origin(inner: T, origin: Uri) -> Self {
|
||||
let inner = tonic::client::Grpc::with_origin(inner, origin);
|
||||
Self { inner }
|
||||
}
|
||||
pub fn with_interceptor<F>(
|
||||
inner: T,
|
||||
interceptor: F,
|
||||
) -> PodResourcesListerClient<InterceptedService<T, F>>
|
||||
where
|
||||
F: tonic::service::Interceptor,
|
||||
T::ResponseBody: Default,
|
||||
T: tonic::codegen::Service<
|
||||
http::Request<tonic::body::Body>,
|
||||
Response = http::Response<
|
||||
<T as tonic::client::GrpcService<tonic::body::Body>>::ResponseBody,
|
||||
>,
|
||||
>,
|
||||
<T as tonic::codegen::Service<
|
||||
http::Request<tonic::body::Body>,
|
||||
>>::Error: Into<StdError> + std::marker::Send + std::marker::Sync,
|
||||
{
|
||||
PodResourcesListerClient::new(InterceptedService::new(inner, interceptor))
|
||||
}
|
||||
/// Compress requests with the given encoding.
|
||||
///
|
||||
/// This requires the server to support it otherwise it might respond with an
|
||||
/// error.
|
||||
#[must_use]
|
||||
pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self {
|
||||
self.inner = self.inner.send_compressed(encoding);
|
||||
self
|
||||
}
|
||||
/// Enable decompressing responses.
|
||||
#[must_use]
|
||||
pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self {
|
||||
self.inner = self.inner.accept_compressed(encoding);
|
||||
self
|
||||
}
|
||||
/// Limits the maximum size of a decoded message.
|
||||
///
|
||||
/// Default: `4MB`
|
||||
#[must_use]
|
||||
pub fn max_decoding_message_size(mut self, limit: usize) -> Self {
|
||||
self.inner = self.inner.max_decoding_message_size(limit);
|
||||
self
|
||||
}
|
||||
/// Limits the maximum size of an encoded message.
|
||||
///
|
||||
/// Default: `usize::MAX`
|
||||
#[must_use]
|
||||
pub fn max_encoding_message_size(mut self, limit: usize) -> Self {
|
||||
self.inner = self.inner.max_encoding_message_size(limit);
|
||||
self
|
||||
}
|
||||
/// / List returns the node resources assigned to pods and containers.
|
||||
pub async fn list(
|
||||
&mut self,
|
||||
request: impl tonic::IntoRequest<super::ListPodResourcesRequest>,
|
||||
) -> std::result::Result<
|
||||
tonic::Response<super::ListPodResourcesResponse>,
|
||||
tonic::Status,
|
||||
> {
|
||||
self.inner
|
||||
.ready()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tonic::Status::unknown(
|
||||
format!("Service was not ready: {}", e.into()),
|
||||
)
|
||||
})?;
|
||||
let codec = tonic_prost::ProstCodec::default();
|
||||
let path = http::uri::PathAndQuery::from_static(
|
||||
"/v1.PodResourcesLister/List",
|
||||
);
|
||||
let mut req = request.into_request();
|
||||
req.extensions_mut()
|
||||
.insert(GrpcMethod::new("v1.PodResourcesLister", "List"));
|
||||
self.inner.unary(req, path, codec).await
|
||||
}
|
||||
/// / GetAllocatableResources returns the node resources that are available for assignment to pods and containers.
|
||||
pub async fn get_allocatable_resources(
|
||||
&mut self,
|
||||
request: impl tonic::IntoRequest<super::AllocatableResourcesRequest>,
|
||||
) -> std::result::Result<
|
||||
tonic::Response<super::AllocatableResourcesResponse>,
|
||||
tonic::Status,
|
||||
> {
|
||||
self.inner
|
||||
.ready()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tonic::Status::unknown(
|
||||
format!("Service was not ready: {}", e.into()),
|
||||
)
|
||||
})?;
|
||||
let codec = tonic_prost::ProstCodec::default();
|
||||
let path = http::uri::PathAndQuery::from_static(
|
||||
"/v1.PodResourcesLister/GetAllocatableResources",
|
||||
);
|
||||
let mut req = request.into_request();
|
||||
req.extensions_mut()
|
||||
.insert(
|
||||
GrpcMethod::new("v1.PodResourcesLister", "GetAllocatableResources"),
|
||||
);
|
||||
self.inner.unary(req, path, codec).await
|
||||
}
|
||||
/// / Get returns the node resources assigned to a specific pod.
|
||||
pub async fn get(
|
||||
&mut self,
|
||||
request: impl tonic::IntoRequest<super::GetPodResourcesRequest>,
|
||||
) -> std::result::Result<
|
||||
tonic::Response<super::GetPodResourcesResponse>,
|
||||
tonic::Status,
|
||||
> {
|
||||
self.inner
|
||||
.ready()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tonic::Status::unknown(
|
||||
format!("Service was not ready: {}", e.into()),
|
||||
)
|
||||
})?;
|
||||
let codec = tonic_prost::ProstCodec::default();
|
||||
let path = http::uri::PathAndQuery::from_static(
|
||||
"/v1.PodResourcesLister/Get",
|
||||
);
|
||||
let mut req = request.into_request();
|
||||
req.extensions_mut().insert(GrpcMethod::new("v1.PodResourcesLister", "Get"));
|
||||
self.inner.unary(req, path, codec).await
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user