diff --git a/src/infragraph/blueprints/fabrics/dragonfly_fabric.py b/src/infragraph/blueprints/fabrics/dragonfly_fabric.py new file mode 100644 index 0000000..9474030 --- /dev/null +++ b/src/infragraph/blueprints/fabrics/dragonfly_fabric.py @@ -0,0 +1,172 @@ +from infragraph import * +from infragraph.infragraph_service import InfraGraphService +from infragraph.infragraph import * + + +class DragonflyFabric(Infrastructure): + """Return a balanced Dragonfly fabric. + + A balanced Dragonfly satisfies a == 2p == 2h, where: + - a = routers per group + - p = host/terminal ports per router + - h = global/inter-group ports per router + Total groups = a*h + 1 (maximum supported by a balanced Dragonfly). + + Each router uses p + (a-1) + h ports: + ports [0, p) -> hosts + ports [p, p + a - 1) -> intra-group peers + ports [p + a - 1, p + a - 1 + h) -> global (inter-group) + + Inputs: + - switch + - host + - a, p, h + - bandwidth + """ + + def __init__( + self, + switch: Device, + host: Device, + a: int, + p: int, + h: int, + bandwidth=None, + ): + super().__init__( + name="dragonfly-fabric", + description="Balanced Dragonfly Fabric", + ) + + # balanced dragonfly constraint + if not (a == 2 * p == 2 * h): + raise ValueError( + f"Balanced Dragonfly requires a == 2p == 2h; got a={a}, p={p}, h={h}" + ) + + switch_port_component = InfraGraphService.get_component(switch, Component.PORT) + host_nic_component = InfraGraphService.get_component(host, Component.NIC) + + ports_per_router = p + (a - 1) + h + if switch_port_component.count < ports_per_router: + raise ValueError( + f"Switch needs at least {ports_per_router} ports " + f"(p + (a-1) + h); got {switch_port_component.count}" + ) + + if p % host_nic_component.count != 0: + raise ValueError( + f"Host NIC count ({host_nic_component.count}) does not divide " + f"terminal ports per router (p={p}) evenly" + ) + + # bandwidth defaults: [host_switch, intra, inter] + if bandwidth is None: + bandwidth = [100, 100, 100] + assert len(bandwidth) == 3, "bandwidth list must have 3 entries" + + # counts + total_groups = a * h + 1 + routers_per_group = a + total_switches = total_groups * routers_per_group + hosts_per_router = p // host_nic_component.count + total_hosts = total_groups * routers_per_group * hosts_per_router + + # port-range bases on each router + intra_port_base = p + global_port_base = p + (a - 1) + + # devices + self.devices.append(host).append(switch) + host_instance = self.instances.add( + name=host.name, device=host.name, count=total_hosts + ) + switch_instance = self.instances.add( + name="router", device=switch.name, count=total_switches + ) + + # links + host_switch_link = self.links.add( + name="host_switch_link", + description="Link between host and router (terminal)", + ) + host_switch_link.physical.bandwidth.gigabits_per_second = bandwidth[0] + + intra_switch_link = self.links.add( + name="intra_switch_link", + description="Link between routers within the same group", + ) + intra_switch_link.physical.bandwidth.gigabits_per_second = bandwidth[1] + + inter_switch_link = self.links.add( + name="inter_switch_link", + description="Link between routers in different groups (global)", + ) + inter_switch_link.physical.bandwidth.gigabits_per_second = bandwidth[2] + + # host and router + for group_idx in range(total_groups): + for local_router in range(routers_per_group): + router_idx = group_idx * routers_per_group + local_router + terminal_port = 0 + for host_in_rack in range(hosts_per_router): + host_idx = router_idx * hosts_per_router + host_in_rack + for nic_idx in range(host_nic_component.count): + edge = self.edges.add( + scheme=InfrastructureEdge.ONE2ONE, + link=host_switch_link.name, + ) + edge.ep1.instance = f"{host_instance.name}[{host_idx}]" + edge.ep1.component = f"{host_nic_component.name}[{nic_idx}]" + edge.ep2.instance = f"{switch_instance.name}[{router_idx}]" + edge.ep2.component = f"{switch_port_component.name}[{terminal_port}]" + terminal_port += 1 + + # router to router intra-group + for group_idx in range(total_groups): + base = group_idx * routers_per_group + intra_cursor = [intra_port_base] * routers_per_group + for i in range(routers_per_group): + for j in range(i + 1, routers_per_group): + edge = self.edges.add( + scheme=InfrastructureEdge.ONE2ONE, + link=intra_switch_link.name, + ) + edge.ep1.instance = f"{switch_instance.name}[{base + i}]" + edge.ep1.component = f"{switch_port_component.name}[{intra_cursor[i]}]" + edge.ep2.instance = f"{switch_instance.name}[{base + j}]" + edge.ep2.component = f"{switch_port_component.name}[{intra_cursor[j]}]" + intra_cursor[i] += 1 + intra_cursor[j] += 1 + + # router to router inter-group + next_router_in_group = [0] * total_groups + next_port_on_router = [0] * total_groups + + for gi in range(total_groups): + for gj in range(gi + 1, total_groups): + ri = gi * routers_per_group + next_router_in_group[gi] + rj = gj * routers_per_group + next_router_in_group[gj] + pi = global_port_base + next_port_on_router[gi] + pj = global_port_base + next_port_on_router[gj] + + edge = self.edges.add( + scheme=InfrastructureEdge.ONE2ONE, + link=inter_switch_link.name, + ) + edge.ep1.instance = f"{switch_instance.name}[{ri}]" + edge.ep1.component = f"{switch_port_component.name}[{pi}]" + edge.ep2.instance = f"{switch_instance.name}[{rj}]" + edge.ep2.component = f"{switch_port_component.name}[{pj}]" + + # if filled all h ports, move to next router + next_port_on_router[gi] += 1 + if next_port_on_router[gi] == h: + next_port_on_router[gi] = 0 + next_router_in_group[gi] += 1 + + # same for gj + next_port_on_router[gj] += 1 + if next_port_on_router[gj] == h: + next_port_on_router[gj] = 0 + next_router_in_group[gj] += 1 \ No newline at end of file diff --git a/src/tests/test_blueprints/test_devices.py b/src/tests/test_blueprints/test_devices.py index 98d065f..72f8da2 100644 --- a/src/tests/test_blueprints/test_devices.py +++ b/src/tests/test_blueprints/test_devices.py @@ -7,6 +7,7 @@ from infragraph.blueprints.devices.ironwood_rack import IronwoodRack from infragraph.blueprints.devices.generic.server import Server from infragraph.blueprints.devices.generic.generic_switch import Switch +from infragraph.blueprints.devices.amd.mi300x import MI300X @pytest.mark.asyncio @@ -19,6 +20,7 @@ Cx5(), NvidiaDGX(), IronwoodRack(), + MI300X() ], ) async def test_devices(count, device): diff --git a/src/tests/test_blueprints/test_dragonfly.py b/src/tests/test_blueprints/test_dragonfly.py new file mode 100644 index 0000000..71971b7 --- /dev/null +++ b/src/tests/test_blueprints/test_dragonfly.py @@ -0,0 +1,107 @@ +import pytest +import networkx +from infragraph.infragraph_service import InfraGraphService +from infragraph.blueprints.devices.generic.generic_switch import Switch +from infragraph.blueprints.devices.generic.server import Server +from infragraph.blueprints.devices.nvidia.dgx import NvidiaDGX +from infragraph.blueprints.devices.amd.mi300x import MI300X +from infragraph.blueprints.fabrics.dragonfly_fabric import DragonflyFabric +from infragraph.visualizer.visualize import run_visualizer + +def print_graph(graph): + # for node, attrs in graph.nodes(data=True): + # print(f"Node: {node}, Attributes: {attrs}") + + # for u, v, attrs in graph.edges(data=True): + # print(f"Edge: ({u}, {v}), Attributes: {attrs}") + pass + +def dump_yaml(clos_fabric, filename): + # import yaml + # with open(filename + ".yaml", "w") as file: + # data = clos_fabric.serialize("dict") + # yaml.dump(data, file, default_flow_style=False, indent=4) + pass + +@pytest.mark.asyncio +async def test_dragonfly_server(): + """ + Generate and validate the dragonfly topology with server as host + """ + server=Server() + switch=Switch(port_count=16) + dragonfly_fabric=DragonflyFabric(switch, server, 4, 2, 2) + assert len(dragonfly_fabric.instances) == 2 + for instance in dragonfly_fabric.instances: + if instance.name == "server": + assert instance.count == 36 + elif instance.name == "router": + assert instance.count == 36 + service = InfraGraphService() + service.set_graph(dragonfly_fabric) + # validations + g = service.get_networkx_graph() + for node, attrs in g.nodes(data=True): + assert attrs, f"Node {node} has empty attributes" + for u, v, attrs in g.edges(data=True): + assert attrs, f"Edge ({u}, {v}) has empty attributes" + #print_graph(g) + #print(networkx.write_network_text(g, vertical_chains=True)) + +@pytest.mark.asyncio +async def test_dragonfly_dgx(): + """ + Generate and validate the dragonfly topology with dgx as host + """ + dgx=NvidiaDGX("dgx_a100") + switch=Switch(port_count=32) + dragonfly_fabric= DragonflyFabric(switch, dgx, 16, 8, 8) + assert len(dragonfly_fabric.instances) == 2 + for instance in dragonfly_fabric.instances: + if instance.name == "dgx_a100": + assert instance.count == 2064 + elif instance.name == "router": + assert instance.count == 2064 + service = InfraGraphService() + service.set_graph(dragonfly_fabric) + # validations + g = service.get_networkx_graph() + for node, attrs in g.nodes(data=True): + assert attrs, f"Node {node} has empty attributes" + for u, v, attrs in g.edges(data=True): + assert attrs, f"Edge ({u}, {v}) has empty attributes" + #dump_yaml(dragonfly_fabric, "dragonfly_fabric_with_dgx") + #print_graph(g) + #print(networkx.write_network_text(g, vertical_chains=True)) + #run_visualizer(infrastructure=dragonfly_fabric, output="/mnt/c/Users/anusghos/git/demo/dragonfly_fabric_with_dgx", + # hosts="dgx", switches="switch") + +@pytest.mark.asyncio +async def test_dragonfly_mi300x(): + """ + Generate and validate the dragonfly topology with mi300x amd device as host + """ + mi300x= MI300X() + switch=Switch(port_count=32) + dragonfly_fabric=DragonflyFabric(switch, mi300x, 16, 8, 8) + assert len(dragonfly_fabric.instances) == 2 + for instance in dragonfly_fabric.instances: + if instance.name == "dgx_a100": + assert instance.count == 2064 + elif instance.name == "router": + assert instance.count == 2064 + service=InfraGraphService() + service.set_graph(dragonfly_fabric) + # validations + g = service.get_networkx_graph() + for node, attrs in g.nodes(data=True): + assert attrs, f"Node {node} has empty attributes" + for u, v, attrs in g.edges(data=True): + assert attrs, f"Edge ({u}, {v}) has empty attributes" + #print_graph(g) + #print(networkx.write_network_text(g, vertical_chains=True) + #run_visualizer(infrastructure=dragonfly_fabric, output="/mnt/c/Users/anusghos/git/demo/dragonfly_fabric_with_mi300x", hosts="mi300x", switches="switch") + + +if __name__ == "__main__": + pytest.main(["-s", __file__]) \ No newline at end of file diff --git a/src/tests/test_blueprints/test_mi300x.py b/src/tests/test_blueprints/test_mi300x.py deleted file mode 100644 index dc82f26..0000000 --- a/src/tests/test_blueprints/test_mi300x.py +++ /dev/null @@ -1,34 +0,0 @@ -import pytest -import networkx -from infragraph.infragraph import Api -from infragraph.infragraph_service import InfraGraphService -from infragraph.blueprints.devices.amd.mi300x import MI300X - - -@pytest.mark.asyncio -@pytest.mark.parametrize("count", [1, 2]) -@pytest.mark.parametrize( - "device", - [ - MI300X(), - ], -) -async def test_mi300x(count, device): - """From an MI300X device, generate a graph and validate the graph. - - - with a count > 1 there should be no connectivity between device instances - """ - device.validate() - infrastructure = Api().infrastructure() - infrastructure.devices.append(device) - infrastructure.instances.add(name=device.name, device=device.name, count=count) - service = InfraGraphService() - service.set_graph(infrastructure) - - g = service.get_networkx_graph() - print(f"\ndevice {device.name} is a {g}") - print(networkx.write_network_text(g, vertical_chains=True)) - - -if __name__ == "__main__": - pytest.main(["-s", __file__]) \ No newline at end of file