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}