Skip to content

feat: Add progress bars for experiment launch timing #38

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 6 commits into
base: main
Choose a base branch
from
Draft

feat: Add progress bars for experiment launch timing #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9ee1d83
Use progress bars when loading FIREWHEEL experiments
mitchnegus Dec 24, 2024
88c4ed7
Refactor methods for adding/updating image cache
mitchnegus Oct 29, 2024
d2a23e0
Allow MC plugins to use an independent console while FIREWHEEL is lau…
mitchnegus Dec 5, 2024
44cfb8a
Enable model component installation objects to use custom consoles
mitchnegus Jan 21, 2025
7258457
Add a progress bar for negative time
mitchnegus Dec 16, 2024
4c74e77
Add tests for the VM utilities
mitchnegus Dec 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 15 additions & 13 deletions src/firewheel/cli/helpers/experiment
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,19 +235,19 @@ def create_cli_parser():
return parser


def print_output(dependency_time=0.0, total_time=0.0, exp_result=None):
def print_output(console, dependency_time=0.0, total_time=0.0, exp_result=None):
"""
Print the output from this Helper in an easy-to-read format.

Args:
console (rich.console.Console): The console for displaying output.
dependency_time (float): The time it took to generate the dependency graph.
total_time (float): The total time for this Helper to run.
exp_result (list): A list of results from the execution of each MC.

Returns:
int: An exit code indicating how many MC failed, or zero if all succeeded.
"""
console = Console()
caption = f"Dependency resolution took [cyan]{dependency_time:.3f}[/ cyan] seconds"

