From 2d4088ad4787516f507d0bdfd293b7db8d5f4aaf Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Wed, 12 Aug 2026 17:51:41 +0200 Subject: [PATCH 1/3] feat: record token counts on tasks Add input_tokens, output_tokens and n_requests to TaskEmissionsData, with energy_per_output_token and emissions_per_request derived from them, so LLM inference can be reported per token instead of per run. Counts are accumulated on the task via tracker.record_tokens(), which can also read them straight from an OpenAI-compatible, Ollama or vLLM response by duck typing, without importing any inference library. Closes #1347 Co-Authored-By: Claude Opus 5 (1M context) --- codecarbon/emissions_tracker.py | 63 ++++++++++- codecarbon/external/task.py | 65 +++++++++++ codecarbon/output_methods/emissions_data.py | 17 +++ docs/reference/output.md | 11 ++ docs/tutorials/python-api.md | 43 ++++++++ tests/test_token_tracking.py | 114 ++++++++++++++++++++ 6 files changed, 312 insertions(+), 1 deletion(-) create mode 100644 tests/test_token_tracking.py diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 96ed00c91..c468f380e 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -27,7 +27,7 @@ from codecarbon.external.logger import logger, set_logger_format, set_logger_level from codecarbon.external.ram import RAM from codecarbon.external.scheduler import PeriodicScheduler -from codecarbon.external.task import Task +from codecarbon.external.task import Task, extract_token_counts from codecarbon.input import DataSource from codecarbon.lock import Lock from codecarbon.output_methods.base_output import BaseOutput, OutputMethod @@ -792,6 +792,41 @@ def start_task(self, task_name=None) -> None: ) self._active_task = task_name + def record_tokens( + self, + input_tokens: int = 0, + output_tokens: int = 0, + n_requests: int = 1, + response=None, + task_name: str = None, + ) -> None: + """ + Record the token counts of one LLM request on a task, so that the task row + carries energy and emissions per token and per request. + + :param input_tokens: Number of prompt tokens of the request. + :param output_tokens: Number of generated tokens of the request. + :param n_requests: Number of requests these counts stand for, default 1. + :param response: Optional response object of an OpenAI compatible client, + Ollama or vLLM, from which the token counts are read. + :param task_name: Task to record on, default the currently active task. + :return: None + """ + task_name = task_name if task_name else self._active_task + task = self._tasks.get(task_name) + if task is None: + logger.warning("record_tokens : No active task to record tokens on.") + return + if response is not None: + extracted_input, extracted_output = extract_token_counts(response) + input_tokens += extracted_input + output_tokens += extracted_output + task.record_tokens( + input_tokens=input_tokens, + output_tokens=output_tokens, + n_requests=n_requests, + ) + def stop_task(self, task_name: str = None) -> EmissionsData: """ Stop tracking a dedicated execution task. Delta energy is computed by task, to isolate its contribution to total @@ -1447,6 +1482,14 @@ class TaskEmissionsTracker: with TaskEmissionsTracker(task_name="Grid search", tracker=tracker): grid = GridSearchCV(estimator=model, param_grid=param_grid) ``` + + For LLM inference, token counts of each request can be recorded on the task: + ```py + with TaskEmissionsTracker(task_name="llama3.1:8b", tracker=tracker) as task: + for prompt in prompts: + response = client.chat.completions.create(...) + task.record_tokens(response=response) + ``` """ def __init__(self, task_name, tracker: EmissionsTracker = None): @@ -1462,6 +1505,24 @@ def __enter__(self): self.tracker.start_task(self.task_name) return self + def record_tokens( + self, + input_tokens: int = 0, + output_tokens: int = 0, + n_requests: int = 1, + response=None, + ) -> None: + """ + Record the token counts of one LLM request on the task under measure. + See `BaseEmissionsTracker.record_tokens`. + """ + self.tracker.record_tokens( + input_tokens=input_tokens, + output_tokens=output_tokens, + n_requests=n_requests, + response=response, + ) + def __exit__(self, exc_type, exc_value, tb) -> None: self.tracker.stop_task() if self.is_default_tracker: diff --git a/codecarbon/external/task.py b/codecarbon/external/task.py index b8945960e..2379afc96 100644 --- a/codecarbon/external/task.py +++ b/codecarbon/external/task.py @@ -4,6 +4,55 @@ from codecarbon.output_methods.emissions_data import EmissionsData, TaskEmissionsData +def _get(response, *names): + """ + Read the first available attribute or mapping key from ``response``. + Returns None if none of ``names`` is present. + """ + for name in names: + if isinstance(response, dict): + value = response.get(name) + else: + value = getattr(response, name, None) + if value is not None: + return value + return None + + +def extract_token_counts(response): + """ + Best effort extraction of (input_tokens, output_tokens) from the response of + an LLM serving stack. Everything is duck-typed, so codecarbon does not import + any inference library: + + - OpenAI compatible clients : ``usage.prompt_tokens`` / ``usage.completion_tokens`` + - Ollama : ``prompt_eval_count`` / ``eval_count`` + - vLLM ``RequestOutput`` : ``prompt_token_ids`` / ``outputs[].token_ids`` + + Counts that cannot be found are reported as 0. + """ + usage = _get(response, "usage") + if usage is not None: + return ( + _get(usage, "prompt_tokens", "input_tokens") or 0, + _get(usage, "completion_tokens", "output_tokens") or 0, + ) + + eval_count = _get(response, "eval_count") + if eval_count is not None: + return _get(response, "prompt_eval_count") or 0, eval_count + + prompt_token_ids = _get(response, "prompt_token_ids") + if prompt_token_ids is not None: + outputs = _get(response, "outputs") or [] + return ( + len(prompt_token_ids), + sum(len(getattr(completion, "token_ids", ())) for completion in outputs), + ) + + return 0, 0 + + class Task: """ A task, used to segregate electrical consumption when executing a treatment. @@ -17,6 +66,19 @@ def __init__(self, task_name): # , task_measure self.task_name: str = task_name self.start_time = time.perf_counter() self.is_active = True + self.input_tokens: int = 0 + self.output_tokens: int = 0 + self.n_requests: int = 0 + + def record_tokens( + self, input_tokens: int = 0, output_tokens: int = 0, n_requests: int = 1 + ) -> None: + """ + Accumulate token counters for this task. + """ + self.input_tokens += input_tokens + self.output_tokens += output_tokens + self.n_requests += n_requests def out(self): return TaskEmissionsData( @@ -52,4 +114,7 @@ def out(self): ram_total_size=self.emissions_data.ram_total_size, tracking_mode=self.emissions_data.tracking_mode, on_cloud=self.emissions_data.on_cloud, + input_tokens=self.input_tokens, + output_tokens=self.output_tokens, + n_requests=self.n_requests, ) diff --git a/codecarbon/output_methods/emissions_data.py b/codecarbon/output_methods/emissions_data.py index 17544aa51..ab0738084 100644 --- a/codecarbon/output_methods/emissions_data.py +++ b/codecarbon/output_methods/emissions_data.py @@ -110,6 +110,23 @@ class TaskEmissionsData: ram_utilization_percent: float = 0.0 ram_used_gb: float = 0.0 on_cloud: str = "N" + input_tokens: int = 0 + output_tokens: int = 0 + n_requests: int = 0 + + @property + def energy_per_output_token(self) -> float: + """Energy in kWh per output token, 0.0 if no output token was recorded.""" + if not self.output_tokens: + return 0.0 + return self.energy_consumed / self.output_tokens + + @property + def emissions_per_request(self) -> float: + """Emissions in kgCO2eq per request, 0.0 if no request was recorded.""" + if not self.n_requests: + return 0.0 + return self.emissions / self.n_requests @property def values(self) -> OrderedDict: diff --git a/docs/reference/output.md b/docs/reference/output.md index 720e0a883..70cb805f7 100644 --- a/docs/reference/output.md +++ b/docs/reference/output.md @@ -65,6 +65,17 @@ The package has an in-built logger that logs data into a CSV file named `emissio | ram_utilization_percent | Average RAM utilization during tracking period (%) | | ram_used_gb | Average RAM used during tracking period (GB) | +Task rows, written to `emissions__.csv`, carry three extra +columns, filled in by `record_tokens()` (see +[LLM inference](../tutorials/python-api.md#llm-inference-energy-per-token)) and +left at `0` otherwise: + +| Field | Description | +|-------|-------------| +| input_tokens | Total prompt tokens recorded on the task | +| output_tokens | Total generated tokens recorded on the task | +| n_requests | Number of requests recorded on the task | + !!! note Developers can enhance the Output interface by implementing a custom class that extends `BaseOutput` at `codecarbon/output.py`. For example, to log into a database. diff --git a/docs/tutorials/python-api.md b/docs/tutorials/python-api.md index 77f3f3b83..7e9c85672 100644 --- a/docs/tutorials/python-api.md +++ b/docs/tutorials/python-api.md @@ -44,6 +44,49 @@ finally: The task manager tracks each sub-task independently. Tasks are not written to disk by default (to reduce overhead), so retrieve results from the `stop_task()` return value. +**Advanced: LLM inference, energy per token** + +Energy per run is not comparable between models, because it depends on how many +prompts you happened to send. Energy per output token is. If the task you are +measuring is LLM inference, record the token counts of each request with +`record_tokens()` and CodeCarbon will report them alongside the energy: + +``` python-skip +from codecarbon import EmissionsTracker +from codecarbon.emissions_tracker import TaskEmissionsTracker + +tracker = EmissionsTracker(project_name="llama3.1-8b-bench") + +with TaskEmissionsTracker(task_name="llama3.1:8b", tracker=tracker) as task: + for prompt in prompts: + response = client.chat.completions.create(model="llama3.1:8b", messages=prompt) + task.record_tokens(response=response) + +tracker.stop() +``` + +`record_tokens(response=...)` reads the counts the serving stack already returns. +It understands OpenAI-compatible `usage` payloads, Ollama's `prompt_eval_count` / +`eval_count`, and vLLM `RequestOutput` objects. Everything is read by duck typing, +so CodeCarbon does not import any inference library. If your client is not one of +these, pass the numbers yourself: + +``` python-skip +task.record_tokens(input_tokens=128, output_tokens=256) +``` + +Counts accumulate over the life of the task, and the resulting `TaskEmissionsData` +exposes `input_tokens`, `output_tokens` and `n_requests`, plus two derived values: +`energy_per_output_token` (kWh per output token) and `emissions_per_request` +(kgCO₂eq per request). + +!!! warning + Measure enough requests for the task to last several `measure_power_secs` + windows. A per-token figure derived from a task shorter than one measurement + window is mostly noise. Under continuous batching, requests overlap and + per-request attribution is not physically meaningful, although the aggregate + per-token figure remains valid. + ### Context Manager Now that you've seen the explicit object approach, let's look at the more idiomatic **context manager** pattern. This is the recommended way for most use cases. diff --git a/tests/test_token_tracking.py b/tests/test_token_tracking.py new file mode 100644 index 000000000..dcf2dce6a --- /dev/null +++ b/tests/test_token_tracking.py @@ -0,0 +1,114 @@ +import os +import shutil +import unittest + +from pandas import read_csv + +from codecarbon import EmissionsTracker +from codecarbon.emissions_tracker import TaskEmissionsTracker +from codecarbon.external.task import extract_token_counts + +OUTPUT_DIR = "test_token_data" + + +class OpenAIUsage: + prompt_tokens = 12 + completion_tokens = 30 + + +class OpenAIResponse: + usage = OpenAIUsage() + + +class VLLMCompletion: + def __init__(self, n): + self.token_ids = list(range(n)) + + +class VLLMRequestOutput: + def __init__(self): + self.prompt_token_ids = list(range(7)) + self.outputs = [VLLMCompletion(4), VLLMCompletion(6)] + + +class TestExtractTokenCounts(unittest.TestCase): + def test_openai_object(self): + self.assertEqual((12, 30), extract_token_counts(OpenAIResponse())) + + def test_openai_dict(self): + response = {"usage": {"prompt_tokens": 3, "completion_tokens": 8}} + self.assertEqual((3, 8), extract_token_counts(response)) + + def test_ollama_dict(self): + response = {"prompt_eval_count": 5, "eval_count": 42, "response": "hi"} + self.assertEqual((5, 42), extract_token_counts(response)) + + def test_vllm_request_output(self): + self.assertEqual((7, 10), extract_token_counts(VLLMRequestOutput())) + + def test_unknown_response_is_zero(self): + self.assertEqual((0, 0), extract_token_counts({"text": "hello"})) + + +class TestTaskEmissionsDataProperties(unittest.TestCase): + def _task_data(self, **kwargs): + tracker = EmissionsTracker(save_to_file=False, allow_multiple_runs=True) + tracker.start_task("properties") + tracker.record_tokens(**kwargs) + tracker.stop_task() + task = tracker._tasks["properties"].out() + tracker.stop() + return task + + def test_zero_counters_do_not_raise(self): + task = self._task_data(n_requests=0) + self.assertEqual(0.0, task.energy_per_output_token) + self.assertEqual(0.0, task.emissions_per_request) + + def test_derived_values(self): + task = self._task_data(output_tokens=100, n_requests=4) + self.assertAlmostEqual(task.energy_consumed / 100, task.energy_per_output_token) + self.assertAlmostEqual(task.emissions / 4, task.emissions_per_request) + + +class TestTokenTracking(unittest.TestCase): + def setUp(self) -> None: + os.makedirs(OUTPUT_DIR, exist_ok=True) + + def tearDown(self) -> None: + shutil.rmtree(OUTPUT_DIR, ignore_errors=True) + + def test_record_tokens_accumulates_over_a_task(self): + tracker = EmissionsTracker(save_to_file=False, allow_multiple_runs=True) + with TaskEmissionsTracker("inference", tracker=tracker) as task: + task.record_tokens(input_tokens=10, output_tokens=20) + task.record_tokens(response=OpenAIResponse()) + task.record_tokens(response={"prompt_eval_count": 1, "eval_count": 2}) + data = tracker._tasks["inference"].out() + tracker.stop() + self.assertEqual(23, data.input_tokens) + self.assertEqual(52, data.output_tokens) + self.assertEqual(3, data.n_requests) + + def test_record_tokens_without_active_task_is_ignored(self): + tracker = EmissionsTracker(save_to_file=False, allow_multiple_runs=True) + tracker.start() + tracker.record_tokens(output_tokens=10) + tracker.stop() + + def test_token_counts_are_written_to_the_task_csv(self): + tracker = EmissionsTracker( + output_dir=OUTPUT_DIR, + experiment_name="tokens", + allow_multiple_runs=True, + ) + with TaskEmissionsTracker("inference", tracker=tracker) as task: + task.record_tokens(input_tokens=4, output_tokens=16) + tracker.stop() + + task_file = [f for f in os.listdir(OUTPUT_DIR) if f.startswith("emissions_")] + df = read_csv(os.path.join(OUTPUT_DIR, task_file[0])) + row = df[df.task_name == "inference"].iloc[0] + self.assertEqual(4, row.input_tokens) + self.assertEqual(16, row.output_tokens) + self.assertEqual(1, row.n_requests) From fbfb0910c44ed0d68a68183bfed062434d652203 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Wed, 12 Aug 2026 19:07:56 +0200 Subject: [PATCH 2/3] docs: make LLM token section a real heading The cross-link from reference/output.md targeted an anchor that never existed: the section was bold text, not a heading. Co-Authored-By: Claude Opus 5 (1M context) --- docs/tutorials/python-api.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/tutorials/python-api.md b/docs/tutorials/python-api.md index 7e9c85672..28d3efb8c 100644 --- a/docs/tutorials/python-api.md +++ b/docs/tutorials/python-api.md @@ -44,7 +44,7 @@ finally: The task manager tracks each sub-task independently. Tasks are not written to disk by default (to reduce overhead), so retrieve results from the `stop_task()` return value. -**Advanced: LLM inference, energy per token** +#### LLM inference: energy per token Energy per run is not comparable between models, because it depends on how many prompts you happened to send. Energy per output token is. If the task you are From cc1a6108bdcbaa01ad2a7de48a7bde40b3010d41 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Wed, 12 Aug 2026 19:49:30 +0200 Subject: [PATCH 3/3] fix: address review nits on task token counts - drop energy_per_output_token / emissions_per_request: properties never reach TaskEmissionsData.values, so they delivered nothing - assert the warning and the untouched counters when no task is active - debug hint when a response carries no usage (streamed OpenAI chunks) Co-Authored-By: Claude Opus 5 (1M context) --- codecarbon/emissions_tracker.py | 9 ++++ codecarbon/output_methods/emissions_data.py | 16 ++----- docs/tutorials/python-api.md | 6 +-- tests/test_token_tracking.py | 47 +++++++++++---------- 4 files changed, 39 insertions(+), 39 deletions(-) diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index c468f380e..30cea99f2 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -819,6 +819,15 @@ def record_tokens( return if response is not None: extracted_input, extracted_output = extract_token_counts(response) + if not extracted_input and not extracted_output: + # Most common cause: a streamed OpenAI chunk, whose `usage` is + # None unless the request passed + # stream_options={"include_usage": True}. + logger.debug( + "record_tokens : No token count found on the given response, " + "recording 0 tokens. For a streamed response, ask your client " + 'for usage data (OpenAI: stream_options={"include_usage": True}).' + ) input_tokens += extracted_input output_tokens += extracted_output task.record_tokens( diff --git a/codecarbon/output_methods/emissions_data.py b/codecarbon/output_methods/emissions_data.py index ab0738084..0e78fa96c 100644 --- a/codecarbon/output_methods/emissions_data.py +++ b/codecarbon/output_methods/emissions_data.py @@ -114,19 +114,9 @@ class TaskEmissionsData: output_tokens: int = 0 n_requests: int = 0 - @property - def energy_per_output_token(self) -> float: - """Energy in kWh per output token, 0.0 if no output token was recorded.""" - if not self.output_tokens: - return 0.0 - return self.energy_consumed / self.output_tokens - - @property - def emissions_per_request(self) -> float: - """Emissions in kgCO2eq per request, 0.0 if no request was recorded.""" - if not self.n_requests: - return 0.0 - return self.emissions / self.n_requests + # Energy per token and emissions per request are deliberately not exposed: + # `values` is built from `__dict__`, so a property reaches no output, and + # both are one division away from the columns above. @property def values(self) -> OrderedDict: diff --git a/docs/tutorials/python-api.md b/docs/tutorials/python-api.md index 28d3efb8c..0879c6a25 100644 --- a/docs/tutorials/python-api.md +++ b/docs/tutorials/python-api.md @@ -76,9 +76,9 @@ task.record_tokens(input_tokens=128, output_tokens=256) ``` Counts accumulate over the life of the task, and the resulting `TaskEmissionsData` -exposes `input_tokens`, `output_tokens` and `n_requests`, plus two derived values: -`energy_per_output_token` (kWh per output token) and `emissions_per_request` -(kgCO₂eq per request). +exposes `input_tokens`, `output_tokens` and `n_requests` — written to the task CSV +alongside `energy_consumed` and `emissions`, so energy per output token and +emissions per request are one division away. !!! warning Measure enough requests for the task to last several `measure_power_secs` diff --git a/tests/test_token_tracking.py b/tests/test_token_tracking.py index dcf2dce6a..a7a6d0091 100644 --- a/tests/test_token_tracking.py +++ b/tests/test_token_tracking.py @@ -1,6 +1,7 @@ import os import shutil import unittest +from unittest import mock from pandas import read_csv @@ -50,27 +51,6 @@ def test_unknown_response_is_zero(self): self.assertEqual((0, 0), extract_token_counts({"text": "hello"})) -class TestTaskEmissionsDataProperties(unittest.TestCase): - def _task_data(self, **kwargs): - tracker = EmissionsTracker(save_to_file=False, allow_multiple_runs=True) - tracker.start_task("properties") - tracker.record_tokens(**kwargs) - tracker.stop_task() - task = tracker._tasks["properties"].out() - tracker.stop() - return task - - def test_zero_counters_do_not_raise(self): - task = self._task_data(n_requests=0) - self.assertEqual(0.0, task.energy_per_output_token) - self.assertEqual(0.0, task.emissions_per_request) - - def test_derived_values(self): - task = self._task_data(output_tokens=100, n_requests=4) - self.assertAlmostEqual(task.energy_consumed / 100, task.energy_per_output_token) - self.assertAlmostEqual(task.emissions / 4, task.emissions_per_request) - - class TestTokenTracking(unittest.TestCase): def setUp(self) -> None: os.makedirs(OUTPUT_DIR, exist_ok=True) @@ -90,12 +70,33 @@ def test_record_tokens_accumulates_over_a_task(self): self.assertEqual(52, data.output_tokens) self.assertEqual(3, data.n_requests) - def test_record_tokens_without_active_task_is_ignored(self): + def test_record_tokens_without_active_task_warns_and_records_nothing(self): tracker = EmissionsTracker(save_to_file=False, allow_multiple_runs=True) tracker.start() - tracker.record_tokens(output_tokens=10) + with mock.patch("codecarbon.emissions_tracker.logger") as mock_logger: + tracker.record_tokens(output_tokens=10) + tracker.stop() + + self.assertEqual({}, tracker._tasks) + mock_logger.warning.assert_called_once() + self.assertIn("No active task", mock_logger.warning.call_args[0][0]) + + def test_response_without_usage_logs_a_debug_hint(self): + tracker = EmissionsTracker(save_to_file=False, allow_multiple_runs=True) + with TaskEmissionsTracker("inference", tracker=tracker) as task: + # A streamed OpenAI chunk carries no usage unless the caller asked + # for stream_options={"include_usage": True}. + with mock.patch("codecarbon.emissions_tracker.logger") as mock_logger: + task.record_tokens(response={"choices": [], "usage": None}) + data = tracker._tasks["inference"].out() tracker.stop() + self.assertEqual( + (0, 0, 1), (data.input_tokens, data.output_tokens, data.n_requests) + ) + mock_logger.debug.assert_called_once() + self.assertIn("include_usage", mock_logger.debug.call_args[0][0]) + def test_token_counts_are_written_to_the_task_csv(self): tracker = EmissionsTracker( output_dir=OUTPUT_DIR,