diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 96ed00c91..30cea99f2 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,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: 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..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..0879c6a25 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. +#### 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. diff --git a/tests/test_token_tracking.py b/tests/test_token_tracking.py new file mode 100644 index 000000000..a7a6d0091 --- /dev/null +++ b/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)