diff --git a/codecarbon/cli/main.py b/codecarbon/cli/main.py index 9031703f6..180c09e86 100644 --- a/codecarbon/cli/main.py +++ b/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) diff --git a/codecarbon/cli/monitor.py b/codecarbon/cli/monitor.py new file mode 100644 index 000000000..77b478d72 --- /dev/null +++ b/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