Subcommands

zeus/monitor/__main__.py

@@ -1,8 +1,8 @@
-
 from __future__ import annotations
 
 import time
 from typing import Optional
+from contextlib import suppress
 
 import rich
 import typer
@@ -14,12 +14,7 @@
 
 @app.command()
 def energy(gpu_indices: Optional[list[int]] = None) -> None:
-    """Measure the time and energy of GPUs during the duration of the CLI program.
-
-    This uses the `ZeusMonitor` class for measurement, ane thus `gpu_indices` respect
-    the `CUDA_VISIBLE_DEVICES` environment variable.
-    For instance, if `CUDA_VISIBLE_DEVICES=2,3`, GPU index `1` passed into `gpu_indices`
-    will be interpreted as CUDA device `3`.
+    """Measure the time and energy of GPUs using the ZeusMonitor.
 
     Args:
         gpu_indices: Indices of GPUs to monitor. Not ommitted, all GPUs will be monitored.
@@ -34,15 +29,25 @@ def energy(gpu_indices: Optional[list[int]] = None) -> None:
 
 
 @app.command()
-def power(gpu_indices: Optional[list[int]] = None) -> None:
+def power(
+    gpu_indices: Optional[list[int]] = None, sleep_ms: int = 100, wait: bool = False
+) -> None:
     """Monitor the power consumption of GPUs during the duration of the CLI program.
 
     Args:
         gpu_indices: Indices of GPUs to monitor. Not ommitted, all GPUs will be monitored.
+        sleep_ms: The number of milliseconds to sleep after each measurement.
+        wait: If True, the program will wait for the user to press Enter before starting.
     """
-    while True:
-        print("Hi")
-        time.sleep(1)
+    try:
+        if wait:
+            input("Press Enter to start monitoring power consumption.")
+        while True:
+            print("Hi")
+            time.sleep(sleep_ms / 1000)
+    except KeyboardInterrupt:
+        raise typer.Exit()
+
 
 if __name__ == "__main__":
     app()