gevulot_rs/models/worker.rs
1//! Worker model and related types for managing worker nodes.
2//!
3//! This module provides the core worker model used throughout the system, including:
4//! - Worker specification and status
5//! - Resource tracking (CPU, GPU, memory, disk)
6//! - Metadata like tags and labels
7//! - Protobuf serialization/deserialization
8
9use super::{
10 metadata::{Label, Metadata},
11 ByteUnit, CoreUnit, DefaultFactorOneMegabyte,
12};
13use crate::proto::gevulot::gevulot;
14use serde::{Deserialize, Serialize};
15
16/// Represents a complete worker definition with metadata, specification and status
17///
18/// # Examples
19///
20/// Creating a basic worker:
21/// ```
22/// use crate::models::Worker;
23/// CoreUnit
24/// let worker = serde_json::from_str::<Worker>(r#"{
25/// "kind": "Worker",
26/// "version": "v0",
27/// "metadata": {
28/// "name": "worker-1",
29/// "tags": ["compute"]
30/// },
31/// "spec": {
32/// "cpus": "8 cores",
33/// "gpus": "1 gpu",
34/// "memory": "16 GiB",
35/// "disk": "100 GiB"
36/// }
37/// }"#).unwrap();
38/// ```
39///
40/// Converting from protobuf:
41/// ```
42/// use crate::proto::gevulot::gevulot;
43/// use crate::models::Worker;
44///
45/// let proto_worker = gevulot::Worker {
46/// metadata: Some(gevulot::Metadata {
47/// name: "worker-1".to_string(),
48/// ..Default::default()
49/// }),
50/// spec: Some(gevulot::WorkerSpec {
51/// cpus: 8,
52/// gpus: 1,
53/// memory: 16000000000,
54/// disk: 100000000000,
55/// }),
56/// ..Default::default()
57/// };
58///
59/// let worker = Worker::from(proto_worker);
60/// ```
61#[derive(Serialize, Deserialize, Debug)]
62pub struct Worker {
63 pub kind: String,
64 pub version: String,
65 pub metadata: Metadata,
66 pub spec: WorkerSpec,
67 pub status: Option<WorkerStatus>,
68}
69
70impl From<gevulot::Worker> for Worker {
71 fn from(proto: gevulot::Worker) -> Self {
72 // Convert protobuf worker to our internal worker model
73 Worker {
74 kind: "Worker".to_string(),
75 version: "v0".to_string(),
76 metadata: Metadata {
77 id: proto.metadata.as_ref().map(|m| m.id.clone()),
78 name: proto
79 .metadata
80 .as_ref()
81 .map(|m| m.name.clone())
82 .unwrap_or_default(),
83 creator: proto.metadata.as_ref().map(|m| m.creator.clone()),
84 description: proto
85 .metadata
86 .as_ref()
87 .map(|m| m.desc.clone())
88 .unwrap_or_default(),
89 tags: proto
90 .metadata
91 .as_ref()
92 .map(|m| m.tags.clone())
93 .unwrap_or_default(),
94 labels: proto
95 .metadata
96 .as_ref()
97 .map(|m| m.labels.clone())
98 .unwrap_or_default()
99 .into_iter()
100 .map(|l| Label {
101 key: l.key,
102 value: l.value,
103 })
104 .collect(),
105 workflow_ref: None,
106 },
107 spec: proto.spec.unwrap().into(),
108 status: proto.status.map(|s| s.into()),
109 }
110 }
111}
112
113/// Specification of worker resources and capabilities
114///
115/// Contains the maximum resources available on this worker:
116/// - CPU cores
117/// - GPU devices
118/// - Memory in bytes
119/// - Disk space in bytes
120#[derive(Serialize, Deserialize, Debug)]
121pub struct WorkerSpec {
122 pub cpus: CoreUnit,
123 pub gpus: CoreUnit,
124 pub memory: ByteUnit<DefaultFactorOneMegabyte>,
125 pub disk: ByteUnit<DefaultFactorOneMegabyte>,
126}
127
128impl From<gevulot::WorkerSpec> for WorkerSpec {
129 fn from(proto: gevulot::WorkerSpec) -> Self {
130 // Convert protobuf spec to internal spec
131 WorkerSpec {
132 cpus: (proto.cpus as i64).into(),
133 gpus: (proto.gpus as i64).into(),
134 memory: (proto.memory as i64).into(),
135 disk: (proto.disk as i64).into(),
136 }
137 }
138}
139
140/// Current status and resource utilization of a worker
141///
142/// Tracks:
143/// - Currently used resources (CPU, GPU, memory, disk)
144/// - When the worker announced it will exit
145#[derive(Serialize, Deserialize, Debug)]
146pub struct WorkerStatus {
147 #[serde(rename = "cpusUsed")]
148 pub cpus_used: CoreUnit,
149 #[serde(rename = "gpusUsed")]
150 pub gpus_used: CoreUnit,
151 #[serde(rename = "memoryUsed")]
152 pub memory_used: ByteUnit<DefaultFactorOneMegabyte>,
153 #[serde(rename = "diskUsed")]
154 pub disk_used: ByteUnit<DefaultFactorOneMegabyte>,
155 #[serde(rename = "exitAnnouncedAt")]
156 pub exit_announced_at: i64,
157}
158
159impl From<gevulot::WorkerStatus> for WorkerStatus {
160 fn from(proto: gevulot::WorkerStatus) -> Self {
161 // Convert protobuf status to internal status
162 WorkerStatus {
163 cpus_used: (proto.cpus_used as i64).into(),
164 gpus_used: (proto.gpus_used as i64).into(),
165 memory_used: (proto.memory_used as i64).into(),
166 disk_used: (proto.disk_used as i64).into(),
167 exit_announced_at: proto.exit_announced_at as i64,
168 }
169 }
170}