drangonfly fabric blueprint

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

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):

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__])