Skip to content

Write Model Updates #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alvaradoo wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Write Model Updates #14

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0457282
working write beta
alvaradoo Dec 17, 2025
e8653a8
changes __write_model & __build_model to skip Device type and model p…
alvaradoo Jan 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 84 additions & 23 deletions src/ahp_graph/SSTGraph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -170,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:
Expand All @@ -183,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:
Expand Down Expand Up @@ -213,6 +221,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.
Expand All @@ -228,10 +237,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.
Expand All @@ -251,13 +260,18 @@ 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)
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,
"params" : self.__encode(d1.attr, True),
"params_global_sets" : global_set,
"params" : self.__encode(params_dict, True),
# "params_global_sets" : global_set,
}
if d1.subs:
item["subcomponents"] = recurseSubcomponents(d1)
Expand All @@ -271,12 +285,21 @@ 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
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,
"params" : self.__encode(d0.attr, True),
"params_global_sets" : global_set,
"params" : self.__encode(params_dict, True),
# "params_global_sets" : global_set,
}
if d0.partition is not None:
component["partition"] = {
Expand Down Expand Up @@ -309,19 +332,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

Expand Down