Skip to content

03 - Session API

Use codegreen.Session from inside a long-running Python program to attribute energy to specific code regions instead of the whole subprocess. Shows the context-manager form and the @codegreen.task decorator side-by-side.

What you learn

  • codegreen.Session("name") as a top-level context.
  • with s.task("region"): to bracket a region.
  • @codegreen.task("name") to bracket a whole function.
  • Where the per-task JSON lands on disk (codegreen_<pid>.json).

Prerequisites

  • pip install codegreen and a working backend (Linux RAPL / macOS IOReport).
  • One Session per process -- the CLI wrapper and this API cannot overlap.

Source

examples/03_session_api/main.py

"""03 Python Session API: context manager + per-task decorator.

`codegreen.Session` is the in-process equivalent of `codegreen run`: useful
when you want to measure specific code blocks inside a long-running script
without spawning a subprocess. The `@codegreen.task` decorator marks
a checkpoint inside an active Session.
"""
import sys

import codegreen


@codegreen.task("hot_loop")
def hot_loop(n: int) -> int:
    acc = 0
    for i in range(n):
        acc += i * i
    return acc


def main() -> int:
    with codegreen.Session("session_api_demo") as s:
        s.start_task("warmup")
        sum(range(100_000))
        s.stop_task("warmup")
        v = hot_loop(2_000_000)
    report = s.stop()
    tasks = report.get("tasks", [])
    total_j = sum(float(t.get("energy_j", 0.0)) for t in tasks)
    print(f"OK: session_api; tasks={len(tasks)}; total_energy_j={total_j:.4f}; hot_loop_result={v}")
    return 0


if __name__ == "__main__":
    sys.exit(main())

Run

python examples/03_session_api/main.py

Expected output

OK: session_api; tasks=2; total_energy_j=0.9320; hot_loop_result=1333332333333000000

A codegreen_<pid>.json is left in the working directory with the full per-task breakdown, per-domain energy, and metadata.

Why it matters

codegreen run measures the whole subprocess, including interpreter startup and imports. For libraries and ML scripts you usually only care about the inner region. Session API brackets just that region, reads the same RAPL/NVML counters in-process, and writes a structured JSON without spawning a child process.

See also: Python API: Session, 05 - GPT-2 staged for the same API on a real ML workload.