From 11f3f20a5e8137433099545b24e2b8cf5542151d Mon Sep 17 00:00:00 2001 From: davidberenstein1957 Date: Sun, 21 Jun 2026 15:05:13 +0100 Subject: [PATCH] perf: speed up CLI monitor startup and fix wrapped commands Lazy-import heavy CLI dependencies, add --log-level to monitor, and strip nested monitor/-- tokens when delegating to wrapped user commands. Co-authored-by: Cursor --- codecarbon/cli/main.py | 40 ++++++++++++++++++---- codecarbon/cli/monitor.py | 9 ++--- docs/reference/cli.md | 2 +- tests/cli/test_cli.py | 17 +++++++-- tests/cli/test_cli_main.py | 70 +++++++++++++++++++++++++++----------- tests/cli/test_monitor.py | 44 +++++++++++++++++++----- 6 files changed, 139 insertions(+), 43 deletions(-) diff --git a/codecarbon/cli/main.py b/codecarbon/cli/main.py index fd6545a3f..c10b32338 100644 --- a/codecarbon/cli/main.py +++ b/codecarbon/cli/main.py @@ -5,15 +5,12 @@ from pathlib import Path from typing import Optional -import questionary -import requests import typer from rich import print from rich.prompt import Confirm from typing_extensions import Annotated from codecarbon import __app_name__, __version__ -from codecarbon.cli.auth import authorize, get_access_token from codecarbon.cli.cli_utils import ( create_new_config_file, get_api_endpoint, @@ -21,10 +18,6 @@ get_existing_exp_id, overwrite_local_config, ) -from codecarbon.cli.monitor import run_and_monitor -from codecarbon.core.api_client import ApiClient, get_datetime_with_timezone -from codecarbon.core.schemas import ExperimentCreate, OrganizationCreate, ProjectCreate -from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker API_URL = os.environ.get("API_URL", "https://dashboard.codecarbon.io/api") @@ -68,6 +61,9 @@ def version( def show_config(path: Path = Path("./.codecarbon.config")) -> None: + from codecarbon.cli.auth import get_access_token + from codecarbon.core.api_client import ApiClient + d = get_config(path) print("Current configuration : \n") print("Config file content : ") @@ -114,6 +110,9 @@ def api_get(): """ ex: test-api """ + from codecarbon.cli.auth import get_access_token + from codecarbon.core.api_client import ApiClient + api_endpoint = get_api_endpoint() api = ApiClient(endpoint_url=api_endpoint) api.set_access_token(get_access_token()) @@ -123,6 +122,9 @@ def api_get(): @codecarbon.command("login", short_help="Login to CodeCarbon") def login(): + from codecarbon.cli.auth import authorize, get_access_token + from codecarbon.core.api_client import ApiClient + authorize() api_endpoint = get_api_endpoint() api = ApiClient(endpoint_url=api_endpoint) @@ -132,6 +134,10 @@ def login(): def get_api_key(project_id: str): + import requests + + from codecarbon.cli.auth import get_access_token + api_endpoint = get_api_endpoint() api_endpoint = api_endpoint.rstrip("/") req = requests.post( @@ -161,6 +167,13 @@ def config(): """ Initialize CodeCarbon, this will prompt you for configuration of Organisation/Team/Project/Experiment. """ + from codecarbon.cli.auth import get_access_token + from codecarbon.core.api_client import ApiClient, get_datetime_with_timezone + from codecarbon.core.schemas import ( + ExperimentCreate, + OrganizationCreate, + ProjectCreate, + ) print("Welcome to CodeCarbon configuration wizard") home = Path.home() @@ -342,6 +355,10 @@ def monitor( str, typer.Option(help="Region/province for offline mode"), ] = None, + log_level: Annotated[ + str, + typer.Option(help="Log level (critical, error, warning, info, debug)"), + ] = "error", ): """Monitor your machine's carbon emissions.""" @@ -349,6 +366,7 @@ def monitor( tracker_args = { "measure_power_secs": measure_power_secs, "api_call_interval": api_call_interval, + "log_level": log_level, } # Set up the tracker arguments based on mode (offline vs online) and validate required args for each mode if offline: @@ -375,8 +393,12 @@ def monitor( tracker_args = {**tracker_args, "save_to_api": api} + from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker + # If extra args are provided (e.g. `codecarbon monitor -- my_script.py`), delegate to `run_and_monitor` if getattr(ctx, "args", None): + from codecarbon.cli.monitor import run_and_monitor + return run_and_monitor(ctx, offline=offline, **tracker_args) # Instantiate the tracker @@ -417,6 +439,8 @@ def detect(): """ Detects hardware and prints information without running any measurements. """ + from codecarbon.emissions_tracker import EmissionsTracker + print("Detecting hardware...") tracker = EmissionsTracker(save_to_file=False) hardware_info = tracker.get_detected_hardware() @@ -438,6 +462,8 @@ def detect(): def questionary_prompt(prompt, list_options, default): + import questionary + value = questionary.select( prompt, list_options, diff --git a/codecarbon/cli/monitor.py b/codecarbon/cli/monitor.py index 98fa4e244..41b3ca353 100644 --- a/codecarbon/cli/monitor.py +++ b/codecarbon/cli/monitor.py @@ -8,8 +8,6 @@ from rich import print from typing_extensions import Annotated -from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker - def run_and_monitor( ctx: typer.Context, @@ -50,12 +48,15 @@ def run_and_monitor( directory. The file path is shown in the final report. """ # Suppress all CodeCarbon logs during execution + from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker from codecarbon.external.logger import set_logger_level set_logger_level(log_level) - # Get the command from remaining args - command = ctx.args + # Get the command from remaining args (strip nested subcommand / `--` leftovers) + command = list(getattr(ctx, "args", None) or []) + while command and command[0] in ("monitor", "--"): + command.pop(0) if not command: print( diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 1e93a588f..71ba05a85 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -44,7 +44,7 @@ Displays real-time emissions data for all processes on your machine. Press `Ctrl | `--no-api` | flag | false | Do not send data to the API (local-only measurement) | | `--offline` | flag | false | Run without internet access | | `--country-iso-code` | string | - | ISO 3166-1 alpha-3 country code (required in offline mode) | -| `--log-level` | choice | INFO | Log level: DEBUG, INFO, WARNING, ERROR | +| `--log-level` | choice | ERROR | Log level: DEBUG, INFO, WARNING, ERROR | **Examples:** ```bash diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 0935cc069..c4f990b4d 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -11,7 +11,7 @@ # MOCK API CLIENT -@patch("codecarbon.cli.main.ApiClient") +@patch("codecarbon.core.api_client.ApiClient") class TestApp(unittest.TestCase): def setUp(self): self.runner = CliRunner() @@ -57,7 +57,7 @@ def test_app(self, MockApiClient): @patch("codecarbon.cli.main.Path.exists") @patch("codecarbon.cli.main.Confirm.ask") @patch("codecarbon.cli.main.questionary_prompt") - @patch("codecarbon.cli.main.get_access_token") + @patch("codecarbon.cli.auth.get_access_token") @patch("typer.prompt") def test_config_no_local_new_all( self, @@ -147,7 +147,7 @@ def side_effect_wrapper(*args, **kwargs): except OSError: pass - @patch("codecarbon.cli.main.get_access_token") + @patch("codecarbon.cli.auth.get_access_token") @patch("codecarbon.cli.main.Path.exists") @patch("codecarbon.cli.main.get_config") @patch("codecarbon.cli.main.questionary_prompt") @@ -186,5 +186,16 @@ def custom_questionary_side_effect(*args, **kwargs): return MagicMock(return_value=default_value) +class TestQuestionaryPrompt(unittest.TestCase): + @patch("questionary.select") + def test_questionary_prompt_returns_selected_value(self, mock_select): + from codecarbon.cli.main import questionary_prompt + + mock_select.return_value.ask.return_value = "selected" + result = questionary_prompt("Pick one", ["a", "b"], "a") + self.assertEqual(result, "selected") + mock_select.assert_called_once_with("Pick one", ["a", "b"], "a") + + if __name__ == "__main__": unittest.main() diff --git a/tests/cli/test_cli_main.py b/tests/cli/test_cli_main.py index 2319dadad..84f42493d 100644 --- a/tests/cli/test_cli_main.py +++ b/tests/cli/test_cli_main.py @@ -3,6 +3,7 @@ from types import SimpleNamespace import pytest +import typer from typer.testing import CliRunner from codecarbon.cli import main as cli_main @@ -34,8 +35,8 @@ def test_version_flag(): def test_api_get_calls_api_and_prints(monkeypatch): runner = CliRunner() - monkeypatch.setattr(cli_main, "ApiClient", FakeApiClient) - monkeypatch.setattr(cli_main, "get_access_token", fake_get_access_token) + monkeypatch.setattr("codecarbon.core.api_client.ApiClient", FakeApiClient) + monkeypatch.setattr("codecarbon.cli.auth.get_access_token", fake_get_access_token) result = runner.invoke(cli_main.codecarbon, ["test-api"]) assert result.exit_code == 0 @@ -51,11 +52,11 @@ def __init__(self, endpoint_url=None): super().__init__(endpoint_url=endpoint_url) runner = CliRunner() - monkeypatch.setattr(cli_main, "ApiClient", CustomApiClient) + monkeypatch.setattr("codecarbon.core.api_client.ApiClient", CustomApiClient) monkeypatch.setattr( cli_main, "get_api_endpoint", lambda: "https://custom.codecarbon.io" ) - monkeypatch.setattr(cli_main, "get_access_token", fake_get_access_token) + monkeypatch.setattr("codecarbon.cli.auth.get_access_token", fake_get_access_token) result = runner.invoke(cli_main.codecarbon, ["test-api"]) assert result.exit_code == 0 @@ -85,7 +86,7 @@ def get_detected_hardware(self): "gpu_ids": None, } - monkeypatch.setattr(cli_main, "EmissionsTracker", FakeTracker) + monkeypatch.setattr("codecarbon.emissions_tracker.EmissionsTracker", FakeTracker) runner = CliRunner() result = runner.invoke(cli_main.codecarbon, ["detect"]) assert result.exit_code == 0 @@ -115,7 +116,7 @@ def set_access_token(self, token): def fake_get_access_token(): raise ValueError("Not able to retrieve the access token, please run login.") - monkeypatch.setattr(cli_main, "ApiClient", FakeApiClient) + monkeypatch.setattr("codecarbon.core.api_client.ApiClient", FakeApiClient) monkeypatch.setattr( cli_main, "get_config", @@ -129,7 +130,7 @@ def fake_get_access_token(): monkeypatch.setattr( cli_main, "get_api_endpoint", lambda path: "https://api.codecarbon.io" ) - monkeypatch.setattr(cli_main, "get_access_token", fake_get_access_token) + monkeypatch.setattr("codecarbon.cli.auth.get_access_token", fake_get_access_token) cli_main.show_config(tmp_path / ".codecarbon.config") captured = capsys.readouterr() @@ -165,16 +166,15 @@ def set_access_token(self, token): def check_auth(self): calls["check_auth"] += 1 - monkeypatch.setattr(cli_main, "ApiClient", FakeApiClient) + monkeypatch.setattr("codecarbon.core.api_client.ApiClient", FakeApiClient) monkeypatch.setattr( - cli_main, - "authorize", + "codecarbon.cli.auth.authorize", lambda: calls.__setitem__("authorize", calls["authorize"] + 1), ) monkeypatch.setattr( cli_main, "get_api_endpoint", lambda: "https://custom-login.codecarbon.io" ) - monkeypatch.setattr(cli_main, "get_access_token", lambda: "login-token") + monkeypatch.setattr("codecarbon.cli.auth.get_access_token", lambda: "login-token") runner = CliRunner() result = runner.invoke(cli_main.codecarbon, ["login"]) @@ -198,8 +198,8 @@ def fake_post(url, json, headers): captured["headers"] = headers return FakeResponse() - monkeypatch.setattr(cli_main, "get_access_token", lambda: "access-token") - monkeypatch.setattr(cli_main.requests, "post", fake_post) + monkeypatch.setattr("codecarbon.cli.auth.get_access_token", lambda: "access-token") + monkeypatch.setattr("requests.post", fake_post) token = cli_main.get_api_key("proj-123") assert token == "project-api-token" @@ -235,8 +235,8 @@ def get_project(self, project_id): def get_experiment(self, experiment_id): return {"id": experiment_id} - monkeypatch.setattr(cli_main, "ApiClient", FakeApiClient) - monkeypatch.setattr(cli_main, "get_access_token", lambda: "fake-token") + monkeypatch.setattr("codecarbon.core.api_client.ApiClient", FakeApiClient) + monkeypatch.setattr("codecarbon.cli.auth.get_access_token", lambda: "fake-token") monkeypatch.setattr( cli_main, "get_api_endpoint", lambda path: "https://api.codecarbon.io" ) @@ -289,7 +289,9 @@ def start(self): def stop(self): return None - monkeypatch.setattr(cli_main, "OfflineEmissionsTracker", FakeOfflineTracker) + monkeypatch.setattr( + "codecarbon.emissions_tracker.OfflineEmissionsTracker", FakeOfflineTracker + ) monkeypatch.setattr(cli_main.signal, "signal", lambda *args, **kwargs: None) runner = CliRunner() @@ -311,7 +313,7 @@ def fake_run_and_monitor(ctx, offline=False, **kwargs): captured["kwargs"] = kwargs return "ok" - monkeypatch.setattr(cli_main, "run_and_monitor", fake_run_and_monitor) + monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor) ctx = SimpleNamespace(args=["python", "-c", "print(1)"]) result = cli_main.monitor( @@ -332,7 +334,7 @@ def fake_run_and_monitor(ctx, offline=False, **kwargs): captured["kwargs"] = kwargs return "ok" - monkeypatch.setattr(cli_main, "run_and_monitor", fake_run_and_monitor) + monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor) monkeypatch.setattr(cli_main, "get_existing_exp_id", lambda: "exp-1") ctx = SimpleNamespace(args=["python", "train.py"]) @@ -350,7 +352,7 @@ def fake_run_and_monitor(ctx, **kwargs): captured["kwargs"] = kwargs return "ok" - monkeypatch.setattr(cli_main, "run_and_monitor", fake_run_and_monitor) + monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor) monkeypatch.setattr(cli_main, "get_existing_exp_id", lambda: "exp-1") ctx = SimpleNamespace(args=["python", "train.py"]) @@ -368,7 +370,7 @@ def fake_run_and_monitor(ctx, offline=False, **kwargs): captured["kwargs"] = kwargs return "ok" - monkeypatch.setattr(cli_main, "run_and_monitor", fake_run_and_monitor) + monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor) monkeypatch.setattr(cli_main, "get_existing_exp_id", lambda: None) ctx = SimpleNamespace(args=["python", "train.py"]) @@ -376,3 +378,31 @@ def fake_run_and_monitor(ctx, offline=False, **kwargs): assert result == "ok" assert captured["offline"] is False assert captured["kwargs"]["save_to_api"] is False + + +def test_monitor_passes_log_level_to_run_and_monitor(monkeypatch): + captured = {} + + def fake_run_and_monitor(ctx, offline=False, **kwargs): + captured["kwargs"] = kwargs + + monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor) + + ctx = SimpleNamespace(args=["echo", "hello"]) + cli_main.monitor( + ctx=ctx, + offline=True, + country_iso_code="FRA", + log_level="debug", + ) + + assert captured["kwargs"]["log_level"] == "debug" + + +def test_monitor_online_requires_experiment_id_for_wrapped_command(monkeypatch): + monkeypatch.setattr(cli_main, "get_existing_exp_id", lambda: None) + + ctx = SimpleNamespace(args=["echo", "hi"]) + with pytest.raises(typer.Exit) as exc_info: + cli_main.monitor(ctx=ctx, offline=False, api=True) + assert exc_info.value.exit_code == 1 diff --git a/tests/cli/test_monitor.py b/tests/cli/test_monitor.py index d4dd718a2..0a9bda365 100644 --- a/tests/cli/test_monitor.py +++ b/tests/cli/test_monitor.py @@ -20,8 +20,15 @@ def stop(self): return 0.123 +def _patch_trackers(monkeypatch, online_cls=FakeTracker, offline_cls=FakeTracker): + monkeypatch.setattr("codecarbon.emissions_tracker.EmissionsTracker", online_cls) + monkeypatch.setattr( + "codecarbon.emissions_tracker.OfflineEmissionsTracker", offline_cls + ) + + def test_run_and_monitor_requires_command(monkeypatch): - monkeypatch.setattr(monitor_module, "EmissionsTracker", FakeTracker) + _patch_trackers(monkeypatch) monkeypatch.setattr(monitor_module, "print", lambda *args, **kwargs: None) with pytest.raises(typer.Exit) as exc_info: @@ -30,12 +37,35 @@ def test_run_and_monitor_requires_command(monkeypatch): assert exc_info.value.exit_code == 1 +def test_run_and_monitor_strips_nested_monitor_prefix(monkeypatch): + captured = {} + + class FakePopen: + def __init__(self, command, text=True): + captured["command"] = command + + def wait(self): + return 0 + + _patch_trackers(monkeypatch) + monkeypatch.setattr(monitor_module.subprocess, "Popen", FakePopen) + monkeypatch.setattr(monitor_module, "print", lambda *args, **kwargs: None) + + with pytest.raises(typer.Exit) as exc_info: + monitor_module.run_and_monitor( + SimpleNamespace(args=["monitor", "--", "echo", "hi"]) + ) + + assert exc_info.value.exit_code == 0 + assert captured["command"] == ["echo", "hi"] + + def test_run_and_monitor_handles_missing_command(monkeypatch): class FakePopen: def __init__(self, command, text=True): raise FileNotFoundError - monkeypatch.setattr(monitor_module, "EmissionsTracker", FakeTracker) + _patch_trackers(monkeypatch) monkeypatch.setattr(monitor_module.subprocess, "Popen", FakePopen) monkeypatch.setattr(monitor_module, "print", lambda *args, **kwargs: None) @@ -50,7 +80,7 @@ class FakePopen: def __init__(self, command, text=True): raise RuntimeError("boom") - monkeypatch.setattr(monitor_module, "EmissionsTracker", FakeTracker) + _patch_trackers(monkeypatch) monkeypatch.setattr(monitor_module.subprocess, "Popen", FakePopen) monkeypatch.setattr(monitor_module, "print", lambda *args, **kwargs: None) @@ -75,8 +105,7 @@ def __init__(self, command, text=True): def wait(self): return 0 - monkeypatch.setattr(monitor_module, "OfflineEmissionsTracker", FakeOfflineTracker) - monkeypatch.setattr(monitor_module, "EmissionsTracker", FakeTracker) + _patch_trackers(monkeypatch, offline_cls=FakeOfflineTracker) monkeypatch.setattr(monitor_module.subprocess, "Popen", FakePopen) monkeypatch.setattr(monitor_module, "print", lambda *args, **kwargs: None) @@ -106,8 +135,7 @@ def __init__(self, command, text=True): def wait(self): return 0 - monkeypatch.setattr(monitor_module, "EmissionsTracker", FakeOnlineTracker) - monkeypatch.setattr(monitor_module, "OfflineEmissionsTracker", FakeTracker) + _patch_trackers(monkeypatch, online_cls=FakeOnlineTracker) monkeypatch.setattr(monitor_module.subprocess, "Popen", FakePopen) monkeypatch.setattr(monitor_module, "print", lambda *args, **kwargs: None) @@ -140,7 +168,7 @@ def terminate(self): def kill(self): process_info["killed"] += 1 - monkeypatch.setattr(monitor_module, "EmissionsTracker", FakeTracker) + _patch_trackers(monkeypatch) monkeypatch.setattr(monitor_module.subprocess, "Popen", FakePopen) monkeypatch.setattr(monitor_module, "print", lambda *args, **kwargs: None)