diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 96ed00c91..c0ade61ca 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -792,6 +792,34 @@ 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, + 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 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 + 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 +1475,17 @@ 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( + input_tokens=response.usage.prompt_tokens, + output_tokens=response.usage.completion_tokens, + ) + ``` """ def __init__(self, task_name, tracker: EmissionsTracker = None): @@ -1462,6 +1501,22 @@ 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, + ) -> 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, + ) + 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..5be5e06b1 100644 --- a/codecarbon/external/task.py +++ b/codecarbon/external/task.py @@ -17,6 +17,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 +65,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..0e78fa96c 100644 --- a/codecarbon/output_methods/emissions_data.py +++ b/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: 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..876e4a501 100644 --- a/docs/tutorials/python-api.md +++ b/docs/tutorials/python-api.md @@ -44,6 +44,43 @@ 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 + +Record the token counts of each request with `record_tokens()`: + +``` 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( + input_tokens=response.usage.prompt_tokens, + output_tokens=response.usage.completion_tokens, + ) + +tracker.stop() +``` + +The counts come from wherever your serving stack reports them: `usage` for +OpenAI-compatible clients, `prompt_eval_count` / `eval_count` for Ollama, and so +on. + +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. diff --git a/tests/test_token_tracking.py b/tests/test_token_tracking.py new file mode 100644 index 000000000..d318bb2fe --- /dev/null +++ b/tests/test_token_tracking.py @@ -0,0 +1,59 @@ +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 + +OUTPUT_DIR = "test_token_data" + + +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(input_tokens=12, output_tokens=30) + task.record_tokens(input_tokens=1, output_tokens=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_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)