558 lines
18 KiB
Rust
558 lines
18 KiB
Rust
// Graph serialization support for unified data+compute graphs
|
|
// Phase 6: Self-Optimizing Platform
|
|
|
|
use indexmap::IndexMap;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
|
|
use crate::{Graph, GraphError, NodeId, Result};
|
|
|
|
/// Serializable graph format for persistence and interchange
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SerializableGraph {
|
|
/// Graph format version for compatibility
|
|
pub version: String,
|
|
/// Graph metadata
|
|
pub metadata: IndexMap<String, String>,
|
|
/// All nodes in the graph
|
|
pub nodes: Vec<SerializableNode>,
|
|
/// All edges in the graph
|
|
pub edges: Vec<SerializableEdge>,
|
|
/// Timestamp of creation/serialization
|
|
pub timestamp: chrono::DateTime<chrono::Utc>,
|
|
}
|
|
|
|
/// Serializable node representation
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SerializableNode {
|
|
/// Node identifier
|
|
pub id: NodeId,
|
|
/// Operation type and configuration
|
|
pub operation: crate::OperationType,
|
|
/// Node metadata
|
|
pub metadata: IndexMap<String, String>,
|
|
}
|
|
|
|
/// Serializable edge representation
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SerializableEdge {
|
|
/// Source node ID
|
|
pub from: NodeId,
|
|
/// Target node ID
|
|
pub to: NodeId,
|
|
/// Edge metadata
|
|
pub metadata: IndexMap<String, String>,
|
|
}
|
|
|
|
impl SerializableGraph {
|
|
/// Current format version
|
|
pub const CURRENT_VERSION: &'static str = "1.0.0";
|
|
|
|
/// Create from graph
|
|
pub fn from_graph(graph: &Graph) -> Self {
|
|
let nodes: Vec<SerializableNode> = graph
|
|
.nodes()
|
|
.map(|node| SerializableNode {
|
|
id: node.id,
|
|
operation: node.operation.clone(),
|
|
metadata: node.metadata.clone(),
|
|
})
|
|
.collect();
|
|
|
|
let mut edges = Vec::new();
|
|
for node in graph.nodes() {
|
|
if let Ok(dependencies) = graph.dependencies(node.id) {
|
|
for dep_id in dependencies {
|
|
edges.push(SerializableEdge {
|
|
from: dep_id,
|
|
to: node.id,
|
|
metadata: IndexMap::new(), // Could be extended with edge metadata
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
Self {
|
|
version: Self::CURRENT_VERSION.to_string(),
|
|
metadata: graph.metadata().clone(),
|
|
nodes,
|
|
edges,
|
|
timestamp: chrono::Utc::now(),
|
|
}
|
|
}
|
|
|
|
/// Convert to graph
|
|
pub fn to_graph(&self) -> Result<Graph> {
|
|
// Validate version compatibility
|
|
if !self.is_compatible_version() {
|
|
return Err(GraphError::serialization_error(format!(
|
|
"Incompatible graph version: {}",
|
|
self.version
|
|
)));
|
|
}
|
|
|
|
let mut builder = crate::GraphBuilder::new();
|
|
let mut node_mapping = HashMap::new();
|
|
|
|
// Add all nodes
|
|
for serializable_node in &self.nodes {
|
|
let new_node_id = match &serializable_node.operation {
|
|
crate::OperationType::Data(data_op) => builder.add_data_op(data_op.clone())?,
|
|
crate::OperationType::Compute(compute_op) => {
|
|
// Need to remap node references
|
|
let remapped_op =
|
|
self.remap_compute_op_for_deserialization(compute_op, &node_mapping)?;
|
|
builder.add_compute_op(remapped_op)?
|
|
}
|
|
};
|
|
node_mapping.insert(serializable_node.id, new_node_id);
|
|
}
|
|
|
|
// Add all edges
|
|
for edge in &self.edges {
|
|
let from_node = node_mapping.get(&edge.from).ok_or_else(|| {
|
|
GraphError::serialization_error(format!("Missing source node: {}", edge.from))
|
|
})?;
|
|
let to_node = node_mapping.get(&edge.to).ok_or_else(|| {
|
|
GraphError::serialization_error(format!("Missing target node: {}", edge.to))
|
|
})?;
|
|
|
|
builder.add_dependency(*to_node, *from_node)?;
|
|
}
|
|
|
|
let mut graph = builder.build()?;
|
|
|
|
// Copy metadata
|
|
for (key, value) in &self.metadata {
|
|
graph.add_metadata(key.clone(), value.clone());
|
|
}
|
|
|
|
Ok(graph)
|
|
}
|
|
|
|
/// Check if version is compatible with current implementation
|
|
fn is_compatible_version(&self) -> bool {
|
|
// Simple version check - in real implementation would have more sophisticated logic
|
|
let version_parts: Vec<&str> = self.version.split('.').collect();
|
|
let current_parts: Vec<&str> = Self::CURRENT_VERSION.split('.').collect();
|
|
|
|
if version_parts.len() >= 2 && current_parts.len() >= 2 {
|
|
// Major version must match, minor version can be different
|
|
version_parts[0] == current_parts[0]
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Remap compute operation node references for deserialization
|
|
fn remap_compute_op_for_deserialization(
|
|
&self,
|
|
op: &crate::ComputeOp,
|
|
node_mapping: &HashMap<NodeId, NodeId>,
|
|
) -> Result<crate::ComputeOp> {
|
|
use crate::ComputeOp;
|
|
|
|
let remap = |id: NodeId| -> Result<NodeId> {
|
|
node_mapping.get(&id).copied().ok_or_else(|| {
|
|
GraphError::serialization_error(format!("Missing node mapping for {id}"))
|
|
})
|
|
};
|
|
|
|
let remapped_op = match op {
|
|
ComputeOp::TensorLoad {
|
|
data_node,
|
|
shape,
|
|
dtype,
|
|
} => ComputeOp::TensorLoad {
|
|
data_node: remap(*data_node)?,
|
|
shape: shape.clone(),
|
|
dtype: dtype.clone(),
|
|
},
|
|
ComputeOp::Add { lhs, rhs } => ComputeOp::Add {
|
|
lhs: remap(*lhs)?,
|
|
rhs: remap(*rhs)?,
|
|
},
|
|
ComputeOp::Mul { lhs, rhs } => ComputeOp::Mul {
|
|
lhs: remap(*lhs)?,
|
|
rhs: remap(*rhs)?,
|
|
},
|
|
ComputeOp::MatMul { lhs, rhs } => ComputeOp::MatMul {
|
|
lhs: remap(*lhs)?,
|
|
rhs: remap(*rhs)?,
|
|
},
|
|
ComputeOp::Conv2d {
|
|
input,
|
|
weight,
|
|
bias,
|
|
stride,
|
|
padding,
|
|
dilation,
|
|
groups,
|
|
} => ComputeOp::Conv2d {
|
|
input: remap(*input)?,
|
|
weight: remap(*weight)?,
|
|
bias: bias.map(remap).transpose()?,
|
|
stride: stride.clone(),
|
|
padding: padding.clone(),
|
|
dilation: dilation.clone(),
|
|
groups: *groups,
|
|
},
|
|
ComputeOp::Activation {
|
|
input,
|
|
activation_type,
|
|
} => ComputeOp::Activation {
|
|
input: remap(*input)?,
|
|
activation_type: activation_type.clone(),
|
|
},
|
|
ComputeOp::Pool2d {
|
|
input,
|
|
pool_type,
|
|
kernel_size,
|
|
stride,
|
|
} => ComputeOp::Pool2d {
|
|
input: remap(*input)?,
|
|
pool_type: pool_type.clone(),
|
|
kernel_size: kernel_size.clone(),
|
|
stride: stride.clone(),
|
|
},
|
|
ComputeOp::BatchNorm {
|
|
input,
|
|
running_mean,
|
|
running_var,
|
|
weight,
|
|
bias,
|
|
training,
|
|
} => ComputeOp::BatchNorm {
|
|
input: remap(*input)?,
|
|
running_mean: remap(*running_mean)?,
|
|
running_var: remap(*running_var)?,
|
|
weight: weight.map(remap).transpose()?,
|
|
bias: bias.map(remap).transpose()?,
|
|
training: *training,
|
|
},
|
|
ComputeOp::Reshape { input, shape } => ComputeOp::Reshape {
|
|
input: remap(*input)?,
|
|
shape: shape.clone(),
|
|
},
|
|
ComputeOp::Transpose { input, dims } => ComputeOp::Transpose {
|
|
input: remap(*input)?,
|
|
dims: dims.clone(),
|
|
},
|
|
};
|
|
|
|
Ok(remapped_op)
|
|
}
|
|
|
|
/// Serialize to JSON string
|
|
pub fn to_json(&self) -> Result<String> {
|
|
serde_json::to_string_pretty(self)
|
|
.map_err(|e| GraphError::serialization_error(e.to_string()))
|
|
}
|
|
|
|
/// Deserialize from JSON string
|
|
pub fn from_json(json: &str) -> Result<Self> {
|
|
serde_json::from_str(json).map_err(|e| GraphError::serialization_error(e.to_string()))
|
|
}
|
|
|
|
/// Serialize to binary format (MessagePack)
|
|
pub fn to_binary(&self) -> Result<Vec<u8>> {
|
|
// For now, use bincode as binary format
|
|
bincode::serialize(self).map_err(|e| GraphError::serialization_error(e.to_string()))
|
|
}
|
|
|
|
/// Deserialize from binary format
|
|
pub fn from_binary(data: &[u8]) -> Result<Self> {
|
|
bincode::deserialize(data).map_err(|e| GraphError::serialization_error(e.to_string()))
|
|
}
|
|
|
|
/// Get format statistics
|
|
pub fn stats(&self) -> SerializationStats {
|
|
SerializationStats {
|
|
node_count: self.nodes.len(),
|
|
edge_count: self.edges.len(),
|
|
data_ops: self
|
|
.nodes
|
|
.iter()
|
|
.filter(|node| node.operation.is_data_op())
|
|
.count(),
|
|
compute_ops: self
|
|
.nodes
|
|
.iter()
|
|
.filter(|node| node.operation.is_compute_op())
|
|
.count(),
|
|
version: self.version.clone(),
|
|
timestamp: self.timestamp,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Statistics about serialized graph
|
|
#[derive(Debug, Clone)]
|
|
pub struct SerializationStats {
|
|
/// Number of nodes
|
|
pub node_count: usize,
|
|
/// Number of edges
|
|
pub edge_count: usize,
|
|
/// Number of data operations
|
|
pub data_ops: usize,
|
|
/// Number of compute operations
|
|
pub compute_ops: usize,
|
|
/// Graph format version
|
|
pub version: String,
|
|
/// Serialization timestamp
|
|
pub timestamp: chrono::DateTime<chrono::Utc>,
|
|
}
|
|
|
|
impl SerializationStats {
|
|
/// Get total operation count
|
|
pub fn total_ops(&self) -> usize {
|
|
self.data_ops + self.compute_ops
|
|
}
|
|
|
|
/// Get data to compute ratio
|
|
pub fn data_compute_ratio(&self) -> f64 {
|
|
if self.compute_ops == 0 {
|
|
f64::INFINITY
|
|
} else {
|
|
self.data_ops as f64 / self.compute_ops as f64
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Graph export formats
|
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
|
pub enum ExportFormat {
|
|
/// JSON format (human readable)
|
|
Json,
|
|
/// Binary format (compact)
|
|
Binary,
|
|
/// GraphViz DOT format (visualization)
|
|
Dot,
|
|
}
|
|
|
|
/// Graph exporter for different formats
|
|
pub struct GraphExporter;
|
|
|
|
impl GraphExporter {
|
|
/// Export graph to specified format
|
|
pub fn export(graph: &Graph, format: ExportFormat) -> Result<Vec<u8>> {
|
|
match format {
|
|
ExportFormat::Json => {
|
|
let json = graph.to_json()?;
|
|
Ok(json.into_bytes())
|
|
}
|
|
ExportFormat::Binary => {
|
|
let serializable = SerializableGraph::from_graph(graph);
|
|
serializable.to_binary()
|
|
}
|
|
ExportFormat::Dot => Self::export_dot(graph),
|
|
}
|
|
}
|
|
|
|
/// Export to GraphViz DOT format for visualization
|
|
fn export_dot(graph: &Graph) -> Result<Vec<u8>> {
|
|
let mut dot = String::new();
|
|
dot.push_str("digraph RustyTorchGraph {\n");
|
|
dot.push_str(" rankdir=TB;\n");
|
|
dot.push_str(" node [shape=box];\n");
|
|
|
|
// Add nodes
|
|
for node in graph.nodes() {
|
|
let label = format!("{}\\n{}", node.id, node.operation.name());
|
|
let color = if node.operation.is_data_op() {
|
|
"lightblue"
|
|
} else {
|
|
"lightgreen"
|
|
};
|
|
|
|
dot.push_str(&format!(
|
|
" \"{}\" [label=\"{}\" fillcolor=\"{}\" style=\"filled\"];\n",
|
|
node.id, label, color
|
|
));
|
|
}
|
|
|
|
// Add edges
|
|
for node in graph.nodes() {
|
|
if let Ok(dependencies) = graph.dependencies(node.id) {
|
|
for dep_id in dependencies {
|
|
dot.push_str(&format!(" \"{}\" -> \"{}\";\n", dep_id, node.id));
|
|
}
|
|
}
|
|
}
|
|
|
|
dot.push_str("}\n");
|
|
Ok(dot.into_bytes())
|
|
}
|
|
|
|
/// Import graph from specified format
|
|
pub fn import(data: &[u8], format: ExportFormat) -> Result<Graph> {
|
|
match format {
|
|
ExportFormat::Json => {
|
|
let json = String::from_utf8(data.to_vec())
|
|
.map_err(|e| GraphError::serialization_error(e.to_string()))?;
|
|
Graph::from_json(&json)
|
|
}
|
|
ExportFormat::Binary => {
|
|
let serializable = SerializableGraph::from_binary(data)?;
|
|
serializable.to_graph()
|
|
}
|
|
ExportFormat::Dot => Err(GraphError::serialization_error(
|
|
"DOT format import not supported".to_string(),
|
|
)),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::{ComputeOp, DataOp, GraphBuilder};
|
|
|
|
#[test]
|
|
fn test_graph_serialization() {
|
|
let mut builder = GraphBuilder::new();
|
|
|
|
let read_node = builder
|
|
.add_data_op(DataOp::Read {
|
|
source: "test.csv".to_string(),
|
|
format: "csv".to_string(),
|
|
schema: None,
|
|
})
|
|
.expect("Failed to add read node");
|
|
|
|
let tensor_node = builder
|
|
.add_compute_op(ComputeOp::TensorLoad {
|
|
data_node: read_node,
|
|
shape: vec![100, 50],
|
|
dtype: "f32".to_string(),
|
|
})
|
|
.expect("Failed to add tensor node");
|
|
|
|
builder
|
|
.add_dependency(tensor_node, read_node)
|
|
.expect("Failed to add dependency");
|
|
|
|
let graph = builder.build().expect("Failed to build graph");
|
|
|
|
// Test JSON serialization
|
|
let json = graph.to_json().expect("Failed to serialize to JSON");
|
|
assert!(json.contains("Read"));
|
|
assert!(json.contains("TensorLoad"));
|
|
|
|
let deserialized = Graph::from_json(&json).expect("Failed to deserialize from JSON");
|
|
assert_eq!(deserialized.node_count(), graph.node_count());
|
|
assert_eq!(deserialized.edge_count(), graph.edge_count());
|
|
}
|
|
|
|
#[test]
|
|
fn test_serializable_graph() {
|
|
let mut builder = GraphBuilder::new();
|
|
|
|
let _read_node = builder
|
|
.add_data_op(DataOp::Read {
|
|
source: "data.parquet".to_string(),
|
|
format: "parquet".to_string(),
|
|
schema: Some("id:i32,value:f64".to_string()),
|
|
})
|
|
.expect("Failed to add read node");
|
|
|
|
let graph = builder.build().expect("Failed to build graph");
|
|
|
|
let serializable = SerializableGraph::from_graph(&graph);
|
|
assert_eq!(serializable.nodes.len(), 1);
|
|
assert_eq!(serializable.edges.len(), 0);
|
|
assert_eq!(serializable.version, SerializableGraph::CURRENT_VERSION);
|
|
|
|
let restored = serializable.to_graph().expect("Failed to restore graph");
|
|
assert_eq!(restored.node_count(), graph.node_count());
|
|
}
|
|
|
|
#[test]
|
|
fn test_export_formats() {
|
|
let mut builder = GraphBuilder::new();
|
|
|
|
let _read_node = builder
|
|
.add_data_op(DataOp::Read {
|
|
source: "test.csv".to_string(),
|
|
format: "csv".to_string(),
|
|
schema: None,
|
|
})
|
|
.expect("Failed to add read node");
|
|
|
|
let graph = builder.build().expect("Failed to build graph");
|
|
|
|
// Test JSON export
|
|
let json_data =
|
|
GraphExporter::export(&graph, ExportFormat::Json).expect("Failed to export to JSON");
|
|
assert!(!json_data.is_empty());
|
|
|
|
// Test binary export
|
|
let binary_data = GraphExporter::export(&graph, ExportFormat::Binary)
|
|
.expect("Failed to export to binary");
|
|
assert!(!binary_data.is_empty());
|
|
|
|
// Test DOT export
|
|
let dot_data =
|
|
GraphExporter::export(&graph, ExportFormat::Dot).expect("Failed to export to DOT");
|
|
let dot_str = String::from_utf8(dot_data).expect("Invalid DOT format");
|
|
assert!(dot_str.contains("digraph"));
|
|
assert!(dot_str.contains("Read"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_version_compatibility() {
|
|
let serializable = SerializableGraph {
|
|
version: "1.0.0".to_string(),
|
|
metadata: IndexMap::new(),
|
|
nodes: Vec::new(),
|
|
edges: Vec::new(),
|
|
timestamp: chrono::Utc::now(),
|
|
};
|
|
|
|
assert!(serializable.is_compatible_version());
|
|
|
|
let incompatible = SerializableGraph {
|
|
version: "2.0.0".to_string(),
|
|
..serializable.clone()
|
|
};
|
|
|
|
assert!(!incompatible.is_compatible_version());
|
|
}
|
|
|
|
#[test]
|
|
fn test_serialization_stats() {
|
|
let mut builder = GraphBuilder::new();
|
|
|
|
let read_node = builder
|
|
.add_data_op(DataOp::Read {
|
|
source: "test.csv".to_string(),
|
|
format: "csv".to_string(),
|
|
schema: None,
|
|
})
|
|
.expect("Failed to add read node");
|
|
|
|
let tensor_node = builder
|
|
.add_compute_op(ComputeOp::TensorLoad {
|
|
data_node: read_node,
|
|
shape: vec![10, 10],
|
|
dtype: "f32".to_string(),
|
|
})
|
|
.expect("Failed to add tensor node");
|
|
|
|
builder
|
|
.add_dependency(tensor_node, read_node)
|
|
.expect("Failed to add dependency");
|
|
|
|
let graph = builder.build().expect("Failed to build graph");
|
|
let serializable = SerializableGraph::from_graph(&graph);
|
|
let stats = serializable.stats();
|
|
|
|
assert_eq!(stats.node_count, 2);
|
|
assert_eq!(stats.edge_count, 1);
|
|
assert_eq!(stats.data_ops, 1);
|
|
assert_eq!(stats.compute_ops, 1);
|
|
assert_eq!(stats.total_ops(), 2);
|
|
assert_eq!(stats.data_compute_ratio(), 1.0);
|
|
}
|
|
}
|