feat: record token counts on tasks for energy-per-token reporting

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,50 @@ 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)
+            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(
+            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 +1491,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 +1514,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:

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,
         )

codecarbon/output_methods/emissions_data.py

@@ -110,6 +110,13 @@ 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
+
+    # 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:

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_<experiment_name>_<run_id>.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.
 

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.
 
+#### 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` — 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`
+    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.

tests/test_token_tracking.py

@@ -0,0 +1,115 @@
+import os
+import shutil
+import unittest
+from unittest import mock
+
+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 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_warns_and_records_nothing(self):
+        tracker = EmissionsTracker(save_to_file=False, allow_multiple_runs=True)
+        tracker.start()
+        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,
+            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)