table = RichDefaultTable(
Expand Down Expand Up @@ -284,13 +284,14 @@ def print_output(dependency_time=0.0, total_time=0.0, exp_result=None):
return exit_code


def run_experiment(mcm, dry_run=False, is_profile=False):
def run_experiment(mcm, console, dry_run=False, is_profile=False):
"""
Execute all model components which have been included within the dependency graph.

Args:
mcm (firewheel.control.model_component_manager.ModelComponentManager): The
ModelComponentManager used to execute the experiment graph.
console (rich.console.Console): The console for displaying output.
dry_run (bool): If this is a dry run. Defaults to False.
is_profile (bool): If the execution should be profiled. Defaults to False.

Expand Down Expand Up @@ -325,11 +326,11 @@ def run_experiment(mcm, dry_run=False, is_profile=False):
exc_str = "".join(
traceback.format_exception(exc_type, exc_value, exc_traceback)
)
Console(theme=cli_output_theme).print(
console.print(
"\nAborting experiment graph processing!"
"\n[error]ERROR:[/error] Invalid arguments to experiment graph plugin."
)
print(f"\nDetails: {exc_str}")
console.print(f"\nDetails: {exc_str}")
sys.exit(errno.EINVAL)
return exp_result

Expand All @@ -339,6 +340,8 @@ def main():
Execute the main logic of this Helper.
"""

console = Console(theme=cli_output_theme)

# Get the start time
total_start = datetime.now()

Expand All @@ -353,11 +356,10 @@ def main():
if vrm_api.get_experiment_launch_time() is not None or mm_api.mm_vms():
active_exp_present = True

# Should we cleanup running experiments?
if cmd_args.force is True or (cmd_args.restart and active_exp_present is True):
print("Cleaning up any running experiments.")
console.print("Cleaning up any running experiments.")
FirewheelCLI().onecmd("restart")
print("All running experiments have been cleared.")
console.print("All running experiments have been cleared.")

# Get the MC names passed in via command line
initial_mc_list = get_mc_list(cmd_args.model_component, install_mcs=cmd_args.no_install)
Expand All @@ -368,9 +370,9 @@ def main():
try:
mcm.build_dependency_graph(initial_mc_list, install_mcs=cmd_args.no_install)
except UnsatisfiableDependenciesError as exp:
print(
"Unable to build model component dependency graph due to the following "
f" error: {exp}\n\nNo Model Components were evaluated and no restart is "
console.print(
"[error]ERROR:[/error] Unable to build model component dependency graph due to the "
f"following error: {exp}\n\nNo Model Components were evaluated and no restart is "
"required. "
)
sys.exit(1)
Expand All @@ -380,14 +382,14 @@ def main():
dependency_time = (de - ds).total_seconds()

# Execute the Model Components (i.e. run the experiment)
exp_result = run_experiment(mcm, cmd_args.dry_run, cmd_args.profile)
exp_result = run_experiment(mcm, console, cmd_args.dry_run, cmd_args.profile)

# Get the total experiment time
total_end = datetime.now()
total_time = (total_end - total_start).total_seconds()

# Print output for the user
exit_code = print_output(dependency_time, total_time, exp_result)
exit_code = print_output(console, dependency_time, total_time, exp_result)

# Exit with the provided exit code
sys.exit(exit_code)
Expand Down
109 changes: 12 additions & 97 deletions src/firewheel/cli/helpers/vm/mix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,110 +33,25 @@ import sys
from time import sleep

from rich.live import Live
from rich.console import Console
from rich.console import Console, Group

import firewheel.vm_resource_manager.api as vm_resource_api
from firewheel.cli.utils import RichDefaultTable
from firewheel.lib.minimega.api import minimegaAPI


def vm_mix():
"""
Create a table showing the mix of VMs in the current experiment.

Returns:
table (firewheel.cli.RichDefaultTable): A table giving the mix of VMs in the experiment.
"""
mm_api = minimegaAPI()
mm_vms = mm_api.mm_vms()

active_exp_present = False
# Check if an experiment is running
if vm_resource_api.get_experiment_launch_time() is not None or mm_api.mm_vms():
active_exp_present = True

mm_state_dict = {}
if active_exp_present:
vm_resource_vms = vm_resource_api.get_vm_states()

for vm_name, vm_dict in mm_vms.items():
mm_state = vm_dict["state"]
vm_resource_state = vm_resource_vms.get(vm_name, "None")

if mm_state not in mm_state_dict:
mm_state_dict[mm_state] = {}
if vm_resource_state not in mm_state_dict[mm_state]:
mm_state_dict[mm_state][vm_resource_state] = []

mm_state_dict[mm_state][vm_resource_state].append(vm_dict["image"])

total_scheduled = 0
rows = []
for mm_state in sorted(mm_state_dict.keys()):
for vm_resource_state in sorted(mm_state_dict[mm_state].keys()):
if vm_resource_state == "None":
continue
os_set = set(mm_state_dict[mm_state][vm_resource_state])
for op_sys in sorted(os_set):
count = mm_state_dict[mm_state][vm_resource_state].count(op_sys)
total_scheduled += count
mm_state_str = mm_state
vm_resource_state_str = vm_resource_state
if "run" in mm_state.lower():
mm_state_str = f"[green]{mm_state}"
if "error" in mm_state.lower():
mm_state_str = f"[red]{mm_state}"
if "configuring" in vm_resource_state.lower():
vm_resource_state_str = f"[yellow]{vm_resource_state}"
if "configured" in vm_resource_state.lower():
vm_resource_state_str = f"[green]{vm_resource_state}"
if "uninitialized" in vm_resource_state.lower():
vm_resource_state_str = f"[yellow]{vm_resource_state}"
if "error" in vm_resource_state.lower():
vm_resource_state_str = f"[red]{vm_resource_state}"
rows.append([op_sys, mm_state_str, vm_resource_state_str, str(count)])

table = build_table(rows, total_scheduled)
return table


def build_table(rows, total_scheduled):
"""
Construct a table given rows of VM information.

Args:
rows (list): A list of information per set of VMs (each list element
representing one row in the output VM mix table).
total_scheduled (int): The total number of VMs scheduled in the experiment.

Returns:
firewheel.cli.utils.RichDefaultTable: A table giving the mix of VMs in the experiment.
"""
table = RichDefaultTable(
title="VM Mix",
show_footer=True,
)
table.add_column("VM Image")
table.add_column("Power State")
table.add_column("VM Resource State", footer="[b]Total Scheduled")
table.add_column("Count", footer=f"[b]{total_scheduled}")
# Add custom rows to the table
if not rows:
rows.append(["N/A", "N/A", "N/A", "0"])
for row in rows:
table.add_row(*row)
return table
from firewheel.cli.vm_utils import VMMixTable, ScheduleProgress


def main():
"""
Provide the primary update functionality for VM Mix.
"""
"""Provide the primary update functionality for VM Mix."""
console = Console()
with Live(vm_mix(), console=console, screen=False, refresh_per_second=1) as live:
mix_table = VMMixTable()
schedule_progress = ScheduleProgress(console=console)
negative_time_progress = schedule_progress.add_negative_time_task()

display_group = Group(mix_table, schedule_progress)
with Live(display_group, console=console, screen=False, refresh_per_second=10) as live:
while True:
display_group.renderables[0] = VMMixTable()
schedule_progress.update_negative_time_task()
live.update(display_group)
sleep(2)
live.update(vm_mix())


if __name__ == "__main__":
Expand Down
Loading
Loading