From 6aa50830d235cfe69168477b84b00b6615162a3e Mon Sep 17 00:00:00 2001
From: Tino Rusch <tino.rusch@gmail.com>
Date: Wed, 18 Dec 2024 16:45:09 +0100
Subject: [PATCH] add comments + make byte values from worker and tasks default
 to a factor of one megabyte when no unit is given

---
 src/models/metadata.rs              |  48 +++++-
 src/models/pin.rs                   | 117 +++++++++++++-
 src/models/serialization_helpers.rs | 238 +++++++++++++++++++++++++++-
 src/models/task.rs                  | 122 +++++++++++++-
 src/models/worker.rs                | 102 ++++++++++--
 src/models/workflow.rs              |  65 ++++++++
 6 files changed, 668 insertions(+), 24 deletions(-)

diff --git a/src/models/metadata.rs b/src/models/metadata.rs
index f31d9ec..75df7b3 100644
--- a/src/models/metadata.rs
+++ b/src/models/metadata.rs
@@ -1,25 +1,70 @@
 use serde::{Deserialize, Serialize};
 use crate::proto::gevulot::gevulot;
 
+/// Metadata represents common metadata fields used across different resource types.
+///
+/// # Examples
+///
+/// ```
+/// use crate::models::Metadata;
+/// use crate::models::Label;
+///
+/// let metadata = Metadata {
+///     id: Some("task-123".to_string()),
+///     name: "my-task".to_string(), 
+///     creator: Some("alice".to_string()),
+///     description: "An example task".to_string(),
+///     tags: vec!["tag1".to_string(), "tag2".to_string()],
+///     labels: vec![
+///         Label {
+///             key: "env".to_string(),
+///             value: "prod".to_string()
+///         }
+///     ],
+///     workflow_ref: None
+/// };
+/// ```
 #[derive(Serialize, Deserialize, Debug, Default)]
 pub struct Metadata {
+    /// Unique identifier for the resource
     pub id: Option<String>,
+    /// Name of the resource
     pub name: String,
+    /// Creator/owner of the resource
     pub creator: Option<String>,
+    /// Detailed description of the resource
     pub description: String,
+    /// List of searchable tags
     pub tags: Vec<String>,
+    /// List of key-value labels
     pub labels: Vec<Label>,
+    /// Reference to a parent workflow (only used in TaskMetadata)
     #[serde(rename = "workflowRef")]
-    pub workflow_ref: Option<String>, // Only used in TaskMetadata
+    pub workflow_ref: Option<String>,
 }
 
