From 04572828a263d07729f907e9b0bb63d1d75fcd50 Mon Sep 17 00:00:00 2001 From: Oliver Alvarado Rodriguez Date: Wed, 17 Dec 2025 22:26:55 +0000 Subject: [PATCH 1/2] working write beta Signed-off-by: Oliver Alvarado Rodriguez --- src/ahp_graph/SSTGraph.py | 93 ++++++++++++++++++++++++++++++--------- 1 file changed, 72 insertions(+), 21 deletions(-) diff --git a/src/ahp_graph/SSTGraph.py b/src/ahp_graph/SSTGraph.py index 96bc588..9d56604 100644 --- a/src/ahp_graph/SSTGraph.py +++ b/src/ahp_graph/SSTGraph.py @@ -98,12 +98,14 @@ def write_json(self, self.__write_model( f"{output}/{filename}", nranks, + rank, program_options) else: (base, ext) = os.path.splitext(f"{output}/{filename}") self.__write_model( base + str(rank) + ext, nranks, + rank, program_options) @staticmethod @@ -213,6 +215,7 @@ def recurseSubcomponents(dev: Device, comp: 'sst.Component') -> None: def __write_model(self, filename: str, nranks: int, + rank: int, program_options: dict = None) -> None: """ Write this DeviceGraph out as JSON. @@ -228,10 +231,10 @@ def __write_model(self, model["program_options"] = dict(program_options) # - # If running in parallel, then set up the SST SELF partitioner. + # If running in parallel, then set up the SST LINEAR partitioner. # if nranks > 1: - model["program_options"]["partitioner"] = "sst.self" + model["program_options"]["partitioner"] = "sst.linear" # # Set up global parameters. @@ -251,13 +254,16 @@ def recurseSubcomponents(dev: Device) -> list: if d1.library is None: raise RuntimeError(f"No library: {d1.name}") - d1.attr.update(type=d1.type, model=d1.model) item = { "slot_name" : n1, "type" : d1.library, "slot_number" : s1, - "params" : self.__encode(d1.attr, True), - "params_global_sets" : global_set, + # omit internal metadata like type/model from params + "params" : self.__encode({ + k: v for (k, v) in d1.attr.items() + if k not in ("type", "model") + }, True), + # "params_global_sets" : global_set, } if d1.subs: item["subcomponents"] = recurseSubcomponents(d1) @@ -271,12 +277,19 @@ def recurseSubcomponents(dev: Device) -> list: components = list() for d0 in self.devices.values(): if d0.subOwner is None and d0.library is not None: - d0.attr.update(type=d0.type, model=d0.model) + # In multi-rank output, only emit components local to this rank + if nranks > 1: + if d0.partition is None or d0.partition[0] != rank: + continue component = { "name" : d0.name, "type" : d0.library, - "params" : self.__encode(d0.attr, True), - "params_global_sets" : global_set, + # omit internal metadata like type/model from params + "params" : self.__encode({ + k: v for (k, v) in d0.attr.items() + if k not in ("type", "model") + }, True), + # "params_global_sets" : global_set, } if d0.partition is not None: component["partition"] = { @@ -309,19 +322,57 @@ def recurseSubcomponents(dev: Device) -> list: name = f'{p0}__{t}__{p1}' else: name = f'{p1}__{t}__{p0}' - links.append({ - "name" : name, - "left" : { - "component" : p0.device.name, - "port" : p0.get_name(), - "latency" : latency - }, - "right" : { - "component" : p1.device.name, - "port" : p1.get_name(), - "latency" : latency - }, - }) + + d0 = p0.device + d1 = p1.device + r0, t0 = (d0.partition or (0, None)) + r1, t1 = (d1.partition or (0, None)) + t0 = 0 if t0 is None else t0 + t1 = 0 if t1 is None else t1 + + # Determine if link is inter-rank relative to this file's rank + is_nonlocal = nranks > 1 and r0 != r1 + + if is_nonlocal: + # Ensure left side is the local endpoint for this rank + if r0 == rank: + left = {"component": d0.name, "port": p0.get_name(), "latency": latency} + right = {"rank": r1, "thread": t1} + elif r1 == rank: + left = {"component": d1.name, "port": p1.get_name(), "latency": latency} + right = {"rank": r0, "thread": t0} + else: + # Neither endpoint is local to this rank; skip emitting + # (should not occur after prune, but guard just in case) + continue + + links.append({ + "name": name, + "noCut": False, + "nonlocal": True, + "left": left, + "right": right, + }) + else: + # Local link (same-rank or single-rank output): keep both endpoints + # If multi-rank, only keep links where both endpoints are on this rank + if nranks > 1 and not (r0 == rank and r1 == rank): + continue + links.append({ + "name": name, + "noCut": False, + "nonlocal": False, + "left": { + "component": d0.name, + "port": p0.get_name(), + "latency": latency, + }, + "right": { + "component": d1.name, + "port": p1.get_name(), + "latency": latency, + }, + }) model["links"] = links From e8653a8491ef4a591502a598a0f852275df41c5e Mon Sep 17 00:00:00 2001 From: Oliver Alvarado Rodriguez Date: Thu, 8 Jan 2026 22:38:31 +0000 Subject: [PATCH 2/2] changes __write_model & __build_model to skip Device type and model params if they are None Signed-off-by: Oliver Alvarado Rodriguez --- src/ahp_graph/SSTGraph.py | 34 ++++++++++++++++++++++------------ 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/src/ahp_graph/SSTGraph.py b/src/ahp_graph/SSTGraph.py index 9d56604..9bc99b8 100644 --- a/src/ahp_graph/SSTGraph.py +++ b/src/ahp_graph/SSTGraph.py @@ -172,7 +172,10 @@ def recurseSubcomponents(dev: Device, comp: 'sst.Component') -> None: c1 = comp.setSubComponent(n1, d1.library) else: c1 = comp.setSubComponent(n1, d1.library, s1) - d1.attr.update(type=d1.type, model=d1.model) + if d1.type is not None: + d1.attr['type'] = d1.type + if d1.model is not None: + d1.attr['model'] = d1.model c1.addParams(self.__encode(d1.attr)) n2c[d1.name] = c1 for key in global_params: @@ -185,7 +188,10 @@ def recurseSubcomponents(dev: Device, comp: 'sst.Component') -> None: for d0 in self.devices.values(): if d0.subOwner is None and d0.library is not None: c0 = sst.Component(d0.name, d0.library) - d0.attr.update(type=d0.type, model=d0.model) + if d0.type is not None: + d0.attr['type'] = d0.type + if d0.model is not None: + d0.attr['model'] = d0.model c0.addParams(self.__encode(d0.attr)) # Set the component partition if we are self-partitioning if self_partition: @@ -254,15 +260,17 @@ def recurseSubcomponents(dev: Device) -> list: if d1.library is None: raise RuntimeError(f"No library: {d1.name}") + params_dict = {k: v for (k, v) in d1.attr.items()} + if d1.type is not None: + params_dict['type'] = d1.type + if d1.model is not None: + params_dict['model'] = d1.model + item = { "slot_name" : n1, "type" : d1.library, "slot_number" : s1, - # omit internal metadata like type/model from params - "params" : self.__encode({ - k: v for (k, v) in d1.attr.items() - if k not in ("type", "model") - }, True), + "params" : self.__encode(params_dict, True), # "params_global_sets" : global_set, } if d1.subs: @@ -281,14 +289,16 @@ def recurseSubcomponents(dev: Device) -> list: if nranks > 1: if d0.partition is None or d0.partition[0] != rank: continue + params_dict = {k: v for (k, v) in d0.attr.items()} + if d0.type is not None: + params_dict['type'] = d0.type + if d0.model is not None: + params_dict['model'] = d0.model + component = { "name" : d0.name, "type" : d0.library, - # omit internal metadata like type/model from params - "params" : self.__encode({ - k: v for (k, v) in d0.attr.items() - if k not in ("type", "model") - }, True), + "params" : self.__encode(params_dict, True), # "params_global_sets" : global_set, } if d0.partition is not None: