Skip to content

Commit e933dea

Browse files
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 <cursoragent@cursor.com>
1 parent 07b82eb commit e933dea

6 files changed

Lines changed: 139 additions & 43 deletions

File tree

codecarbon/cli/main.py

Lines changed: 33 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,26 +5,19 @@
55
from pathlib import Path
66
from typing import Optional
77

8-
import questionary
9-
import requests
108
import typer
119
from rich import print
1210
from rich.prompt import Confirm
1311
from typing_extensions import Annotated
1412

1513
from codecarbon import __app_name__, __version__
16-
from codecarbon.cli.auth import authorize, get_access_token
1714
from codecarbon.cli.cli_utils import (
1815
create_new_config_file,
1916
get_api_endpoint,
2017
get_config,
2118
get_existing_exp_id,
2219
overwrite_local_config,
2320
)
24-
from codecarbon.cli.monitor import run_and_monitor
25-
from codecarbon.core.api_client import ApiClient, get_datetime_with_timezone
26-
from codecarbon.core.schemas import ExperimentCreate, OrganizationCreate, ProjectCreate
27-
from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker
2821

2922
API_URL = os.environ.get("API_URL", "https://dashboard.codecarbon.io/api")
3023

@@ -68,6 +61,9 @@ def version(
6861

6962

7063
def show_config(path: Path = Path("./.codecarbon.config")) -> None:
64+
from codecarbon.cli.auth import get_access_token
65+
from codecarbon.core.api_client import ApiClient
66+
7167
d = get_config(path)
7268
print("Current configuration : \n")
7369
print("Config file content : ")
@@ -114,6 +110,9 @@ def api_get():
114110
"""
115111
ex: test-api
116112
"""
113+
from codecarbon.cli.auth import get_access_token
114+
from codecarbon.core.api_client import ApiClient
115+
117116
api_endpoint = get_api_endpoint()
118117
api = ApiClient(endpoint_url=api_endpoint)
119118
api.set_access_token(get_access_token())
@@ -123,6 +122,9 @@ def api_get():
123122

124123
@codecarbon.command("login", short_help="Login to CodeCarbon")
125124
def login():
125+
from codecarbon.cli.auth import authorize, get_access_token
126+
from codecarbon.core.api_client import ApiClient
127+
126128
authorize()
127129
api_endpoint = get_api_endpoint()
128130
api = ApiClient(endpoint_url=api_endpoint)
@@ -132,6 +134,10 @@ def login():
132134

133135

134136
def get_api_key(project_id: str):
137+
import requests
138+
139+
from codecarbon.cli.auth import get_access_token
140+
135141
api_endpoint = get_api_endpoint()
136142
api_endpoint = api_endpoint.rstrip("/")
137143
req = requests.post(
@@ -161,6 +167,13 @@ def config():
161167
"""
162168
Initialize CodeCarbon, this will prompt you for configuration of Organisation/Team/Project/Experiment.
163169
"""
170+
from codecarbon.cli.auth import get_access_token
171+
from codecarbon.core.api_client import ApiClient, get_datetime_with_timezone
172+
from codecarbon.core.schemas import (
173+
ExperimentCreate,
174+
OrganizationCreate,
175+
ProjectCreate,
176+
)
164177

165178
print("Welcome to CodeCarbon configuration wizard")
166179
home = Path.home()
@@ -342,13 +355,18 @@ def monitor(
342355
str,
343356
typer.Option(help="Region/province for offline mode"),
344357
] = None,
358+
log_level: Annotated[
359+
str,
360+
typer.Option(help="Log level (critical, error, warning, info, debug)"),
361+
] = "error",
345362
):
346363
"""Monitor your machine's carbon emissions."""
347364

348365
# Shared tracker args so monitor and run_and_monitor behave the same
349366
tracker_args = {
350367
"measure_power_secs": measure_power_secs,
351368
"api_call_interval": api_call_interval,
369+
"log_level": log_level,
352370
}
353371
# Set up the tracker arguments based on mode (offline vs online) and validate required args for each mode
354372
if offline:
@@ -375,8 +393,12 @@ def monitor(
375393

376394
tracker_args = {**tracker_args, "save_to_api": api}
377395

396+
from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker
397+
378398
# If extra args are provided (e.g. `codecarbon monitor -- my_script.py`), delegate to `run_and_monitor`
379399
if getattr(ctx, "args", None):
400+
from codecarbon.cli.monitor import run_and_monitor
401+
380402
return run_and_monitor(ctx, offline=offline, **tracker_args)
381403

382404
# Instantiate the tracker
@@ -417,6 +439,8 @@ def detect():
417439
"""
418440
Detects hardware and prints information without running any measurements.
419441
"""
442+
from codecarbon.emissions_tracker import EmissionsTracker
443+
420444
print("Detecting hardware...")
421445
tracker = EmissionsTracker(save_to_file=False)
422446
hardware_info = tracker.get_detected_hardware()
@@ -438,6 +462,8 @@ def detect():
438462

439463

440464
def questionary_prompt(prompt, list_options, default):
465+
import questionary
466+
441467
value = questionary.select(
442468
prompt,
443469
list_options,

codecarbon/cli/monitor.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@
88
from rich import print
99
from typing_extensions import Annotated
1010

11-
from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker
12-
1311

1412
def run_and_monitor(
1513
ctx: typer.Context,
@@ -50,12 +48,15 @@ def run_and_monitor(
5048
directory. The file path is shown in the final report.
5149
"""
5250
# Suppress all CodeCarbon logs during execution
51+
from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker
5352
from codecarbon.external.logger import set_logger_level
5453

5554
set_logger_level(log_level)
5655

57-
# Get the command from remaining args
58-
command = ctx.args
56+
# Get the command from remaining args (strip nested subcommand / `--` leftovers)
57+
command = list(getattr(ctx, "args", None) or [])
58+
while command and command[0] in ("monitor", "--"):
59+
command.pop(0)
5960

6061
if not command:
6162
print(

docs/reference/cli.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ Displays real-time emissions data for all processes on your machine. Press `Ctrl
4444
| `--no-api` | flag | false | Do not send data to the API (local-only measurement) |
4545
| `--offline` | flag | false | Run without internet access |
4646
| `--country-iso-code` | string | - | ISO 3166-1 alpha-3 country code (required in offline mode) |
47-
| `--log-level` | choice | INFO | Log level: DEBUG, INFO, WARNING, ERROR |
47+
| `--log-level` | choice | ERROR | Log level: DEBUG, INFO, WARNING, ERROR |
4848

4949
**Examples:**
5050
```bash

tests/cli/test_cli.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
# MOCK API CLIENT
1212

1313

14-
@patch("codecarbon.cli.main.ApiClient")
14+
@patch("codecarbon.core.api_client.ApiClient")
1515
class TestApp(unittest.TestCase):
1616
def setUp(self):
1717
self.runner = CliRunner()
@@ -57,7 +57,7 @@ def test_app(self, MockApiClient):
5757
@patch("codecarbon.cli.main.Path.exists")
5858
@patch("codecarbon.cli.main.Confirm.ask")
5959
@patch("codecarbon.cli.main.questionary_prompt")
60-
@patch("codecarbon.cli.main.get_access_token")
60+
@patch("codecarbon.cli.auth.get_access_token")
6161
@patch("typer.prompt")
6262
def test_config_no_local_new_all(
6363
self,
@@ -147,7 +147,7 @@ def side_effect_wrapper(*args, **kwargs):
147147
except OSError:
148148
pass
149149

150-
@patch("codecarbon.cli.main.get_access_token")
150+
@patch("codecarbon.cli.auth.get_access_token")
151151
@patch("codecarbon.cli.main.Path.exists")
152152
@patch("codecarbon.cli.main.get_config")
153153
@patch("codecarbon.cli.main.questionary_prompt")
@@ -186,5 +186,16 @@ def custom_questionary_side_effect(*args, **kwargs):
186186
return MagicMock(return_value=default_value)
187187

188188

189+
class TestQuestionaryPrompt(unittest.TestCase):
190+
@patch("questionary.select")
191+
def test_questionary_prompt_returns_selected_value(self, mock_select):
192+
from codecarbon.cli.main import questionary_prompt
193+
194+
mock_select.return_value.ask.return_value = "selected"
195+
result = questionary_prompt("Pick one", ["a", "b"], "a")
196+
self.assertEqual(result, "selected")
197+
mock_select.assert_called_once_with("Pick one", ["a", "b"], "a")
198+
199+
189200
if __name__ == "__main__":
190201
unittest.main()

tests/cli/test_cli_main.py

Lines changed: 50 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from types import SimpleNamespace
44

55
import pytest
6+
import typer
67
from typer.testing import CliRunner
78

89
from codecarbon.cli import main as cli_main
@@ -34,8 +35,8 @@ def test_version_flag():
3435

3536
def test_api_get_calls_api_and_prints(monkeypatch):
3637
runner = CliRunner()
37-
monkeypatch.setattr(cli_main, "ApiClient", FakeApiClient)
38-
monkeypatch.setattr(cli_main, "get_access_token", fake_get_access_token)
38+
monkeypatch.setattr("codecarbon.core.api_client.ApiClient", FakeApiClient)
39+
monkeypatch.setattr("codecarbon.cli.auth.get_access_token", fake_get_access_token)
3940

4041
result = runner.invoke(cli_main.codecarbon, ["test-api"])
4142
assert result.exit_code == 0
@@ -51,11 +52,11 @@ def __init__(self, endpoint_url=None):
5152
super().__init__(endpoint_url=endpoint_url)
5253

5354
runner = CliRunner()
54-
monkeypatch.setattr(cli_main, "ApiClient", CustomApiClient)
55+
monkeypatch.setattr("codecarbon.core.api_client.ApiClient", CustomApiClient)
5556
monkeypatch.setattr(
5657
cli_main, "get_api_endpoint", lambda: "https://custom.codecarbon.io"
5758
)
58-
monkeypatch.setattr(cli_main, "get_access_token", fake_get_access_token)
59+
monkeypatch.setattr("codecarbon.cli.auth.get_access_token", fake_get_access_token)
5960

6061
result = runner.invoke(cli_main.codecarbon, ["test-api"])
6162
assert result.exit_code == 0
@@ -85,7 +86,7 @@ def get_detected_hardware(self):
8586
"gpu_ids": None,
8687
}
8788

88-
monkeypatch.setattr(cli_main, "EmissionsTracker", FakeTracker)
89+
monkeypatch.setattr("codecarbon.emissions_tracker.EmissionsTracker", FakeTracker)
8990
runner = CliRunner()
9091
result = runner.invoke(cli_main.codecarbon, ["detect"])
9192
assert result.exit_code == 0
@@ -115,7 +116,7 @@ def set_access_token(self, token):
115116
def fake_get_access_token():
116117
raise ValueError("Not able to retrieve the access token, please run login.")
117118

118-
monkeypatch.setattr(cli_main, "ApiClient", FakeApiClient)
119+
monkeypatch.setattr("codecarbon.core.api_client.ApiClient", FakeApiClient)
119120
monkeypatch.setattr(
120121
cli_main,
121122
"get_config",
@@ -129,7 +130,7 @@ def fake_get_access_token():
129130
monkeypatch.setattr(
130131
cli_main, "get_api_endpoint", lambda path: "https://api.codecarbon.io"
131132
)
132-
monkeypatch.setattr(cli_main, "get_access_token", fake_get_access_token)
133+
monkeypatch.setattr("codecarbon.cli.auth.get_access_token", fake_get_access_token)
133134

134135
cli_main.show_config(tmp_path / ".codecarbon.config")
135136
captured = capsys.readouterr()
@@ -165,16 +166,15 @@ def set_access_token(self, token):
165166
def check_auth(self):
166167
calls["check_auth"] += 1
167168

168-
monkeypatch.setattr(cli_main, "ApiClient", FakeApiClient)
169+
monkeypatch.setattr("codecarbon.core.api_client.ApiClient", FakeApiClient)
169170
monkeypatch.setattr(
170-
cli_main,
171-
"authorize",
171+
"codecarbon.cli.auth.authorize",
172172
lambda: calls.__setitem__("authorize", calls["authorize"] + 1),
173173
)
174174
monkeypatch.setattr(
175175
cli_main, "get_api_endpoint", lambda: "https://custom-login.codecarbon.io"
176176
)
177-
monkeypatch.setattr(cli_main, "get_access_token", lambda: "login-token")
177+
monkeypatch.setattr("codecarbon.cli.auth.get_access_token", lambda: "login-token")
178178

179179
runner = CliRunner()
180180
result = runner.invoke(cli_main.codecarbon, ["login"])
@@ -198,8 +198,8 @@ def fake_post(url, json, headers):
198198
captured["headers"] = headers
199199
return FakeResponse()
200200

201-
monkeypatch.setattr(cli_main, "get_access_token", lambda: "access-token")
202-
monkeypatch.setattr(cli_main.requests, "post", fake_post)
201+
monkeypatch.setattr("codecarbon.cli.auth.get_access_token", lambda: "access-token")
202+
monkeypatch.setattr("requests.post", fake_post)
203203

204204
token = cli_main.get_api_key("proj-123")
205205
assert token == "project-api-token"
@@ -235,8 +235,8 @@ def get_project(self, project_id):
235235
def get_experiment(self, experiment_id):
236236
return {"id": experiment_id}
237237

238-
monkeypatch.setattr(cli_main, "ApiClient", FakeApiClient)
239-
monkeypatch.setattr(cli_main, "get_access_token", lambda: "fake-token")
238+
monkeypatch.setattr("codecarbon.core.api_client.ApiClient", FakeApiClient)
239+
monkeypatch.setattr("codecarbon.cli.auth.get_access_token", lambda: "fake-token")
240240
monkeypatch.setattr(
241241
cli_main, "get_api_endpoint", lambda path: "https://api.codecarbon.io"
242242
)
@@ -289,7 +289,9 @@ def start(self):
289289
def stop(self):
290290
return None
291291

292-
monkeypatch.setattr(cli_main, "OfflineEmissionsTracker", FakeOfflineTracker)
292+
monkeypatch.setattr(
293+
"codecarbon.emissions_tracker.OfflineEmissionsTracker", FakeOfflineTracker
294+
)
293295
monkeypatch.setattr(cli_main.signal, "signal", lambda *args, **kwargs: None)
294296

295297
runner = CliRunner()
@@ -311,7 +313,7 @@ def fake_run_and_monitor(ctx, offline=False, **kwargs):
311313
captured["kwargs"] = kwargs
312314
return "ok"
313315

314-
monkeypatch.setattr(cli_main, "run_and_monitor", fake_run_and_monitor)
316+
monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor)
315317

316318
ctx = SimpleNamespace(args=["python", "-c", "print(1)"])
317319
result = cli_main.monitor(
@@ -332,7 +334,7 @@ def fake_run_and_monitor(ctx, offline=False, **kwargs):
332334
captured["kwargs"] = kwargs
333335
return "ok"
334336

335-
monkeypatch.setattr(cli_main, "run_and_monitor", fake_run_and_monitor)
337+
monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor)
336338
monkeypatch.setattr(cli_main, "get_existing_exp_id", lambda: "exp-1")
337339

338340
ctx = SimpleNamespace(args=["python", "train.py"])
@@ -350,7 +352,7 @@ def fake_run_and_monitor(ctx, **kwargs):
350352
captured["kwargs"] = kwargs
351353
return "ok"
352354

353-
monkeypatch.setattr(cli_main, "run_and_monitor", fake_run_and_monitor)
355+
monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor)
354356
monkeypatch.setattr(cli_main, "get_existing_exp_id", lambda: "exp-1")
355357

356358
ctx = SimpleNamespace(args=["python", "train.py"])
@@ -368,11 +370,39 @@ def fake_run_and_monitor(ctx, offline=False, **kwargs):
368370
captured["kwargs"] = kwargs
369371
return "ok"
370372

371-
monkeypatch.setattr(cli_main, "run_and_monitor", fake_run_and_monitor)
373+
monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor)
372374
monkeypatch.setattr(cli_main, "get_existing_exp_id", lambda: None)
373375

374376
ctx = SimpleNamespace(args=["python", "train.py"])
375377
result = cli_main.monitor(ctx=ctx, api=False)
376378
assert result == "ok"
377379
assert captured["offline"] is False
378380
assert captured["kwargs"]["save_to_api"] is False
381+
382+
383+
def test_monitor_passes_log_level_to_run_and_monitor(monkeypatch):
384+
captured = {}
385+
386+
def fake_run_and_monitor(ctx, offline=False, **kwargs):
387+
captured["kwargs"] = kwargs
388+
389+
monkeypatch.setattr("codecarbon.cli.monitor.run_and_monitor", fake_run_and_monitor)
390+
391+
ctx = SimpleNamespace(args=["echo", "hello"])
392+
cli_main.monitor(
393+
ctx=ctx,
394+
offline=True,
395+
country_iso_code="FRA",
396+
log_level="debug",
397+
)
398+
399+
assert captured["kwargs"]["log_level"] == "debug"
400+
401+
402+
def test_monitor_online_requires_experiment_id_for_wrapped_command(monkeypatch):
403+
monkeypatch.setattr(cli_main, "get_existing_exp_id", lambda: None)
404+
405+
ctx = SimpleNamespace(args=["echo", "hi"])
406+
with pytest.raises(typer.Exit) as exc_info:
407+
cli_main.monitor(ctx=ctx, offline=False, api=True)
408+
assert exc_info.value.exit_code == 1

0 commit comments

Comments
 (0)