#615: spec: add new methods to adding communication edges to spec

src/lbaf/Utils/lbsJSONSpecFileMaker.py

@@ -1,11 +1,12 @@
 import os
 import yaml
 from typing import Set, List, Union
 
 from lbaf.Execution.lbsPhaseSpecification import (
     PhaseSpecification, SharedBlockSpecification,
     RankSpecification, TaskSpecification,
-    PhaseSpecificationNormalizer
+    PhaseSpecificationNormalizer,
+    CommunicationSpecification
 )
 
 class JSONSpecFileMaker:
@@ -26,20 +27,22 @@
             "shared": 0,
             "rank": 0,
             "phase": 0,
+            "comm": 0,
         }
 
         self.id_sets = {
             "seq": set(),
             "shared": set(),
             "rank": set(),
             "phase": set(),
+            "comm": set(),
         }
 
 
     ###########################################################################
     ## Private generators
 
     def generateNextID_(self, key: str) -> int:
         """
         Generic function to generate and update the next available ID for a given category.
         """
@@ -53,7 +56,7 @@
         self.current_ids[key] += 1
         return next_id
 
     def checkID_(self, try_id: int, key: str) -> int:
         if try_id < 0:
             try_id = self.generateNextID_(key)
         elif id in self.id_sets[key]:
@@ -64,11 +67,11 @@
     ###########################################################################
     ## Private functions for assertions
 
     def assertIDExists_(self, test_id: int, id_type: str):
         assert test_id in self.id_sets[id_type], \
                f"Task {test_id} has not been created yet. Use createTask() method."
 
     def assertAllTasksHaveBeenAssigned_(self):
         for t in self.tasks:
             assert t in self.assignments, \
                    f"Task {t} has not been assigned. Call " \
@@ -79,7 +82,7 @@
     ###########################################################################
     ## Public Functions for creating JSON fields
 
     def createTask(self,
             time: float,
             seq_id: int = -1,
             collection_id: int = 7,
@@ -102,10 +105,10 @@
         return task
 
     def createSharedBlock(self,
             bytes: float,
             tasks: Union[List[int], List[TaskSpecification]],
             rank_id: int = -1,
             id: int = -1):
         shared_id = self.checkID_(id, "shared")
         shared_block = SharedBlockSpecification({
             "size": bytes,
@@ -123,13 +126,22 @@
                     task_specs[t["seq_id"]] = t
                 else:
                     raise RuntimeError(
                         f"tasks must be a list of ints (IDs) or TaskSpecifications")
         shared_block["tasks"] = task_specs
         self.shared_blocks[shared_id] = shared_block
         return shared_block
 
+    def createComm(self, to_obj : int, from_obj : int, size : float) -> CommunicationSpecification:
+        new_comm = CommunicationSpecification({
+            'size': size,
+            'from': from_obj,
+            'to': to_obj
+        })
+        comm_id = self.checkID_(-1, "comm")
+        self.comms[comm_id] = new_comm
+
     def createRank(self,
             id: int = -1,
             tasks: Union[List[int], List[TaskSpecification]] = None,
             user_defined: dict = None
         ) -> RankSpecification:
@@ -192,13 +204,13 @@
     ###########################################################################
     ## Standard Getters
 
     def getRank(self, id: int) -> RankSpecification:
         return self.ranks[id]
 
     def getTask(self, seq_id: int) -> TaskSpecification:
         return self.tasks[seq_id]
 
     def getSharedBlock(self, id: int) -> SharedBlockSpecification:
         return self.shared_blocks[id]
 
     ###########################################################################
@@ -220,5 +232,18 @@
         # Write out the spec
         path = os.path.join(os.getcwd(), "spec.yaml") if path is None else path
         os.makedirs(os.path.dirname(path), exist_ok=True)
         with open(path, 'w') as output_file:
             yaml.dump(spec, output_file, default_flow_style=False)
+
+    def makeSpecification(self, phase_id):
+        self.assertAllTasksHaveBeenAssigned_()
+        phase = PhaseSpecification({
+            "tasks": self.tasks,
+            "shared_blocks": self.shared_blocks,
+            "communications": self.comms,
+            "ranks": self.ranks,
+            "id": phase_id
+        })
+        norm = PhaseSpecificationNormalizer()
+        spec = norm.normalize(phase)
+        return spec