From b0b68e68888e921aab0695a3b36ed95ef44c7265 Mon Sep 17 00:00:00 2001 From: Tim Stavenger Date: Wed, 24 Jun 2026 08:38:01 -0700 Subject: [PATCH] have claude sonnet 4.5 add edge attributes --- examples/link_attributes_example.py | 106 ++++++++++++++++++++++++++++ pyproject.toml | 5 ++ src/ahp_graph/DeviceGraph.py | 38 +++++++--- src/ahp_graph/SSTGraph.py | 18 ++--- tests/test_DeviceGraph.py | 43 ++++++++++- 5 files changed, 192 insertions(+), 18 deletions(-) create mode 100644 examples/link_attributes_example.py diff --git a/examples/link_attributes_example.py b/examples/link_attributes_example.py new file mode 100644 index 0000000..7ce0c05 --- /dev/null +++ b/examples/link_attributes_example.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +""" +Example demonstrating arbitrary key/value attributes on links. + +This example shows how to add custom attributes to links in addition to +the standard latency attribute. This provides flexibility for modeling +various link properties such as bandwidth, protocol, reliability, etc. +""" + +import sys +sys.path.insert(0, '../src') + +from ahp_graph.Device import * +from ahp_graph.DeviceGraph import * + + +class NetworkSwitch(Device): + """A simple network switch device.""" + library = 'network.Switch' + portinfo = PortInfo() + portinfo.add('port', limit=None, required=False) # Multi-port + + def __init__(self, name: str): + super().__init__(name, attr={'switch_type': '10G'}) + + +class Server(Device): + """A server device.""" + library = 'compute.Server' + portinfo = PortInfo() + portinfo.add('nic') # Network interface + + def __init__(self, name: str): + super().__init__(name, attr={'cores': 32, 'memory': '256GB'}) + + +def main(): + """Create a simple network topology with diverse link attributes.""" + print("Creating network topology with link attributes...\n") + + graph = DeviceGraph() + + # Create devices + switch1 = NetworkSwitch('switch1') + switch2 = NetworkSwitch('switch2') + server1 = Server('server1') + server2 = Server('server2') + + # Link servers to switches with different link properties + # Server 1 has a high-speed, low-latency connection + graph.link( + server1.nic, + switch1.port('port', 0), + attr={ + 'latency': '1ns', + 'bandwidth': '100GB/s', + 'protocol': 'InfiniBand', + 'mtu': 9000, + 'reliability': 0.99999 + } + ) + + # Server 2 has a standard Ethernet connection + graph.link( + server2.nic, + switch1.port('port', 1), + attr={ + 'latency': '10ns', + 'bandwidth': '10GB/s', + 'protocol': 'Ethernet', + 'mtu': 1500, + 'reliability': 0.9999 + } + ) + + # Inter-switch link with high bandwidth + graph.link( + switch1.port('port', 2), + switch2.port('port', 0), + attr={ + 'latency': '5ns', + 'bandwidth': '400GB/s', + 'protocol': 'Ethernet', + 'trunk': True, + 'vlan_ids': [100, 200, 300] + } + ) + + # Display the graph + print(graph) + print("\n" + "="*80 + "\n") + + # Demonstrate accessing link attributes + print("Link Attributes Summary:") + print("-" * 80) + for (p0, p1), attr in graph.links.items(): + print(f"\nLink: {p0} <--> {p1}") + for key, value in sorted(attr.items()): + print(f" {key}: {value}") + + print("\n" + "="*80) + print("Example complete! Link attributes provide rich metadata for modeling.") + + +if __name__ == '__main__': + main() diff --git a/pyproject.toml b/pyproject.toml index 393c5ae..50ddc7a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,3 +22,8 @@ classifiers = [ [project.urls] "Homepage" = "https://github.com/lpsmodsim/ahp_graph" + +[tool.pytest.ini_options] +pythonpath = [ + ".", "src" +] \ No newline at end of file diff --git a/src/ahp_graph/DeviceGraph.py b/src/ahp_graph/DeviceGraph.py index 824a650..d8c3642 100644 --- a/src/ahp_graph/DeviceGraph.py +++ b/src/ahp_graph/DeviceGraph.py @@ -76,7 +76,9 @@ def __repr__(self) -> str: for device in self.devices.values(): lines.append(str(device)) for p0, p1 in self.links: - lines.append(f"{p0} <--{self.links[(p0, p1)]}--> {p1}") + attr = self.links[(p0, p1)] + attr_str = ', '.join(f"{k}={v}" for k, v in sorted(attr.items())) + lines.append(f"{p0} <--[{attr_str}]--> {p1}") return "\n".join(lines) def _link_other_port(self, p0: DevicePort, p1: DevicePort) -> None: @@ -91,27 +93,47 @@ def _link_other_port(self, p0: DevicePort, p1: DevicePort) -> None: p1.link = p2 self.ports.remove(p0) self.ports.add(p1) - latency = self.links.pop(_orderedtuple(p0, p2)) + attr = self.links.pop(_orderedtuple(p0, p2)) # add the other device to the graph if p1.device.name not in self.devices: self.add(p1.device) - self.links[_orderedtuple(p1, p2)] = latency + self.links[_orderedtuple(p1, p2)] = attr if self.expand_new_links is not None: self.expand_new_links.append((p1,p2)) def link(self, p0: DevicePort, p1: DevicePort, - latency: str = '0s') -> None: + latency: str = None, attr: dict = None) -> None: """ - Link two DevicePorts with latency if provided. + Link two DevicePorts with optional attributes. Links are bidirectional and the key is a frozenset of the two DevicePorts. Duplicate links (links between the same DevicePorts) are not permitted. Keep in mind that a unique DevicePort is created for each port number in a multi-port style port. If the link types to not match, then throw an exception. Devices that are linked - to will be added to the graph automatically. Latency is expressed - as a string with time units (ps, ns, us...) + to will be added to the graph automatically. + + Args: + p0: First DevicePort to link + p1: Second DevicePort to link + latency: Optional latency string (e.g., '0s', '10ns'). If provided, + added to attr as 'latency' key. For backward compatibility. + attr: Optional dictionary of link attributes. If not provided, + defaults to {'latency': '0s'} or {'latency': latency} if + latency is specified. """ + # Process attributes + if attr is None: + attr = {} + else: + attr = dict(attr) # Make a copy to avoid modifying caller's dict + + # Handle backward compatibility: if latency is provided, add it to attr + if latency is not None: + attr['latency'] = latency + elif 'latency' not in attr: + attr['latency'] = '0s' + if callable(p0) or callable(p1): raise RuntimeError(f"{p0} or {p1} is callable. This probably means" f" you have a multi port and didn't pick a port" @@ -150,7 +172,7 @@ def link(self, p0: DevicePort, p1: DevicePort, p0.link = p1 p1.link = p0 key = _orderedtuple(p0, p1) - self.links[key] = latency + self.links[key] = attr if self.expand_new_links is not None: self.expand_new_links.append(key) diff --git a/src/ahp_graph/SSTGraph.py b/src/ahp_graph/SSTGraph.py index 9b00934..7022b4d 100644 --- a/src/ahp_graph/SSTGraph.py +++ b/src/ahp_graph/SSTGraph.py @@ -198,18 +198,19 @@ def recurseSubcomponents(dev: Device, comp: 'sst.Component') -> None: recurseSubcomponents(d0, c0) # Second, link the component ports using graph links - for ((p0,p1),t) in self.links.items(): + for ((p0,p1),attr) in self.links.items(): if p0.device.library is not None \ and p1.device.library is not None: c0 = n2c[p0.device.name] c1 = n2c[p1.device.name] s0 = p0.get_name() s1 = p1.get_name() + latency_str = attr.get('latency', '0s') if str(p0) < str(p1): - link = sst.Link(f'{p0}__{t}__{p1}') + link = sst.Link(f'{p0}__{latency_str}__{p1}') else: - link = sst.Link(f'{p1}__{t}__{p0}') - latency = t if t != '0s' else '1ps' + link = sst.Link(f'{p1}__{latency_str}__{p0}') + latency = latency_str if latency_str != '0s' else '1ps' link.connect((c0, s0, latency), (c1, s1, latency)) def __write_model(self, @@ -307,17 +308,18 @@ def recurseSubcomponents(dev: Device) -> list: # Now define the links between components. # links = list() - for ((p0,p1),t) in self.links.items(): + for ((p0,p1),attr) in self.links.items(): if p0.device.library is None: raise RuntimeError(f"No SST library: {p0.device.name}") if p1.device.library is None: raise RuntimeError(f"No SST library: {p1.device.name}") - latency = t if t != '0s' else '1ps' + latency_str = attr.get('latency', '0s') + latency = latency_str if latency_str != '0s' else '1ps' if str(p0) < str(p1): - name = f'{p0}__{t}__{p1}' + name = f'{p0}__{latency_str}__{p1}' else: - name = f'{p1}__{t}__{p0}' + name = f'{p1}__{latency_str}__{p0}' d0 = p0.device d1 = p1.device diff --git a/tests/test_DeviceGraph.py b/tests/test_DeviceGraph.py index 01091a2..3de16dc 100644 --- a/tests/test_DeviceGraph.py +++ b/tests/test_DeviceGraph.py @@ -106,7 +106,11 @@ def test_link() -> None: graph.link(lptd.input, ptd0.optional) # type: ignore[arg-type] assert len(graph.devices) == 3, 'linking add submodule parent' assert ltd in graph.devices, 'submodule parent included' - assert graph.links[frozenset({lptd.input, ptd0.optional})] == '0s', 'default latency' # type: ignore[arg-type] + # Find the link key (links are keyed by ordered tuples) + link_key = [k for k in graph.links.keys() if lptd.input in k and ptd0.optional in k][0] + link_attr = graph.links[link_key] + assert isinstance(link_attr, dict), 'link attributes are dict' + assert link_attr['latency'] == '0s', 'default latency' linkAgain = None try: graph.link(ptd0.optional, lptd.input) # type: ignore[arg-type] @@ -136,7 +140,42 @@ def test_link() -> None: changeSinglePortLink = True assert changeSinglePortLink, 'linking from a single port again' graph.link(ptd0.limit(0), ptd1.limit(0), '123ns') # type: ignore[operator] - assert graph.links[frozenset({ptd0.limit(0), ptd1.limit(0)})] == '123ns', 'latency' # type: ignore[operator] + # Find the link key (links are keyed by ordered tuples) + link_key = [k for k in graph.links.keys() if ptd0.limit(0) in k and ptd1.limit(0) in k][0] + link_attr = graph.links[link_key] + assert link_attr['latency'] == '123ns', 'latency from positional param' + + +def test_link_attributes() -> None: + """Test of arbitrary link attributes in a DeviceGraph.""" + graph = DeviceGraph() + + ptd0 = PortTestDevice('0') + ptd1 = PortTestDevice('1') + + # Test link with attr dict including latency + graph.link(ptd0.default, ptd1.default, attr={'latency': '10ns', 'bandwidth': '100GB/s'}) # type: ignore[arg-type] + # Find the link key (links are keyed by ordered tuples of ports) + link_key = [k for k in graph.links.keys() if ptd0.default in k and ptd1.default in k][0] + link_attr = graph.links[link_key] + assert link_attr['latency'] == '10ns', 'latency in attr dict' + assert link_attr['bandwidth'] == '100GB/s', 'custom attribute' + + # Test link with latency param and additional attr + graph.link(ptd0.ptype, ptd1.ptype, latency='5ns', attr={'bandwidth': '50GB/s', 'protocol': 'PCIe'}) # type: ignore[arg-type] + link_key = [k for k in graph.links.keys() if ptd0.ptype in k and ptd1.ptype in k][0] + link_attr = graph.links[link_key] + assert link_attr['latency'] == '5ns', 'latency param overrides attr dict' + assert link_attr['bandwidth'] == '50GB/s', 'custom attribute with latency param' + assert link_attr['protocol'] == 'PCIe', 'multiple custom attributes' + + # Test link with only custom attributes (no latency specified) + graph.link(ptd0.no_limit(0), ptd1.no_limit(0), attr={'bandwidth': '200GB/s', 'protocol': 'NVLink'}) # type: ignore[operator] + link_key = [k for k in graph.links.keys() if ptd0.no_limit(0) in k and ptd1.no_limit(0) in k][0] + link_attr = graph.links[link_key] + assert link_attr['latency'] == '0s', 'default latency when not in attr' + assert link_attr['bandwidth'] == '200GB/s', 'custom attribute without latency' + assert link_attr['protocol'] == 'NVLink', 'protocol attribute' def test_verifyLinks() -> None: