Generalize for N partitions, specified on commmand line. See multi-part-sim.py

choices = ['random', 'benchmark', 'peak', 'idle']

parser.add_argument('-w', '--workload', type=str, choices=choices, default=choices[0], help='Type of synthetic workload')

choices = ['layout1', 'layout2']

parser.add_argument('-x', '--partitions', nargs='+', default=None, help='List of machine configurations to use, e.g., -x setonix-cpu setonix-gpu')

parser.add_argument('--layout', type=str, choices=choices, default=choices[0], help='Layout of UI')

args = parser.parse_args()

args_dict = vars(args)

if args.verbose: print(jobs)

# Create generator objects for both partitions

gen1 = layout_manager1.run_nonblocking(jobs1, timesteps=timesteps)

gen2 = layout_manager2.run_nonblocking(jobs2, timesteps=timesteps)

gen1 = layout_manager1.run_stepwise(jobs1, timesteps=timesteps)

gen2 = layout_manager2.run_stepwise(jobs2, timesteps=timesteps)

# Step through both generators in lockstep

for _ in range(timesteps):

    next(gen1)  # Advance first scheduler

    next(gen2)  # Advance second scheduler

#for _ in range(timesteps):

#    next(gen1)  # Advance first scheduler

#    next(gen2)  # Advance second scheduler

for timestep in range(timesteps):

    # Advance generators

    next(gen1)

    next(gen2)

    # Timestep

    print(f"[DEBUG] Timestep: {timestep}")

    # Queue lengths

    print(f"[DEBUG] setonix-cpu Queue: {len(layout_manager1.scheduler.queue)}")

    print(f"[DEBUG] setonix-gpu Queue: {len(layout_manager2.scheduler.queue)}")

    # System utilization

    sys_util1 = layout_manager1.scheduler.sys_util_history[-1][1] if layout_manager1.scheduler.sys_util_history else 0.0

    sys_util2 = layout_manager2.scheduler.sys_util_history[-1][1] if layout_manager2.scheduler.sys_util_history else 0.0

    print(f"[DEBUG] setonix-cpu Util: {sys_util1:.2f}%")

    print(f"[DEBUG] setonix-gpu Util: {sys_util2:.2f}%")

from raps.helpers import check_python_version

check_python_version()

import random

import sys

from args import args

from raps.config import ConfigManager

from raps.policy import PolicyType

from raps.ui import LayoutManager

from raps.scheduler import Scheduler

from raps.flops import FLOPSManager

from raps.power import PowerManager, compute_node_power

from raps.workload import Workload

from raps.utils import convert_to_seconds

# Load configurations for each partition

partition_names = args.partitions

configs = [ConfigManager(system_name=partition).get_config() for partition in partition_names]

args_dicts = [{**vars(args), 'config': config} for config in configs]

# Initialize Workload with all configurations

wl = Workload(*configs)

# Generate jobs based on workload type

jobs = getattr(wl, args.workload)(num_jobs=args.numjobs)

# Group jobs by partition

jobs_by_partition = {partition: [] for partition in wl.partitions}

for job in jobs:

    jobs_by_partition[job['partition']].append(job)

# Initialize layout managers for each partition

layout_managers = {}

for i, config in enumerate(configs):

    pm = PowerManager(compute_node_power, **configs[i])

    fm = FLOPSManager(**args_dicts[i])

    scheduler = Scheduler(power_manager=pm, flops_manager=fm, cooling_model=None, **args_dicts[i])

    layout_managers[config['system_name']] = LayoutManager(args.layout, scheduler=scheduler, debug=args.debug, **config)

# Set simulation timesteps

if args.time:

    timesteps = convert_to_seconds(args.time)

else:

    timesteps = 88200  # Default to 24 hours

# Create generators for each layout manager

generators = {name: lm.run_stepwise(jobs_by_partition[name], timesteps=timesteps)

              for name, lm in layout_managers.items()}

# Step through all generators in lockstep

for timestep in range(timesteps):

    for name, gen in generators.items():

        next(gen)  # Advance each generator

    # Print debug info every UI_UPDATE_FREQ

    if timestep % configs[0]['UI_UPDATE_FREQ'] == 0:  # Assuming same frequency for all partitions

        for name, lm in layout_managers.items():

            sys_util = lm.scheduler.sys_util_history[-1] if lm.scheduler.sys_util_history else 0.0

            print(f"[DEBUG] {name} - Timestep {timestep} - Jobs in queue: {len(lm.scheduler.queue)} - Utilization: {sys_util[1]:.2f}%")

print("Simulation complete.")

        self.timesteps = timesteps

        if self.debug: 

            limits = self.get_gauge_limits()

            print(f"[DEBUG] Limits for {self.config['system_name']}:", limits)

        for timestep in range(timesteps):

            # Print the current timestep for this partition

            if timestep % self.config['UI_UPDATE_FREQ'] == 0:

                sys_util = self.sys_util_history[-1][1] if self.sys_util_history else 0

                print(f"[DEBUG] {self.config['system_name']} - Timestep {timestep} - Jobs in queue: {len(self.queue)} - Utilization: {sys_util:.1f}% - Power: {self.sys_power:.1f} kW")

            while self.current_time >= last_submit_time and jobs:

                job = jobs.pop(0)

                print(f"[DEBUG] {self.config['system_name']} - Scheduling job {job['id']} at timestep {timestep}")

                self.schedule([job])

                if jobs:

                    last_submit_time = job['submit_time']

                self.update(data)

                self.render()

    def run_nonblocking(self, jobs, timesteps):

    def run_stepwise(self, jobs, timesteps):

        """ Prepares the UI and returns a generator for the simulation """

        return self.scheduler.run_simulation(jobs, timesteps)