Monitor subprocess

codecarbon/cli/main.py

@@ -1,14 +1,14 @@
+import shlex
 import sys
-import time
 
 import click
 
-from codecarbon import EmissionsTracker
 from codecarbon.cli.cli_utils import (
     get_api_endpoint,
     get_existing_local_exp_id,
     write_local_exp_id,
 )
+from codecarbon.cli.monitor import monitor_infinite_loop, monitor_subprocess
 from codecarbon.core.api_client import ApiClient, get_datetime_with_timezone
 from codecarbon.core.schemas import ExperimentCreate
 
@@ -60,7 +60,23 @@ def init():
 
 
 @codecarbon.command(
-    "monitor", short_help="Run an infinite loop to monitor this machine."
+    "monitor",
+    short_help="Run an infinite loop to monitor this machine.",
+    help=f"""
+    Monitor the environmental impact of programs.
+
+    This command has two usages:
+
+    Examples:
+
+    \b
+    # Monitor the system until you press CTRL-C
+    {sys.argv[0]} monitor
+
+    \b
+    # Run and monitor a subprocess called "recon-all" and forward arguments to "recon-all":
+    {sys.argv[0]} monitor -- recon-all -s sub-101 -i sub-101_ses-BL_T1w.nii.gz -all
+    """,
 )
 @click.option(
     "--measure_power_secs", default=10, help="Interval between two measures. (10)"
@@ -73,17 +89,18 @@ def init():
 @click.option(
     "--api/--no-api", default=True, help="Choose to call Code Carbon API or not. (yes)"
 )
-def monitor(measure_power_secs, api_call_interval, api):
+@click.argument("cmd", nargs=-1)
+def monitor(measure_power_secs, api_call_interval, api, cmd):
     experiment_id = get_existing_local_exp_id()
     if api and experiment_id is None:
         click.echo("ERROR: No experiment id, call 'codecarbon init' first.")
         sys.exit(1)
-    click.echo("CodeCarbon is going in an infinite loop to monitor this machine.")
-    with EmissionsTracker(
-        measure_power_secs=measure_power_secs,
-        api_call_interval=api_call_interval,
-        save_to_api=api,
-    ):
-        # Infinite loop
-        while True:
-            time.sleep(300)
+
+    if cmd:
+        click.echo("CodeCarbon is going to run and monitor this command:")
+        click.echo(f"\n\t{shlex.join(cmd)}\n")
+        rc = monitor_subprocess(measure_power_secs, api_call_interval, api, cmd)
+        sys.exit(rc)
+    else:
+        click.echo("CodeCarbon is going in an infinite loop to monitor this machine.")
+        monitor_infinite_loop(measure_power_secs, api_call_interval, api)

codecarbon/cli/monitor.py

@@ -0,0 +1,42 @@
+"""
+Implementations of the ``codecarbon monitor`` subcommand.
+"""
+import os
+import subprocess as sp
+import time
+from typing import Sequence
+
+from codecarbon import EmissionsTracker
+
+
+def monitor_infinite_loop(measure_power_secs, api_call_interval, api):
+    """
+    Monitor all activity on the system until a user presses CTRL-C.
+    """
+    with EmissionsTracker(
+        measure_power_secs=measure_power_secs,
+        api_call_interval=api_call_interval,
+        save_to_api=api,
+        tracking_mode="machine",
+    ):
+        while True:
+            time.sleep(300)
+
+
+def monitor_subprocess(
+    measure_power_secs, api_call_interval, api, cmd: Sequence[str | os.PathLike]
+) -> int:
+    """
+    Run and monitor a subprocess.
+
+    :return: return code of the subprocess.
+    """
+
+    with EmissionsTracker(
+        measure_power_secs=measure_power_secs,
+        api_call_interval=api_call_interval,
+        save_to_api=api,
+        tracking_mode="process",
+    ):
+        proc = sp.run(cmd)
+    return proc.returncode