+/// Label represents a key-value pair used for resource classification and filtering.
+///
+/// # Examples
+///
+/// ```
+/// use crate::models::Label;
+///
+/// let label = Label {
+///     key: "environment".to_string(),
+///     value: "production".to_string()
+/// };
+/// ```
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Label {
+    /// The label key
     pub key: String,
+    /// The label value
     pub value: String,
 }
 
 impl From<gevulot::Label> for Label {
+    /// Converts a protobuf Label into our domain Label
     fn from(proto: gevulot::Label) -> Self {
         Label {
             key: proto.key,
@@ -29,6 +74,7 @@ impl From<gevulot::Label> for Label {
 }
 
 impl From<Label> for gevulot::Label {
+    /// Converts our domain Label into a protobuf Label
     fn from(val: Label) -> Self {
         gevulot::Label {
             key: val.key,
diff --git a/src/models/pin.rs b/src/models/pin.rs
index b0fa2cf..c59b5e9 100644
--- a/src/models/pin.rs
+++ b/src/models/pin.rs
@@ -1,3 +1,9 @@
+//! Pin module provides functionality for managing pinned data in the system.
+//! 
+//! A Pin represents data that should be stored and maintained by workers in the network.
+//! It includes specifications for storage duration, size, redundancy and can reference
+//! data either by CID or fallback URLs.
+
 use super::{
     metadata::{Label, Metadata},
     serialization_helpers::{ByteUnit, DefaultFactorOne, TimeUnit},
@@ -5,6 +11,59 @@ use super::{
 use crate::proto::gevulot::gevulot;
 use serde::{Deserialize, Serialize};
 
+/// Represents a Pin resource for storing data in the network
+///
+/// A Pin defines what data should be stored, for how long, and with what redundancy level.
+/// The data can be referenced either by CID or fallback URLs.
+///
+/// # Examples
+///
+/// Creating a Pin with CID:
+/// ```
+/// use crate::models::{Pin, PinSpec, Metadata};
+/// 
+/// let pin = Pin {
+///     kind: "Pin".to_string(),
+///     version: "v0".to_string(),
+///     metadata: Metadata {
+///         name: "my-data".to_string(),
+///         ..Default::default()
+///     },
+///     spec: PinSpec {
+///         cid: Some("QmExample123".to_string()),
+///         bytes: "1GB".parse().unwrap(),
+///         time: "24h".parse().unwrap(),
+///         redundancy: 3,
+///         fallback_urls: None,
+///     },
+///     status: None,
+/// };
+/// ```
+///
+/// Creating a Pin with fallback URLs:
+/// ```
+/// use crate::models::{Pin, PinSpec, Metadata};
+///
+/// let pin = Pin {
+///     kind: "Pin".to_string(),
+///     version: "v0".to_string(), 
+///     metadata: Metadata {
+///         name: "my-backup".to_string(),
+///         ..Default::default()
+///     },
+///     spec: PinSpec {
+///         cid: None,
+///         bytes: "500MB".parse().unwrap(),
+///         time: "7d".parse().unwrap(),
+///         redundancy: 2,
+///         fallback_urls: Some(vec![
+///             "https://example.com/backup1".to_string(),
+///             "https://backup.example.com/data".to_string()
+///         ]),
+///     },
+///     status: None,
+/// };
+/// ```
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Pin {
     pub kind: String,
@@ -44,7 +103,7 @@ impl From<gevulot::Pin> for Pin {
                     .as_ref()
                     .map(|m| m.tags.clone())
                     .unwrap_or_default(),
-                labels: proto
+                labels: proto 
                     .metadata
                     .as_ref()
                     .map(|m| m.labels.clone())
@@ -63,6 +122,24 @@ impl From<gevulot::Pin> for Pin {
     }
 }
 
+/// Specification for a Pin resource
+///
+/// Defines the key parameters for pinning data including size, duration and redundancy.
+/// Either a CID or fallback URLs must be specified.
+///
+/// # Examples
+///
+/// ```
+/// use crate::models::PinSpec;
+///
+/// let spec = PinSpec {
+///     cid: Some("QmExample123".to_string()),
+///     bytes: "1GB".parse().unwrap(),
+///     time: "24h".parse().unwrap(),
+///     redundancy: 3,
+///     fallback_urls: None,
+/// };
+/// ```
 #[derive(Serialize, Debug)]
 pub struct PinSpec {
     #[serde(default)]
@@ -136,6 +213,28 @@ impl From<gevulot::PinSpec> for PinSpec {
     }
 }
 
+/// Status information for a Pin
+///
+/// Tracks which workers are assigned to store the data and their acknowledgments.
+///
+/// # Examples
+///
+/// ```
+/// use crate::models::{PinStatus, PinAck};
+///
+/// let status = PinStatus {
+///     assigned_workers: vec!["worker1".to_string(), "worker2".to_string()],
+///     worker_acks: vec![
+///         PinAck {
+///             worker: "worker1".to_string(),
+///             block_height: 1000,
+///             success: true,
+///             error: None,
+///         }
+///     ],
+///     cid: Some("QmExample123".to_string()),
+/// };
+/// ```
 #[derive(Serialize, Deserialize, Debug)]
 pub struct PinStatus {
     #[serde(rename = "assignedWorkers", default)]
@@ -168,6 +267,22 @@ impl From<gevulot::PinStatus> for PinStatus {
     }
 }
 
+/// Acknowledgment from a worker about pinning data
+///
+/// Contains information about whether the pinning was successful and any errors encountered.
+///
+/// # Examples
+///
+/// ```
+/// use crate::models::PinAck;
+///
+/// let ack = PinAck {
+///     worker: "worker1".to_string(),
+///     block_height: 1000,
+///     success: true,
+///     error: None,
+/// };
+/// ```
 #[derive(Serialize, Deserialize, Debug)]
 pub struct PinAck {
     pub worker: String,
diff --git a/src/models/serialization_helpers.rs b/src/models/serialization_helpers.rs
index f2b710f..8b0bd83 100644
--- a/src/models/serialization_helpers.rs
+++ b/src/models/serialization_helpers.rs
@@ -1,30 +1,68 @@
+//! Serialization helpers for handling byte sizes, CPU/GPU cores, and time durations.
+//! 
+//! This module provides types and traits to handle parsing and serialization of:
+//! - Byte sizes (e.g. "500MB", "1.5GB")
+//! - CPU/GPU core counts (e.g. "2 cores", "500mcpu") 
+//! - Time durations (e.g. "24h", "7d")
+
 use std::str::FromStr;
 
 use bytesize::ByteSize;
 use serde::{Deserialize, Serialize};
 
+/// Trait for specifying default multiplication factors for byte units
 pub trait DefaultFactor {
     const FACTOR: i64;
 }
 
+/// Default factor of 1 (no multiplication)
 #[derive(Debug)]
 pub struct DefaultFactorOne;
 impl DefaultFactor for DefaultFactorOne {
     const FACTOR: i64 = 1;
 }
 
+/// Default factor of 1KB (1024 bytes)
+#[derive(Debug)]
+pub struct DefaultFactorOneKilobyte;
+impl DefaultFactor for DefaultFactorOneKilobyte {
+    const FACTOR: i64 = 1024;
+}
+
+/// Default factor of 1MB (1024 * 1024 bytes)
 #[derive(Debug)]
 pub struct DefaultFactorOneMegabyte;
 impl DefaultFactor for DefaultFactorOneMegabyte {
     const FACTOR: i64 = 1024 * 1024;
 }
 
+/// Default factor of 1GB (1024 * 1024 * 1024 bytes)
 #[derive(Debug)]
 pub struct DefaultFactorOneGigabyte;
 impl DefaultFactor for DefaultFactorOneGigabyte {
     const FACTOR: i64 = 1024 * 1024 * 1024;
 }
 
+/// Type for handling byte sizes with configurable default factors
+/// 
+/// Supports three formats:
+/// - Raw numbers (e.g. 1024)
+/// - String representations (e.g. "500MB", "1.5GB")
+/// - Default factor based on type parameter
+///
+/// # Examples
+///
+/// ```rust
+/// use crate::models::serialization_helpers::{ByteUnit, DefaultFactorOne};
+/// 
+/// // Parse from string
+/// let bytes: ByteUnit = "500MB".parse().unwrap();
+/// assert_eq!(bytes.bytes().unwrap(), 500 * 1024 * 1024);
+///
+/// // Use raw number
+/// let bytes = ByteUnit::<DefaultFactorOneMegabyte>::from(1);
+/// assert_eq!(bytes.bytes().unwrap(), 1 * 1024 * 1024);
+/// ```
 #[derive(Debug, Serialize, Deserialize, Eq)]
 #[serde(untagged)]
 pub enum ByteUnit<D: DefaultFactor = DefaultFactorOne>{
@@ -35,10 +73,16 @@ pub enum ByteUnit<D: DefaultFactor = DefaultFactorOne>{
 }
 
 impl<D: DefaultFactor> ByteUnit<D> {
+    /// Convert to number of bytes
     pub fn bytes(&self) -> Result<i64, String> {
         match self {
-            ByteUnit::Number(n) => Ok(*n),
-            ByteUnit::String(s) => s.parse::<ByteSize>().map(|b| b.0 as i64),
+            ByteUnit::Number(n) => Ok(*n * D::FACTOR),
+            ByteUnit::String(s) => {
+                if s.chars().all(|c| c.is_ascii_digit()) {
+                    return Ok(s.parse::<i64>().map_err(|e| e.to_string())? * D::FACTOR);
+                }
+                s.parse::<ByteSize>().map(|b| b.0 as i64)
+            },
             ByteUnit::Factor(_) => Ok(D::FACTOR),
         }
     }
@@ -65,6 +109,27 @@ impl<D: DefaultFactor> From<i64> for ByteUnit<D> {
     }
 }
 
+/// Type for handling CPU/GPU core counts
+///
+/// Supports formats like:
+/// - Raw numbers (interpreted as cores, e.g. 2 = 2000 millicores)
+/// - String representations with units:
+///   - Cores: "2 cores", "2 cpus", "2 gpus"
+///   - Millicores: "500mcpu", "500mgpu", "500mcore"
+///
+/// # Examples
+///
+/// ```rust
+/// use crate::models::serialization_helpers::CoreUnit;
+///
+/// // Parse core counts
+/// let cores: CoreUnit = "2 cores".parse().unwrap();
+/// assert_eq!(cores.millicores().unwrap(), 2000);
+///
+/// // Parse millicores
+/// let cores: CoreUnit = "500mcpu".parse().unwrap();
+/// assert_eq!(cores.millicores().unwrap(), 500);
+/// ```
 #[derive(Debug, Serialize, Deserialize, Eq)]
 #[serde(untagged)]
 pub enum CoreUnit {
@@ -73,13 +138,18 @@ pub enum CoreUnit {
 }
 
 impl CoreUnit {
+    /// Convert to millicores (1 core = 1000 millicores)
     pub fn millicores(&self) -> Result<i64, String> {
         match self {
-            CoreUnit::Number(n) => Ok(*n * 1000), // Attention! default factor without unit is 1000, i.e. cpus: 1 results in 1000mcpu
+            CoreUnit::Number(n) => Ok(*n * 1000), // Default factor without unit is 1000
             CoreUnit::String(s) => {
+                // Extract numeric part
                 let numeric: String = s.chars().take_while(|c| c.is_digit(10)).collect();
+                // Extract and normalize unit part
                 let unit = s[numeric.len()..].to_lowercase().replace(" ", "");
                 let base: i64 = numeric.parse().map_err(|e| format!("Invalid number: {}", e))?;
+                
+                // Convert based on unit, using 1000 millicores = 1 core
                 Ok(base
                     * match unit.as_str() {
                         "cpu" | "cpus" => 1000, 
@@ -88,7 +158,7 @@ impl CoreUnit {
                         "mcpu" | "mcpus" | "millicpu" | "millicpus" => 1,
                         "mgpu" | "mgpus" | "milligpu" | "milligpus" => 1,
                         "mcore" | "mcores" | "millicore" | "millicores" => 1,
-                        "" => 1000, // Attention! default factor is 1000, i.e. cpus: 1 results in 1000mcpu
+                        "" => 1000, // Default to cores if no unit specified
                         _ => return Err(format!("Invalid unit: {}", unit)),
                     })
             }
@@ -118,6 +188,25 @@ impl From<i64> for CoreUnit {
     }
 }
 
+/// Type for handling time durations
+///
+/// Supports:
+/// - Raw numbers (interpreted as seconds)
+/// - Human readable durations (e.g. "24h", "7d", "1y")
+///
+/// # Examples
+///
+/// ```rust
+/// use crate::models::serialization_helpers::TimeUnit;
+///
+/// // Parse durations
+/// let time: TimeUnit = "24h".parse().unwrap();
+/// assert_eq!(time.seconds().unwrap(), 24 * 60 * 60);
+///
+/// // Use raw seconds
+/// let time = TimeUnit::from(3600);
+/// assert_eq!(time.seconds().unwrap(), 3600);
+/// ```
 #[derive(Debug, Serialize, Deserialize, Eq)]
 #[serde(untagged)]
 pub enum TimeUnit {
@@ -126,6 +215,7 @@ pub enum TimeUnit {
 }
 
 impl TimeUnit {
+    /// Convert to seconds
     pub fn seconds(&self) -> Result<i64, String> {
         match self {
             TimeUnit::Number(n) => Ok(*n),
@@ -157,4 +247,144 @@ impl From<i64> for TimeUnit {
     fn from(n: i64) -> Self {
         TimeUnit::Number(n)
     }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json;
+
+    #[test]
+    fn test_byte_unit() {
+        let bytes: ByteUnit = "500MiB".parse().unwrap();
+        assert_eq!(bytes.bytes().unwrap(), 500 * 1024 * 1024);
+    }
+
+    #[test]
+    fn test_byte_unit_default_factor() {
+        let bytes: ByteUnit = 1.into();
+        assert_eq!(bytes.bytes(), Ok(1), "Default factor should be 1 (no multiplication)");
+        let bytes: ByteUnit<DefaultFactorOneKilobyte> = 1.into();
+        assert_eq!(bytes.bytes(), Ok(1 * 1024), "1 with KB factor should be 1024 bytes");
+        let bytes: ByteUnit<DefaultFactorOneMegabyte> = 1.into();
+        assert_eq!(bytes.bytes(), Ok(1 * 1024 * 1024), "1 with MB factor should be 1024*1024 bytes");
+        let bytes: ByteUnit<DefaultFactorOneGigabyte> = 1.into();
+        assert_eq!(bytes.bytes(), Ok(1 * 1024 * 1024 * 1024), "1 with GB factor should be 1024*1024*1024 bytes");
+
+        let bytes: ByteUnit = "1".parse().unwrap();
+        assert_eq!(bytes.bytes(), Ok(1), "String '1' with default factor should be 1 byte");
+        let bytes: ByteUnit<DefaultFactorOneKilobyte> = "1".parse().unwrap();
+        assert_eq!(bytes.bytes(), Ok(1 * 1024), "String '1' with KB factor should be 1024 bytes");
+        let bytes: ByteUnit<DefaultFactorOneMegabyte> = "1".parse().unwrap();
+        assert_eq!(bytes.bytes(), Ok(1 * 1024 * 1024), "String '1' with MB factor should be 1024*1024 bytes");
+        let bytes: ByteUnit<DefaultFactorOneGigabyte> = "1".parse().unwrap();
+        assert_eq!(bytes.bytes(), Ok(1 * 1024 * 1024 * 1024), "String '1' with GB factor should be 1024*1024*1024 bytes");
+
+        let bytes: ByteUnit<DefaultFactorOneKilobyte> = "1MiB".parse().unwrap();
+        assert_eq!(bytes.bytes(), Ok(1 * 1024 * 1024), "Explicit unit (1MiB) should override default factor (KB)");
+    }
+
+    #[test]
+    fn test_time_unit() {
+        let time: TimeUnit = "24h".parse().unwrap();
+        assert_eq!(time.seconds().unwrap(), 24 * 60 * 60);
+
+        let time: TimeUnit = "7d".parse().unwrap();
+        assert_eq!(time.seconds().unwrap(), 7 * 24 * 60 * 60);
+
+        let time: TimeUnit = "30m".parse().unwrap();
+        assert_eq!(time.seconds().unwrap(), 30 * 60);
+
+        let time: TimeUnit = "90s".parse().unwrap();
+        assert_eq!(time.seconds().unwrap(), 90);
+
+        let time: TimeUnit = 3600.into();
+        assert_eq!(time.seconds().unwrap(), 3600);
+    }
+
+    #[test]
+    fn test_core_unit() {
+        let cores: CoreUnit = "2 cores".parse().unwrap();
+        assert_eq!(cores.millicores().unwrap(), 2000);
+
+        let cores: CoreUnit = "500mcpu".parse().unwrap();
+        assert_eq!(cores.millicores().unwrap(), 500);
+
+        let cores: CoreUnit = "1.5 cpus".parse().unwrap();
+        assert_eq!(cores.millicores().unwrap(), 1500);
+
+        let cores: CoreUnit = "2 gpus".parse().unwrap();
+        assert_eq!(cores.millicores().unwrap(), 2000);
+
+        let cores: CoreUnit = "750mgpu".parse().unwrap();
+        assert_eq!(cores.millicores().unwrap(), 750);
+
+        let cores: CoreUnit = 2.into();
+        assert_eq!(cores.millicores().unwrap(), 2000);
+    }
+
+    #[test]
+    fn test_invalid_formats() {
+        assert!("invalid".parse::<ByteUnit>().is_err());
+        assert!("invalid".parse::<TimeUnit>().is_err());
+        assert!("invalid".parse::<CoreUnit>().is_err());
+
+        assert!("-1GB".parse::<ByteUnit>().is_err());
+        assert!("-24h".parse::<TimeUnit>().is_err());
+        assert!("-2cores".parse::<CoreUnit>().is_err());
+    }
+
+    #[test]
+    fn test_json_serialization() {
+        // Test ByteUnit serialization
+        let bytes: ByteUnit = "500MiB".parse().unwrap();
+        let json = serde_json::to_string(&bytes).unwrap();
+        assert_eq!(json, "\"500MiB\"");
+
+        let bytes: ByteUnit = 1024.into();
+        let json = serde_json::to_string(&bytes).unwrap();
+        assert_eq!(json, "1024");
+
+        // Test TimeUnit serialization
+        let time: TimeUnit = "24h".parse().unwrap();
+        let json = serde_json::to_string(&time).unwrap();
+        assert_eq!(json, "\"24h\"");
+
+        let time: TimeUnit = 3600.into();
+        let json = serde_json::to_string(&time).unwrap();
+        assert_eq!(json, "3600");
+
+        // Test CoreUnit serialization
+        let cores: CoreUnit = "2 cores".parse().unwrap();
+        let json = serde_json::to_string(&cores).unwrap();
+        assert_eq!(json, "\"2 cores\"");
+
+        let cores: CoreUnit = 2.into();
+        let json = serde_json::to_string(&cores).unwrap();
+        assert_eq!(json, "2");
+    }
+
+    #[test]
+    fn test_json_deserialization() {
+        // Test ByteUnit deserialization
+        let bytes: ByteUnit = serde_json::from_str("\"500MiB\"").unwrap();
+        assert_eq!(bytes.bytes().unwrap(), 500 * 1024 * 1024);
+
+        let bytes: ByteUnit = serde_json::from_str("1024").unwrap();
+        assert_eq!(bytes.bytes().unwrap(), 1024);
+
+        // Test TimeUnit deserialization
+        let time: TimeUnit = serde_json::from_str("\"24h\"").unwrap();
+        assert_eq!(time.seconds().unwrap(), 24 * 60 * 60);
+
+        let time: TimeUnit = serde_json::from_str("3600").unwrap();
+        assert_eq!(time.seconds().unwrap(), 3600);
+
+        // Test CoreUnit deserialization
+        let cores: CoreUnit = serde_json::from_str("\"2 cores\"").unwrap();
+        assert_eq!(cores.millicores().unwrap(), 2000);
+
+        let cores: CoreUnit = serde_json::from_str("2").unwrap();
+        assert_eq!(cores.millicores().unwrap(), 2000);
+    }
 }
\ No newline at end of file
diff --git a/src/models/task.rs b/src/models/task.rs
index 5c30d43..a000919 100644
--- a/src/models/task.rs
+++ b/src/models/task.rs
@@ -1,22 +1,89 @@
+//! Task model and related types for managing task execution.
+//! 
+//! This module provides the core task model used throughout the system, including:
+//! - Task specification and status
+//! - Resource requirements (CPU, GPU, memory, time)
+//! - Input/output context handling
+//! - Environment variables
+//! - Metadata like tags and labels
+
 use crate::proto::gevulot::gevulot;
 use serde::{Deserialize, Serialize};
 
+use super::serialization_helpers::DefaultFactorOneMegabyte;
+
+/// Represents a complete task definition with metadata, specification and status
+///
+/// # Examples
+///
+/// Creating a basic task:
+/// ```
+/// use crate::models::Task;
+/// 
+/// let task = serde_json::from_str::<Task>(r#"{
+///     "kind": "Task",
+///     "version": "v0",
+///     "spec": {
+///         "image": "ubuntu:latest",
+///         "command": ["echo", "hello"],
+///         "resources": {
+///             "cpus": "1cpu",
+///             "gpus": "0gpu", 
+///             "memory": "512mb",
+///             "time": "1h"
+///         }
+///     }
+/// }"#).unwrap();
+/// ```
+///
+/// Task with input/output contexts:
+/// ```
+/// let task = serde_json::from_str::<Task>(r#"{
+///     "kind": "Task", 
+///     "version": "v0",
+///     "spec": {
+///         "image": "processor:v1",
+///         "inputContexts": [{
+///             "source": "input-data",
+///             "target": "/data"
+///         }],
+///         "outputContexts": [{
+///             "source": "/results",
+///             "retentionPeriod": 86400
+///         }],
+///         "resources": {
+///             "cpus": "2cpu",
+///             "gpus": "1gpu",
+///             "memory": "4gb",
+///             "time": "2h"
+///         }
+///     }
+/// }"#).unwrap();
+/// ```
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Task {
+    // The kind is always "Task" - used for type identification in serialized form
     pub kind: String,
+    // API version, currently "v0"
     pub version: String,
+    // Task metadata like name, description, tags etc
     #[serde(default)]
     pub metadata: crate::models::Metadata,
+    // Core task specification
     pub spec: TaskSpec,
+    // Optional runtime status
     pub status: Option<TaskStatus>,
 }
 
+// Conversion from protobuf Task message
 impl From<gevulot::Task> for Task {
     fn from(proto: gevulot::Task) -> Self {
+        // Extract workflow reference if present in spec
         let workflow_ref = match proto.spec.as_ref() {
             Some(spec) if !spec.workflow_ref.is_empty() => Some(spec.workflow_ref.clone()),
             _ => None,
         };
+        
         Task {
             kind: "Task".to_string(),
             version: "v0".to_string(),
@@ -57,26 +124,54 @@ impl From<gevulot::Task> for Task {
     }
 }
 
+/// Task specification containing all execution parameters
+///
+/// # Examples
+///
+/// Basic spec with just image and resources:
+/// ```
+/// use crate::models::TaskSpec;
+///
+/// let spec = serde_json::from_str::<TaskSpec>(r#"{
+///     "image": "ubuntu:latest",
+///     "resources": {
+///         "cpus": "1cpu",
+///         "gpus": "0gpu",
+///         "memory": "512mb",
+///         "time": "1h"
+///     }
+/// }"#).unwrap();
+/// ```
 #[derive(Serialize, Deserialize, Debug)]
 pub struct TaskSpec {
+    // Container image to run
     pub image: String,
+    // Optional command to override image entrypoint
     #[serde(default)]
     pub command: Vec<String>,
+    // Optional arguments for the command
     #[serde(default)]
     pub args: Vec<String>,
+    // Environment variables to set in container
     #[serde(default)]
     pub env: Vec<TaskEnv>,
+    // Input data contexts to mount
     #[serde(rename = "inputContexts", default)]
     pub input_contexts: Vec<InputContext>,
+    // Output data contexts to capture
     #[serde(rename = "outputContexts", default)]
     pub output_contexts: Vec<OutputContext>,
+    // Resource requirements
     pub resources: TaskResources,
+    // Whether to store stdout stream
     #[serde(rename = "storeStdout", default)]
     pub store_stdout: bool,
+    // Whether to store stderr stream
     #[serde(rename = "storeStderr", default)]
     pub store_stderr: bool,
 }
 
+// Conversion from protobuf TaskSpec message
 impl From<gevulot::TaskSpec> for TaskSpec {
     fn from(proto: gevulot::TaskSpec) -> Self {
         TaskSpec {
@@ -119,57 +214,79 @@ impl From<gevulot::TaskSpec> for TaskSpec {
     }
 }
 
+/// Environment variable definition for task container
 #[derive(Serialize, Deserialize, Debug)]
 pub struct TaskEnv {
     pub name: String,
     pub value: String,
 }
 
+/// Input context for mounting data into task container
 #[derive(Serialize, Deserialize, Debug)]
 pub struct InputContext {
+    // Source data identifier
     pub source: String,
+    // Target mount path in container
     pub target: String,
 }
 
+/// Output context for capturing data from task container
 #[derive(Serialize, Deserialize, Debug)]
 pub struct OutputContext {
+    // Source path in container to capture
     pub source: String,
+    // How long to retain the output data
     #[serde(rename = "retentionPeriod")]
     pub retention_period: i64,
 }
 
+/// Resource requirements for task execution
 #[derive(Serialize, Deserialize, Debug)]
 pub struct TaskResources {
+    // CPU cores required (supports units like "2cpu", "500mcpu")
     pub cpus: crate::models::CoreUnit,
+    // GPU cores required (supports units like "1gpu", "500mgpu") 
     pub gpus: crate::models::CoreUnit,
-    pub memory: crate::models::ByteUnit,
+    // Memory required (supports units like "1gb", "512mb")
+    pub memory: crate::models::ByteUnit<DefaultFactorOneMegabyte>, // when no unit is specified, we assume MiB
+    // Time limit (supports units like "1h", "30m")
     pub time: crate::models::TimeUnit,
 }
 
+/// Runtime status of a task
 #[derive(Serialize, Deserialize, Debug)]
 pub struct TaskStatus {
+    // Current state (Pending, Running, Done, Failed etc)
     pub state: String,
+    // Timestamps for task lifecycle
     #[serde(rename = "createdAt")]
     pub created_at: i64,
     #[serde(rename = "startedAt")]
     pub started_at: i64,
     #[serde(rename = "completedAt")]
     pub completed_at: i64,
+    // Workers assigned/active for this task
     #[serde(rename = "assignedWorkers")]
     pub assigned_workers: Vec<String>,
     #[serde(rename = "activeWorker")]
     pub active_worker: String,
+    // Exit code if task completed
     #[serde(rename = "exitCode")]
     pub exit_code: Option<i64>,
+    // Output context identifiers
     #[serde(rename = "outputContexts")]
     pub output_contexts: Vec<String>,
+    // Captured output streams if enabled
     pub stdout: Option<String>,
     pub stderr: Option<String>,
+    // Error message if failed
     pub error: Option<String>,
 }
 
+// Conversion from protobuf TaskStatus message
 impl From<gevulot::TaskStatus> for TaskStatus {
     fn from(proto: gevulot::TaskStatus) -> Self {
+        // Map numeric state to string representation
         let mut exit_code = None;
         let state = match proto.state {
             0 => "Pending".to_string(),
@@ -185,6 +302,8 @@ impl From<gevulot::TaskStatus> for TaskStatus {
             }
             _ => "Unknown".to_string(),
         };
+
+        // Convert empty strings to None
         let error = if proto.error.is_empty() {
             None
         } else {
@@ -200,6 +319,7 @@ impl From<gevulot::TaskStatus> for TaskStatus {
         } else {
             Some(proto.stderr)
         };
+
         TaskStatus {
             state,
             created_at: proto.created_at as i64,
diff --git a/src/models/worker.rs b/src/models/worker.rs
index 22115fe..cbd7334 100644
--- a/src/models/worker.rs
+++ b/src/models/worker.rs
@@ -1,7 +1,60 @@
+//! Worker model and related types for managing worker nodes.
+//! 
+//! This module provides the core worker model used throughout the system, including:
+//! - Worker specification and status
+//! - Resource tracking (CPU, GPU, memory, disk)
+//! - Metadata like tags and labels
+//! - Protobuf serialization/deserialization
+
 use serde::{Deserialize, Serialize};
 use crate::proto::gevulot::gevulot;
-use super::metadata::{Label, Metadata};
+use super::{metadata::{Label, Metadata}, ByteUnit, CoreUnit, DefaultFactorOneMegabyte};
 
+/// Represents a complete worker definition with metadata, specification and status
+///
+/// # Examples
+///
+/// Creating a basic worker:
+/// ```
+/// use crate::models::Worker;
+/// CoreUnit
+/// let worker = serde_json::from_str::<Worker>(r#"{
+///     "kind": "Worker",
+///     "version": "v0",
+///     "metadata": {
+///         "name": "worker-1",
+///         "tags": ["compute"]
+///     },
+///     "spec": {
+///         "cpus": "8 cores",
+///         "gpus": "1 gpu",
+///         "memory": "16 GiB",
+///         "disk": "100 GiB"
+///     }
+/// }"#).unwrap();
+/// ```
+///
+/// Converting from protobuf:
+/// ```
+/// use crate::proto::gevulot::gevulot;
+/// use crate::models::Worker;
+///
+/// let proto_worker = gevulot::Worker {
+///     metadata: Some(gevulot::Metadata {
+///         name: "worker-1".to_string(),
+///         ..Default::default()
+///     }),
+///     spec: Some(gevulot::WorkerSpec {
+///         cpus: 8,
+///         gpus: 1,
+///         memory: 16000000000,
+///         disk: 100000000000,
+///     }),
+///     ..Default::default()
+/// };
+///
+/// let worker = Worker::from(proto_worker);
+/// ```
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Worker {
     pub kind: String,
@@ -13,6 +66,7 @@ pub struct Worker {
 
 impl From<gevulot::Worker> for Worker {
     fn from(proto: gevulot::Worker) -> Self {
+        // Convert protobuf worker to our internal worker model
         Worker {
             kind: "Worker".to_string(),
             version: "v0".to_string(),
@@ -53,46 +107,60 @@ impl From<gevulot::Worker> for Worker {
     }
 }
 
+/// Specification of worker resources and capabilities
+///
+/// Contains the maximum resources available on this worker:
+/// - CPU cores
+/// - GPU devices  
+/// - Memory in bytes
+/// - Disk space in bytes
 #[derive(Serialize, Deserialize, Debug)]
 pub struct WorkerSpec {
-    pub cpus: i64,
-    pub gpus: i64,
-    pub memory: i64,
-    pub disk: i64,
+    pub cpus: CoreUnit,
+    pub gpus: CoreUnit,
+    pub memory: ByteUnit<DefaultFactorOneMegabyte>,
+    pub disk: ByteUnit<DefaultFactorOneMegabyte>,
 }
 
 impl From<gevulot::WorkerSpec> for WorkerSpec {
     fn from(proto: gevulot::WorkerSpec) -> Self {
+        // Convert protobuf spec to internal spec
         WorkerSpec {
-            cpus: proto.cpus as i64,
-            gpus: proto.gpus as i64,
-            memory: proto.memory as i64,
-            disk: proto.disk as i64,
+            cpus: (proto.cpus as i64).into(),
+            gpus: (proto.gpus as i64).into(),
+            memory: (proto.memory as i64).into(),
+            disk: (proto.disk as i64).into(),
         }
     }
 }
 
+/// Current status and resource utilization of a worker
+///
+/// Tracks:
+/// - Currently used resources (CPU, GPU, memory, disk)
+/// - When the worker announced it will exit
 #[derive(Serialize, Deserialize, Debug)]
 pub struct WorkerStatus {
     #[serde(rename = "cpusUsed")]
-    pub cpus_used: i64,
+    pub cpus_used: CoreUnit,
     #[serde(rename = "gpusUsed")]
-    pub gpus_used: i64,
+    pub gpus_used: CoreUnit,
     #[serde(rename = "memoryUsed")]
-    pub memory_used: i64,
+    pub memory_used: ByteUnit<DefaultFactorOneMegabyte>,
     #[serde(rename = "diskUsed")]
-    pub disk_used: i64,
+    pub disk_used: ByteUnit<DefaultFactorOneMegabyte>,
     #[serde(rename = "exitAnnouncedAt")]
     pub exit_announced_at: i64,
 }
 
 impl From<gevulot::WorkerStatus> for WorkerStatus {
     fn from(proto: gevulot::WorkerStatus) -> Self {
+        // Convert protobuf status to internal status
         WorkerStatus {
-            cpus_used: proto.cpus_used as i64,
-            gpus_used: proto.gpus_used as i64,
-            memory_used: proto.memory_used as i64,
-            disk_used: proto.disk_used as i64,
+            cpus_used: (proto.cpus_used as i64).into(),
+            gpus_used: (proto.gpus_used as i64).into(),
+            memory_used: (proto.memory_used as i64).into(),
+            disk_used: (proto.disk_used as i64).into(),
             exit_announced_at: proto.exit_announced_at as i64,
         }
     }
diff --git a/src/models/workflow.rs b/src/models/workflow.rs
index 0ddd2cb..75ad4eb 100644
--- a/src/models/workflow.rs
+++ b/src/models/workflow.rs
@@ -1,7 +1,46 @@
+//! Workflow model and related types for managing workflow execution.
+//! 
+//! This module provides the core workflow model used throughout the system, including:
+//! - Workflow specification with stages and tasks
+//! - Status tracking for workflow execution
+//! - Metadata like tags and labels
+//! - Protobuf serialization/deserialization
+
 use crate::proto::gevulot::gevulot;
 use serde::{Deserialize, Serialize};
 use super::{Label, Metadata, TaskSpec};
 
+/// Represents a complete workflow definition with metadata, specification and status
+///
+/// A workflow consists of one or more stages that are executed sequentially. Each stage
+/// contains one or more tasks that can be executed in parallel.
+///
+/// # Examples
+///
+/// Creating a basic workflow:
+/// ```
+/// use crate::models::Workflow;
+///
+/// let workflow = serde_json::from_str::<Workflow>(r#"{
+///     "kind": "Workflow",
+///     "version": "v0", 
+///     "metadata": {
+///         "name": "my-workflow",
+///         "tags": ["compute"]
+///     },
+///     "spec": {
+///         "stages": [{
+///             "tasks": [{
+///                 "image": "alpine",
+///                 "resources": {
+///                     "cpus": "1cpu",
+///                     "memory": "1GiB"
+///                 }
+///             }]
+///         }]
+///     }
+/// }"#).unwrap();
+/// ```
 #[derive(Serialize, Deserialize, Debug)]
 pub struct Workflow {
     pub kind: String,
@@ -12,8 +51,10 @@ pub struct Workflow {
     pub status: Option<WorkflowStatus>,
 }
 
+// Converts a protobuf workflow message into our internal Workflow model
 impl From<gevulot::Workflow> for Workflow {
     fn from(proto: gevulot::Workflow) -> Self {
+        // Create a new workflow, carefully mapping all protobuf fields to our model
         Workflow {
             kind: "Workflow".to_string(),
             version: "v0".to_string(),
@@ -54,18 +95,28 @@ impl From<gevulot::Workflow> for Workflow {
     }
 }
 
+/// Represents a single stage in a workflow containing one or more tasks
+///
+/// Tasks within a stage can be executed in parallel. The workflow will only
+/// proceed to the next stage once all tasks in the current stage are complete.
 #[derive(Serialize, Deserialize, Debug)]
 pub struct WorkflowStage {
     pub tasks: Vec<TaskSpec>,
 }
 
+/// Specification for a workflow defining its stages and tasks
+///
+/// The stages are executed sequentially, with tasks in each stage potentially
+/// running in parallel depending on available resources.
 #[derive(Serialize, Deserialize, Debug)]
 pub struct WorkflowSpec {
     pub stages: Vec<WorkflowStage>,
 }
 
+// Converts a protobuf workflow spec into our internal WorkflowSpec model
 impl From<gevulot::WorkflowSpec> for WorkflowSpec {
     fn from(proto: gevulot::WorkflowSpec) -> Self {
+        // Map each protobuf stage to our stage model, converting tasks as well
         WorkflowSpec {
             stages: proto
                 .stages
@@ -78,6 +129,9 @@ impl From<gevulot::WorkflowSpec> for WorkflowSpec {
     }
 }
 
+/// Status information for a single stage in a workflow
+///
+/// Tracks which tasks have been created and how many have completed.
 #[derive(Serialize, Deserialize, Debug)]
 pub struct WorkflowStageStatus {
     #[serde(rename = "taskIds")]
@@ -86,6 +140,12 @@ pub struct WorkflowStageStatus {
     pub finished_tasks: u64,
 }
 
+/// Current status of a workflow's execution
+///
+/// Tracks:
+/// - Overall workflow state (Pending, Running, Done, Failed)
+/// - Which stage is currently executing
+/// - Status of each stage including task completion
 #[derive(Serialize, Deserialize, Debug)]
 pub struct WorkflowStatus {
     pub state: String,
@@ -94,9 +154,11 @@ pub struct WorkflowStatus {
     pub stages: Vec<WorkflowStageStatus>,
 }
 
+// Converts a protobuf workflow status into our internal WorkflowStatus model
 impl From<gevulot::WorkflowStatus> for WorkflowStatus {
     fn from(proto: gevulot::WorkflowStatus) -> Self {
         WorkflowStatus {
+            // Map numeric states to human readable strings
             state: match proto.state {
                 0 => "Pending".to_string(),
                 1 => "Running".to_string(), 
@@ -117,6 +179,7 @@ impl From<gevulot::WorkflowStatus> for WorkflowStatus {
     }
 }
 
+// Unit tests to verify workflow serialization/deserialization and field mapping
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -124,6 +187,7 @@ mod tests {
 
     #[test]
     fn test_parse_workflow() {
+        // Test parsing a complete workflow JSON with all fields populated
         let workflow = serde_json::from_value::<Workflow>(json!({
             "kind": "Workflow",
             "version": "v0",
@@ -204,6 +268,7 @@ mod tests {
 
     #[test]
     fn test_parse_workflow_with_minimum() {
+        // Test parsing a minimal workflow JSON with only required fields
         let workflow = serde_json::from_value::<Workflow>(json!({
             "kind": "Workflow",
             "version": "v0",