Merge pull request #117 from Janekdererste/partitioning-options

introduced new metis partitioning options

Cargo.toml

@@ -21,7 +21,7 @@ mpi = { git = "https://github.com/rsmpi/rsmpi", rev = "1622101c9fc3a78620e9649f7
 prost = "0.11.5"
 # bytes replaces the signature of std::io::Cursor to implemnt Buf somehow...
 bytes = "1.3.0"
-serial_test = "1.0.0"
+serial_test = "2.0.0"
 wait-timeout = "0.2.0"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3.0", features = ["json", "fmt", "std", "registry"] }

src/bin/partition_network.rs

@@ -3,7 +3,7 @@ use std::path::PathBuf;
 use clap::{arg, Parser};
 use tracing::info;
 
-use rust_q_sim::simulation::config::PartitionMethod;
+use rust_q_sim::simulation::config::{MetisOptions, PartitionMethod};
 use rust_q_sim::simulation::network::global_network::Network;
 
 fn main() {
@@ -32,7 +32,11 @@ fn main() {
         args.in_path, args.num_parts
     );
 
-    let net1 = Network::from_file(&args.in_path, args.num_parts, PartitionMethod::Metis);
+    let net1 = Network::from_file(
+        &args.in_path,
+        args.num_parts,
+        PartitionMethod::Metis(MetisOptions::default()),
+    );
     info!(
         "Network is loaded with {} links and {} nodes.",
         net1.links.len(),

src/simulation/config.rs

@@ -210,15 +210,94 @@ pub enum RoutingMode {
     UsePlans,
 }
 
-#[derive(PartialEq, Debug, ValueEnum, Clone, Copy, Serialize, Deserialize)]
+#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
 pub enum PartitionMethod {
-    Metis,
+    Metis(MetisOptions),
     None,
 }
 
+#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
+pub struct MetisOptions {
+    pub vertex_weight: Vec<VertexWeight>,
+    #[serde(default = "edge_weight_constant")]
+    pub edge_weight: EdgeWeight,
+    #[serde(default = "f32_value_0_03")]
+    pub imbalance_factor: f32,
+    #[serde(default = "u32_value_100")]
+    pub iteration_number: u32,
+}
+
+#[derive(PartialEq, Debug, ValueEnum, Clone, Copy, Serialize, Deserialize)]
+pub enum VertexWeight {
+    InLinkCapacity,
+    InLinkCount,
+    Constant,
+}
+
+#[derive(PartialEq, Debug, ValueEnum, Clone, Copy, Serialize, Deserialize)]
+pub enum EdgeWeight {
+    Capacity,
+    Constant,
+}
+
+impl Default for MetisOptions {
+    fn default() -> Self {
+        MetisOptions {
+            vertex_weight: vec![],
+            edge_weight: EdgeWeight::Constant,
+            imbalance_factor: 0.03,
+            iteration_number: 10,
+        }
+    }
+}
+
+impl MetisOptions {
+    pub fn set_imbalance_factor(mut self, imbalance_factor: f32) -> Self {
+        self.imbalance_factor = imbalance_factor;
+        self
+    }
+
+    pub fn add_vertex_weight(mut self, vertex_weight: VertexWeight) -> Self {
+        self.vertex_weight.push(vertex_weight);
+        self
+    }
+
+    pub fn set_edge_weight(mut self, edge_weight: EdgeWeight) -> Self {
+        self.edge_weight = edge_weight;
+        self
+    }
+
+    pub fn set_iteration_number(mut self, iteration_number: u32) -> Self {
+        self.iteration_number = iteration_number;
+        self
+    }
+
+    pub fn ufactor(&self) -> i32 {
+        let val = (self.imbalance_factor * 1000.) as i32;
+        if val <= 0 {
+            return 1;
+        };
+        val
+    }
+}
+
+fn f32_value_0_03() -> f32 {
+    0.03
+}
+
+fn edge_weight_constant() -> EdgeWeight {
+    EdgeWeight::Constant
+}
+
+fn u32_value_100() -> u32 {
+    100
+}
+
 #[cfg(test)]
 mod tests {
-    use crate::simulation::config::{Config, PartitionMethod, Partitioning};
+    use crate::simulation::config::{
+        Config, EdgeWeight, MetisOptions, PartitionMethod, Partitioning, VertexWeight,
+    };
 
     #[test]
     fn read_from_yaml() {
@@ -227,7 +306,12 @@ mod tests {
         };
         let partitioning = Partitioning {
             num_parts: 1,
-            method: PartitionMethod::Metis,
+            method: PartitionMethod::Metis(MetisOptions {
+                vertex_weight: vec![VertexWeight::InLinkCount, VertexWeight::InLinkCapacity],
+                edge_weight: EdgeWeight::Constant,
+                imbalance_factor: 1.02,
+                iteration_number: 100,
+            }),
         };
         config
             .modules
@@ -238,7 +322,92 @@ mod tests {
 
         println!("{yaml}");
 
-        let _parsed_config: Config = serde_yaml::from_str(&yaml).expect("failed to parse config");
-        println!("done.")
+        let parsed_config: Config = serde_yaml::from_str(&yaml).expect("failed to parse config");
+        println!("done.");
+
+        assert_eq!(parsed_config.partitioning().num_parts, 1);
+        assert_eq!(
+            parsed_config.partitioning().method,
+            PartitionMethod::Metis(MetisOptions {
+                vertex_weight: vec![VertexWeight::InLinkCount, VertexWeight::InLinkCapacity],
+                edge_weight: EdgeWeight::Constant,
+                imbalance_factor: 1.02,
+                iteration_number: 100,
+            })
+        );
+    }
+
+    #[test]
+    fn read_none_partitioning() {
+        let yaml = r#"
+        modules: 
+          partitioning:
+            type: Partitioning
+            num_parts: 1
+            method: None
+        "#;
+        let parsed_config: Config = serde_yaml::from_str(&yaml).expect("failed to parse config");
+        assert_eq!(parsed_config.partitioning().num_parts, 1);
+        assert_eq!(parsed_config.partitioning().method, PartitionMethod::None);
+    }
+
+    #[test]
+    fn read_metis_partitioning() {
+        let yaml = r#"
+        modules:
+          partitioning:
+            type: Partitioning
+            num_parts: 1
+            method: !Metis
+              vertex_weight: 
+              - InLinkCount
+              imbalance_factor: 1.1
+              edge_weight: Capacity
+        "#;
+        let parsed_config: Config = serde_yaml::from_str(yaml).expect("failed to parse config");
+        assert_eq!(parsed_config.partitioning().num_parts, 1);
+        assert_eq!(
+            parsed_config.partitioning().method,
+            PartitionMethod::Metis(MetisOptions {
+                vertex_weight: vec![VertexWeight::InLinkCount],
+                edge_weight: EdgeWeight::Capacity,
+                imbalance_factor: 1.1,
+                iteration_number: 100
+            })
+        );
+    }
+
+    #[test]
+    fn test_imbalance_factor() {
+        assert_eq!(
+            MetisOptions::default().set_imbalance_factor(0.03).ufactor(),
+            30
+        );
+        assert_eq!(
+            MetisOptions::default()
+                .set_imbalance_factor(0.001)
+                .ufactor(),
+            1
+        );
+        assert_eq!(
+            MetisOptions::default()
+                .set_imbalance_factor(0.00001)
+                .ufactor(),
+            1
+        );
+        assert_eq!(
+            MetisOptions::default()
+                .set_imbalance_factor(0.00000)
+                .ufactor(),
+            1
+        );
+        assert_eq!(
+            MetisOptions::default().set_imbalance_factor(-1.).ufactor(),
+            1
+        );
+        assert_eq!(
+            MetisOptions::default().set_imbalance_factor(1.1).ufactor(),
+            1100
+        );
     }
 }

src/simulation/controller.rs

@@ -186,7 +186,7 @@ fn try_join(mut handles: IntMap<u32, JoinHandle<()>>) {
 
 pub fn partition_input(config: &Config) {
     id::load_from_file(&PathBuf::from(config.proto_files().ids));
-    let net = if config.partitioning().method == PartitionMethod::Metis {
+    let net = if let PartitionMethod::Metis(_) = config.partitioning().method {
         info!("Config param Partition method was set to metis. Loading input network, running metis conversion and then store it into output folder");
         partition_network(config)
     } else {

src/simulation/network/global_network.rs

@@ -127,8 +127,8 @@ impl Network {
 
     fn partition_network(network: &mut Network, partition_method: PartitionMethod, num_parts: u32) {
         match partition_method {
-            PartitionMethod::Metis => {
-                let partitions = metis_partitioning::partition(network, num_parts);
+            PartitionMethod::Metis(options) => {
+                let partitions = metis_partitioning::partition(network, num_parts, options);
                 for node in network.nodes.iter_mut() {
                     let partition = partitions[node.id.internal() as usize] as u32;
                     node.partition = partition;
@@ -232,7 +232,7 @@ impl Link {
 
 #[cfg(test)]
 mod tests {
-    use crate::simulation::config::PartitionMethod;
+    use crate::simulation::config::{EdgeWeight, MetisOptions, PartitionMethod};
     use crate::simulation::id::Id;
 
     use super::{Link, Network, Node};
@@ -308,11 +308,13 @@ mod tests {
     }
 
     #[test]
+    #[serial_test::serial]
     fn from_file() {
         let network = Network::from_file(
             "./assets/equil/equil-network.xml",
             2,
-            PartitionMethod::Metis,
+            //I don't know, why "edge_weight = true" sets 1 as partition for all nodes.
+            PartitionMethod::Metis(MetisOptions::default().set_edge_weight(EdgeWeight::Constant)),
         );
 
         // check partitioning