Commit ad2b93e7 authored by Brewer, Wes's avatar Brewer, Wes
Browse files

Refactor network plotting to use NetworkModel.plot_topology method

parent 3194fa6a

raps/engine.py

+1 −3
Original line number Diff line number Diff line
@@ -330,9 +330,7 @@ class Engine: 
            available_nodes = self.resource_manager.available_nodes

            self.network_model = NetworkModel(

                available_nodes=available_nodes,

                config=self.config,

                plot=self.sim_config.plot,

                output_dir=self.sim_config.get_output(),

                config=self.config

            )

        else:

            self.network_model = None

raps/network/__init__.py

+21 −16
Original line number Diff line number Diff line 
import os

import warnings



from .base import (

    all_to_all_paths,
@@ -44,9 +45,6 @@ __all__ = [ 
class NetworkModel:

    def __init__(self, *, available_nodes, config, **kwargs):

        self.config = config

        self.output_dir = kwargs.get('output_dir')

        if self.output_dir:

            self.output_dir.mkdir(parents=True, exist_ok=True)

        self.topology = config.get("TOPOLOGY")

        self.max_link_bw = config.get("NETWORK_MAX_BW", 1e9)  # default safeguard

        self.real_to_fat_idx = kwargs.get("real_to_fat_idx", {})
@@ -57,9 +55,6 @@ class NetworkModel: 
            self.net_graph = build_fattree(self.fattree_k, total_nodes)

            # TODO: future testing of subsampling feature

            #self.net_graph = subsample_hosts(self.net_graph, num_hosts=4626)

            if self.output_dir:

                save_path = os.path.join(self.output_dir, "net-fat-tree.png")

                plot_fattree_hierarchy(self.net_graph, k=self.fattree_k, save_path=save_path)



        elif self.topology == "torus3d":

            dims = (
@@ -73,12 +68,6 @@ class NetworkModel: 
            # Build the graph and metadata

            self.net_graph, self.meta = build_torus3d(dims, wrap, hosts_per_router=hosts_per_router)



            if self.output_dir:

                save_path = os.path.join(self.output_dir, "net-torus2d.png")

                plot_torus2d(self.net_graph, save_path=save_path)

                save_path = os.path.join(self.output_dir, "net-torus3d.png")

                plot_torus3d(self.net_graph, save_path=save_path)



            # Deterministic numeric → host mapping

            X, Y, Z = self.meta["dims"]

            self.id_to_host = {}
@@ -109,10 +98,6 @@ class NetworkModel: 
            self.real_to_fat_idx = build_dragonfly_idx_map(D, A, P, total_real_nodes)

            print(f"[DEBUG] Dragonfly mapping: {len(self.real_to_fat_idx)} entries")



            if self.output_dir:

                save_path = os.path.join(self.output_dir, "net-dragonfly.png")

                plot_dragonfly(self.net_graph, save_path=save_path)



        elif self.topology == "capacity":

            # Capacity-only model: no explicit graph

            self.net_graph = None
@@ -174,3 +159,23 @@ class NetworkModel: 
            raise ValueError(f"Unsupported topology: {self.topology}")



        return net_util, net_cong, net_tx, net_rx, max_throughput



    def plot_topology(self, output_dir):

        """Plot network topology - save as png file in output_dir."""

        if output_dir:

            if self.topology == "fat-tree":

                save_path = output_dir / "net-fat-tree.png"

                plot_fattree_hierarchy(self.net_graph, k=self.fattree_k, save_path=save_path)

            elif self.topology == "dragonfly":

                save_path = output_dir / "net-dragonfly.png"

                plot_dragonfly(self.net_graph, save_path=save_path)

            elif self.topology == "torus3d":

                save_path = output_dir / "net-torus2d.png"

                plot_torus2d(self.net_graph, save_path=save_path)

                save_path = output_dir / "net-torus3d.png"

                plot_torus3d(self.net_graph, save_path=save_path)

            else:

                warnings.warn(

                    f"plotting not supported for {self.topology} topology",

                    UserWarning

                )

raps/run_sim.py

+3 −0
Original line number Diff line number Diff line
@@ -137,6 +137,9 @@ def run_sim(sim_config: SingleSimConfig): 
            else:

                print('Cooling model not enabled... skipping output of plot')



        if 'net' in sim_config.plot:

            engine.network_model.plot_topology(out)



        if 'temp' in sim_config.plot:

            if engine.cooling_model:

                ylabel = 'Tr_pri_Out[1]'

raps/sim_config.py

+1 −1
Original line number Diff line number Diff line
@@ -327,7 +327,7 @@ class SimConfig(RAPSBaseModel, abc.ABC): 
        if self.jobsize_is_power_of is not None and self.jobsize_is_of_degree is not None:

            raise ValueError("jobsize_is_power_of and jobsize_is_of_degree are mutually exclusive")



        if self.plot and not self.output:

        if self.plot and self.output == "none":

            raise ValueError("plot requires an output directory to be set")



        if self.live and not self.replay and self.time is None: