Feature proposal: record token counts on tasks for energy-per-token reporting #1347

This is a feature proposal for discussion before implementation.

The user problem

When people measure a local LLM with CodeCarbon, the number they get back is energy per script run. That number is not comparable to anything: it depends entirely on how many prompts happened to be sent during the run. The number people actually want — and the one that can go in a model card, a benchmark table, or a quantization decision — is energy per output token, or CO2 per request.

CodeCarbon already measures the hard part (the power draw over a task). The only missing ingredient is the token counter, and that is something only the caller knows. Today users have to keep that bookkeeping entirely outside CodeCarbon and do the division themselves, and nothing lands in the task CSV.

Proposed design

Small and contained, reusing the existing task machinery rather than adding a subsystem.

1. Three optional fields on TaskEmissionsData (defaults 0, so existing rows and readers are unaffected): input_tokens, output_tokens, n_requests. Plus two derived properties, energy_per_output_token (kWh) and emissions_per_request (kgCO2eq), computed rather than stored so they cannot go stale, and returning 0.0 instead of raising when the counter is zero.

2. Counters carried on Task and passed through Task.out(), so they flow to every output handler through the existing BaseOutput.task_out() path and appear as three new CSV columns.

3. One recording entry point on the tracker:

tracker.record_tokens(input_tokens=..., output_tokens=..., n_requests=1)
tracker.record_tokens(response=resp)  # duck-typed extraction

The response= form reads the counts that serving stacks already return — usage.prompt_tokens / usage.completion_tokens for OpenAI-compatible clients, prompt_eval_count / eval_count for Ollama, prompt_token_ids / outputs[].token_ids for vLLM RequestOutput. All of it is getattr / dict.get, so CodeCarbon imports nothing new and works with any client following those shapes.

4. The same method exposed on TaskEmissionsTracker, so the context-manager form reads naturally:

with TaskEmissionsTracker("llama3.1:8b", tracker) as task:
    for prompt in prompts:
        resp = client.chat.completions.create(...)
        task.record_tokens(response=resp)

Counters accumulate across the life of one task rather than opening a task per request, because a single request is typically far shorter than measure_power_secs and per-request tasks would be dominated by measurement noise.

Why it fits the existing extension points

start_task / stop_task, Task, TaskEmissionsData and task_out() already exist and already do the energy isolation and the routing to output handlers. This change adds fields to a dataclass and one method; it introduces no new module, no new dependency, and no new measurement path.

Scope boundary

Deliberately out of scope:

• No prefill/decode split. Attributing those separately needs sub-second sampling that the current scheduler cannot deliver.
• No per-request task granularity. Accumulation only, for the noise reason above.
• No vendor-specific or model-specific machinery, and no optional dependency on any serving library — extraction stays duck-typed.
• Token counts stay task-level and do not propagate to run-level EmissionsData, since a run may mix inference with other work.
• Under continuous batching (e.g. vLLM), requests overlap and per-request attribution is not physically meaningful; aggregate per-token figures remain valid and the limitation is documented rather than papered over.

Happy to adjust the surface before implementing.