fix(optimizer): minor mypy and formatting fixes

discopop-project · Dec 13, 2023 · 4bb27ed · 4bb27ed
1 parent 4b134d2
commit 4bb27ed
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 12 deletions.
diff --git a/discopop_library/discopop_optimizer/classes/system/Network.py b/discopop_library/discopop_optimizer/classes/system/Network.py
@@ -25,7 +25,7 @@ def __init__(self):
 
     def add_connection(self, source: Device, target: Device, transfer_speed: Expr, initialization_delay: Expr):
         if source == target:
-            transfer_speed = Integer(10000*1000) # 10000 GB/s = 10000 * 1000 MB/s
+            transfer_speed = Integer(10000 * 1000)  # 10000 GB/s = 10000 * 1000 MB/s
         self.__transfer_speeds[(source, target)] = transfer_speed
         self.__transfer_initialization_costs[(source, target)] = initialization_delay
 

diff --git a/discopop_library/discopop_optimizer/classes/system/System.py b/discopop_library/discopop_optimizer/classes/system/System.py
@@ -64,17 +64,26 @@ def __build_from_configuration_file(self, arguments: OptimizerArguments):
                 raise ValueError("Unknown device type: " + str(device_dict["device_type"]) + " in: " + str(device_dict))
         # register host device for device -> device transfers
         self.__network.set_host_device(self.get_device(self.__host_device_id))
-        
+
         # register connections
         for device_dict in system_configuration["devices"]:
             if "transfer_speeds" in device_dict:
                 # register connection with given transfer speed
-                self.__network.add_connection(source = self.get_device(self.__host_device_id), target = self.get_device(device_dict["device_id"]), transfer_speed=Float(device_dict["transfer_speeds"]["H2D_MB/s"]), initialization_delay=Float(device_dict["transfer_init_delays"]["average"]))  # H2D
-                self.__network.add_connection(source = self.get_device(device_dict["device_id"]), target = self.get_device(self.__host_device_id), transfer_speed=Float(device_dict["transfer_speeds"]["D2H_MB/s"]), initialization_delay=Float(device_dict["transfer_init_delays"]["average"]))  # D2H                
+                self.__network.add_connection(
+                    source=self.get_device(self.__host_device_id),
+                    target=self.get_device(device_dict["device_id"]),
+                    transfer_speed=Float(device_dict["transfer_speeds"]["H2D_MB/s"]),
+                    initialization_delay=Float(device_dict["transfer_init_delays"]["average"]),
+                )  # H2D
+                self.__network.add_connection(
+                    source=self.get_device(device_dict["device_id"]),
+                    target=self.get_device(self.__host_device_id),
+                    transfer_speed=Float(device_dict["transfer_speeds"]["D2H_MB/s"]),
+                    initialization_delay=Float(device_dict["transfer_init_delays"]["average"]),
+                )  # D2H
             else:
                 # no transfer speed information exists
-                pass      
-
+                pass
 
     def __build_CPU(self, device_configuration: Dict[str, Any]):
         cpu = CPU(

diff --git a/discopop_library/discopop_optimizer/classes/system/system_utils.py b/discopop_library/discopop_optimizer/classes/system/system_utils.py
@@ -8,7 +8,7 @@
 
 import json
 import os
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Union
 
 from sympy import Integer
 
@@ -23,17 +23,17 @@ def generate_default_system_configuration(file_path: str):
         return
     # build default system configuration
     system_configuration: Dict[str, Any] = dict()
-    devices: List[Dict[str, float | int | Dict[str, float]]] = []
+    devices: List[Dict[str, Union[float, int, Dict[str, float]]]] = []
     # configure host device
-    host_device: Dict[str, float | int | Dict[str, float]] = {
+    host_device: Dict[str, Union[float, int, Dict[str, float]]] = {
         "device_id": 0,
         "device_type": DeviceTypeEnum.CPU,
         "frequency": 3000000000,
         "processors": 16,
         "threads": 16,
     }
     # configure gpu_1
-    gpu_1: Dict[str, float | int | Dict[str, float]] = {
+    gpu_1: Dict[str, Union[float, int, Dict[str, float]]] = {
         "compute_init_delays": {"target_teams_distribute_parallel_for": 0.01},
         "device_id": 1,
         "device_type": DeviceTypeEnum.GPU,
@@ -50,15 +50,20 @@ def generate_default_system_configuration(file_path: str):
         "transfer_speeds": {"D2H_MB/s": 1800, "H2D_MB/s": 3600},  # MB/s
     }
     # configure gpu_2
-    gpu_2: Dict[str, float | int | Dict[str, float]] = {
+    gpu_2: Dict[str, Union[float, int, Dict[str, float]]] = {
         "compute_init_delays": {"target_teams_distribute_parallel_for": 0.005},
         "device_id": 2,
         "device_type": DeviceTypeEnum.GPU,
         "frequency": 512000000,
         "processors": 128,
         "teams": 3200,
         "threads": 3200,
-        "transfer_init_delays": {"target_data_update": 7.1e-05, "target_enter_data": 0.0002, "target_exit_data": 5e-06, "average": 6e-05},
+        "transfer_init_delays": {
+            "target_data_update": 7.1e-05,
+            "target_enter_data": 0.0002,
+            "target_exit_data": 5e-06,
+            "average": 6e-05,
+        },
         "transfer_speeds": {"D2H_MB/s": 1900, "H2D_MB/s": 4200},  # MB/s
     }
     # assemble system_configuration