diff --git a/.gitignore b/.gitignore index 7f6619314..d9d5de839 100644 --- a/.gitignore +++ b/.gitignore @@ -136,3 +136,12 @@ tests/test_data/rapl/* credentials* .codecarbon.config* scripts/agent-vm.personal.config.sh + +# Added by ggshield +.cache_ggshield + +# Added by ggshield +.cache_ggshield + +# Added by ggshield +.cache_ggshield diff --git a/codecarbon/core/util.py b/codecarbon/core/util.py index 3bb0ca39c..744b2e3e5 100644 --- a/codecarbon/core/util.py +++ b/codecarbon/core/util.py @@ -8,6 +8,7 @@ from pathlib import Path from typing import Optional, Union +import cpuinfo import psutil from codecarbon.external.logger import logger @@ -75,8 +76,6 @@ def backup(file_path: Union[str, Path], ext: Optional[str] = ".bak") -> None: @lru_cache(maxsize=1) def detect_cpu_model() -> Optional[str]: - import cpuinfo - cpu_info = cpuinfo.get_cpu_info() if cpu_info: cpu_model_detected = cpu_info.get("brand_raw", "") diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 270723466..d37603838 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -3,26 +3,28 @@ OfflineEmissionsTracker, context manager and decorator @track_emissions """ -from __future__ import annotations - import dataclasses import os import platform import re +import threading import time import uuid import warnings from abc import ABC, abstractmethod from datetime import datetime from functools import wraps -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional, Union import psutil from codecarbon._version import __version__ from codecarbon.core.config import get_hierarchical_config, normalize_gpu_ids +from codecarbon.core.emissions import Emissions +from codecarbon.core.resource_tracker import ResourceTracker from codecarbon.core.units import Energy, Power, Time, Water from codecarbon.core.util import count_cpus, count_physical_cpus, suppress +from codecarbon.external.geography import CloudMetadata, GeoMetadata from codecarbon.external.hardware import CPU, GPU, AppleSiliconChip from codecarbon.external.logger import logger, set_logger_format, set_logger_level from codecarbon.external.ram import RAM @@ -30,12 +32,18 @@ from codecarbon.external.task import Task from codecarbon.input import DataSource from codecarbon.lock import Lock -from codecarbon.output_methods.base_output import BaseOutput, OutputMethod -from codecarbon.output_methods.emissions_data import EmissionsData - -if TYPE_CHECKING: - from codecarbon.external.geography import CloudMetadata, GeoMetadata - from codecarbon.output_methods.logger import LoggerOutput +from codecarbon.output import ( + BaseOutput, + BoAmpsOutput, + CodeCarbonAPIOutput, + EmissionsData, + FileOutput, + HTTPOutput, + LogfireOutput, + LoggerOutput, + OutputMethod, + PrometheusOutput, +) # /!\ Warning: current implementation prevents the user from setting any value to None # from the script call @@ -56,6 +64,21 @@ _sentinel = object() +@dataclasses.dataclass(frozen=True) +class HttpRequestBaseline: + """Per-request totals snapshot for FastAPI middleware (lifespan tracker).""" + + task_name: str + started_at: float + duration_at_start: float + emissions: float + cpu_energy: float + gpu_energy: float + ram_energy: float + energy_consumed: float + water_consumed: float + + class BaseEmissionsTracker(ABC): """ Primary abstraction with Emissions Tracking functionality. @@ -296,16 +319,11 @@ def _initialize_runtime_state(self) -> None: self._tasks: Dict[str, Task] = {} self._active_task: Optional[str] = None self._active_task_emissions_at_start: Optional[EmissionsData] = None + self._http_task_lock = threading.Lock() + self._measure_lock = threading.Lock() + self._cached_cloud_metadata: Optional[CloudMetadata] = None + self._http_emissions_template: Optional[EmissionsData] = None self._hardware = [] - self._hardware_initialized = False - - def _ensure_hardware_ready(self) -> None: - if self._hardware_initialized: - return - self._populate_system_metadata() - self._initialize_hardware_tracking() - self._hardware_initialized = True - self._log_tracker_metadata() def _populate_system_metadata(self) -> None: self._conf["os"] = platform.platform() @@ -314,8 +332,6 @@ def _populate_system_metadata(self) -> None: self._conf["cpu_physical_count"] = count_physical_cpus() def _initialize_hardware_tracking(self) -> None: - from codecarbon.core.resource_tracker import ResourceTracker - resource_tracker = ResourceTracker(self) resource_tracker.set_CPU_GPU_ram_tracking() self._conf["hardware"] = [item.description() for item in self._hardware] @@ -353,38 +369,21 @@ def _initialize_scheduler_state(self) -> None: def _initialize_emissions_context(self) -> None: self._data_source = DataSource() - self._geo = None - self._emissions = None + cloud: CloudMetadata = self._get_cloud_metadata() + self._geo = self._get_geo_metadata() + + if cloud.is_on_private_infra: + self._conf["longitude"] = self._geo.longitude + self._conf["latitude"] = self._geo.latitude - def _ensure_cloud_conf(self) -> None: - if self._conf.get("_cloud_conf_initialized"): - return - cloud = self._get_cloud_metadata() self._conf["region"] = cloud.region self._conf["provider"] = cloud.provider - self._conf["_cloud_conf_initialized"] = True - - def _ensure_emissions_engine(self) -> None: - if self._emissions is not None: - return - from codecarbon.core.emissions import Emissions - - self._emissions = Emissions( + self._emissions: Emissions = Emissions( self._data_source, self._electricitymaps_api_token, force_carbon_intensity_g_co2e_kwh=self.force_carbon_intensity_g_co2e_kwh, ) - def _ensure_geo_metadata(self) -> None: - """Load geo metadata on first use to avoid blocking tracker construction.""" - if self._geo is not None: - return - self._geo = self._get_geo_metadata() - cloud: CloudMetadata = self._get_cloud_metadata() - if cloud.is_on_private_infra: - self._conf["longitude"] = self._geo.longitude - self._conf["latitude"] = self._geo.latitude - def __init__( self, project_name: Optional[str] = _sentinel, @@ -599,6 +598,9 @@ def __init__( set_logger_level(self._log_level) set_logger_format(self._logger_preamble) self._initialize_runtime_state() + self._populate_system_metadata() + self._initialize_hardware_tracking() + self._log_tracker_metadata() self._initialize_scheduler_state() self._initialize_emissions_context() self._init_output_methods(api_key=self._api_key) @@ -609,18 +611,6 @@ def _init_output_methods(self, *, api_key: str = None): """ methods = set(self._output_methods) if self._output_methods else set() - if not methods and not self._emissions_endpoint: - self.run_id = uuid.uuid4() - return - - from codecarbon.output_methods.boamps import BoAmpsOutput - from codecarbon.output_methods.file import FileOutput - from codecarbon.output_methods.http import CodeCarbonAPIOutput, HTTPOutput - from codecarbon.output_methods.metrics.logfire import LogfireOutput - from codecarbon.output_methods.metrics.prometheus import PrometheusOutput - - methods = set(self._output_methods) if self._output_methods else set() - if OutputMethod.CSV in methods: self._output_handlers.append( FileOutput( @@ -671,7 +661,6 @@ def get_detected_hardware(self) -> Dict[str, Any]: Get the detected hardware. :return: A dictionary containing hardware data. """ - self._ensure_hardware_ready() hardware_info = { "ram_total_size": self._conf.get("ram_total_size"), "cpu_count": self._conf.get("cpu_count"), @@ -703,11 +692,15 @@ def start(self) -> None: "Another instance of codecarbon is already running. Exiting." ) return + try: + _ = self._emissions + except AttributeError: + logger.error("Tracker not initialized. Please check the logs.") + return if self._start_time is not None: logger.warning("Already started tracking") return - self._ensure_hardware_ready() self._last_measured_time = self._start_time = time.perf_counter() # Clear utilization history for fresh measurements @@ -721,9 +714,7 @@ def start(self) -> None: hardware.start() self._scheduler.start() - if self._output_handlers: - self._scheduler_monitor_power.start() - self._measure_power_and_energy() + self._scheduler_monitor_power.start() def start_task(self, task_name=None) -> None: """ @@ -741,13 +732,11 @@ def start_task(self, task_name=None) -> None: ) return try: - self._ensure_emissions_engine() - except Exception: + _ = self._emissions + except AttributeError: logger.error("Tracker not initialized. Please check the logs.") return - self._ensure_hardware_ready() - # Stop scheduler as we do not want it to interfere with the task measurement if self._scheduler: self._scheduler.stop() @@ -800,9 +789,12 @@ def stop_task(self, task_name: str = None) -> EmissionsData: self._scheduler_monitor_power.stop() task_name = task_name if task_name else self._active_task - if self._tasks.get(task_name) is None: + task = self._tasks.get(task_name) + if task is None: logger.warning("stop_task : No active task to stop.") return None + if not task.is_active and task.emissions_data is not None: + return task.emissions_data self._measure_power_and_energy() emissions_data = ( self._prepare_emissions_data() @@ -846,8 +838,124 @@ def stop_task(self, task_name: str = None) -> EmissionsData: self._active_task = None self._active_task_emissions_at_start = None # Clear task-specific start data + if self._scheduler is not None and self._scheduler._stopped: + if self._start_time is not None: + self._scheduler.start() + return task_emission_data + def _resolve_http_task_name(self, task_name: str) -> str: + """Return a unique task name for HTTP request tracking.""" + if not task_name: + task_name = uuid.uuid4().__str__() + if task_name in self._tasks: + task_name += "_" + uuid.uuid4().__str__() + return task_name + + def mark_http_request_start(self, task_name: str) -> HttpRequestBaseline: + """Snapshot cumulative totals at request start (FastAPI lifespan path). + + Use with :meth:`finish_http_request` while the main tracker scheduler keeps + running. Avoids per-request scheduler and hardware restarts from + :meth:`start_task`. + + Args: + task_name: Logical name for this HTTP request (e.g. route key). + + Returns: + Baseline to pass to :meth:`finish_http_request`. + + Raises: + RuntimeError: If the tracker has not been started with :meth:`start`. + """ + if self._start_time is None: + raise RuntimeError("EmissionsTracker.start() must run before HTTP requests") + with self._http_task_lock: + resolved = self._resolve_http_task_name(task_name) + self._tasks[resolved] = Task(task_name=resolved) + duration_at_start = time.perf_counter() - self._start_time + return HttpRequestBaseline( + task_name=resolved, + started_at=time.perf_counter(), + duration_at_start=duration_at_start, + emissions=self._total_emissions, + cpu_energy=self._total_cpu_energy.kWh, + gpu_energy=self._total_gpu_energy.kWh, + ram_energy=self._total_ram_energy.kWh, + energy_consumed=self._total_energy.kWh, + water_consumed=self._total_water.litres, + ) + + def _http_finalize_measure_threshold(self) -> float: + return min(1.0, self._measure_power_secs / 4) + + def _maybe_measure_power_and_energy(self) -> None: + """Sample hardware only when totals may be stale (HTTP finalize path).""" + with self._measure_lock: + if ( + time.perf_counter() - self._last_measured_time + >= self._http_finalize_measure_threshold() + ): + self._run_power_measurement() + + def finish_http_request( + self, baseline: HttpRequestBaseline + ) -> Optional[EmissionsData]: + """Compute per-request emissions from a :meth:`mark_http_request_start` baseline. + + Args: + baseline: Value returned by :meth:`mark_http_request_start`. + + Returns: + Request-scoped :class:`~codecarbon.output.EmissionsData`, or ``None`` if + the task record is missing. + """ + self._maybe_measure_power_and_energy() + with self._http_task_lock: + task = self._tasks.get(baseline.task_name) + if task is None: + logger.warning( + "finish_http_request: unknown task %s", baseline.task_name + ) + return None + emissions_at_stop = self._prepare_http_request_emissions_data() + previous = dataclasses.replace(emissions_at_stop) + previous.emissions = baseline.emissions + previous.cpu_energy = baseline.cpu_energy + previous.gpu_energy = baseline.gpu_energy + previous.ram_energy = baseline.ram_energy + previous.energy_consumed = baseline.energy_consumed + previous.water_consumed = baseline.water_consumed + previous.duration = baseline.duration_at_start + + task_emission_data = dataclasses.replace(emissions_at_stop) + request_duration = time.perf_counter() - baseline.started_at + task_emission_data.duration = Time.from_seconds(request_duration).seconds + task_emission_data.compute_delta_emission(previous) + + task.emissions_data = task_emission_data + task.is_active = False + return task_emission_data + + def persist_completed_task(self, task_name: str) -> None: + """Push a finished task's emissions to API handlers (e.g. after ``stop_task``). + + Args: + task_name: Name of the task that was stopped with :meth:`stop_task`. + """ + if not self._save_to_api: + return + task = self._tasks.get(task_name) + if task is None or task.is_active or task.emissions_data is None: + return + if task.uploaded_to_api: + return + task_payload = [task.out()] + for handler in self._output_handlers: + if isinstance(handler, CodeCarbonAPIOutput): + handler.task_out(task_payload, self._experiment_name) + task.uploaded_to_api = True + @suppress(Exception) def flush(self) -> Optional[float]: """ @@ -870,7 +978,7 @@ def flush(self) -> Optional[float]: # Run to calculate the power used from last # scheduled measurement to shutdown - self._measure_power_and_energy_if_stale() + self._measure_power_and_energy() emissions_data = self._prepare_emissions_data() emissions_data_delta = self._compute_emissions_delta(emissions_data) @@ -912,13 +1020,21 @@ def stop(self) -> Optional[float]: self._scheduler_monitor_power = None else: logger.warning("Tracker already stopped !") - for task_name in self._tasks: - if self._tasks[task_name].is_active: + for task_name in list(self._tasks): + task = self._tasks[task_name] + if not task.is_active: + continue + if ( + self._active_task == task_name + and self._active_task_emissions_at_start is not None + ): self.stop_task(task_name=task_name) + else: + task.is_active = False # Run to calculate the power used from last # scheduled measurement to shutdown # or if scheduler interval was longer than the run - self._measure_power_and_energy_if_stale() + self._measure_power_and_energy() emissions_data = self._prepare_emissions_data() emissions_data_delta = self._compute_emissions_delta(emissions_data) @@ -944,24 +1060,102 @@ def _persist_data( experiment_name=None, ): task_emissions_data = [] + api_task_emissions_data = [] for task in self._tasks: - task_emissions_data.append(self._tasks[task].out()) + task_entry = self._tasks[task].out() + task_emissions_data.append(task_entry) + if not self._tasks[task].uploaded_to_api: + api_task_emissions_data.append(task_entry) for handler in self._output_handlers: handler.out(total_emissions, delta_emissions) if len(task_emissions_data) > 0: - handler.task_out(task_emissions_data, experiment_name) + if isinstance(handler, CodeCarbonAPIOutput): + if api_task_emissions_data: + handler.task_out(api_task_emissions_data, experiment_name) + for task_obj in self._tasks.values(): + if not task_obj.is_active and task_obj.emissions_data: + task_obj.uploaded_to_api = True + else: + handler.task_out(task_emissions_data, experiment_name) + + def _cached_cloud(self) -> CloudMetadata: + if self._cached_cloud_metadata is None: + self._cached_cloud_metadata = self._get_cloud_metadata() + return self._cached_cloud_metadata + + def _average_power_values(self) -> tuple[float, float, float]: + if self._power_measurement_count > 0: + return ( + self._cpu_power_sum / self._power_measurement_count, + self._gpu_power_sum / self._power_measurement_count, + self._ram_power_sum / self._power_measurement_count, + ) + return self._cpu_power.W, self._gpu_power.W, self._ram_power.W + + def _utilization_averages(self) -> tuple[float, float, float, float]: + cpu_util = ( + sum(self._cpu_utilization_history) / len(self._cpu_utilization_history) + if self._cpu_utilization_history + else 0 + ) + gpu_util = ( + sum(self._gpu_utilization_history) / len(self._gpu_utilization_history) + if self._gpu_utilization_history + else 0 + ) + ram_util = ( + sum(self._ram_utilization_history) / len(self._ram_utilization_history) + if self._ram_utilization_history + else 0 + ) + ram_used = ( + sum(self._ram_used_history) / len(self._ram_used_history) + if self._ram_used_history + else 0 + ) + return cpu_util, gpu_util, ram_util, ram_used + + def _prepare_http_request_emissions_data(self) -> EmissionsData: + """Build emissions snapshot for HTTP finalize with cached static metadata.""" + if self._http_emissions_template is None: + snapshot = self._prepare_emissions_data() + self._http_emissions_template = dataclasses.replace(snapshot) + return dataclasses.replace(snapshot) + + self._update_emissions() + duration = Time.from_seconds(time.perf_counter() - self._start_time) + emissions = self._total_emissions + avg_cpu_power, avg_gpu_power, avg_ram_power = self._average_power_values() + cpu_util, gpu_util, ram_util, ram_used = self._utilization_averages() + return dataclasses.replace( + self._http_emissions_template, + timestamp=datetime.now().strftime("%Y-%m-%dT%H:%M:%S"), + duration=duration.seconds, + emissions=emissions, + emissions_rate=emissions / duration.seconds if duration.seconds else 0, + cpu_utilization_percent=cpu_util, + gpu_utilization_percent=gpu_util, + ram_utilization_percent=ram_util, + ram_used_gb=ram_used, + cpu_power=avg_cpu_power, + gpu_power=avg_gpu_power, + ram_power=avg_ram_power, + cpu_energy=self._total_cpu_energy.kWh, + gpu_energy=self._total_gpu_energy.kWh, + ram_energy=self._total_ram_energy.kWh, + energy_consumed=self._total_energy.kWh, + water_consumed=self._total_water.litres, + ) def _update_emissions(self) -> None: """ Compute emissions for the energy consumed since the last update and add them to the total emissions. """ - self._ensure_geo_metadata() - self._ensure_emissions_engine() delta_energy = self._total_energy - self._last_energy_covered if delta_energy.kWh > 0: - cloud: CloudMetadata = self._get_cloud_metadata() + cloud: CloudMetadata = self._cached_cloud() if cloud.is_on_private_infra: delta_emissions = self._emissions.get_private_infra_emissions( delta_energy, self._geo @@ -979,8 +1173,7 @@ def _prepare_emissions_data(self) -> EmissionsData: :return: EmissionsData object with the total emissions data. """ self._update_emissions() - self._ensure_cloud_conf() - cloud = self._get_cloud_metadata() + cloud: CloudMetadata = self._cached_cloud() duration: Time = Time.from_seconds(time.perf_counter() - self._start_time) emissions = self._total_emissions @@ -1024,22 +1217,8 @@ def _prepare_emissions_data(self) -> EmissionsData: cloud_provider = cloud.provider cloud_region = cloud.region - # Calculate average power values across all measurements - avg_cpu_power = ( - self._cpu_power_sum / self._power_measurement_count - if self._power_measurement_count > 0 - else self._cpu_power.W - ) - avg_gpu_power = ( - self._gpu_power_sum / self._power_measurement_count - if self._power_measurement_count > 0 - else self._gpu_power.W - ) - avg_ram_power = ( - self._ram_power_sum / self._power_measurement_count - if self._power_measurement_count > 0 - else self._ram_power.W - ) + avg_cpu_power, avg_gpu_power, avg_ram_power = self._average_power_values() + cpu_util, gpu_util, ram_util, ram_used = self._utilization_averages() total_emissions = EmissionsData( timestamp=datetime.now().strftime("%Y-%m-%dT%H:%M:%S"), @@ -1049,26 +1228,10 @@ def _prepare_emissions_data(self) -> EmissionsData: duration=duration.seconds, emissions=emissions, # kg emissions_rate=emissions / duration.seconds, # kg/s - cpu_utilization_percent=( - sum(self._cpu_utilization_history) / len(self._cpu_utilization_history) - if self._cpu_utilization_history - else 0 - ), - gpu_utilization_percent=( - sum(self._gpu_utilization_history) / len(self._gpu_utilization_history) - if self._gpu_utilization_history - else 0 - ), - ram_utilization_percent=( - sum(self._ram_utilization_history) / len(self._ram_utilization_history) - if self._ram_utilization_history - else 0 - ), - ram_used_gb=( - sum(self._ram_used_history) / len(self._ram_used_history) - if self._ram_used_history - else 0 - ), + cpu_utilization_percent=cpu_util, + gpu_utilization_percent=gpu_util, + ram_utilization_percent=ram_util, + ram_used_gb=ram_used, cpu_power=avg_cpu_power, gpu_power=avg_gpu_power, ram_power=avg_ram_power, @@ -1147,8 +1310,7 @@ def _monitor_power(self) -> None: self._ram_utilization_history.append(psutil.virtual_memory().percent) self._ram_used_history.append(psutil.virtual_memory().used / (1024**3)) - # Collect GPU utilization metrics (lightweight path — skips - # heavyweight calls like process lists, memory, temperature). + # Collect GPU utilization metrics (lightweight path — skips heavyweight calls). for hardware in self._hardware: if isinstance(hardware, GPU): gpu_ids_to_monitor = hardware.gpu_ids @@ -1239,17 +1401,16 @@ def _do_measurements(self) -> None: f"{self._total_energy.kWh:.6f} kWh of electricity and {self._total_water.litres:.6f} L of water were used since the beginning." ) - def _measure_power_and_energy_if_stale(self, min_interval_s: float = 0.05) -> None: - """Measure only if the last sample is older than ``min_interval_s``.""" - if time.perf_counter() - self._last_measured_time >= min_interval_s: - self._measure_power_and_energy() - def _measure_power_and_energy(self) -> None: """ A function that is periodically run by the `BackgroundScheduler` every `self._measure_power_secs` seconds. :return: None """ + with self._measure_lock: + self._run_power_measurement() + + def _run_power_measurement(self) -> None: try: last_duration = time.perf_counter() - self._last_measured_time except AttributeError as e: @@ -1273,7 +1434,6 @@ def _measure_power_and_energy(self) -> None: self._do_measurements() self._last_measured_time = time.perf_counter() self._measure_occurrence += 1 - # Special case: metrics and api calls are sent every `api_call_interval` measures if ( self._api_call_interval != -1 and len(self._output_handlers) > 0 @@ -1357,48 +1517,40 @@ def __init__( "Cloud Region must be provided " + " if cloud provider is set" ) + df = DataSource().get_cloud_emissions_data() + if ( + len( + df.loc[ + (df["provider"] == self._cloud_provider) + & (df["region"] == self._cloud_region) + ] + ) + == 0 + ): + logger.error( + "Cloud Provider/Region " + f"{self._cloud_provider} {self._cloud_region} " + "not found in cloud emissions data." + ) + if self._country_iso_code: + try: + self._country_name: str = DataSource().get_global_energy_mix_data()[ + self._country_iso_code + ]["country_name"] + except KeyError as e: + logger.error( + "Does not support country" + + f" with ISO code {self._country_iso_code} " + f"Exception occurred {e}" + ) + if self._country_2letter_iso_code: assert isinstance(self._country_2letter_iso_code, str) self._country_2letter_iso_code: str = self._country_2letter_iso_code.upper() super().__init__(*args, **kwargs) - def _resolve_offline_country_name(self) -> None: - if self._country_name is not None or not self._country_iso_code: - return - try: - self._country_name = DataSource().get_global_energy_mix_data()[ - self._country_iso_code - ]["country_name"] - except KeyError as e: - logger.error( - "Does not support country" + f" with ISO code {self._country_iso_code} " - f"Exception occurred {e}" - ) - - def _validate_offline_cloud_provider(self) -> None: - if not self._cloud_provider: - return - df = DataSource().get_cloud_emissions_data() - if ( - len( - df.loc[ - (df["provider"] == self._cloud_provider) - & (df["region"] == self._cloud_region) - ] - ) - == 0 - ): - logger.error( - "Cloud Provider/Region " - f"{self._cloud_provider} {self._cloud_region} " - "not found in cloud emissions data." - ) - def _get_geo_metadata(self) -> GeoMetadata: - from codecarbon.external.geography import GeoMetadata - - self._resolve_offline_country_name() return GeoMetadata( country_iso_code=self._country_iso_code, country_name=self._country_name, @@ -1407,9 +1559,6 @@ def _get_geo_metadata(self) -> GeoMetadata: ) def _get_cloud_metadata(self) -> CloudMetadata: - from codecarbon.external.geography import CloudMetadata - - self._validate_offline_cloud_provider() if self._cloud is None: self._cloud = CloudMetadata( provider=self._cloud_provider, region=self._cloud_region @@ -1424,13 +1573,9 @@ class EmissionsTracker(BaseEmissionsTracker): """ def _get_geo_metadata(self) -> GeoMetadata: - from codecarbon.external.geography import GeoMetadata - return GeoMetadata.from_geo_js(self._data_source.geo_js_url) def _get_cloud_metadata(self) -> CloudMetadata: - from codecarbon.external.geography import CloudMetadata - if self._cloud is None: self._cloud = CloudMetadata.from_utils() return self._cloud diff --git a/codecarbon/external/task.py b/codecarbon/external/task.py index b8945960e..705cf6160 100644 --- a/codecarbon/external/task.py +++ b/codecarbon/external/task.py @@ -1,7 +1,7 @@ import time from uuid import uuid4 -from codecarbon.output_methods.emissions_data import EmissionsData, TaskEmissionsData +from codecarbon.output import EmissionsData, TaskEmissionsData class Task: @@ -17,6 +17,7 @@ def __init__(self, task_name): # , task_measure self.task_name: str = task_name self.start_time = time.perf_counter() self.is_active = True + self.uploaded_to_api = False def out(self): return TaskEmissionsData( diff --git a/codecarbon/input.py b/codecarbon/input.py index 4ed23db2b..93a96c988 100644 --- a/codecarbon/input.py +++ b/codecarbon/input.py @@ -1,21 +1,19 @@ """ App configuration and static reference data loading. -Static CSV/JSON reference data is loaded lazily on first DataSource access -to keep `import codecarbon` fast for measurement startup. +Data files are static reference data that never change during runtime. +They are loaded once at module import to avoid repeated file I/O on the hot path +(start_task/stop_task calls for instance). """ -from __future__ import annotations - import atexit import json from contextlib import ExitStack from importlib.resources import as_file as importlib_resources_as_file from importlib.resources import files as importlib_resources_files -from typing import TYPE_CHECKING, Any, Dict +from typing import Any, Dict -if TYPE_CHECKING: - import pandas as pd +import pandas as pd _CACHE: Dict[str, Any] = {} _MODULE_NAME = "codecarbon" @@ -37,8 +35,6 @@ def _load_static_data() -> None: Called once when codecarbon is imported. All data loaded here is immutable and shared across all tracker instances. """ - import pandas as pd - # Global energy mix - used for emissions calculations path = _get_resource_path("data/private_infra/global_energy_mix.json") with open(path) as f: @@ -63,16 +59,8 @@ def _load_static_data() -> None: _CACHE["nordic_country_energy_mix"] = json.load(f) -_STATIC_DATA_LOADED = False - - -def _ensure_static_data_loaded() -> None: - """Load immutable reference data on first use instead of at import.""" - global _STATIC_DATA_LOADED - if _STATIC_DATA_LOADED: - return - _load_static_data() - _STATIC_DATA_LOADED = True +# Load static data at module import +_load_static_data() class DataSource: @@ -142,17 +130,15 @@ def cpu_power_path(self): def get_global_energy_mix_data(self) -> Dict: """ Returns Global Energy Mix Data. - Data is loaded on first access and cached for all tracker instances. + Data is pre-loaded at module import for performance. """ - _ensure_static_data_loaded() return _CACHE["global_energy_mix"] def get_cloud_emissions_data(self) -> pd.DataFrame: """ Returns Cloud Regions Impact Data. - Data is loaded on first access and cached for all tracker instances. + Data is pre-loaded at module import for performance. """ - _ensure_static_data_loaded() return _CACHE["cloud_emissions"] def get_country_emissions_data(self, country_iso_code: str) -> Dict: @@ -190,25 +176,22 @@ def get_country_energy_mix_data(self, country_iso_code: str) -> Dict: def get_carbon_intensity_per_source_data(self) -> Dict: """ Returns Carbon intensity per source. In gCO2.eq/kWh. - Data is loaded on first access and cached for all tracker instances. + Data is pre-loaded at module import for performance. """ - _ensure_static_data_loaded() return _CACHE["carbon_intensity_per_source"] def get_cpu_power_data(self) -> pd.DataFrame: """ Returns CPU power Data. - Data is loaded on first access and cached for all tracker instances. + Data is pre-loaded at module import for performance. """ - _ensure_static_data_loaded() return _CACHE["cpu_power"] def get_nordic_country_energy_mix_data(self) -> Dict: """ Returns Nordic Country Energy Mix Data. - Data is loaded on first access and cached for all tracker instances. + Data is cached on first access per country. """ - _ensure_static_data_loaded() return _CACHE["nordic_country_energy_mix"] diff --git a/codecarbon/integrations/__init__.py b/codecarbon/integrations/__init__.py new file mode 100644 index 000000000..9c5777a96 --- /dev/null +++ b/codecarbon/integrations/__init__.py @@ -0,0 +1 @@ +"""Optional integrations for frameworks and platforms.""" diff --git a/codecarbon/integrations/fastapi/__init__.py b/codecarbon/integrations/fastapi/__init__.py new file mode 100644 index 000000000..8a3aa36c8 --- /dev/null +++ b/codecarbon/integrations/fastapi/__init__.py @@ -0,0 +1,27 @@ +"""FastAPI integration: middleware and lifespan helpers.""" + +try: + from codecarbon.integrations.fastapi.lifespan import ( + compose_lifespans, + create_codecarbon_lifespan, + ) + from codecarbon.integrations.fastapi.middleware import ( + CodeCarbonMiddleware, + add_codecarbon_middleware, + log_request_complete, + shutdown_codecarbon_middleware, + ) +except ImportError as exc: + raise ImportError( + "CodeCarbon FastAPI integration requires Starlette (installed with FastAPI). " + "Install optional dependencies with: pip install 'codecarbon[fastapi]'" + ) from exc + +__all__ = [ + "CodeCarbonMiddleware", + "add_codecarbon_middleware", + "compose_lifespans", + "create_codecarbon_lifespan", + "log_request_complete", + "shutdown_codecarbon_middleware", +] diff --git a/codecarbon/integrations/fastapi/_routing.py b/codecarbon/integrations/fastapi/_routing.py new file mode 100644 index 000000000..2bf385d00 --- /dev/null +++ b/codecarbon/integrations/fastapi/_routing.py @@ -0,0 +1,131 @@ +"""Route naming and endpoint filter helpers for FastAPI/Starlette.""" + +from collections.abc import Iterable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from starlette.requests import Request + +DEFAULT_EXCLUDE: frozenset[str] = frozenset( + { + "/docs", + "/redoc", + "/openapi.json", + "/health", + "/healthz", + "/ready", + "/live", + } +) + +HTTP_METHODS = frozenset( + {"GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS", "TRACE", "CONNECT"} +) + + +def get_endpoint_path(request: "Request") -> str: + """Return the mounted route template or the raw URL path. + + Args: + request: Current Starlette/FastAPI request. + + Returns: + Route template such as ``/items/{item_id}``, or ``request.url.path``. + """ + route = request.scope.get("route") + if route is not None: + return route.path + return request.url.path + + +def build_endpoint_key(request: "Request") -> str: + """Build a stable endpoint identifier such as ``GET /predict``. + + Args: + request: Current Starlette/FastAPI request. + + Returns: + HTTP method plus route template or URL path. + """ + return f"{request.method} {get_endpoint_path(request)}" + + +def is_method_pattern(pattern: str) -> bool: + """Return True when ``pattern`` is ``METHOD /path``.""" + method, _, path = pattern.partition(" ") + return method in HTTP_METHODS and path.startswith("/") + + +def matches_filter_pattern( + pattern: str, + endpoint_key: str, + endpoint_path: str, + url_path: str, + *, + exclude: bool, +) -> bool: + """Return True when an include or exclude pattern matches the request.""" + if is_method_pattern(pattern): + return endpoint_key == pattern + if not pattern.startswith("/"): + return endpoint_key == pattern + if exclude: + return ( + url_path == pattern + or url_path.startswith(f"{pattern}/") + or endpoint_path == pattern + ) + return endpoint_path == pattern + + +def should_track_request( + request: "Request", + include: Iterable[str] | None, + exclude: Iterable[str], +) -> bool: + """Return True when the request should be measured. + + Patterns use one of two forms: + + * ``METHOD /route/template`` — one HTTP method on one route (e.g. ``GET /predict``) + * ``/route/template`` — any method on that route, or a URL path prefix when excluding + + Args: + request: Current Starlette/FastAPI request. + include: When set, only matching endpoints are tracked. + exclude: Endpoints or URL prefixes to skip. + + Returns: + True when CodeCarbon should track this request. + """ + url_path = request.url.path + if include is None: + needs_full_match = any(is_method_pattern(pattern) for pattern in exclude) + if not needs_full_match: + for pattern in exclude: + if url_path == pattern or url_path.startswith(f"{pattern}/"): + return False + return True + endpoint_key = build_endpoint_key(request) + endpoint_path = get_endpoint_path(request) + for pattern in exclude: + if matches_filter_pattern( + pattern, + endpoint_key, + endpoint_path, + url_path, + exclude=True, + ): + return False + if include is None: + return True + return any( + matches_filter_pattern( + pattern, + endpoint_key, + endpoint_path, + url_path, + exclude=False, + ) + for pattern in include + ) diff --git a/codecarbon/integrations/fastapi/lifespan.py b/codecarbon/integrations/fastapi/lifespan.py new file mode 100644 index 000000000..6a03084fd --- /dev/null +++ b/codecarbon/integrations/fastapi/lifespan.py @@ -0,0 +1,72 @@ +"""Lifespan helpers for sharing one ``EmissionsTracker`` across requests.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Callable +from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager +from typing import Any + +from codecarbon import EmissionsTracker +from codecarbon.integrations.fastapi.middleware import shutdown_codecarbon_middleware + + +@asynccontextmanager +async def create_codecarbon_lifespan( + app: Any, + *, + project_name: str = "codecarbon-fastapi", + **tracker_kwargs: Any, +) -> AsyncIterator[None]: + """Start a tracker for the app lifetime and expose it on ``app.state``. + + Args: + app: Starlette/FastAPI application with ``state`` namespace. + project_name: ``project_name`` for :class:`~codecarbon.EmissionsTracker`. + **tracker_kwargs: Extra constructor kwargs for the tracker. + + Yields: + ``None`` while the app runs. + """ + merged = dict(tracker_kwargs) + merged.setdefault("allow_multiple_runs", True) + tracker = EmissionsTracker(project_name=project_name, **merged) + tracker.start() + app.state.codecarbon_tracker = tracker + try: + yield + finally: + shutdown_codecarbon_middleware(app, wait=True) + tracker.stop() + app.state.codecarbon_tracker = None + + +def compose_lifespans( + *factories: Callable[[Any], AbstractAsyncContextManager[Any]], +) -> Callable[[Any], AbstractAsyncContextManager[None]]: + """Nest multiple lifespan context managers into one FastAPI lifespan. + + FastAPI accepts a single ``lifespan`` handler. Use this helper to stack + CodeCarbon with database, cache, or other startup/shutdown contexts:: + + app = FastAPI( + lifespan=compose_lifespans( + lambda a: create_codecarbon_lifespan(a, project_name="my-api"), + lambda a: db_lifespan(a), + ) + ) + + Args: + *factories: Callables that take the app and return an async context manager. + + Returns: + A lifespan callable suitable for ``FastAPI(lifespan=...)``. + """ + + @asynccontextmanager + async def lifespan(app: Any) -> AsyncIterator[None]: + async with AsyncExitStack() as stack: + for factory in factories: + await stack.enter_async_context(factory(app)) + yield + + return lifespan diff --git a/codecarbon/integrations/fastapi/middleware.py b/codecarbon/integrations/fastapi/middleware.py new file mode 100644 index 000000000..967b334b2 --- /dev/null +++ b/codecarbon/integrations/fastapi/middleware.py @@ -0,0 +1,578 @@ +"""FastAPI/Starlette middleware for per-request emissions tracking.""" + +from __future__ import annotations + +import asyncio +import collections +import threading +from collections.abc import Awaitable, Callable, Iterable, Sequence +from concurrent import futures +from typing import Any + +from starlette.requests import Request +from starlette.responses import Response +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from codecarbon import EmissionsTracker +from codecarbon.emissions_tracker import HttpRequestBaseline +from codecarbon.external.logger import logger +from codecarbon.integrations.fastapi._routing import ( + DEFAULT_EXCLUDE, + build_endpoint_key, + should_track_request, +) +from codecarbon.output_methods.emissions_data import EmissionsData + +DEFAULT_TRACKER_KWARGS: dict[str, Any] = { + "save_to_file": False, + "save_to_api": False, + "save_to_logger": False, +} + +# ponytail: local map only; full preset taxonomy if headers become a public API +_HEADER_UNITS: dict[str, str] = { + "emissions": "kg", + "emissions_rate": "kg-per-s", + "duration": "s", + "energy_consumed": "kwh", + "cpu_energy": "kwh", + "gpu_energy": "kwh", + "ram_energy": "kwh", + "cpu_power": "w", + "gpu_power": "w", + "ram_power": "w", +} + +_Job = tuple[Callable[..., Any], tuple[Any, ...], futures.Future[Any]] + + +def _codecarbon_header_name(field: str) -> str: + unit = _HEADER_UNITS.get(field, "") + title = "-".join(part.capitalize() for part in field.split("_")) + suffix = f"-{unit}" if unit else "" + return f"X-CodeCarbon-{title}{suffix}" + + +def _resolve_header_fields( + response_headers: bool | Sequence[str] | None, +) -> tuple[str, ...]: + if not response_headers: + return () + if response_headers is True: + return ("emissions",) + return tuple(response_headers) + + +def _inject_emission_headers( + message: Message, + emissions_data: EmissionsData | None, + fields: Sequence[str], +) -> Message: + if not fields or emissions_data is None: + return message + headers = list(message.get("headers", [])) + for field in fields: + if not hasattr(emissions_data, field): + continue + name = _codecarbon_header_name(field) + value = str(getattr(emissions_data, field)) + headers.append((name.encode("latin-1"), value.encode("latin-1"))) + return {**message, "headers": headers} + + +class _TrackerRunner: + """Single tracker thread: request-path jobs first, then pending finalization.""" + + REQUEST = 0 + FINALIZE = 1 + + def __init__(self, thread_name: str = "codecarbon-tracker") -> None: + self._request_jobs: collections.deque[_Job] = collections.deque() + self._finalize_jobs: collections.deque[_Job] = collections.deque() + self._cond = threading.Condition() + self._closed = False + self._thread = threading.Thread( + target=self._worker, name=thread_name, daemon=True + ) + self._thread.start() + + def _run_job(self, job: _Job) -> None: + func, args, future = job + if future.cancelled(): + return + try: + result = func(*args) + except Exception as exc: + try: + future.set_exception(exc) + except futures.InvalidStateError: + pass + return + try: + future.set_result(result) + except futures.InvalidStateError: + pass + + def _worker(self) -> None: + while True: + with self._cond: + while ( + not self._closed + and not self._request_jobs + and not self._finalize_jobs + ): + self._cond.wait() + if self._closed and not self._request_jobs and not self._finalize_jobs: + return + if self._request_jobs: + job = self._request_jobs.popleft() + lane = self.REQUEST + else: + job = self._finalize_jobs.popleft() + lane = self.FINALIZE + self._run_job(job) + if lane == self.REQUEST: + while True: + with self._cond: + if self._request_jobs: + break + if not self._finalize_jobs: + break + finalize_job = self._finalize_jobs.popleft() + self._run_job(finalize_job) + + def submit( + self, lane: int, func: Callable[..., Any], *args: Any + ) -> futures.Future[Any]: + if self._closed: + raise RuntimeError("cannot schedule tracker work after shutdown") + future: futures.Future[Any] = futures.Future() + job = (func, args, future) + with self._cond: + if lane == self.REQUEST: + self._request_jobs.append(job) + else: + self._finalize_jobs.append(job) + self._cond.notify() + return future + + def submit_request( + self, func: Callable[..., Any], *args: Any + ) -> futures.Future[Any]: + return self.submit(self.REQUEST, func, *args) + + async def run_async(self, lane: int, func: Callable[..., Any], *args: Any) -> Any: + return await asyncio.wrap_future(self.submit(lane, func, *args)) + + def shutdown(self, *, wait: bool = True) -> None: + if self._closed: + return + with self._cond: + self._closed = True + self._cond.notify_all() + if wait: + self._thread.join() + + +def log_request_complete( + request: Request, + response: Response, + emissions_data: EmissionsData | None, + task_name: str, +) -> None: + """Default ``on_request_complete`` handler; logs via the ``codecarbon`` logger.""" + emissions = getattr(emissions_data, "emissions", None) if emissions_data else None + logger.info( + "CodeCarbon %s: emissions=%s kg CO2 status=%s", + task_name, + emissions, + response.status_code, + ) + + +class CodeCarbonMiddleware: + """ASGI middleware using a shared tracker and deferred per-request measurement.""" + + def __init__( + self, + app: ASGIApp, + *, + project_name: str = "codecarbon-fastapi", + include: Iterable[str] | None = None, + exclude: Iterable[str] | None = None, + task_name_formatter: Callable[[Request], str] | None = None, + on_request_complete: Callable[..., Any] | None = log_request_complete, + response_headers: bool | Sequence[str] | None = None, + include_background_tasks: bool = True, + tracker_kwargs: dict[str, Any] | None = None, + **emissions_tracker_kwargs: Any, + ) -> None: + """Configure middleware. + + Args: + app: Inner ASGI application. + project_name: ``project_name`` passed to :class:`~codecarbon.EmissionsTracker`. + include: When set, only matching endpoints are tracked (e.g. ``GET /predict``). + exclude: Endpoints or URL prefixes to skip. Defaults to common docs and health routes. + task_name_formatter: Overrides default route-based task naming. + on_request_complete: Callback ``(request, response, emissions_data | None, task_name)``. + Defaults to :func:`log_request_complete`; pass ``None`` to disable logging. + response_headers: When set, measure before ``http.response.start`` and inject + ``X-CodeCarbon-*`` headers (``True`` → ``emissions`` only, or a field list). + Adds sampling latency to the client response path. + include_background_tasks: When ``True`` (default), finalize after the ASGI call + returns so FastAPI/Starlette ``BackgroundTasks`` are included. When ``False``, + finalize at end of response body (excludes post-body background work). + tracker_kwargs: Baseline kwargs merged into the tracker constructor. + **emissions_tracker_kwargs: Additional :class:`~codecarbon.EmissionsTracker` kwargs. + """ + self.app = app + self.project_name = project_name + self.include = set(include) if include is not None else None + self.exclude = set(exclude if exclude is not None else DEFAULT_EXCLUDE) + self.task_name_formatter = task_name_formatter + self.on_request_complete = on_request_complete + self.header_fields = _resolve_header_fields(response_headers) + self.include_background_tasks = include_background_tasks + merged: dict[str, Any] = dict(DEFAULT_TRACKER_KWARGS) + merged.update(tracker_kwargs or {}) + merged.update(emissions_tracker_kwargs) + merged.setdefault("allow_multiple_runs", True) + self.tracker_kwargs = merged + self._app_tracker: EmissionsTracker | None = None + self._tracker_init_lock = threading.Lock() + self._tracker_runner = _TrackerRunner() + + def shutdown_tracker_executor(self, *, wait: bool = True) -> None: + """Shut down the tracker background thread (idempotent). + + Args: + wait: When ``True``, block until queued tracker work finishes. + """ + self._tracker_runner.shutdown(wait=wait) + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + """ASGI entrypoint.""" + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + request = Request(scope, receive) + if not should_track_request(request, self.include, self.exclude): + await self.app(scope, receive, send) + return + + task_name = self._task_name(request) + tracker, baseline = await self._run_begin_request(request, task_name) + await self._handle_tracked( + scope, receive, send, request, tracker, task_name, baseline + ) + + def _task_name(self, request: Request) -> str: + if self.task_name_formatter is not None: + return self.task_name_formatter(request) + return build_endpoint_key(request) + + async def _run_begin_request( + self, request: Request, task_name: str + ) -> tuple[EmissionsTracker, HttpRequestBaseline | None]: + return await self._tracker_runner.run_async( + _TrackerRunner.REQUEST, self._begin_request, request, task_name + ) + + async def _run_finalize_tracker(self, func: Callable[..., Any], *args: Any) -> Any: + return await self._tracker_runner.run_async( + _TrackerRunner.FINALIZE, func, *args + ) + + def _create_and_start_tracker(self) -> EmissionsTracker: + tracker = EmissionsTracker( + project_name=self.project_name, **self.tracker_kwargs + ) + tracker.start() + return tracker + + def _lifespan_tracker(self, request: Request) -> EmissionsTracker | None: + return getattr(request.app.state, "codecarbon_tracker", None) + + def _tracker_running(self, tracker: EmissionsTracker) -> bool: + return getattr(tracker, "_start_time", None) is not None + + def _begin_request( + self, request: Request, task_name: str + ) -> tuple[EmissionsTracker, HttpRequestBaseline | None]: + tracker = self._lifespan_tracker(request) + if tracker is None: + with self._tracker_init_lock: + if self._app_tracker is None: + self._app_tracker = self._create_and_start_tracker() + tracker = self._app_tracker + if self._tracker_running(tracker): + baseline = tracker.mark_http_request_start(task_name) + return tracker, baseline + tracker.start_task(task_name) + return tracker, None + + def _finalize_on_worker( + self, + tracker: EmissionsTracker, + task_name: str, + request: Request, + response: Response, + run_callback: bool, + baseline: HttpRequestBaseline | None, + ) -> EmissionsData | None: + if baseline is not None: + emissions_data = tracker.finish_http_request(baseline) + resolved_task = baseline.task_name + else: + active_task = getattr(tracker, "_active_task", None) + resolved_task = active_task if isinstance(active_task, str) else task_name + emissions_data = tracker.stop_task(resolved_task) + tracker.persist_completed_task(resolved_task) + if run_callback: + self._run_request_complete(request, response, emissions_data, resolved_task) + return emissions_data + + def _run_request_complete( + self, + request: Request, + response: Response | None, + emissions_data: EmissionsData | None, + task_name: str, + ) -> None: + if self.on_request_complete is None or response is None: + return + self.on_request_complete(request, response, emissions_data, task_name) + + def _schedule_finalize(self, coro: Awaitable[None]) -> None: + async def _run() -> None: + try: + await coro + except Exception: + logger.exception("CodeCarbon deferred measurement failed") + + asyncio.create_task(_run()) + + async def _finalize_after_response( + self, + tracker: EmissionsTracker, + task_name: str, + request: Request, + response: Response, + baseline: HttpRequestBaseline | None, + *, + run_callback: bool, + ) -> EmissionsData | None: + return await self._run_finalize_tracker( + self._finalize_on_worker, + tracker, + task_name, + request, + response, + run_callback, + baseline, + ) + + async def _handle_tracked( + self, + scope: Scope, + receive: Receive, + send: Send, + request: Request, + tracker: EmissionsTracker, + task_name: str, + baseline: HttpRequestBaseline | None, + ) -> None: + if self.header_fields: + await self._handle_tracked_sync_headers( + scope, receive, send, request, tracker, task_name, baseline + ) + return + if self.include_background_tasks: + await self._handle_tracked_after_app( + scope, receive, send, request, tracker, task_name, baseline + ) + return + await self._handle_tracked_end_of_body( + scope, receive, send, request, tracker, task_name, baseline + ) + + async def _handle_tracked_after_app( + self, + scope: Scope, + receive: Receive, + send: Send, + request: Request, + tracker: EmissionsTracker, + task_name: str, + baseline: HttpRequestBaseline | None, + ) -> None: + status_code = 500 + + async def send_wrapper(message: Message) -> None: + nonlocal status_code + if message["type"] == "http.response.start": + status_code = message["status"] + await send(message) + + error: BaseException | None = None + try: + await self.app(scope, receive, send_wrapper) + except BaseException as exc: + error = exc + finally: + response = Response(status_code=status_code) + self._schedule_finalize( + self._finalize_after_response( + tracker, + task_name, + request, + response, + baseline, + run_callback=error is None, + ) + ) + if error is not None: + raise error + + async def _handle_tracked_end_of_body( + self, + scope: Scope, + receive: Receive, + send: Send, + request: Request, + tracker: EmissionsTracker, + task_name: str, + baseline: HttpRequestBaseline | None, + ) -> None: + status_code = 500 + finalized = False + + def _kick_finalize(*, run_callback: bool) -> None: + nonlocal finalized + if finalized: + return + finalized = True + response = Response(status_code=status_code) + self._schedule_finalize( + self._finalize_after_response( + tracker, + task_name, + request, + response, + baseline, + run_callback=run_callback, + ) + ) + + async def send_wrapper(message: Message) -> None: + nonlocal status_code + if message["type"] == "http.response.start": + status_code = message["status"] + await send(message) + if message["type"] == "http.response.body" and not message.get( + "more_body", False + ): + _kick_finalize(run_callback=True) + + error: BaseException | None = None + try: + await self.app(scope, receive, send_wrapper) + except BaseException as exc: + error = exc + finally: + _kick_finalize(run_callback=error is None) + if error is not None: + raise error + + async def _handle_tracked_sync_headers( + self, + scope: Scope, + receive: Receive, + send: Send, + request: Request, + tracker: EmissionsTracker, + task_name: str, + baseline: HttpRequestBaseline | None, + ) -> None: + status_code = 500 + finalized = False + + async def send_wrapper(message: Message) -> None: + nonlocal status_code, finalized + if message["type"] != "http.response.start": + await send(message) + return + status_code = message["status"] + response = Response(status_code=status_code) + emissions_data = await self._finalize_after_response( + tracker, + task_name, + request, + response, + baseline, + run_callback=True, + ) + finalized = True + await send( + _inject_emission_headers(message, emissions_data, self.header_fields) + ) + + error: BaseException | None = None + try: + await self.app(scope, receive, send_wrapper) + except BaseException as exc: + error = exc + finally: + if not finalized: + response = Response(status_code=status_code) + self._schedule_finalize( + self._finalize_after_response( + tracker, + task_name, + request, + response, + baseline, + run_callback=error is None, + ) + ) + if error is not None: + raise error + + +def shutdown_codecarbon_middleware(app: Any, *, wait: bool = True) -> None: + """Shut down the middleware tracker background thread registered on ``app``. + + Args: + app: Application that called :func:`add_codecarbon_middleware`. + wait: Passed to :meth:`CodeCarbonMiddleware.shutdown_tracker_executor`. + """ + middleware = getattr(app.state, "codecarbon_middleware", None) + if middleware is not None: + middleware.shutdown_tracker_executor(wait=wait) + + +def add_codecarbon_middleware(app: Any, **kwargs: Any) -> None: + """Register :class:`CodeCarbonMiddleware` on a FastAPI or Starlette app. + + Registers the instance on ``app.state.codecarbon_middleware`` so + :func:`create_codecarbon_lifespan` or :func:`shutdown_codecarbon_middleware` + can shut down the tracker background thread on teardown. + + Args: + app: Application instance with ``add_middleware``. + **kwargs: Forwarded to :class:`CodeCarbonMiddleware`. + """ + registered: list[CodeCarbonMiddleware] = [] + + class _RegisteredCodeCarbonMiddleware(CodeCarbonMiddleware): + def __init__(self, asgi_app: ASGIApp, **kw: Any) -> None: + super().__init__(asgi_app, **kw) + registered.clear() + registered.append(self) + + app.add_middleware(_RegisteredCodeCarbonMiddleware, **kwargs) + app.build_middleware_stack() + if registered: + app.state.codecarbon_middleware = registered[0] diff --git a/codecarbon/output_methods/http.py b/codecarbon/output_methods/http.py index e0ff710b1..27b08dc6d 100644 --- a/codecarbon/output_methods/http.py +++ b/codecarbon/output_methods/http.py @@ -6,7 +6,7 @@ from codecarbon.core.api_client import ApiClient from codecarbon.external.logger import logger from codecarbon.output_methods.base_output import BaseOutput -from codecarbon.output_methods.emissions_data import EmissionsData +from codecarbon.output_methods.emissions_data import EmissionsData, TaskEmissionsData class HTTPOutput(BaseOutput): @@ -74,3 +74,7 @@ def live_out(self, _, delta: EmissionsData): def out(self, _, delta: EmissionsData): self._emit(delta) + + def task_out(self, data: list[TaskEmissionsData], experiment_name: str) -> None: + for task_data in data: + self._emit(task_data) diff --git a/docs/how-to/fastapi.md b/docs/how-to/fastapi.md new file mode 100644 index 000000000..6f87d87c7 --- /dev/null +++ b/docs/how-to/fastapi.md @@ -0,0 +1,225 @@ +# FastAPI middleware + +Track HTTP request carbon emissions for a [FastAPI](https://fastapi.tiangolo.com/) (or Starlette) app. Install the optional integration extra, register the middleware, and each route is measured without per-handler boilerplate. + +## Install + +```console +pip install "codecarbon[fastapi]" +``` + +With uv: + +```console +uv add "codecarbon[fastapi]" +``` + +## Basic usage + +```python +from fastapi import FastAPI +from codecarbon.integrations.fastapi import add_codecarbon_middleware + +app = FastAPI() +add_codecarbon_middleware(app, project_name="my-api") +``` + +By default, measurement runs after the response is sent (clients are not blocked on hardware sampling), and emissions are logged on the `codecarbon` logger. Pass `on_request_complete=None` to turn logging off, or supply your own callback. + +A minimal runnable app lives at [`examples/fastapi_middleware.py`](https://github.com/mlco2/codecarbon/blob/master/examples/fastapi_middleware.py). For a production-style setup with a Hugging Face embedder and safe concurrent requests, see [`examples/fastapi_embedder.py`](https://github.com/mlco2/codecarbon/blob/master/examples/fastapi_embedder.py). Run it with: + +```console +uv run --extra fastapi uvicorn examples.fastapi_middleware:app --reload +``` + +Then open or `curl` `http://127.0.0.1:8000/predict` and check application logs for per-request emissions. + +## Lifespan (recommended) + +Start one shared `EmissionsTracker` at boot and flush on shutdown: + +```python +from contextlib import asynccontextmanager + +from fastapi import FastAPI +from codecarbon.integrations.fastapi import add_codecarbon_middleware, create_codecarbon_lifespan + + +@asynccontextmanager +async def lifespan(app: FastAPI): + async with create_codecarbon_lifespan(app, project_name="my-api"): + yield + + +app = FastAPI(lifespan=lifespan) +add_codecarbon_middleware(app) +``` + +`create_codecarbon_lifespan` puts a shared tracker on `app.state` for the middleware to reuse and stops it cleanly on shutdown. If you skip lifespan, call `shutdown_codecarbon_middleware(app)` before exit. + +### Combining with other lifespans + +FastAPI accepts only one `lifespan` handler. Nest CodeCarbon with your own startup/shutdown using `compose_lifespans`: + +```python +from contextlib import asynccontextmanager + +from fastapi import FastAPI +from codecarbon.integrations.fastapi import ( + add_codecarbon_middleware, + compose_lifespans, + create_codecarbon_lifespan, +) + + +@asynccontextmanager +async def db_lifespan(app: FastAPI): + app.state.db = "connected" + try: + yield + finally: + app.state.db = None + + +app = FastAPI( + lifespan=compose_lifespans( + lambda a: create_codecarbon_lifespan(a, project_name="my-api"), + db_lifespan, + ) +) +add_codecarbon_middleware(app) +``` + +## Measurement model + +- **Default (deferred):** response is sent first; finalize runs on a dedicated tracker worker thread; sampling is synchronous on that worker before `on_request_complete`. +- **`response_headers=...`:** measure before `http.response.start` and inject `X-CodeCarbon-*` headers (adds sampling latency on the client path). Header values cover work up to response start, not post-body background tasks. +- With `create_codecarbon_lifespan`, concurrent requests on the same route get unique internal task IDs via `HttpRequestBaseline`. + +## Cloud API + +Use **global config only** (`~/.codecarbon.config`). Do not add a repo-local `./.codecarbon.config`, or it will override these values when you run from the project directory. + +```ini +[codecarbon] +api_endpoint = https://api.codecarbon.io +project_id = 00000000-0000-0000-0000-000000000001 +experiment_id = 00000000-0000-0000-0000-000000000002 +``` + +Run `codecarbon login` to store your `api_key` in the same file (never commit it). + +To upload emissions to the dashboard, enable `save_to_api` (IDs are read from global config unless overridden in code): + +```python +add_codecarbon_middleware( + app, + tracker_kwargs={"save_to_api": True}, +) +``` + +One **run** is created per app process when the shared tracker starts; each measured request uploads one emission after the response. See [Use the Cloud API & Dashboard](cloud-api.md). + +Verify logging, CSV, and API locally: + +```console +CODECARBON_ALLOW_MULTIPLE_RUNS=True uv run --extra fastapi \ + python scripts/verify_fastapi_middleware_outputs.py --save-to-api +``` + +## Performance overview + +By default, CodeCarbon measures **after** the response is sent. Clients see only a small amount of middleware bookkeeping; hardware sampling and logging run on a background tracker worker. + +| Mode | Client path | When to use | +|------|-------------|-------------| +| Deferred + logging (default) | Response first, then measure | Most APIs | +| Deferred, `on_request_complete=None` | Response first, measure without log | Lowest overhead while still tracking | +| `response_headers=True` | Measure **before** `http.response.start` | Clients need `X-CodeCarbon-*` headers | +| `create_codecarbon_lifespan` | Same as above + one shared tracker | Production (recommended) | + +### Measured overhead (HF embedder, live tracker) + +Benchmarks use [`scripts/benchmark_fastapi_middleware.py`](https://github.com/mlco2/codecarbon/blob/master/scripts/benchmark_fastapi_middleware.py) with a **live** `EmissionsTracker`, [`paraphrase-MiniLM-L3-v2`](https://huggingface.co/sentence-transformers/paraphrase-MiniLM-L3-v2), uvicorn, `create_codecarbon_lifespan`, 50 timed requests after 5 warmup, concurrency 4. + +Run: + +```console +CODECARBON_ALLOW_MULTIPLE_RUNS=True uv run --extra fastapi --with uvicorn \ + --with sentence-transformers --with torch \ + python scripts/benchmark_fastapi_middleware.py --realistic --with-headers +``` + +Measured on **Darwin arm64**, Python 3.12 (**2026-07-29**): + +| Setup | Mean response time | vs baseline | +|--------|-------------------:|------------:| +| No middleware | 42 ms | — | +| Deferred, logging off | 30 ms | ~same order as baseline | +| Deferred + logging (default) | 32 ms | ~same order as baseline | +| Sync headers (`response_headers=True`) | 52 ms | **~+24%** | + +**What this means** + +- **Deferred (default):** response is sent before finalize; client latency stays in the same ballpark as inference under concurrency (not hundreds of ms). +- **Request path:** mark runs on the tracker REQUEST lane; finalize reuses cached metadata and skips redundant power samples when the scheduler is fresh. +- **Sync headers:** measure before `http.response.start`; latency includes sample time on the client path when a fresh hardware read is needed. +- **`save_to_api=True`:** uploads after the response; network time is not on the HTTP critical path in deferred mode. + +Prefer deferred + logging/API unless clients need response headers. + +## `include` and `exclude` + +Two filters control which requests are measured. Both accept the same pattern forms: + +| Pattern | Meaning | +|---------|---------| +| `GET /predict` | One HTTP method on one route | +| `/predict` | Any method on that route (`include`), or skip that route/URL prefix (`exclude`) | + +- **`exclude`** — skip matching requests. Defaults to docs and health paths (`/docs`, `/health`, …). Pass your own list to replace the default. +- **`include`** — when set, only matching endpoints are tracked (allowlist). + +```python +add_codecarbon_middleware( + app, + include=["GET /predict", "POST /train"], + exclude=["GET /admin", "/internal"], +) +``` + +## `response_headers`, `include_background_tasks`, `task_name_formatter`, `on_request_complete` + +- **`response_headers`** — `True` (emissions only) or a list of field names (`emissions`, `duration`, `energy_consumed`, …). Measures before the response starts and sets `X-CodeCarbon-*` headers. Default `None` / off (deferred, no headers). +- **`include_background_tasks`** — default `True`: FastAPI/Starlette `BackgroundTasks` on the response are included. Set `False` to finalize at end of body and exclude post-body background work. +- **`task_name_formatter`** — optional `(Request) -> str`; default is `METHOD /route/template`. Concurrent requests on the same route still get unique internal task IDs with `create_codecarbon_lifespan`. +- **`on_request_complete`** — optional callback; default logs via `log_request_complete`; `None` disables it. + +```python +add_codecarbon_middleware( + app, + response_headers=["emissions", "energy_consumed", "duration"], + include_background_tasks=False, +) +``` + +## Middleware order + +Per [FastAPI middleware order](https://fastapi.tiangolo.com/tutorial/middleware/), the **last** middleware added is **outermost** on the request path (runs first on the way in). Add CodeCarbon **after** other middleware so it wraps inner layers and includes work done by inner middleware and route handlers: + +```python +from starlette.middleware.cors import CORSMiddleware + +app.add_middleware(CORSMiddleware, ...) +add_codecarbon_middleware(app) # outermost on request → measures the full stack below +``` + +## Limitations (v1) + +- **WebSockets** are not supported. The middleware ignores non-`http` scopes and does not wrap connect/disconnect or messages. +- **Background tasks:** by default, FastAPI/Starlette `BackgroundTasks` / `Response.background` **are included** (they finish before deferred finalize). Use `include_background_tasks=False` to measure only through the response body. Fire-and-forget `asyncio.create_task`, unjoined threads, and external queues (Celery, RQ, …) are **not** tracked. +- **Response headers** require `response_headers=...` (sync measure). Deferred mode cannot attach real emissions to headers because values are computed after the response is sent. + +## Per-endpoint tracking + +For a single route or fine-grained control without global middleware, use the [`@track_emissions` decorator](../reference/api.md#track_emissions-decorator) (same parameters as `EmissionsTracker`). diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 71ba05a85..1e93a588f 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -44,7 +44,7 @@ Displays real-time emissions data for all processes on your machine. Press `Ctrl | `--no-api` | flag | false | Do not send data to the API (local-only measurement) | | `--offline` | flag | false | Run without internet access | | `--country-iso-code` | string | - | ISO 3166-1 alpha-3 country code (required in offline mode) | -| `--log-level` | choice | ERROR | Log level: DEBUG, INFO, WARNING, ERROR | +| `--log-level` | choice | INFO | Log level: DEBUG, INFO, WARNING, ERROR | **Examples:** ```bash diff --git a/examples/fastapi_concurrency.py b/examples/fastapi_concurrency.py new file mode 100644 index 000000000..0f724dedf --- /dev/null +++ b/examples/fastapi_concurrency.py @@ -0,0 +1,174 @@ +"""Concurrent FastAPI + live tracker: what breaks and what works. + +Root cause of ``_active_task_emissions_at_start was None`` under load +-------------------------------------------------------------------- +``EmissionsTracker.start_task`` / ``stop_task`` assume **one** active task +(``_active_task``, ``_active_task_emissions_at_start``). Concurrent HTTP +requests that call ``start_task`` while another request is in flight either +bail out ("A task is already under measure") or corrupt shared state; the +next ``stop_task`` then logs the error and reports zero delta. + +The middleware avoids this whenever the tracker is already running +(``tracker.start()`` was called) by using per-request baselines:: + + mark_http_request_start("GET /embed") -> HttpRequestBaseline + ... handle request ... + finish_http_request(baseline) + +That path is concurrency-safe (unique internal task names, locks). + +**Recommended:** ``create_codecarbon_lifespan`` (see ``fastapi_embedder.py``). +**Also OK:** middleware-only lazy ``tracker.start()`` — same mark/finish path +after the fix in ``CodeCarbonMiddleware._begin_request``. + +Run the embedder with lifespan (production-style):: + + uv run --extra fastapi --with uvicorn --with sentence-transformers --with torch \\ + uvicorn examples.fastapi_concurrency:app_lifespan --host 127.0.0.1 --port 8000 + +Run the minimal lazy-tracker variant (no lifespan):: + + uv run --extra fastapi --with uvicorn --with sentence-transformers --with torch \\ + uvicorn examples.fastapi_concurrency:app_lazy --host 127.0.0.1 --port 8001 + +Load test (20 requests, concurrency 4):: + + uv run --extra fastapi --with httpx python examples/fastapi_concurrency.py \\ + --url http://127.0.0.1:8000/embed +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import sys +from contextlib import asynccontextmanager +from functools import lru_cache +from typing import Any + +import httpx +from fastapi import FastAPI +from sentence_transformers import SentenceTransformer + +from codecarbon.integrations.fastapi import add_codecarbon_middleware, create_codecarbon_lifespan + +MODEL_ID = "sentence-transformers/paraphrase-MiniLM-L3-v2" +SAMPLE_TEXT = "CodeCarbon measures the carbon footprint of machine learning workloads." + +_TRACKER_KWARGS = { + "save_to_file": False, + "save_to_api": False, + "save_to_logger": False, + "allow_multiple_runs": True, + "measure_power_secs": 2, +} + + +@lru_cache(maxsize=1) +def _load_model() -> SentenceTransformer: + return SentenceTransformer(MODEL_ID) + + +def _build_routes(application: FastAPI) -> None: + @application.get("/embed") + def embed(text: str = SAMPLE_TEXT) -> dict[str, Any]: + vector = _load_model().encode(text) + return {"dimensions": int(vector.shape[0]), "model": MODEL_ID} + + +def _wire_middleware(application: FastAPI) -> None: + add_codecarbon_middleware( + application, + project_name="fastapi-concurrency", + tracker_kwargs=_TRACKER_KWARGS, + on_request_complete=None, + ) + + +@asynccontextmanager +async def _lifespan(application: FastAPI): + _load_model() + async with create_codecarbon_lifespan( + application, + project_name="fastapi-concurrency", + **_TRACKER_KWARGS, + ): + yield + + +app_lifespan = FastAPI(title="CodeCarbon concurrency (lifespan)", lifespan=_lifespan) +_build_routes(app_lifespan) +_wire_middleware(app_lifespan) + +app_lazy = FastAPI(title="CodeCarbon concurrency (lazy tracker)") +_build_routes(app_lazy) +_wire_middleware(app_lazy) + + +@app_lazy.on_event("startup") +def _warmup_lazy_model() -> None: + _load_model() + + +async def _load_test(base_url: str, *, requests: int, concurrency: int) -> None: + sem = asyncio.Semaphore(concurrency) + + async def one(client: httpx.AsyncClient) -> None: + async with sem: + response = await client.get("/embed", params={"text": SAMPLE_TEXT}) + response.raise_for_status() + + async with httpx.AsyncClient(base_url=base_url, timeout=60.0) as client: + await asyncio.gather(*(one(client) for _ in range(requests))) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--url", + default="http://127.0.0.1:8000", + help="Base URL of a running app (default: lifespan app on :8000)", + ) + parser.add_argument("--requests", type=int, default=20) + parser.add_argument("--concurrency", type=int, default=4) + args = parser.parse_args(argv) + + captured: list[str] = [] + + class _ErrorCapture(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if ( + record.levelno >= logging.ERROR + and "_active_task_emissions_at_start" in record.getMessage() + ): + captured.append(record.getMessage()) + + codecarbon_logger = logging.getLogger("codecarbon") + codecarbon_logger.addHandler(_ErrorCapture()) + try: + asyncio.run( + _load_test( + args.url.rstrip("/"), + requests=args.requests, + concurrency=args.concurrency, + ) + ) + finally: + codecarbon_logger.removeHandler(_ErrorCapture()) + + if captured: + print(f"FAIL: {len(captured)} concurrency error(s)", file=sys.stderr) + for message in captured[:5]: + print(f" {message}", file=sys.stderr) + return 1 + + print( + f"OK: {args.requests} requests at concurrency {args.concurrency} " + f"— no _active_task_emissions_at_start errors" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/fastapi_embedder.py b/examples/fastapi_embedder.py new file mode 100644 index 000000000..883406676 --- /dev/null +++ b/examples/fastapi_embedder.py @@ -0,0 +1,77 @@ +"""FastAPI embedder API with CodeCarbon middleware (production-style setup). + +Uses ``create_codecarbon_lifespan`` so one shared tracker handles concurrent +requests safely via ``mark_http_request_start`` / ``finish_http_request``. + +Run:: + + uv run --extra fastapi --with uvicorn --with sentence-transformers --with torch \\ + uvicorn examples.fastapi_embedder:app --host 127.0.0.1 --port 8000 + +Load test (optional):: + + uv run --extra fastapi --with httpx python -c " + import asyncio, httpx + async def main(): + sem = asyncio.Semaphore(4) + async def one(): + async with sem: + async with httpx.AsyncClient(timeout=60) as c: + r = await c.get('http://127.0.0.1:8000/embed', params={'text': 'hello'}) + r.raise_for_status() + await asyncio.gather(*[one() for _ in range(20)]) + asyncio.run(main()) + " +""" + +from __future__ import annotations + +from contextlib import asynccontextmanager +from functools import lru_cache +from typing import Any + +from fastapi import FastAPI +from sentence_transformers import SentenceTransformer + +from codecarbon.integrations.fastapi import add_codecarbon_middleware, create_codecarbon_lifespan + +MODEL_ID = "sentence-transformers/paraphrase-MiniLM-L3-v2" +SAMPLE_TEXT = "CodeCarbon measures the carbon footprint of machine learning workloads." + +_tracker_kwargs = { + "save_to_file": False, + "save_to_api": False, + "save_to_logger": False, + "allow_multiple_runs": True, + "measure_power_secs": 2, +} + + +@lru_cache(maxsize=1) +def _load_model() -> SentenceTransformer: + return SentenceTransformer(MODEL_ID) + + +@asynccontextmanager +async def lifespan(app: FastAPI): + _load_model() + async with create_codecarbon_lifespan( + app, + project_name="fastapi-embedder", + **_tracker_kwargs, + ): + yield + + +app = FastAPI(title="CodeCarbon embedder demo", lifespan=lifespan) +add_codecarbon_middleware( + app, + project_name="fastapi-embedder", + tracker_kwargs=_tracker_kwargs, +) + + +@app.get("/embed") +def embed(text: str = SAMPLE_TEXT) -> dict[str, Any]: + vector = _load_model().encode(text) + return {"dimensions": int(vector.shape[0]), "model": MODEL_ID} diff --git a/examples/fastapi_middleware.py b/examples/fastapi_middleware.py new file mode 100644 index 000000000..097c6adb3 --- /dev/null +++ b/examples/fastapi_middleware.py @@ -0,0 +1,56 @@ +"""Minimal FastAPI app with CodeCarbon middleware.""" + +from contextlib import asynccontextmanager +from pathlib import Path + +from fastapi import FastAPI + +from codecarbon.integrations.fastapi import ( + add_codecarbon_middleware, + create_codecarbon_lifespan, +) + +_OUTPUT_DIR = Path(__file__).resolve().parent / "output" +_OUTPUT_DIR.mkdir(exist_ok=True) + +# api_key, experiment_id, project_id: read from ~/.codecarbon.config (not repo .codecarbon.config). +_tracker_kwargs = { + "save_to_file": True, + "save_to_api": True, + "save_to_logger": False, + "log_level": "info", + "output_dir": str(_OUTPUT_DIR), + "allow_multiple_runs": True, +} + + +@asynccontextmanager +async def lifespan(app: FastAPI): + async with create_codecarbon_lifespan( + app, + project_name="fastapi-demo", + **_tracker_kwargs, + ): + yield + + +app = FastAPI(title="CodeCarbon FastAPI demo", lifespan=lifespan) +add_codecarbon_middleware( + app, + project_name="fastapi-demo", + tracker_kwargs=_tracker_kwargs, +) + + +@app.get("/predict") +def predict(text: str = "hello"): + return {"text": text, "label": "demo"} + + +# Stack other startup with compose_lifespans — see docs/how-to/fastapi.md +# Per-request: codecarbon logger (INFO) after each response. +# CSV: examples/output/emissions.csv on shutdown. +# Run: +# CODECARBON_ALLOW_MULTIPLE_RUNS=True uv run --extra fastapi --with uvicorn \ +# uvicorn examples.fastapi_middleware:app --reload +# curl 'http://127.0.0.1:8000/predict?text=hello' diff --git a/mkdocs.yml b/mkdocs.yml index 19c451331..fb9114b63 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -149,6 +149,7 @@ nav: - How-to Guides: - Examples: how-to/examples.md - Configure CodeCarbon: how-to/configuration.md + - FastAPI middleware: how-to/fastapi.md - Compare Model Efficiency: tutorials/comparing-model-efficiency.md - Dashboard & Visualization: - Use the Cloud API & Dashboard: how-to/cloud-api.md diff --git a/pyproject.toml b/pyproject.toml index d4999aba8..37c235ffc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -96,6 +96,8 @@ dev = [ "jsonschema", # For BoAmps schema validation tests "mktestdocs", # For testing documentation code blocks "scikit-learn", # For documentation examples and tests + "fastapi>=0.100", + "httpx", ] doc = [ "requests", @@ -120,6 +122,10 @@ viz-legacy = [ "dash_bootstrap_components > 1.0.0", "fire", ] +fastapi = [ + "fastapi>=0.100", + "httpx", +] [project.scripts] carbonboard = "codecarbon.viz.carbonboard:main" @@ -161,6 +167,9 @@ version_pattern = "MAJOR.MINOR.PATCH[_TAGNUM]" [tool.pytest.ini_options] pythonpath = "." +markers = [ + "no_immediate_finalize: disable deferred-finalize autouse fixture for a test", +] [tool.coverage.run] source = [ diff --git a/scripts/benchmark_fastapi_middleware.py b/scripts/benchmark_fastapi_middleware.py new file mode 100644 index 000000000..2965ac632 --- /dev/null +++ b/scripts/benchmark_fastapi_middleware.py @@ -0,0 +1,1232 @@ +"""Benchmark FastAPI middleware overhead with a realistic ML inference workload. + +Run from repo root: + + uv run --extra fastapi --with uvicorn --with sentence-transformers --with torch \\ + python scripts/benchmark_fastapi_middleware.py + +Uses async HTTP clients (``httpx.AsyncClient``). Reports 95% bootstrap CIs on mean +latency. Verifies default middleware emits one ``codecarbon`` log line per request. + +Optional ``--with-save-to-api`` adds a scenario with ``save_to_api=True`` and +``api_call_interval=1`` (API ``live_out`` after each task measurement). Mocked runs +add ``--api-delay-ms`` sleep on ``stop_task``; ``--real-tracker`` patches +``ApiClient`` instead of calling the network. + +Use ``--quick`` for in-process ASGI (no uvicorn per scenario), noop workload, and +normal-approx CIs. ML workloads are preloaded once across scenarios when using HF. +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("CODECARBON_LOG_LEVEL", "ERROR") + +import argparse # noqa: E402 +import asyncio # noqa: E402 +import logging # noqa: E402 +import platform # noqa: E402 +import random # noqa: E402 +import statistics # noqa: E402 +import sys # noqa: E402 +import threading # noqa: E402 +import time # noqa: E402 +from contextlib import asynccontextmanager # noqa: E402 +from dataclasses import dataclass # noqa: E402 +from typing import Any # noqa: E402 +from unittest.mock import MagicMock, patch # noqa: E402 + +import httpx # noqa: E402 +from fastapi import FastAPI # noqa: E402 + +import codecarbon.integrations.fastapi.middleware as cc_fastapi_middleware # noqa: E402 +from codecarbon.external.logger import logger as codecarbon_logger # noqa: E402 +from codecarbon.integrations.fastapi import ( # noqa: E402 + add_codecarbon_middleware, + shutdown_codecarbon_middleware, +) + +DEFAULT_MEASUREMENT_DELAY_S = 0.02 +WARMUP_REQUESTS = 50 +BENCHMARK_REQUESTS = 300 +QUICK_WARMUP_REQUESTS = 5 +QUICK_BENCHMARK_REQUESTS = 50 +QUICK_SECONDARY_WARMUP = 2 +SMOKE_WARMUP_REQUESTS = 2 +SMOKE_BENCHMARK_REQUESTS = 20 +SMOKE_INFERENCE_DELAY_MS = 15.0 +QUICK_LOGGING_SAMPLE = 10 +CONCURRENCY = 8 +BOOTSTRAP_SAMPLES = 2000 +QUICK_BOOTSTRAP_SAMPLES = 200 +FINALIZE_DRAIN_MULTIPLIER = 4 +QUICK_INFERENCE_DELAY_MS = 25.0 +CONFIDENCE_LEVEL = 0.95 +FASTAPI_BENCHMARK_PROJECT_ID = "25bf2346-49de-4658-911e-4c9003000e13" +FASTAPI_BENCHMARK_EXPERIMENT_ID = "d2d69403-1373-42b4-a2c1-09589aed4801" +REALISTIC_BENCHMARK_REQUESTS = 50 +REALISTIC_WARMUP_REQUESTS = 5 +REALISTIC_CONCURRENCY = 4 +TRACKER_KWARGS = {"save_to_file": False, "save_to_api": False} +TRACKER_KWARGS_SAVE_TO_API = { + "save_to_file": False, + "save_to_api": True, + "save_to_logger": False, + "api_call_interval": 1, + "experiment_id": FASTAPI_BENCHMARK_EXPERIMENT_ID, +} +DEFAULT_EMBEDDER_MODEL = "sentence-transformers/paraphrase-MiniLM-L3-v2" +DEFAULT_CLASSIFIER_MODEL = "distilbert-base-uncased-finetuned-sst-2-english" +SAMPLE_TEXT = "CodeCarbon measures the carbon footprint of machine learning workloads." + + +@dataclass(frozen=True) +class BenchmarkResult: + """Aggregated HTTP benchmark metrics for one configuration.""" + + name: str + requests: int + concurrency: int + mean_ms: float + ci_low_ms: float + ci_high_ms: float + median_ms: float + p95_ms: float + requests_per_sec: float + overhead_pct: float | None + codecarbon_log_lines: int | None = None + + +def _mock_emissions_data(measurement_delay_s: float) -> MagicMock: + return MagicMock( + emissions=0.001, + duration=measurement_delay_s, + energy_consumed=0.002, + emissions_rate=0.002, + ) + + +def _install_tracker_patch( + measurement_delay_s: float, + *, + api_delay_state: dict[str, float] | None = None, + api_delay_s: float = 0.0, +) -> Any: + delays = api_delay_state if api_delay_state is not None else {"api": api_delay_s} + + def _stop() -> float: + time.sleep(measurement_delay_s) + return 0.001 + + def _stop_task(_name: str) -> MagicMock: + time.sleep(measurement_delay_s) + if delays.get("api", 0.0) > 0: + time.sleep(delays["api"]) + return _mock_emissions_data(measurement_delay_s) + + tracker = MagicMock() + tracker.start.return_value = None + tracker.stop.side_effect = _stop + tracker.start_task.return_value = None + tracker.stop_task.side_effect = _stop_task + tracker.persist_completed_task.return_value = None + tracker.final_emissions_data = _mock_emissions_data(measurement_delay_s) + return patch.object(cc_fastapi_middleware, "EmissionsTracker", return_value=tracker) + + +def _config_ids() -> tuple[str, str]: + """Read project_id and experiment_id from hierarchical config when present.""" + from codecarbon.core.config import get_hierarchical_config + + section = get_hierarchical_config() + project_id = section.get("project_id") or FASTAPI_BENCHMARK_PROJECT_ID + experiment_id = section.get("experiment_id") or FASTAPI_BENCHMARK_EXPERIMENT_ID + return project_id, experiment_id + + +def _install_api_client_patch(api_delay_s: float) -> Any: + """Avoid network I/O while exercising ``save_to_api`` output handlers.""" + + import uuid + + from codecarbon.core import api_client as api_client_module + + def _create_run(self: Any, experiment_id: str) -> None: + self.run_id = str(uuid.uuid4()) + + def _add_emission(self: Any, carbon_emission: dict) -> bool: + time.sleep(api_delay_s) + return True + + return patch.multiple( + api_client_module.ApiClient, + _create_run=_create_run, + add_emission=_add_emission, + ) + + +_Z_95 = 1.96 + + +def bootstrap_mean_ci( + latencies_ms: list[float], + *, + samples: int = BOOTSTRAP_SAMPLES, + confidence: float = CONFIDENCE_LEVEL, +) -> tuple[float, float, float]: + """Return mean and two-sided bootstrap CI bounds for mean latency.""" + if not latencies_ms: + return 0.0, 0.0, 0.0 + n = len(latencies_ms) + boot_means = [ + statistics.mean(random.choices(latencies_ms, k=n)) for _ in range(samples) + ] + boot_means.sort() + alpha = (1.0 - confidence) / 2.0 + low_index = max(0, int(alpha * samples) - 1) + high_index = min(samples - 1, int((1.0 - alpha) * samples)) + return ( + statistics.mean(latencies_ms), + boot_means[low_index], + boot_means[high_index], + ) + + +def normal_mean_ci(latencies_ms: list[float]) -> tuple[float, float, float]: + """Approximate 95% CI for the mean (faster than bootstrap for --quick).""" + if not latencies_ms: + return 0.0, 0.0, 0.0 + n = len(latencies_ms) + mean = statistics.mean(latencies_ms) + if n < 2: + return mean, mean, mean + margin = _Z_95 * statistics.stdev(latencies_ms) / (n**0.5) + return mean, mean - margin, mean + margin + + +def summarize_latencies( + latencies_ms: list[float], + *, + bootstrap_samples: int, + use_normal_ci: bool, +) -> tuple[float, float, float, float, float]: + """Return mean, CI low/high, median, and p95.""" + if use_normal_ci: + mean_ms, ci_low_ms, ci_high_ms = normal_mean_ci(latencies_ms) + else: + mean_ms, ci_low_ms, ci_high_ms = bootstrap_mean_ci( + latencies_ms, samples=bootstrap_samples + ) + return ( + mean_ms, + ci_low_ms, + ci_high_ms, + statistics.median(latencies_ms), + _percentile(latencies_ms, 0.95), + ) + + +class InferenceWorkload: + """Runs a small Hugging Face model once per request.""" + + def __init__( + self, + workload: str, + model_id: str, + *, + inference_delay_s: float = 0.0, + ) -> None: + self.workload = workload + self.model_id = model_id + self.inference_delay_s = inference_delay_s + self._embedder: Any = None + self._classifier: Any = None + self._loaded = False + + def ensure_loaded(self) -> None: + """Load the model at most once (shared across benchmark scenarios).""" + if self._loaded: + return + self.load() + self._loaded = True + + def load(self) -> None: + """Load the model into memory.""" + if self.workload == "noop": + self._loaded = True + return + if self.workload == "hf-embedder": + from sentence_transformers import SentenceTransformer + + self._embedder = SentenceTransformer(self.model_id) + self._loaded = True + return + if self.workload == "hf-classifier": + from transformers import pipeline + + self._classifier = pipeline( + "sentiment-analysis", + model=self.model_id, + device=-1, + ) + self._loaded = True + return + raise ValueError(f"Unknown workload: {self.workload}") + + def run(self, text: str = SAMPLE_TEXT) -> dict[str, Any]: + """Execute one inference and return a small JSON-serializable payload.""" + if self.inference_delay_s > 0: + time.sleep(self.inference_delay_s) + if self.workload == "noop": + return {"ok": True} + if self.workload == "hf-embedder": + vector = self._embedder.encode(text) + return {"dimensions": int(vector.shape[0])} + if self.workload == "hf-classifier": + result = self._classifier(text[:512])[0] + return {"label": result["label"], "score": float(result["score"])} + raise ValueError(f"Unknown workload: {self.workload}") + + +def build_app( # noqa: C901 + mode: str, + workload: InferenceWorkload, + *, + project_name: str = FASTAPI_BENCHMARK_PROJECT_ID, + experiment_id: str = FASTAPI_BENCHMARK_EXPERIMENT_ID, + real_tracker: bool = False, +) -> FastAPI: + """Build a FastAPI app for the given benchmark mode.""" + codecarbon_modes = { + "deferred_no_logging", + "deferred_logging", + "deferred_save_to_api", + "sync_headers", + } + if real_tracker and mode in codecarbon_modes: + from codecarbon.integrations.fastapi import create_codecarbon_lifespan + + tracker_kwargs = ( + TRACKER_KWARGS_SAVE_TO_API + if mode == "deferred_save_to_api" + else TRACKER_KWARGS + ) + if mode == "deferred_save_to_api": + tracker_kwargs = {**tracker_kwargs, "experiment_id": experiment_id} + + @asynccontextmanager + async def lifespan(_app: FastAPI): + workload.ensure_loaded() + async with create_codecarbon_lifespan( + _app, + project_name=project_name, + allow_multiple_runs=True, + **tracker_kwargs, + ): + yield + + else: + + @asynccontextmanager + async def lifespan(_app: FastAPI): + workload.ensure_loaded() + yield + + application = FastAPI(lifespan=lifespan) + + @application.get("/predict") + def predict(text: str = SAMPLE_TEXT) -> dict[str, Any]: + return workload.run(text) + + if mode == "baseline": + return application + + if mode == "noop_middleware": + + class _NoopMiddleware: + def __init__(self, app: Any) -> None: + self.app = app + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + await self.app(scope, receive, send) + + application.add_middleware(_NoopMiddleware) + return application + + if mode == "logfire_instrumented": + try: + import logfire + except ImportError as exc: + raise ImportError( + "Logfire scenario requires logfire. Install with: " + "uv run --with 'logfire[fastapi]' ..." + ) from exc + try: + logfire.configure(send_to_logfire=False) + logfire.instrument_fastapi(application) + except RuntimeError as exc: + raise RuntimeError( + "Logfire FastAPI instrumentation requires " + "`opentelemetry-instrumentation-fastapi`. Install with: " + "uv run --with 'logfire[fastapi]' ..." + ) from exc + return application + + kwargs: dict[str, Any] = { + "tracker_kwargs": TRACKER_KWARGS, + "exclude": [], + } + if mode == "deferred_no_logging": + kwargs["on_request_complete"] = None + elif mode == "deferred_logging": + pass + elif mode == "deferred_save_to_api": + kwargs["tracker_kwargs"] = { + **TRACKER_KWARGS_SAVE_TO_API, + "experiment_id": experiment_id, + } + kwargs["on_request_complete"] = None + elif mode == "sync_headers": + kwargs["response_headers"] = True + kwargs["on_request_complete"] = None + else: + raise ValueError(f"Unknown mode: {mode}") + + add_codecarbon_middleware(application, project_name=project_name, **kwargs) + return application + + +def _percentile(values: list[float], pct: float) -> float: + ordered = sorted(values) + index = max(0, min(len(ordered) - 1, int(len(ordered) * pct) - 1)) + return ordered[index] + + +class _CodeCarbonLogCounter(logging.Handler): + """Count ``codecarbon`` INFO lines emitted during a benchmark scenario.""" + + def __init__(self) -> None: + super().__init__(level=logging.INFO) + self.emissions_lines = 0 + + def emit(self, record: logging.LogRecord) -> None: + if record.name != codecarbon_logger.name: + return + if record.levelno < logging.INFO: + return + message = record.getMessage() + if message.startswith("CodeCarbon ") and "emissions=" in message: + self.emissions_lines += 1 + + +async def _run_load_async( + client: httpx.AsyncClient, + url: str, + requests: int, + concurrency: int, +) -> list[float]: + """Issue concurrent async GET requests and return client-side latencies (ms).""" + semaphore = asyncio.Semaphore(concurrency) + + async def _get() -> float: + async with semaphore: + start = time.perf_counter() + response = await client.get(url, timeout=120.0) + response.raise_for_status() + return (time.perf_counter() - start) * 1000 + + return list(await asyncio.gather(*(_get() for _ in range(requests)))) + + +async def _wait_for_deferred_finalize( + measurement_delay_s: float, + *, + requests: int, + concurrency: int, +) -> None: + """Yield until deferred finalize tasks are likely submitted.""" + waves = max(1, (requests + concurrency - 1) // concurrency) + estimate_s = measurement_delay_s * min(waves, 4) + await asyncio.sleep(min(0.06, max(0.01, estimate_s))) + + +def _drain_middleware(app: FastAPI) -> None: + """Wait for deferred tracker work before tearing down an in-process app.""" + shutdown_codecarbon_middleware(app, wait=True) + + +def _summarize( + name: str, + latencies_ms: list[float], + concurrency: int, + baseline_mean_ms: float | None, + *, + bootstrap_samples: int, + use_normal_ci: bool, + codecarbon_log_lines: int | None = None, +) -> BenchmarkResult: + total_s = sum(latencies_ms) / 1000 + mean_ms, ci_low_ms, ci_high_ms, median_ms, p95_ms = summarize_latencies( + latencies_ms, + bootstrap_samples=bootstrap_samples, + use_normal_ci=use_normal_ci, + ) + overhead = None + if baseline_mean_ms and baseline_mean_ms > 0: + overhead = ((mean_ms - baseline_mean_ms) / baseline_mean_ms) * 100 + return BenchmarkResult( + name=name, + requests=len(latencies_ms), + concurrency=concurrency, + mean_ms=mean_ms, + ci_low_ms=ci_low_ms, + ci_high_ms=ci_high_ms, + median_ms=median_ms, + p95_ms=p95_ms, + requests_per_sec=len(latencies_ms) / total_s if total_s else 0.0, + overhead_pct=overhead, + codecarbon_log_lines=codecarbon_log_lines, + ) + + +async def _wait_for_server_async( + client: httpx.AsyncClient, url: str, timeout_s: float = 120.0 +) -> None: + deadline = time.perf_counter() + timeout_s + while time.perf_counter() < deadline: + try: + response = await client.get(url, timeout=30.0) + response.raise_for_status() + return + except (httpx.HTTPError, OSError): + await asyncio.sleep(0.02) + raise RuntimeError(f"Server at {url} did not become ready") + + +async def _run_scenario_in_process( + mode: str, + display_name: str, + requests: int, + warmup: int, + concurrency: int, + workload: InferenceWorkload, + measurement_delay_s: float, + *, + real_tracker: bool, + bootstrap_samples: int, + use_normal_ci: bool, + verify_logging: bool, + logging_sample: int | None, + experiment_id: str, + project_name: str, +) -> BenchmarkResult: + """Benchmark one configuration in-process via ASGI transport.""" + app = build_app( + mode, + workload, + project_name=project_name, + experiment_id=experiment_id, + real_tracker=real_tracker, + ) + workload.ensure_loaded() + log_counter: _CodeCarbonLogCounter | None = None + logging_level_restore: int | None = None + predict_url = "http://benchmark/predict" + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, timeout=120.0) as client: + if warmup > 0: + await _run_load_async(client, predict_url, warmup, concurrency) + if verify_logging and mode == "deferred_logging": + log_counter = _CodeCarbonLogCounter() + logging_level_restore = codecarbon_logger.level + codecarbon_logger.setLevel(logging.INFO) + codecarbon_logger.addHandler(log_counter) + latencies = await _run_load_async(client, predict_url, requests, concurrency) + if mode != "baseline": + drain_s = 0.5 if real_tracker else measurement_delay_s + await _wait_for_deferred_finalize( + drain_s, requests=requests, concurrency=concurrency + ) + if log_counter is not None: + expected_logs = logging_sample or requests + deadline = time.perf_counter() + min( + 2.0, + measurement_delay_s * (requests / max(concurrency, 1) + 2) + 0.25, + ) + while ( + log_counter.emissions_lines < expected_logs + and time.perf_counter() < deadline + ): + await asyncio.sleep(0.005) + log_lines = log_counter.emissions_lines if log_counter is not None else None + if mode != "baseline": + _drain_middleware(app) + if log_counter is not None: + codecarbon_logger.removeHandler(log_counter) + if logging_level_restore is not None: + codecarbon_logger.setLevel(logging_level_restore) + return _summarize( + display_name, + latencies, + concurrency, + None, + bootstrap_samples=bootstrap_samples, + use_normal_ci=use_normal_ci, + codecarbon_log_lines=log_lines, + ) + + +async def _run_scenario_network( + mode: str, + display_name: str, + port: int, + requests: int, + warmup: int, + concurrency: int, + measurement_delay_s: float, + workload: InferenceWorkload, + real_tracker: bool, + *, + bootstrap_samples: int, + use_normal_ci: bool, + verify_logging: bool, + api_delay_s: float = 0.0, + experiment_id: str = FASTAPI_BENCHMARK_EXPERIMENT_ID, + project_name: str = FASTAPI_BENCHMARK_PROJECT_ID, +) -> BenchmarkResult: + import uvicorn + + app = build_app( + mode, + workload, + project_name=project_name, + experiment_id=experiment_id, + real_tracker=real_tracker, + ) + api_patcher = None + uses_save_to_api = mode == "deferred_save_to_api" + if uses_save_to_api and not real_tracker: + api_patcher = _install_api_client_patch(api_delay_s) + api_patcher.start() + + config = uvicorn.Config( + app, host="127.0.0.1", port=port, log_level="error", access_log=False + ) + server = uvicorn.Server(config) + + def _serve() -> None: + server.run() + + thread = threading.Thread(target=_serve, daemon=True) + thread.start() + predict_url = f"http://127.0.0.1:{port}/predict" + log_counter: _CodeCarbonLogCounter | None = None + logging_level_restore: int | None = None + try: + async with httpx.AsyncClient() as client: + await _wait_for_server_async(client, predict_url) + if warmup > 0: + await _run_load_async(client, predict_url, warmup, concurrency) + if mode != "baseline": + finalize_drain_s = ( + 3.0 + if real_tracker + else measurement_delay_s * FINALIZE_DRAIN_MULTIPLIER + ) + time.sleep(finalize_drain_s) + if verify_logging and mode == "deferred_logging": + log_counter = _CodeCarbonLogCounter() + logging_level_restore = codecarbon_logger.level + codecarbon_logger.setLevel(logging.INFO) + codecarbon_logger.addHandler(log_counter) + latencies = await _run_load_async( + client, predict_url, requests, concurrency + ) + if mode != "baseline": + time.sleep( + 3.0 + if real_tracker + else measurement_delay_s * FINALIZE_DRAIN_MULTIPLIER + ) + log_lines = log_counter.emissions_lines if log_counter is not None else None + if log_counter is not None: + codecarbon_logger.removeHandler(log_counter) + if logging_level_restore is not None: + codecarbon_logger.setLevel(logging_level_restore) + return _summarize( + display_name, + latencies, + concurrency, + None, + bootstrap_samples=bootstrap_samples, + use_normal_ci=use_normal_ci, + codecarbon_log_lines=log_lines, + ) + finally: + server.should_exit = True + thread.join(timeout=3.0) + if api_patcher is not None: + api_patcher.stop() + + +def _format_results( + results: list[BenchmarkResult], + *, + workload: str, + model_id: str, + real_tracker: bool, + measurement_delay_ms: float | None, + api_delay_ms: float | None, + with_save_to_api: bool, + experiment_id: str, + project_id: str, + bootstrap_samples: int, + use_normal_ci: bool, + in_process: bool, + logging_verified: bool | None, +) -> str: + confidence_pct = int(CONFIDENCE_LEVEL * 100) + ci_method = ( + f"{confidence_pct}% normal approx" + if use_normal_ci + else f"{confidence_pct}% bootstrap ({bootstrap_samples} resamples)" + ) + transport = "in-process ASGI" if in_process else "HTTP (uvicorn)" + lines = [ + f"Platform: {platform.system()} {platform.release()} ({platform.machine()})", + f"Python: {sys.version.split()[0]}", + f"Workload: {workload} ({model_id})", + f"Transport: {transport}", + "HTTP client: async (httpx.AsyncClient)", + f"EmissionsTracker: {'live' if real_tracker else f'mocked ({measurement_delay_ms:.0f} ms stop delay)'}", + f"save_to_api scenario: {'yes (api_call_interval=1)' if with_save_to_api else 'no'}", + f"project_id: {project_id}", + ( + f"experiment_id (save_to_api): {experiment_id}" + if with_save_to_api + else "experiment_id (save_to_api): n/a" + ), + ( + f"Mocked API upload delay: {api_delay_ms:.0f} ms" + if with_save_to_api and api_delay_ms is not None + else "Mocked API upload delay: n/a" + ), + "Middleware: default deferred measurement", + f"Logger namespace: {codecarbon_logger.name}", + f"Requests per scenario: {results[0].requests} (warmup excluded), " + f"concurrency: {results[0].concurrency}", + f"Mean CI: {ci_method}", + "", + f"| Configuration | Mean (ms) | {confidence_pct}% CI (ms) | Median (ms) | " + f"p95 (ms) | req/s | vs baseline |", + "|---|---:|---|---:|---:|---:|---:|---:|", + ] + for result in results: + ci_cell = f"[{result.ci_low_ms:.1f}, {result.ci_high_ms:.1f}]" + overhead = result.overhead_pct + if overhead is None: + overhead_str = "—" + elif overhead >= 0: + overhead_str = f"+{overhead:.1f}%" + else: + overhead_str = f"{overhead:.1f}%" + lines.append( + f"| {result.name} | {result.mean_ms:.2f} | {ci_cell} | " + f"{result.median_ms:.2f} | {result.p95_ms:.2f} | " + f"{result.requests_per_sec:.1f} | {overhead_str} |" + ) + if logging_verified is not None: + status = "yes" if logging_verified else "no" + lines.append("") + lines.append( + f"CodeCarbon per-request log lines (default middleware): verified={status}" + ) + return "\n".join(lines) + + +SCENARIO_KEYS = { + "no_logging": ("deferred_no_logging", "Deferred, no logging"), + "logging": ("deferred_logging", "Deferred + logging (default)"), + "save_to_api": ("deferred_save_to_api", "Deferred + save_to_api (no logging)"), + "headers": ("sync_headers", "Sync response_headers=True"), + "noop_middleware": ("noop_middleware", "Empty ASGI middleware (stack cost)"), + "logfire": ("logfire_instrumented", "Logfire instrumentation only"), +} + + +async def _run_benchmarks_async( + *, + requests: int, + warmup: int, + secondary_warmup: int, + concurrency: int, + measurement_delay_s: float, + workload_name: str, + model_id: str, + real_tracker: bool, + bootstrap_samples: int, + use_normal_ci: bool, + verify_logging: bool, + logging_sample: int | None, + with_save_to_api: bool, + scenario_keys: list[str] | None, + api_delay_s: float, + experiment_id: str, + project_id: str, + inference_delay_s: float, + in_process: bool, +) -> tuple[list[BenchmarkResult], bool | None]: + """Run baseline and middleware scenarios.""" + workload = InferenceWorkload( + workload_name, model_id, inference_delay_s=inference_delay_s + ) + if workload_name != "noop": + print(f"Preloading workload {workload_name} ({model_id})...", flush=True) + workload.ensure_loaded() + + api_delay_state = {"api": 0.0} + tracker_patcher: Any | None = None + api_patcher: Any | None = None + + async def _run_one( + mode: str, + label: str, + *, + port: int | None, + scenario_warmup: int, + ) -> BenchmarkResult: + if in_process: + return await _run_scenario_in_process( + mode, + label, + requests, + scenario_warmup, + concurrency, + workload, + measurement_delay_s, + real_tracker=real_tracker, + bootstrap_samples=bootstrap_samples, + use_normal_ci=use_normal_ci, + verify_logging=verify_logging, + logging_sample=logging_sample, + experiment_id=experiment_id, + project_name=project_id, + ) + assert port is not None + return await _run_scenario_network( + mode, + label, + port, + requests, + scenario_warmup, + concurrency, + measurement_delay_s, + workload, + real_tracker, + bootstrap_samples=bootstrap_samples, + use_normal_ci=use_normal_ci, + verify_logging=verify_logging, + api_delay_s=api_delay_state["api"], + experiment_id=experiment_id, + project_name=project_id, + ) + + baseline = await _run_one( + "baseline", + "No middleware (baseline)", + port=8765 if not in_process else None, + scenario_warmup=warmup, + ) + + scenarios: list[tuple[str, str]] = [] + selected = scenario_keys or ["no_logging", "logging"] + if with_save_to_api and "save_to_api" not in selected: + selected = [*selected, "save_to_api"] + for key in selected: + if key not in SCENARIO_KEYS: + raise ValueError( + f"Unknown scenario {key!r}; choose from {sorted(SCENARIO_KEYS)}" + ) + scenarios.append(SCENARIO_KEYS[key]) + + if not real_tracker: + tracker_patcher = _install_tracker_patch( + measurement_delay_s, api_delay_state=api_delay_state + ) + tracker_patcher.start() + + results: list[BenchmarkResult] = [baseline] + logging_result: BenchmarkResult | None = None + try: + for index, (mode, label) in enumerate(scenarios): + api_delay_state["api"] = ( + api_delay_s if mode == "deferred_save_to_api" else 0.0 + ) + middleware_warmup = ( + secondary_warmup + if secondary_warmup > 0 + else min(10, warmup) if in_process else warmup + ) + result = await _run_one( + mode, + label, + port=None if in_process else 8766 + index, + scenario_warmup=middleware_warmup if in_process else warmup, + ) + if mode == "deferred_logging": + logging_result = result + results.append(result) + finally: + if api_patcher is not None: + api_patcher.stop() + if tracker_patcher is not None: + tracker_patcher.stop() + + baseline_mean = baseline.mean_ms + enriched: list[BenchmarkResult] = [ + BenchmarkResult( + name=baseline.name, + requests=baseline.requests, + concurrency=baseline.concurrency, + mean_ms=baseline.mean_ms, + ci_low_ms=baseline.ci_low_ms, + ci_high_ms=baseline.ci_high_ms, + median_ms=baseline.median_ms, + p95_ms=baseline.p95_ms, + requests_per_sec=baseline.requests_per_sec, + overhead_pct=None, + ) + ] + for result in results[1:]: + enriched.append( + BenchmarkResult( + name=result.name, + requests=result.requests, + concurrency=result.concurrency, + mean_ms=result.mean_ms, + ci_low_ms=result.ci_low_ms, + ci_high_ms=result.ci_high_ms, + median_ms=result.median_ms, + p95_ms=result.p95_ms, + requests_per_sec=result.requests_per_sec, + overhead_pct=((result.mean_ms - baseline_mean) / baseline_mean * 100), + codecarbon_log_lines=result.codecarbon_log_lines, + ) + ) + + logging_verified: bool | None = None + if logging_result is not None and logging_result.codecarbon_log_lines is not None: + expected_logs = logging_sample or logging_result.requests + logging_verified = logging_result.codecarbon_log_lines >= expected_logs + return enriched, logging_verified + + +def run_benchmarks( + *, + requests: int = BENCHMARK_REQUESTS, + warmup: int = WARMUP_REQUESTS, + secondary_warmup: int = 0, + concurrency: int = CONCURRENCY, + measurement_delay_s: float = DEFAULT_MEASUREMENT_DELAY_S, + workload_name: str, + model_id: str, + real_tracker: bool, + bootstrap_samples: int, + use_normal_ci: bool, + verify_logging: bool, + logging_sample: int | None, + with_save_to_api: bool, + scenario_keys: list[str] | None, + api_delay_s: float, + experiment_id: str, + project_id: str, + inference_delay_s: float, + in_process: bool, +) -> tuple[list[BenchmarkResult], bool | None]: + """Run all scenarios under one asyncio event loop.""" + return asyncio.run( + _run_benchmarks_async( + requests=requests, + warmup=warmup, + secondary_warmup=secondary_warmup, + concurrency=concurrency, + measurement_delay_s=measurement_delay_s, + workload_name=workload_name, + model_id=model_id, + real_tracker=real_tracker, + bootstrap_samples=bootstrap_samples, + use_normal_ci=use_normal_ci, + verify_logging=verify_logging, + logging_sample=logging_sample, + with_save_to_api=with_save_to_api, + scenario_keys=scenario_keys, + api_delay_s=api_delay_s, + experiment_id=experiment_id, + project_id=project_id, + inference_delay_s=inference_delay_s, + in_process=in_process, + ) + ) + + +def _resolve_model_id(workload: str, model_id: str | None) -> str: + if model_id: + return model_id + if workload == "hf-embedder": + return DEFAULT_EMBEDDER_MODEL + if workload == "hf-classifier": + return DEFAULT_CLASSIFIER_MODEL + return "n/a" + + +def main() -> None: # noqa: C901 + """CLI entrypoint.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--requests", type=int, default=BENCHMARK_REQUESTS) + parser.add_argument("--warmup", type=int, default=WARMUP_REQUESTS) + parser.add_argument("--concurrency", type=int, default=CONCURRENCY) + parser.add_argument( + "--bootstrap-samples", + type=int, + default=BOOTSTRAP_SAMPLES, + help="Bootstrap resamples for mean latency CI", + ) + parser.add_argument( + "--workload", + choices=("noop", "hf-embedder", "hf-classifier"), + default="hf-embedder", + ) + parser.add_argument("--model", default=None, help="Hugging Face model id override") + parser.add_argument( + "--real-tracker", + action="store_true", + help="Use a live EmissionsTracker instead of a mocked stop() delay", + ) + parser.add_argument( + "--realistic", + action="store_true", + help=( + "Live tracker + hf-embedder + uvicorn HTTP: " + f"{REALISTIC_BENCHMARK_REQUESTS} requests, concurrency {REALISTIC_CONCURRENCY}" + ), + ) + parser.add_argument( + "--no-verify-logging", + action="store_true", + help="Skip counting codecarbon logger lines after the default scenario", + ) + parser.add_argument( + "--measurement-delay-ms", + type=float, + default=DEFAULT_MEASUREMENT_DELAY_S * 1000, + help="Mocked tracker stop() duration when --real-tracker is not set", + ) + parser.add_argument( + "--with-save-to-api", + action="store_true", + help="Add a scenario with save_to_api=True and api_call_interval=1", + ) + parser.add_argument( + "--with-logfire", + action="store_true", + help="Add noop middleware and Logfire instrumentation comparison scenarios", + ) + parser.add_argument( + "--with-headers", + action="store_true", + help="Add sync response_headers=True scenario (measure on request path)", + ) + parser.add_argument( + "--project-id", + default=FASTAPI_BENCHMARK_PROJECT_ID, + help="CodeCarbon project UUID (middleware project_name for tracked scenarios)", + ) + parser.add_argument( + "--experiment-id", + default=FASTAPI_BENCHMARK_EXPERIMENT_ID, + help="CodeCarbon experiment UUID for the save_to_api scenario", + ) + parser.add_argument( + "--api-delay-ms", + type=float, + default=None, + help="Simulated API upload latency (defaults to --measurement-delay-ms)", + ) + parser.add_argument( + "--smoke", + action="store_true", + help=( + "Fastest run: in-process ASGI, 20 requests, skips log verify, " + "no_logging+logging only" + ), + ) + parser.add_argument( + "--quick", + action="store_true", + help=( + "Fast run: in-process ASGI, noop + 25 ms simulated inference, " + "50 timed requests, normal-approx CI" + ), + ) + parser.add_argument( + "--in-process", + action="store_true", + help="Benchmark via httpx ASGI transport (no uvicorn TCP per scenario)", + ) + parser.add_argument( + "--network", + action="store_true", + help="Force uvicorn HTTP even when --quick is set", + ) + parser.add_argument( + "--inference-delay-ms", + type=float, + default=0.0, + help="Optional sleep per /predict request (useful with --workload noop)", + ) + parser.add_argument( + "--logging-sample", + type=int, + default=None, + help="Verify at least N log lines (default: all requests; quick uses 10)", + ) + parser.add_argument( + "--scenarios", + default=None, + help="Comma-separated middleware scenarios: no_logging, logging, save_to_api, " + "headers, noop_middleware, logfire", + ) + args = parser.parse_args() + if args.realistic: + args.real_tracker = True + args.network = True + args.quick = False + args.workload = "hf-embedder" + if args.requests == BENCHMARK_REQUESTS: + args.requests = REALISTIC_BENCHMARK_REQUESTS + if args.warmup == WARMUP_REQUESTS: + args.warmup = REALISTIC_WARMUP_REQUESTS + if args.concurrency == CONCURRENCY: + args.concurrency = REALISTIC_CONCURRENCY + config_project, config_experiment = _config_ids() + if args.project_id == FASTAPI_BENCHMARK_PROJECT_ID: + args.project_id = config_project + if args.experiment_id == FASTAPI_BENCHMARK_EXPERIMENT_ID: + args.experiment_id = config_experiment + os.environ.setdefault("CODECARBON_ALLOW_MULTIPLE_RUNS", "True") + scenario_keys = ( + [part.strip() for part in args.scenarios.split(",") if part.strip()] + if args.scenarios + else None + ) + if args.with_logfire: + extras = ["noop_middleware", "logfire"] + if scenario_keys is None: + scenario_keys = ["no_logging", "logging", *extras] + else: + for key in extras: + if key not in scenario_keys: + scenario_keys.append(key) + if args.with_headers: + if scenario_keys is None: + scenario_keys = ["no_logging", "logging", "headers"] + elif "headers" not in scenario_keys: + scenario_keys.append("headers") + use_normal_ci = False + secondary_warmup = 0 + logging_sample = args.logging_sample + if args.smoke: + args.quick = True + if args.requests == BENCHMARK_REQUESTS: + args.requests = SMOKE_BENCHMARK_REQUESTS + if args.warmup == WARMUP_REQUESTS: + args.warmup = SMOKE_WARMUP_REQUESTS + if args.inference_delay_ms == 0.0: + args.inference_delay_ms = SMOKE_INFERENCE_DELAY_MS + args.no_verify_logging = True + if scenario_keys is None: + scenario_keys = ["no_logging", "logging"] + if args.quick: + if args.workload == "hf-embedder": + args.workload = "noop" + if args.requests == BENCHMARK_REQUESTS: + args.requests = QUICK_BENCHMARK_REQUESTS + if args.warmup == WARMUP_REQUESTS: + args.warmup = QUICK_WARMUP_REQUESTS + if args.bootstrap_samples == BOOTSTRAP_SAMPLES: + args.bootstrap_samples = QUICK_BOOTSTRAP_SAMPLES + if args.inference_delay_ms == 0.0: + args.inference_delay_ms = QUICK_INFERENCE_DELAY_MS + use_normal_ci = True + secondary_warmup = QUICK_SECONDARY_WARMUP + if logging_sample is None and not args.no_verify_logging: + logging_sample = QUICK_LOGGING_SAMPLE + in_process = (args.in_process or args.quick) and not args.network + if in_process and not args.quick and args.bootstrap_samples == BOOTSTRAP_SAMPLES: + use_normal_ci = False + model_id = _resolve_model_id(args.workload, args.model) + measurement_delay_s = args.measurement_delay_ms / 1000 + api_delay_ms = ( + args.api_delay_ms + if args.api_delay_ms is not None + else args.measurement_delay_ms + ) + api_delay_s = api_delay_ms / 1000 + inference_delay_s = args.inference_delay_ms / 1000 + + previous_log_level = codecarbon_logger.level + codecarbon_logger.setLevel(logging.WARNING) + + results, logging_verified = run_benchmarks( + requests=args.requests, + warmup=args.warmup, + secondary_warmup=secondary_warmup, + concurrency=args.concurrency, + measurement_delay_s=measurement_delay_s, + workload_name=args.workload, + model_id=model_id, + real_tracker=args.real_tracker, + bootstrap_samples=args.bootstrap_samples, + use_normal_ci=use_normal_ci, + verify_logging=not args.no_verify_logging, + logging_sample=logging_sample, + with_save_to_api=args.with_save_to_api, + scenario_keys=scenario_keys, + api_delay_s=api_delay_s, + experiment_id=args.experiment_id, + project_id=args.project_id, + inference_delay_s=inference_delay_s, + in_process=in_process, + ) + codecarbon_logger.setLevel(previous_log_level) + delay_label = None if args.real_tracker else args.measurement_delay_ms + print( + _format_results( + results, + workload=args.workload, + model_id=model_id, + real_tracker=args.real_tracker, + measurement_delay_ms=delay_label or 0.0, + api_delay_ms=api_delay_ms if args.with_save_to_api else None, + with_save_to_api=args.with_save_to_api, + experiment_id=args.experiment_id, + project_id=args.project_id, + bootstrap_samples=args.bootstrap_samples, + use_normal_ci=use_normal_ci, + in_process=in_process, + logging_verified=logging_verified, + ) + ) + if logging_verified is False: + logging_result = results[-1] + print( + f"\nWARNING: expected at least {logging_sample or logging_result.requests} " + f"CodeCarbon log lines, got {logging_result.codecarbon_log_lines}", + file=sys.stderr, + ) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark_gpu_monitoring.py b/scripts/benchmark_gpu_monitoring.py new file mode 100755 index 000000000..23d642916 --- /dev/null +++ b/scripts/benchmark_gpu_monitoring.py @@ -0,0 +1,644 @@ +#!/usr/bin/env python3 +""" +Benchmark: GPU monitoring overhead — heavyweight get_gpu_details vs lightweight get_gpu_utilization_list. + +Measures how many unnecessary NVML calls the per-second _monitor_power() hot path +makes on multi-GPU systems, and the latency difference between the old full-detail +path and the new lightweight utilization-only path. + +Usage: + # Quick run (default) + uv run python scripts/benchmark_gpu_monitoring.py + + # Full benchmark with subprocess cold-start samples + uv run python scripts/benchmark_gpu_monitoring.py all + + # Simulated multi-GPU scale (no real GPU needed) + uv run python scripts/benchmark_gpu_monitoring.py all --simulate-gpus 8 + +Methodology: + - Cold metrics: spawn fresh Python subprocesses, each performing full GPU init + - Warm metrics: repeat calls in the same process after warm-up + - p50 (median) reported across multiple samples + - NVML call counts derived from source code audit (gpu_nvidia.py + gpu_device.py) + - On real NVIDIA hardware: wall-clock timing of actual NVML calls + - On non-NVIDIA hardware: mock NVML with realistic simulated call latencies +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import sys +import time +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +RESULTS_DIR = REPO_ROOT / ".context" +RESULTS_DIR.mkdir(parents=True, exist_ok=True) +DEFAULT_RESULTS = RESULTS_DIR / "gpu-benchmark-results.jsonl" + +# NVML call categories based on source audit (gpu_nvidia.py + gpu_device.py) +# _monitor_power() calls get_gpu_details() every 1s but only uses gpu_utilization +NVML_CALLS_HEAVY = [ + "nvmlDeviceGetMemoryInfo", # → free_memory, total_memory, used_memory — DISCARDED + "nvmlDeviceGetTemperature", # → temperature — DISCARDED + "nvmlDeviceGetPowerUsage", # → power_usage — DISCARDED + "nvmlDeviceGetTotalEnergyConsumption", # → total_energy_consumption — DISCARDED + "nvmlDeviceGetUtilizationRates", # → gpu_utilization — USED + "nvmlDeviceGetComputeMode", # → compute_mode — DISCARDED + "nvmlDeviceGetComputeRunningProcesses", # → compute_processes — DISCARDED (most expensive) + "nvmlDeviceGetGraphicsRunningProcesses", # → graphics_processes — DISCARDED (most expensive) +] + +NVML_CALLS_LIGHTWEIGHT = [ + "nvmlDeviceGetUtilizationRates", # ← the only call we need for utilization +] + +# Simulated per-call latencies (microseconds) for non-GPU systems. +# Based on typical NVML overheads reported in NVIDIA docs & community benchmarks. +# Process enumeration (GetComputeRunningProcesses) is the most expensive because +# it iterates active GPU processes and collects PID-level info. +SIMULATED_LATENCY_US: dict[str, float] = { + "nvmlDeviceGetMemoryInfo": 50, + "nvmlDeviceGetTemperature": 40, + "nvmlDeviceGetPowerUsage": 45, + "nvmlDeviceGetTotalEnergyConsumption": 40, + "nvmlDeviceGetUtilizationRates": 50, + "nvmlDeviceGetComputeMode": 35, + "nvmlDeviceGetComputeRunningProcesses": 500, # ← expensive: process enumeration + "nvmlDeviceGetGraphicsRunningProcesses": 500, # ← expensive: process enumeration + "nvmlDeviceGetName": 40, + "nvmlDeviceGetUUID": 35, + "nvmlDeviceGetEnforcedPowerLimit": 40, +} + + +@dataclass +class LatencyStats: + count: int = 0 + min_ms: float = 0.0 + max_ms: float = 0.0 + mean_ms: float = 0.0 + p50_ms: float = 0.0 + p95_ms: float = 0.0 + + +@dataclass +class NvmlCallBreakdown: + call_name: str + latency_us: float + used_by_monitor: bool + + +@dataclass +class GpuDetailMethodBenchmark: + method: str # "get_gpu_details" or "get_gpu_utilization_list" + gpu_count: int + nvml_calls_per_second: int + nvml_calls_unused_per_second: int + latency_per_call_ms: LatencyStats + latency_per_second_ms: float # projected = per_gpu * gpu_count + + +@dataclass +class MonitoringOverheadProjection: + metric: str + heavy_path: float + lightweight_path: float + savings: float + unit: str + + +@dataclass +class BenchmarkReport: + timestamp: str + mode: str + hostname: str + gpu_backend: str + gpu_count_real: int + simulated: bool + call_breakdown: list[dict] + method_benchmarks: list[dict] + projections: list[dict] + result: str = "" + + +def _now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _percentile(sorted_values: list[float], pct: float) -> float: + if not sorted_values: + return 0.0 + if len(sorted_values) == 1: + return sorted_values[0] + k = (len(sorted_values) - 1) * (pct / 100.0) + f = int(k) + c = min(f + 1, len(sorted_values) - 1) + if f == c: + return sorted_values[f] + return sorted_values[f] + (sorted_values[c] - sorted_values[f]) * (k - f) + + +def compute_stats(values_ms: list[float]) -> LatencyStats: + if not values_ms: + return LatencyStats(count=0) + s = sorted(values_ms) + return LatencyStats( + count=len(s), + min_ms=s[0], + max_ms=s[-1], + mean_ms=statistics.mean(s), + p50_ms=_percentile(s, 50), + p95_ms=_percentile(s, 95), + ) + + +def _detect_gpu_backend() -> tuple[str, int]: + """Detect real GPU backend and count. Returns (backend_name, count).""" + try: + from codecarbon.core.gpu import AMDSMI_AVAILABLE, PYNVML_AVAILABLE + + if PYNVML_AVAILABLE: + from codecarbon.core import gpu_nvidia + + count = gpu_nvidia.pynvml.nvmlDeviceGetCount() + return ("nvidia", count) + if AMDSMI_AVAILABLE: + return ("amd", 0) # count not trivial + except Exception: + pass + return ("none", 0) + + +def _collect_call_breakdown() -> list[dict]: + """Return the per-NVML-call breakdown showing what's used vs discarded.""" + results = [] + for call in NVML_CALLS_HEAVY: + results.append( + { + "call_name": call, + "used_by_monitor": call == "nvmlDeviceGetUtilizationRates", + "simulated_latency_us": SIMULATED_LATENCY_US.get(call, 50), + } + ) + return results + + +def _mock_time_for_call(call_name: str) -> None: + """Sleep to simulate NVML call latency when no real GPU is available.""" + time.sleep(SIMULATED_LATENCY_US.get(call_name, 50) / 1_000_000) + + +class MockNvidiaGPUDevice: + """A lightweight mock that simulates NVML call latencies. + + Used on non-NVIDIA systems so the benchmark can still measure + relative overhead and project multi-GPU scaling. + """ + + def __init__(self, gpu_index: int): + self.gpu_index = gpu_index + + def get_gpu_details(self) -> dict: + _mock_time_for_call("nvmlDeviceGetMemoryInfo") + _mock_time_for_call("nvmlDeviceGetTemperature") + _mock_time_for_call("nvmlDeviceGetPowerUsage") + _mock_time_for_call("nvmlDeviceGetTotalEnergyConsumption") + _mock_time_for_call("nvmlDeviceGetUtilizationRates") + _mock_time_for_call("nvmlDeviceGetComputeMode") + _mock_time_for_call("nvmlDeviceGetComputeRunningProcesses") + _mock_time_for_call("nvmlDeviceGetGraphicsRunningProcesses") + return {"gpu_index": self.gpu_index, "gpu_utilization": 50} + + def get_gpu_utilization_lightweight(self) -> dict: + _mock_time_for_call("nvmlDeviceGetUtilizationRates") + return {"gpu_index": self.gpu_index, "gpu_utilization": 50} + + +def _benchmark_method( + devices: list, + method_name: str, + samples: int = 200, + warmup: int = 20, +) -> LatencyStats: + """Benchmark a GPU method. Returns latency stats in milliseconds.""" + for _ in range(warmup): + if method_name == "get_gpu_details": + [d.get_gpu_details() for d in devices] + else: + [d.get_gpu_utilization_lightweight() for d in devices] + + timings = [] + for _ in range(samples): + t0 = time.perf_counter() + if method_name == "get_gpu_details": + [d.get_gpu_details() for d in devices] + else: + [d.get_gpu_utilization_lightweight() for d in devices] + elapsed_ms = (time.perf_counter() - t0) * 1000 + timings.append(elapsed_ms) + + return compute_stats(timings) + + +def _benchmark_real_gpu(gpu_count: int) -> tuple[list[dict], list[dict]]: + """Benchmark using real GPU hardware via AllGPUDevices.""" + sys.path.insert(0, str(REPO_ROOT)) + from codecarbon.core.gpu import AllGPUDevices + + devices = AllGPUDevices() + actual_count = devices.device_count + + heavy_stats = _benchmark_method(devices.devices, "get_gpu_details") + light_stats = _benchmark_method(devices.devices, "get_gpu_utilization_lightweight") + + method_benchmarks = [ + { + "method": "get_gpu_details", + "gpu_count": actual_count, + "nvml_calls_per_second": len(NVML_CALLS_HEAVY) * actual_count, + "nvml_calls_unused_per_second": (len(NVML_CALLS_HEAVY) - 1) * actual_count, + "latency_per_call_ms": asdict(heavy_stats), + "latency_per_second_ms": heavy_stats.p50_ms, + }, + { + "method": "get_gpu_utilization_list", + "gpu_count": actual_count, + "nvml_calls_per_second": len(NVML_CALLS_LIGHTWEIGHT) * actual_count, + "nvml_calls_unused_per_second": 0, + "latency_per_call_ms": asdict(light_stats), + "latency_per_second_ms": light_stats.p50_ms, + }, + ] + + # Scale projections for multi-GPU + for simulated_count in [1, 4, 8]: + scale = simulated_count / actual_count if actual_count else 1 + method_benchmarks.append( + { + "method": f"get_gpu_details (projected {simulated_count} GPU)", + "gpu_count": simulated_count, + "nvml_calls_per_second": len(NVML_CALLS_HEAVY) * simulated_count, + "nvml_calls_unused_per_second": (len(NVML_CALLS_HEAVY) - 1) + * simulated_count, + "latency_per_call_ms": asdict(heavy_stats), + "latency_per_second_ms": heavy_stats.p50_ms * scale, + } + ) + method_benchmarks.append( + { + "method": f"get_gpu_utilization_list (projected {simulated_count} GPU)", + "gpu_count": simulated_count, + "nvml_calls_per_second": len(NVML_CALLS_LIGHTWEIGHT) * simulated_count, + "nvml_calls_unused_per_second": 0, + "latency_per_call_ms": asdict(light_stats), + "latency_per_second_ms": light_stats.p50_ms * scale, + } + ) + + return method_benchmarks, [] + + +def _benchmark_simulated_gpu(simulate_gpus: int) -> tuple[list[dict], list[dict]]: + """Benchmark using mock devices with simulated NVML latencies.""" + devices = [MockNvidiaGPUDevice(i) for i in range(simulate_gpus)] + + heavy_stats = _benchmark_method(devices, "get_gpu_details") + light_stats = _benchmark_method(devices, "get_gpu_utilization_lightweight") + + method_benchmarks = [ + { + "method": "get_gpu_details", + "gpu_count": simulate_gpus, + "nvml_calls_per_second": len(NVML_CALLS_HEAVY) * simulate_gpus, + "nvml_calls_unused_per_second": (len(NVML_CALLS_HEAVY) - 1) * simulate_gpus, + "latency_per_call_ms": asdict(heavy_stats), + "latency_per_second_ms": heavy_stats.p50_ms, + }, + { + "method": "get_gpu_utilization_list", + "gpu_count": simulate_gpus, + "nvml_calls_per_second": len(NVML_CALLS_LIGHTWEIGHT) * simulate_gpus, + "nvml_calls_unused_per_second": 0, + "latency_per_call_ms": asdict(light_stats), + "latency_per_second_ms": light_stats.p50_ms, + }, + ] + + return method_benchmarks, [] + + +def _compute_projections(method_benchmarks: list[dict]) -> list[dict]: + """Compute time-savings projections from benchmark results.""" + heavy = next( + (m for m in method_benchmarks if m["method"] == "get_gpu_details"), None + ) + light = next( + (m for m in method_benchmarks if m["method"] == "get_gpu_utilization_list"), + None, + ) + if not heavy or not light: + return [] + + heavy_per_sec = heavy["latency_per_second_ms"] + light_per_sec = light["latency_per_second_ms"] + savings_per_sec = heavy_per_sec - light_per_sec + + gpu_count = heavy["gpu_count"] + + return [ + { + "metric": "Per-second monitoring overhead", + "heavy_path_ms": heavy_per_sec, + "lightweight_path_ms": light_per_sec, + "savings_ms": savings_per_sec, + "savings_pct": ( + round((savings_per_sec / heavy_per_sec) * 100, 1) + if heavy_per_sec + else 0 + ), + "unit": "ms/s", + }, + { + "metric": "Per-minute monitoring overhead", + "heavy_path_ms": heavy_per_sec * 60, + "lightweight_path_ms": light_per_sec * 60, + "savings_ms": savings_per_sec * 60, + "savings_pct": ( + round((savings_per_sec / heavy_per_sec) * 100, 1) + if heavy_per_sec + else 0 + ), + "unit": "ms/min", + }, + { + "metric": "Per-hour monitoring overhead", + "heavy_path_ms": heavy_per_sec * 3600, + "lightweight_path_ms": light_per_sec * 3600, + "savings_ms": savings_per_sec * 3600, + "savings_pct": ( + round((savings_per_sec / heavy_per_sec) * 100, 1) + if heavy_per_sec + else 0 + ), + "unit": "ms/hr", + }, + { + "metric": "Per-day monitoring overhead (24h)", + "heavy_path_ms": heavy_per_sec * 86400, + "lightweight_path_ms": light_per_sec * 86400, + "savings_ms": savings_per_sec * 86400, + "savings_pct": ( + round((savings_per_sec / heavy_per_sec) * 100, 1) + if heavy_per_sec + else 0 + ), + "unit": "ms/day", + }, + { + "metric": "Unnecessary NVML calls per second", + "heavy_path_value": heavy["nvml_calls_unused_per_second"], + "lightweight_path_value": 0, + "savings_value": heavy["nvml_calls_unused_per_second"], + "unit": "calls/s", + }, + { + "metric": f"Unnecessary NVML calls per hour (on {gpu_count} GPU{'s' if gpu_count != 1 else ''})", + "heavy_path_value": heavy["nvml_calls_unused_per_second"] * 3600, + "lightweight_path_value": 0, + "savings_value": heavy["nvml_calls_unused_per_second"] * 3600, + "unit": "calls/hr", + }, + ] + + +def run_all(simulate_gpus: int | None = None) -> BenchmarkReport: + backend, real_count = _detect_gpu_backend() + simulated = backend == "none" and simulate_gpus is not None + + if backend != "none" and real_count > 0: + gpu_backend = f"nvidia ({real_count} GPU{'s' if real_count != 1 else ''})" + method_bms, _ = _benchmark_real_gpu(real_count) + elif simulate_gpus: + gpu_backend = ( + f"simulated ({simulate_gpus} GPU{'s' if simulate_gpus != 1 else ''})" + ) + method_bms, _ = _benchmark_simulated_gpu(simulate_gpus) + else: + gpu_backend = "none (no GPU available, use --simulate-gpus N)" + method_bms = [] + + projections = _compute_projections(method_bms) if method_bms else [] + + call_breakdown = _collect_call_breakdown() + + return BenchmarkReport( + timestamp=_now_iso(), + mode="all", + hostname=os.uname().nodename, + gpu_backend=gpu_backend, + gpu_count_real=real_count, + simulated=simulated, + call_breakdown=call_breakdown, + method_benchmarks=method_bms, + projections=projections, + ) + + +def print_report(report: BenchmarkReport) -> None: + sep = "─" * 72 + + print(f"\n{' GPU Monitoring Overhead Benchmark ':=^72}") + print(f" Host: {report.hostname}") + print(f" GPU backend: {report.gpu_backend}") + print(f" Simulated: {report.simulated}") + print(f" Timestamp: {report.timestamp}") + + if report.simulated: + print(f"\n{' ⚠ SIMULATED — No real GPU detected ':=^72}") + print(" Call latencies are estimated (see SIMULATED_LATENCY_US in script).") + print(" Run this on an NVIDIA GPU machine for real hardware measurements.") + + # NVML call breakdown + print(f"\n{sep}") + print(f"{' NVML Call Breakdown (per GPU, per call to get_gpu_details) ':=^72}") + print(f"{'NVML Call':40s} {'Latency (µs)':15s} {'Used by monitor':20s}") + print("-" * 72) + for cb in report.call_breakdown: + used = "YES" if cb["used_by_monitor"] else "" + print( + f"{cb['call_name']:40s} {cb['simulated_latency_us']:>10.0f} µs {used:20s}" + ) + + unused = sum(1 for cb in report.call_breakdown if not cb["used_by_monitor"]) + total = len(report.call_breakdown) + print(f"\n → {unused}/{total} NVML calls DISCARDED by _monitor_power()") + print(f" → Only 1/{total} calls actually used (gpu_utilization)") + + # Method benchmarks + if report.method_benchmarks: + print(f"\n{sep}") + print(f"{' Method Latency Benchmarks ':=^72}") + print( + f"{'Method':50s} {'p50':>8s} {'mean':>8s} {'p95':>8s} {'NVML calls/s':>14s}" + ) + print("-" * 72) + for mb in report.method_benchmarks: + lat = mb["latency_per_call_ms"] + print( + f"{mb['method']:50s} " + f"{lat['p50_ms']:>7.2f}ms {lat['mean_ms']:>7.2f}ms {lat['p95_ms']:>7.2f}ms " + f"{mb['nvml_calls_per_second']:>8d}/s" + ) + + # Projections + if report.projections: + print(f"\n{sep}") + print(f"{' Projected Savings (heavyweight → lightweight) ':=^72}") + print(f"{'Metric':50s} {'Heavy':>12s} {'Light':>12s} {'Savings':>12s}") + print("-" * 72) + for p in report.projections: + if "savings_pct" in p: + print( + f"{p['metric']:50s} " + f"{p['heavy_path_ms']:>8.1f}ms {p['lightweight_path_ms']:>8.1f}ms " + f"{p['savings_ms']:>8.1f}ms ({p['savings_pct']}%)" + ) + else: + print( + f"{p['metric']:50s} " + f"{p['heavy_path_value']:>12,d} {p['lightweight_path_value']:>12,d} " + f"{p['savings_value']:>12,d}" + ) + + print(f"\n{sep}") + print(f"{' Summary ':=^72}") + if report.projections: + hourly = next( + ( + p + for p in report.projections + if p["metric"] == "Per-hour monitoring overhead" + ), + None, + ) + daily = next( + ( + p + for p in report.projections + if p["metric"] == "Per-day monitoring overhead (24h)" + ), + None, + ) + nvml_daily = next( + (p for p in report.projections if "NVML calls per hour" in p["metric"]), + None, + ) + if hourly: + print( + f" Each second of monitoring saves {hourly['savings_ms'] / 3600:.3f} ms" + ) + print( + f" Per hour of continuous monitoring saves {hourly['savings_ms'] / 1000:.1f} s" + ) + if daily: + print( + f" Per 24h day of monitoring saves {daily['savings_ms'] / 1000:.0f} s ({daily['savings_ms'] / 60000:.1f} min)" + ) + if nvml_daily: + print( + f" Unnecessary NVML calls per 24h: {nvml_daily['savings_value'] * 24:,d}" + ) + print(f"{'=' * 72}\n") + + +def run_cold_subprocess(simulate_gpus: int | None = None) -> BenchmarkReport: + """Spawn a fresh subprocess to measure cold-start GPU detection overhead.""" + cmd = [ + sys.executable, + __file__, + "cold", + "--json", + ] + if simulate_gpus: + cmd.extend(["--simulate-gpus", str(simulate_gpus)]) + env = os.environ.copy() + t0 = time.perf_counter() + proc = subprocess.run(cmd, capture_output=True, text=True, timeout=60, env=env) + elapsed_ms = (time.perf_counter() - t0) * 1000 + if proc.returncode != 0: + print(f"Subprocess failed: {proc.stderr[:500]}") + return BenchmarkReport( + timestamp=_now_iso(), + mode="cold_subprocess", + hostname=os.uname().nodename, + gpu_backend="error", + gpu_count_real=0, + simulated=False, + call_breakdown=[], + method_benchmarks=[], + projections=[], + result="error", + ) + report = json.loads(proc.stdout) + report["mode"] = "cold_subprocess" + report["result"] = f"cold_subprocess_overhead_ms={elapsed_ms:.1f}" + return BenchmarkReport(**report) + + +def main() -> None: + p = argparse.ArgumentParser(description="GPU monitoring overhead benchmark") + p.add_argument("mode", nargs="?", default="quick", choices=["quick", "all", "cold"]) + p.add_argument( + "--simulate-gpus", + type=int, + default=None, + help="Simulate N GPUs (default: auto-detect)", + ) + p.add_argument( + "--json", action="store_true", help="Output JSON (for subprocess consumption)" + ) + p.add_argument("--results-file", type=Path, default=DEFAULT_RESULTS) + args = p.parse_args() + + if args.mode == "quick": + report = run_all(args.simulate_gpus) + print_report(report) + + elif args.mode == "all": + report = run_all(args.simulate_gpus) + if args.json: + print(json.dumps(asdict(report), default=str)) + else: + print_report(report) + + # Also run cold subprocess if not already in one + if not args.json and not os.environ.get("_BENCHMARK_CHILD"): + print("\n--- Cold subprocess benchmark ---") + cold_report = run_cold_subprocess(args.simulate_gpus) + print(f"Cold subprocess overhead: {cold_report.result}") + + elif args.mode == "cold": + os.environ["_BENCHMARK_CHILD"] = "1" + report = run_all(args.simulate_gpus) + if args.json: + print(json.dumps(asdict(report), default=str)) + else: + print_report(report) + + # Append to results file + if not args.json and args.mode != "cold": + with open(args.results_file, "a") as f: + f.write(json.dumps(asdict(report), default=str) + "\n") + print(f"→ Results appended to {args.results_file}") + + +if __name__ == "__main__": + main() diff --git a/scripts/repro_fastapi_concurrency.py b/scripts/repro_fastapi_concurrency.py new file mode 100644 index 000000000..6f0f662ef --- /dev/null +++ b/scripts/repro_fastapi_concurrency.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +"""Repro: live tracker + concurrent FastAPI requests. + +Uses the embedder from examples/fastapi_concurrency.py. Pass --lazy to exercise +middleware-only tracker startup (no lifespan). + + uv run --extra fastapi --with uvicorn --with httpx --with sentence-transformers --with torch \\ + python scripts/repro_fastapi_concurrency.py + + uv run ... python scripts/repro_fastapi_concurrency.py --lazy +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import sys +from pathlib import Path + +import httpx +import uvicorn + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from examples.fastapi_concurrency import SAMPLE_TEXT, app_lazy, app_lifespan + +logging.basicConfig(level=logging.ERROR) +codecarbon_logger = logging.getLogger("codecarbon") +errors: list[str] = [] + + +class _ErrorCapture(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if record.levelno >= logging.ERROR and "_active_task_emissions_at_start" in record.getMessage(): + errors.append(record.getMessage()) + + +codecarbon_logger.addHandler(_ErrorCapture()) + + +async def fire_requests(base_url: str, n: int, concurrency: int) -> None: + sem = asyncio.Semaphore(concurrency) + + async with httpx.AsyncClient(base_url=base_url, timeout=120.0) as client: + await client.get("/embed", params={"text": SAMPLE_TEXT}) + + async def one() -> None: + async with sem: + response = await client.get("/embed", params={"text": SAMPLE_TEXT}) + response.raise_for_status() + + await asyncio.gather(*(one() for _ in range(n))) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--lazy", + action="store_true", + help="Use app_lazy (middleware creates tracker) instead of lifespan app", + ) + parser.add_argument("--requests", type=int, default=20) + parser.add_argument("--concurrency", type=int, default=4) + args = parser.parse_args() + + app = app_lazy if args.lazy else app_lifespan + mode = "lazy" if args.lazy else "lifespan" + host, port = "127.0.0.1", 0 + config = uvicorn.Config(app, host=host, port=port, log_level="error") + server = uvicorn.Server(config) + + async def run() -> None: + serve_task = asyncio.create_task(server.serve()) + while not server.started: + await asyncio.sleep(0.05) + bound_port = server.servers[0].sockets[0].getsockname()[1] + base_url = f"http://{host}:{bound_port}" + await fire_requests( + base_url, + n=args.requests, + concurrency=args.concurrency, + ) + server.should_exit = True + await serve_task + + asyncio.run(run()) + + if errors: + print( + f"FAIL [{mode}]: {len(errors)} _active_task_emissions_at_start error(s)", + file=sys.stderr, + ) + for msg in errors[:5]: + print(f" {msg}", file=sys.stderr) + return 1 + print(f"OK [{mode}]: no concurrency errors ({args.requests} req, c={args.concurrency})") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/verify_fastapi_middleware_outputs.py b/scripts/verify_fastapi_middleware_outputs.py new file mode 100644 index 000000000..edd026daa --- /dev/null +++ b/scripts/verify_fastapi_middleware_outputs.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +"""Verify FastAPI middleware logging, CSV, and optional API upload. + +Per-request emissions appear in logs via ``on_request_complete`` (default). +CSV rows and API ``add_emission`` calls are written when the shared tracker +stops (use ``create_codecarbon_lifespan``), not after each ``stop_task``. + +Examples: + uv run --extra fastapi python scripts/verify_fastapi_middleware_outputs.py + uv run --extra fastapi python scripts/verify_fastapi_middleware_outputs.py --save-to-api +""" + +from __future__ import annotations + +import argparse +import logging +import sys +import tempfile +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any + +import requests +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import codecarbon.integrations.fastapi.middleware as cc_fastapi_middleware +from codecarbon.core.api_client import ApiClient +from codecarbon.core.config import get_hierarchical_config +from codecarbon.integrations.fastapi import ( + add_codecarbon_middleware, + create_codecarbon_lifespan, +) +from codecarbon.integrations.fastapi.middleware import log_request_complete + + +class _LogCounter(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.INFO) + self.request_log_lines = 0 + + def emit(self, record: logging.LogRecord) -> None: + if record.name != "codecarbon": + return + message = record.getMessage() + if message.startswith("CodeCarbon ") and "emissions=" in message: + self.request_log_lines += 1 + + +def _build_app( + *, + output_dir: Path, + save_to_api: bool, + project_name: str, +) -> FastAPI: + tracker_kwargs: dict[str, Any] = { + "save_to_file": True, + "save_to_api": save_to_api, + "save_to_logger": False, + "output_dir": str(output_dir), + "measure_power_secs": 2, + "api_call_interval": 1, + "allow_multiple_runs": True, + } + + @asynccontextmanager + async def lifespan(application: FastAPI): + async with create_codecarbon_lifespan( + application, + project_name=project_name, + **tracker_kwargs, + ): + yield + + application = FastAPI(lifespan=lifespan) + add_codecarbon_middleware( + application, + project_name=project_name, + tracker_kwargs=tracker_kwargs, + on_request_complete=log_request_complete, + ) + + @application.get("/predict") + def predict(text: str = "hello") -> dict[str, str]: + return {"text": text, "label": "demo"} + + return application + + +def _count_run_emissions(api: ApiClient, run_id: str) -> int: + url = f"{api.url}/runs/{run_id}/emissions" + response = requests.get(url, headers=api._get_headers(), timeout=15) + if response.status_code != 200: + return 0 + payload = response.json() + items = payload.get("items") or payload.get("data") or [] + if isinstance(items, list): + return len(items) + return 0 + + +def _get_api_client_from_config() -> ApiClient | None: + conf = get_hierarchical_config() + section = conf.get("codecarbon", conf) + api_key = section.get("api_key") or section.get("api_token") + experiment_id = section.get("experiment_id") + endpoint = section.get("api_endpoint") or "https://api.codecarbon.io" + if not api_key or not experiment_id: + return None + return ApiClient( + endpoint_url=endpoint, + experiment_id=experiment_id, + api_key=api_key, + conf=conf, + create_run_automatically=False, + ) + + +def main(argv: list[str] | None = None) -> int: # noqa: C901 + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--save-to-api", + action="store_true", + help="Enable save_to_api using ~/.codecarbon.config (requires api_key).", + ) + parser.add_argument( + "--requests", + type=int, + default=3, + help="Number of GET /predict calls (default: 3).", + ) + args = parser.parse_args(argv) + + save_to_api = args.save_to_api + if save_to_api: + api_probe = _get_api_client_from_config() + if api_probe is None: + print( + "ERROR: --save-to-api needs api_key and experiment_id in " + "~/.codecarbon.config", + file=sys.stderr, + ) + return 1 + if api_probe.check_auth() is None: + print( + "WARN: /auth/check failed; continuing (upload probe uses run emissions)." + ) + + log_counter = _LogCounter() + cc_fastapi_middleware.logger.addHandler(log_counter) + + failures: list[str] = [] + try: + with tempfile.TemporaryDirectory(prefix="cc-fastapi-verify-") as tmp: + output_dir = Path(tmp) + app = _build_app( + output_dir=output_dir, + save_to_api=save_to_api, + project_name="fastapi-verify", + ) + run_id: str | None = None + with TestClient(app) as client: + for _ in range(args.requests): + response = client.get("/predict", params={"text": "verify"}) + if response.status_code != 200: + failures.append( + f"predict returned status {response.status_code}" + ) + break + tracker = getattr(app.state, "codecarbon_tracker", None) + if tracker is not None: + for handler in tracker._output_handlers: + handler_run_id = getattr(handler, "run_id", None) + if handler_run_id: + run_id = handler_run_id + break + + if log_counter.request_log_lines < args.requests: + failures.append( + f"expected {args.requests} per-request log lines, got " + f"{log_counter.request_log_lines}" + ) + else: + print( + f"OK: {log_counter.request_log_lines} per-request log line(s) " + "(on_request_complete)" + ) + + emissions_csv = output_dir / "emissions.csv" + if not emissions_csv.is_file() or emissions_csv.stat().st_size == 0: + failures.append( + f"missing or empty CSV at {emissions_csv} (written on tracker.stop)" + ) + else: + line_count = len(emissions_csv.read_text().splitlines()) + print( + f"OK: CSV {emissions_csv} ({line_count} line(s) including header)" + ) + + task_csvs = list(output_dir.glob("emissions_*.csv")) + if task_csvs: + print(f"OK: task CSV(s): {', '.join(p.name for p in task_csvs)}") + else: + print( + "NOTE: no per-task CSV (emissions__.csv); " + "run-level emissions.csv is the main artifact on stop" + ) + + if save_to_api: + api = _get_api_client_from_config() + if api is None or run_id is None: + failures.append("could not resolve API client or run_id after stop") + else: + count = _count_run_emissions(api, run_id) + if count < 1: + failures.append( + f"no emissions listed for run {run_id} at " + f"{api.url}/runs/.../emissions" + ) + else: + print(f"OK: API run {run_id} has {count} emission record(s)") + finally: + cc_fastapi_middleware.logger.removeHandler(log_counter) + + if failures: + for msg in failures: + print(f"FAIL: {msg}", file=sys.stderr) + return 1 + + print("All checks passed.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/integrations/test_fastapi_import.py b/tests/integrations/test_fastapi_import.py new file mode 100644 index 000000000..5289310ec --- /dev/null +++ b/tests/integrations/test_fastapi_import.py @@ -0,0 +1,51 @@ +"""Import surface for the optional FastAPI integration package.""" + +import builtins +import importlib +import sys + +import pytest + + +def test_fastapi_integration_importable() -> None: + """Public helpers are importable without instantiating middleware.""" + from codecarbon.integrations.fastapi import ( + CodeCarbonMiddleware, + add_codecarbon_middleware, + create_codecarbon_lifespan, + log_request_complete, + shutdown_codecarbon_middleware, + ) + + assert CodeCarbonMiddleware is not None + assert callable(add_codecarbon_middleware) + assert callable(create_codecarbon_lifespan) + assert callable(log_request_complete) + assert callable(shutdown_codecarbon_middleware) + + +def test_missing_starlette_shows_helpful_error(monkeypatch: pytest.MonkeyPatch) -> None: + """Middleware import surfaces an actionable hint without Starlette/FastAPI.""" + for key in list(sys.modules): + if key.startswith("starlette") or key.startswith( + "codecarbon.integrations.fastapi" + ): + del sys.modules[key] + + real_import = builtins.__import__ + + def mock_import( + name: str, + globals: dict | None = None, + locals: dict | None = None, + fromlist: tuple[str, ...] = (), + level: int = 0, + ): + root = name.split(".", 1)[0] + if root in ("starlette", "fastapi"): + raise ImportError("no starlette") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", mock_import) + with pytest.raises(ImportError, match=r"pip install .*codecarbon\[fastapi\]"): + importlib.import_module("codecarbon.integrations.fastapi.middleware") diff --git a/tests/integrations/test_fastapi_lifespan.py b/tests/integrations/test_fastapi_lifespan.py new file mode 100644 index 000000000..eae624d28 --- /dev/null +++ b/tests/integrations/test_fastapi_lifespan.py @@ -0,0 +1,29 @@ +import asyncio +from unittest.mock import MagicMock, patch + +import pytest +from fastapi import FastAPI + +import codecarbon.integrations.fastapi.lifespan as cc_fastapi_lifespan +from codecarbon.integrations.fastapi.lifespan import create_codecarbon_lifespan + + +@pytest.fixture +def app(): + return FastAPI() + + +@patch.object(cc_fastapi_lifespan, "EmissionsTracker") +def test_lifespan_stops_tracker_on_shutdown(MockTracker, app): + tracker = MagicMock() + MockTracker.return_value = tracker + + async def run(): + async with create_codecarbon_lifespan(app, project_name="api"): + assert app.state.codecarbon_tracker is tracker + tracker.start.assert_called_once() + + asyncio.run(run()) + + tracker.stop.assert_called_once() + assert app.state.codecarbon_tracker is None diff --git a/tests/integrations/test_fastapi_middleware.py b/tests/integrations/test_fastapi_middleware.py new file mode 100644 index 000000000..b1d021f73 --- /dev/null +++ b/tests/integrations/test_fastapi_middleware.py @@ -0,0 +1,1051 @@ +import asyncio +import logging +from concurrent import futures +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import codecarbon.integrations.fastapi.lifespan as cc_fastapi_lifespan +import codecarbon.integrations.fastapi.middleware as cc_fastapi_middleware +from codecarbon.external.logger import logger as codecarbon_logger +from codecarbon.integrations.fastapi import ( + add_codecarbon_middleware, + create_codecarbon_lifespan, + shutdown_codecarbon_middleware, +) +from codecarbon.integrations.fastapi.middleware import log_request_complete + + +def _configure_mock_running_tracker( + tracker_instance: MagicMock, + *, + task_name: str = "GET /predict", + emissions: float = 0.001, +) -> MagicMock: + """Mock a started tracker that uses mark/finish HTTP paths (concurrency-safe).""" + baseline = MagicMock(task_name=task_name) + + def mark_started() -> None: + tracker_instance._start_time = 1.0 + + tracker_instance.start.side_effect = mark_started + tracker_instance._start_time = None + tracker_instance.mark_http_request_start.return_value = baseline + tracker_instance.finish_http_request.return_value = MagicMock(emissions=emissions) + return baseline + + +def _run_finalize_immediately(coro: Any) -> None: + def run_in_thread() -> None: + loop = asyncio.new_event_loop() + try: + loop.run_until_complete(coro) + finally: + loop.close() + + futures.ThreadPoolExecutor(max_workers=1).submit(run_in_thread).result() + + +@pytest.fixture(autouse=True) +def finalize_deferred_immediately(request): + if request.node.get_closest_marker("no_immediate_finalize"): + yield + return + with patch.object( + cc_fastapi_middleware.CodeCarbonMiddleware, + "_schedule_finalize", + side_effect=_run_finalize_immediately, + ): + yield + + +@pytest.fixture +def app(): + application = FastAPI() + + @application.get("/items/{item_id}") + def get_item(item_id: int): + return {"item_id": item_id} + + @application.get("/health") + def health(): + return {"ok": True} + + add_codecarbon_middleware(application, project_name="test-api") + return application + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_middleware_tracks_routed_request(MockTracker, app) -> None: + tracker_instance = MockTracker.return_value + _configure_mock_running_tracker(tracker_instance, task_name="GET /items/7") + + response = TestClient(app).get("/items/7") + + assert response.status_code == 200 + MockTracker.assert_called_once() + tracker_instance.start.assert_called_once() + tracker_instance.mark_http_request_start.assert_called_once() + tracker_instance.finish_http_request.assert_called_once() + tracker_instance.persist_completed_task.assert_called_once() + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_middleware_skips_excluded_paths(MockTracker, app) -> None: + response = TestClient(app).get("/health") + assert response.status_code == 200 + MockTracker.assert_not_called() + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_middleware_on_request_complete_callback(MockTracker) -> None: + application = FastAPI() + completed = [] + + @application.get("/predict") + def predict(): + return {"ok": True} + + add_codecarbon_middleware( + application, + on_request_complete=lambda request, response, data, task_name: completed.append( + (request.url.path, response.status_code, data, task_name) + ), + ) + tracker_instance = MockTracker.return_value + emissions = MagicMock(emissions=0.001) + baseline = _configure_mock_running_tracker( + tracker_instance, task_name="GET /predict", emissions=0.001 + ) + tracker_instance.finish_http_request.return_value = emissions + + response = TestClient(application).get("/predict") + assert response.status_code == 200 + assert completed == [("/predict", 200, emissions, baseline.task_name)] + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_middleware_uses_lifespan_tracker(MockTracker) -> None: + application = FastAPI() + tracker_instance = MagicMock() + tracker_instance._start_time = 1.0 + baseline = MagicMock(task_name="GET /predict") + emissions = MagicMock(emissions=0.003) + tracker_instance.mark_http_request_start.return_value = baseline + tracker_instance.finish_http_request.return_value = emissions + application.state.codecarbon_tracker = tracker_instance + completed = [] + + @application.get("/predict") + def predict(): + return {"ok": True} + + add_codecarbon_middleware( + application, + on_request_complete=lambda request, response, data, task_name: completed.append( + (request.url.path, data, task_name) + ), + ) + + response = TestClient(application).get("/predict") + assert response.status_code == 200 + MockTracker.assert_not_called() + tracker_instance.mark_http_request_start.assert_called_once_with("GET /predict") + tracker_instance.finish_http_request.assert_called_once_with(baseline) + tracker_instance.persist_completed_task.assert_called_once_with("GET /predict") + assert completed == [("/predict", emissions, "GET /predict")] + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_middleware_skips_callback_when_handler_raises(MockTracker) -> None: + application = FastAPI() + tracker_instance = MagicMock() + tracker_instance._start_time = 1.0 + baseline = MagicMock(task_name="GET /fail") + tracker_instance.mark_http_request_start.return_value = baseline + tracker_instance.finish_http_request.return_value = MagicMock(emissions=0.001) + application.state.codecarbon_tracker = tracker_instance + completed = [] + + @application.get("/fail") + def fail(): + raise RuntimeError("boom") + + add_codecarbon_middleware( + application, + on_request_complete=lambda *args: completed.append(args), + ) + + with pytest.raises(RuntimeError, match="boom"): + TestClient(application, raise_server_exceptions=True).get("/fail") + + assert completed == [] + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_middleware_lazy_tracker(MockTracker) -> None: + application = FastAPI() + tracker_instance = MagicMock() + _configure_mock_running_tracker( + tracker_instance, task_name="GET /run", emissions=0.005 + ) + MockTracker.return_value = tracker_instance + + @application.get("/run") + def run(): + return {"ok": True} + + add_codecarbon_middleware(application) + + response = TestClient(application).get("/run") + assert response.status_code == 200 + MockTracker.assert_called_once() + tracker_instance.start.assert_called_once() + tracker_instance.mark_http_request_start.assert_called_once_with("GET /run") + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_middleware_no_logging_when_callback_disabled(MockTracker) -> None: + application = FastAPI() + + @application.get("/predict") + def predict(): + return {"ok": True} + + add_codecarbon_middleware(application, on_request_complete=None) + _configure_mock_running_tracker(MockTracker.return_value) + + with patch.object(cc_fastapi_middleware.logger, "info") as mock_info: + response = TestClient(application).get("/predict") + + assert response.status_code == 200 + mock_info.assert_not_called() + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_middleware_include_endpoints_allowlist(MockTracker) -> None: + application = FastAPI() + + @application.get("/predict") + def predict(): + return {"ok": True} + + @application.get("/metrics") + def metrics(): + return {"ok": True} + + add_codecarbon_middleware(application, include=["GET /predict"]) + _configure_mock_running_tracker(MockTracker.return_value) + + client = TestClient(application) + assert client.get("/predict").status_code == 200 + assert client.get("/metrics").status_code == 200 + MockTracker.assert_called_once() + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_middleware_exclude_endpoints(MockTracker) -> None: + application = FastAPI() + + @application.get("/predict") + def predict(): + return {"tracked": True} + + @application.get("/admin") + def admin(): + return {"admin": True} + + add_codecarbon_middleware(application, exclude=["GET /admin"]) + _configure_mock_running_tracker(MockTracker.return_value) + + client = TestClient(application) + client.get("/predict") + client.get("/admin") + MockTracker.assert_called_once() + + +def test_log_request_complete_uses_codecarbon_logger() -> None: + request = MagicMock(url=MagicMock(path="/predict")) + response = MagicMock(status_code=200) + emissions = MagicMock(emissions=0.0012) + counter = _CodeCarbonLogCapture() + previous_level = codecarbon_logger.level + + codecarbon_logger.setLevel(logging.INFO) + cc_fastapi_middleware.logger.addHandler(counter) + try: + log_request_complete(request, response, emissions, "GET /predict") + finally: + cc_fastapi_middleware.logger.removeHandler(counter) + codecarbon_logger.setLevel(previous_level) + + assert codecarbon_logger.name == "codecarbon" + assert counter.emissions_lines == 1 + + +class _CodeCarbonLogCapture(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.INFO) + self.emissions_lines = 0 + + def emit(self, record: logging.LogRecord) -> None: + if record.name != "codecarbon": + return + message = record.getMessage() + if message.startswith("CodeCarbon ") and "emissions=" in message: + self.emissions_lines += 1 + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +@patch.object(cc_fastapi_middleware.logger, "info") +def test_middleware_default_logs_after_request(mock_logger_info, MockTracker) -> None: + application = FastAPI() + _configure_mock_running_tracker(MockTracker.return_value) + + @application.get("/predict") + def predict(): + return {"ok": True} + + add_codecarbon_middleware(application, project_name="test-api") + response = TestClient(application).get("/predict") + + assert response.status_code == 200 + mock_logger_info.assert_called_once() + + +def test_add_codecarbon_middleware_registers_instance_on_app_state() -> None: + application = FastAPI() + add_codecarbon_middleware(application, project_name="shutdown-test") + middleware = application.state.codecarbon_middleware + middleware.shutdown_tracker_executor() + with pytest.raises(RuntimeError, match="shutdown"): + middleware._tracker_runner.submit_request(lambda: None) + + +def test_shutdown_codecarbon_middleware_helper() -> None: + application = FastAPI() + add_codecarbon_middleware(application, project_name="shutdown-test") + shutdown_codecarbon_middleware(application) + middleware = application.state.codecarbon_middleware + with pytest.raises(RuntimeError, match="shutdown"): + middleware._tracker_runner.submit_request(lambda: None) + + +@patch.object(cc_fastapi_lifespan, "EmissionsTracker") +def test_create_codecarbon_lifespan_shuts_down_middleware_executor( + MockTracker: MagicMock, +) -> None: + MockTracker.return_value = MagicMock() + + @asynccontextmanager + async def lifespan(application: FastAPI): + async with create_codecarbon_lifespan( + application, project_name="lifespan-test" + ): + yield + + application = FastAPI(lifespan=lifespan) + add_codecarbon_middleware(application, project_name="lifespan-test") + + with TestClient(application): + pass + + middleware = application.state.codecarbon_middleware + with pytest.raises(RuntimeError, match="shutdown"): + middleware._tracker_runner.submit_request(lambda: None) + + +def test_middleware_real_tracker_logs_and_csv_on_lifespan_stop(tmp_path: Path) -> None: + tracker_kwargs = { + "save_to_file": True, + "save_to_api": False, + "save_to_logger": False, + "output_dir": str(tmp_path), + "measure_power_secs": 10, + "allow_multiple_runs": True, + } + + @asynccontextmanager + async def lifespan(application: FastAPI): + async with create_codecarbon_lifespan( + application, + project_name="outputs-test", + **tracker_kwargs, + ): + yield + + application = FastAPI(lifespan=lifespan) + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + add_codecarbon_middleware( + application, + project_name="outputs-test", + tracker_kwargs=tracker_kwargs, + ) + log_counter = _CodeCarbonLogCapture() + cc_fastapi_middleware.logger.addHandler(log_counter) + try: + with TestClient(application) as client: + assert client.get("/predict").status_code == 200 + assert client.get("/predict").status_code == 200 + finally: + cc_fastapi_middleware.logger.removeHandler(log_counter) + + assert log_counter.emissions_lines == 2 + emissions_csv = tmp_path / "emissions.csv" + assert emissions_csv.is_file() + assert emissions_csv.stat().st_size > 0 + + +@patch("codecarbon.output_methods.http.ApiClient") +def test_middleware_real_tracker_calls_api_per_request( + MockApiClient, tmp_path: Path +) -> None: + mock_api = MockApiClient.return_value + mock_api.run_id = "test-run-id" + mock_api.add_emission.return_value = True + tracker_kwargs = { + "save_to_file": False, + "save_to_api": True, + "save_to_logger": False, + "output_dir": str(tmp_path), + "experiment_id": "00000000-0000-0000-0000-000000000001", + "api_key": "test-key", + "measure_power_secs": 10, + "allow_multiple_runs": True, + } + + @asynccontextmanager + async def lifespan(application: FastAPI): + async with create_codecarbon_lifespan( + application, + project_name="api-outputs-test", + **tracker_kwargs, + ): + yield + + application = FastAPI(lifespan=lifespan) + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + add_codecarbon_middleware( + application, + project_name="api-outputs-test", + tracker_kwargs=tracker_kwargs, + on_request_complete=None, + ) + with TestClient(application) as client: + assert client.get("/predict").status_code == 200 + assert client.get("/predict").status_code == 200 + + assert mock_api.add_emission.call_count >= 2 + + +def test_finalize_measures_before_on_request_complete() -> None: + order: list[str] = [] + + @asynccontextmanager + async def lifespan(application: FastAPI): + async with create_codecarbon_lifespan( + application, + project_name="measure-order", + save_to_file=False, + save_to_api=False, + allow_multiple_runs=True, + measure_power_secs=10, + ): + yield + + application = FastAPI(lifespan=lifespan) + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + def on_complete(request, response, emissions_data, task_name) -> None: + order.append("callback") + + add_codecarbon_middleware( + application, + project_name="measure-order", + on_request_complete=on_complete, + tracker_kwargs={ + "save_to_file": False, + "save_to_api": False, + "allow_multiple_runs": True, + "measure_power_secs": 10, + }, + ) + + with TestClient(application) as client: + tracker = application.state.codecarbon_tracker + tracker._last_measured_time = 0.0 + original = tracker._run_power_measurement + + def wrapped() -> None: + order.append("measure") + return original() + + with patch.object(tracker, "_run_power_measurement", side_effect=wrapped): + assert client.get("/predict").status_code == 200 + + assert order == ["measure", "callback"] + + +def test_concurrent_same_route_gets_distinct_task_names() -> None: + task_names: list[str] = [] + + @asynccontextmanager + async def lifespan(application: FastAPI): + async with create_codecarbon_lifespan( + application, + project_name="concurrent-test", + save_to_file=False, + save_to_api=False, + allow_multiple_runs=True, + measure_power_secs=10, + ): + yield + + application = FastAPI(lifespan=lifespan) + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + def on_complete(request, response, emissions_data, task_name) -> None: + task_names.append(task_name) + + add_codecarbon_middleware( + application, + project_name="concurrent-test", + on_request_complete=on_complete, + tracker_kwargs={ + "save_to_file": False, + "save_to_api": False, + "allow_multiple_runs": True, + "measure_power_secs": 10, + }, + ) + + with TestClient(application) as client: + tracker = application.state.codecarbon_tracker + baselines = [ + tracker.mark_http_request_start("GET /predict"), + tracker.mark_http_request_start("GET /predict"), + ] + assert baselines[0].task_name != baselines[1].task_name + assert baselines[0].task_name.startswith("GET /predict") + assert "GET /predict" in baselines[1].task_name + for baseline in baselines: + tracker.finish_http_request(baseline) + + assert client.get("/predict").status_code == 200 + assert client.get("/predict").status_code == 200 + + assert len(task_names) == 2 + assert all(name.startswith("GET /predict") for name in task_names) + + +def test_concurrent_live_tracker_no_stop_task_errors() -> None: + """Concurrent HTTP requests must not trigger stop_task on mark_http_request tasks.""" + import logging + + from starlette.testclient import TestClient + + error_messages: list[str] = [] + + class _ErrorHandler(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if ( + record.levelno >= logging.ERROR + and "_active_task_emissions_at_start" in record.getMessage() + ): + error_messages.append(record.getMessage()) + + handler = _ErrorHandler() + cc_fastapi_middleware.logger.addHandler(handler) + try: + + @asynccontextmanager + async def lifespan(application: FastAPI): + async with create_codecarbon_lifespan( + application, + project_name="concurrent-live", + save_to_file=False, + save_to_api=False, + allow_multiple_runs=True, + measure_power_secs=2, + ): + yield + + application = FastAPI(lifespan=lifespan) + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + add_codecarbon_middleware( + application, + project_name="concurrent-live", + on_request_complete=None, + tracker_kwargs={ + "save_to_file": False, + "save_to_api": False, + "allow_multiple_runs": True, + "measure_power_secs": 2, + }, + ) + + with TestClient(application) as client: + for _ in range(12): + assert client.get("/predict").status_code == 200 + + assert error_messages == [] + finally: + cc_fastapi_middleware.logger.removeHandler(handler) + + +def test_concurrent_lazy_tracker_without_lifespan() -> None: + """Lazy-started tracker must use mark/finish, not start_task/stop_task, under load.""" + import concurrent.futures + import logging + + error_messages: list[str] = [] + + class _ErrorHandler(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if ( + record.levelno >= logging.ERROR + and "_active_task_emissions_at_start" in record.getMessage() + ): + error_messages.append(record.getMessage()) + + handler = _ErrorHandler() + cc_fastapi_middleware.logger.addHandler(handler) + try: + application = FastAPI() + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + add_codecarbon_middleware( + application, + project_name="lazy-concurrent", + on_request_complete=None, + tracker_kwargs={ + "save_to_file": False, + "save_to_api": False, + "allow_multiple_runs": True, + "measure_power_secs": 2, + }, + ) + + with TestClient(application) as client: + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: + futures = [pool.submit(client.get, "/predict") for _ in range(16)] + for future in futures: + assert future.result().status_code == 200 + + assert error_messages == [] + finally: + cc_fastapi_middleware.logger.removeHandler(handler) + + +def test_compose_lifespans_stacks_contexts() -> None: + from codecarbon.integrations.fastapi import compose_lifespans + + events: list[str] = [] + + @asynccontextmanager + async def other(app: FastAPI): + events.append("other-enter") + app.state.other = True + try: + yield + finally: + events.append("other-exit") + + application = FastAPI( + lifespan=compose_lifespans( + lambda a: create_codecarbon_lifespan( + a, + project_name="compose-test", + save_to_file=False, + save_to_api=False, + allow_multiple_runs=True, + ), + other, + ) + ) + add_codecarbon_middleware( + application, + project_name="compose-test", + on_request_complete=None, + tracker_kwargs={"save_to_file": False, "save_to_api": False}, + ) + + with TestClient(application) as client: + assert application.state.other is True + assert application.state.codecarbon_tracker is not None + assert client.get("/docs").status_code == 200 + + assert events == ["other-enter", "other-exit"] + assert application.state.codecarbon_tracker is None + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_response_headers_sync_mode_injects_emissions_header(MockTracker) -> None: + _configure_mock_running_tracker( + MockTracker.return_value, emissions=0.0012 + ) + MockTracker.return_value.finish_http_request.return_value = MagicMock( + emissions=0.0012 + ) + application = FastAPI() + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + add_codecarbon_middleware( + application, + project_name="headers-test", + response_headers=True, + on_request_complete=None, + ) + response = TestClient(application).get("/predict") + assert response.status_code == 200 + assert response.headers.get("X-CodeCarbon-Emissions-kg") == "0.0012" + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_default_mode_has_no_emission_headers(MockTracker) -> None: + _configure_mock_running_tracker(MockTracker.return_value, emissions=0.0012) + application = FastAPI() + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + add_codecarbon_middleware( + application, project_name="no-headers", on_request_complete=None + ) + response = TestClient(application).get("/predict") + assert "X-CodeCarbon-Emissions-kg" not in response.headers + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_include_background_tasks_false_finalizes_before_background( + MockTracker, +) -> None: + from fastapi import BackgroundTasks + + order: list[str] = [] + mock_tracker = MockTracker.return_value + _configure_mock_running_tracker(mock_tracker) + + def finish_http_request(*args: Any, **kwargs: Any) -> MagicMock: + order.append("finalize") + return MagicMock(emissions=0.001) + + mock_tracker.finish_http_request.side_effect = finish_http_request + application = FastAPI() + + @application.get("/predict") + def predict_with_bg(background_tasks: BackgroundTasks) -> dict[str, bool]: + def work() -> None: + order.append("background") + + background_tasks.add_task(work) + return {"ok": True} + + add_codecarbon_middleware( + application, + project_name="bg-false", + include_background_tasks=False, + on_request_complete=None, + ) + assert TestClient(application).get("/predict").status_code == 200 + assert "finalize" in order + assert "background" in order + assert order.index("finalize") < order.index("background") + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_include_background_tasks_true_finalizes_after_background(MockTracker) -> None: + from fastapi import BackgroundTasks + + order: list[str] = [] + mock_tracker = MockTracker.return_value + _configure_mock_running_tracker(mock_tracker) + + def finish_http_request(*args: Any, **kwargs: Any) -> MagicMock: + order.append("finalize") + return MagicMock(emissions=0.001) + + mock_tracker.finish_http_request.side_effect = finish_http_request + application = FastAPI() + + @application.get("/predict") + def predict_with_bg(background_tasks: BackgroundTasks) -> dict[str, bool]: + def work() -> None: + order.append("background") + + background_tasks.add_task(work) + return {"ok": True} + + add_codecarbon_middleware( + application, + project_name="bg-true", + include_background_tasks=True, + on_request_complete=None, + ) + assert TestClient(application).get("/predict").status_code == 200 + assert "background" in order + assert "finalize" in order + assert order.index("background") < order.index("finalize") + + +def test_resolve_header_fields_and_header_names() -> None: + from codecarbon.integrations.fastapi.middleware import ( + _codecarbon_header_name, + _inject_emission_headers, + _resolve_header_fields, + ) + + assert _resolve_header_fields(None) == () + assert _resolve_header_fields(True) == ("emissions",) + assert _resolve_header_fields(["emissions", "duration"]) == ( + "emissions", + "duration", + ) + assert ( + _codecarbon_header_name("energy_consumed") == "X-CodeCarbon-Energy-Consumed-kwh" + ) + + message = {"type": "http.response.start", "headers": []} + assert _inject_emission_headers(message, None, ["emissions"]) is message + + emissions = MagicMock(spec=["emissions", "duration"]) + emissions.emissions = 0.0012 + emissions.duration = 1.5 + injected = _inject_emission_headers( + message, emissions, ["emissions", "unknown_field", "duration"] + ) + header_names = {name.decode() for name, _ in injected["headers"]} + assert header_names == { + "X-CodeCarbon-Emissions-kg", + "X-CodeCarbon-Duration-s", + } + + +def test_tracker_runner_handles_cancelled_and_failed_jobs() -> None: + from concurrent import futures + + runner = cc_fastapi_middleware._TrackerRunner() + cancelled = runner.submit_request(lambda: 1) + cancelled.cancel() + runner.shutdown() + assert cancelled.cancelled() + + runner = cc_fastapi_middleware._TrackerRunner() + + def boom() -> None: + raise ValueError("tracker failed") + + with pytest.raises(ValueError, match="tracker failed"): + runner.submit_request(boom).result(timeout=2) + runner.shutdown() + + done = futures.Future() + done.set_result(1) + runner = cc_fastapi_middleware._TrackerRunner() + runner._run_job((lambda: 99, (), done)) + + def raise_runtime() -> None: + raise RuntimeError("x") + + already_done = futures.Future() + already_done.set_result(1) + runner._run_job((raise_runtime, (), already_done)) + runner.shutdown() + + +def test_tracker_runner_finalize_lane_and_no_wait_shutdown() -> None: + runner = cc_fastapi_middleware._TrackerRunner() + assert runner.submit(runner.FINALIZE, lambda: 42).result(timeout=2) == 42 + runner.shutdown(wait=False) + runner.shutdown() + + +def test_tracker_runner_drains_finalize_after_request_job() -> None: + order: list[str] = [] + runner = cc_fastapi_middleware._TrackerRunner() + + def request_job() -> None: + order.append("request") + + def finalize_job() -> None: + order.append("finalize") + + runner.submit_request(request_job) + runner.submit(runner.FINALIZE, finalize_job) + runner.shutdown() + assert order == ["request", "finalize"] + + +def test_tracker_runner_prioritizes_new_requests_over_finalize_drain() -> None: + import threading + + order: list[str] = [] + gate = threading.Event() + runner = cc_fastapi_middleware._TrackerRunner() + + def slow_request() -> None: + gate.wait(timeout=2) + order.append("request1") + + def finalize_job() -> None: + order.append("finalize") + + def second_request() -> None: + order.append("request2") + + runner.submit_request(slow_request) + runner.submit(runner.FINALIZE, finalize_job) + runner.submit_request(second_request) + gate.set() + runner.shutdown() + assert order.index("request1") < order.index("request2") + assert order.index("request2") < order.index("finalize") + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_task_name_formatter(MockTracker) -> None: + application = FastAPI() + _configure_mock_running_tracker(MockTracker.return_value, task_name="custom-/predict") + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + add_codecarbon_middleware( + application, + task_name_formatter=lambda request: f"custom-{request.url.path}", + on_request_complete=None, + ) + assert TestClient(application).get("/predict").status_code == 200 + MockTracker.return_value.mark_http_request_start.assert_called_once_with( + "custom-/predict" + ) + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_response_headers_custom_field_list(MockTracker) -> None: + application = FastAPI() + emissions = MagicMock(emissions=0.0012, duration=1.5) + _configure_mock_running_tracker(MockTracker.return_value) + MockTracker.return_value.finish_http_request.return_value = emissions + + @application.get("/predict") + def predict() -> dict[str, bool]: + return {"ok": True} + + add_codecarbon_middleware( + application, + response_headers=["emissions", "duration"], + on_request_complete=None, + ) + response = TestClient(application).get("/predict") + assert response.headers.get("X-CodeCarbon-Emissions-kg") == "0.0012" + assert response.headers.get("X-CodeCarbon-Duration-s") == "1.5" + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_websocket_scope_is_not_tracked(MockTracker) -> None: + from unittest.mock import AsyncMock + + inner = AsyncMock() + + async def run() -> None: + middleware = cc_fastapi_middleware.CodeCarbonMiddleware(inner) + await middleware({"type": "websocket"}, MagicMock(), MagicMock()) + + asyncio.run(run()) + inner.assert_awaited_once() + MockTracker.assert_not_called() + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_end_of_body_mode_reraises_handler_error(MockTracker) -> None: + application = FastAPI() + _configure_mock_running_tracker(MockTracker.return_value) + + @application.get("/fail") + def fail() -> None: + raise RuntimeError("boom") + + add_codecarbon_middleware( + application, + include_background_tasks=False, + on_request_complete=None, + ) + with pytest.raises(RuntimeError, match="boom"): + TestClient(application, raise_server_exceptions=True).get("/fail") + + +@patch.object(cc_fastapi_middleware, "EmissionsTracker") +def test_sync_headers_mode_reraises_handler_error(MockTracker) -> None: + application = FastAPI() + _configure_mock_running_tracker(MockTracker.return_value) + + @application.get("/fail") + def fail() -> None: + raise RuntimeError("boom") + + add_codecarbon_middleware( + application, + response_headers=True, + on_request_complete=None, + ) + with pytest.raises(RuntimeError, match="boom"): + TestClient(application, raise_server_exceptions=True).get("/fail") + + +@pytest.mark.no_immediate_finalize +def test_schedule_finalize_logs_measurement_failure() -> None: + application = FastAPI() + middleware = cc_fastapi_middleware.CodeCarbonMiddleware(application) + + async def fail() -> None: + raise RuntimeError("measurement failed") + + async def run() -> None: + scheduled: list[asyncio.Task[None]] = [] + + def track_create_task(coro: Any) -> asyncio.Task[None]: + task = asyncio.get_running_loop().create_task(coro) + scheduled.append(task) + return task + + with patch("asyncio.create_task", side_effect=track_create_task): + with patch.object( + cc_fastapi_middleware.logger, "exception" + ) as mock_exception: + middleware._schedule_finalize(fail()) + await asyncio.gather(*scheduled) + mock_exception.assert_called_once() + + asyncio.run(run()) diff --git a/tests/integrations/test_fastapi_routing.py b/tests/integrations/test_fastapi_routing.py new file mode 100644 index 000000000..176bbffc7 --- /dev/null +++ b/tests/integrations/test_fastapi_routing.py @@ -0,0 +1,112 @@ +"""Tests for route naming and endpoint filter helpers.""" + +from unittest.mock import MagicMock + +from codecarbon.integrations.fastapi._routing import ( + build_endpoint_key, + should_track_request, +) + + +def _mock_request(method: str, route_path: str | None, url_path: str) -> MagicMock: + request = MagicMock() + request.method = method + request.url.path = url_path + if route_path is None: + request.scope = {} + else: + route = MagicMock() + route.path = route_path + request.scope = {"route": route} + return request + + +def test_build_endpoint_key_uses_route_template() -> None: + request = _mock_request("GET", "/predict", "/predict") + assert build_endpoint_key(request) == "GET /predict" + + +def test_should_track_request_exclude_path_prefix() -> None: + request = _mock_request("GET", "/docs", "/docs/oauth2-redirect") + assert should_track_request(request, None, ["/docs"]) is False + + +def test_should_track_request_exclude_by_method_and_path() -> None: + request = _mock_request("GET", "/predict", "/predict") + assert should_track_request(request, None, ["GET /predict"]) is False + assert should_track_request(request, None, ["POST /predict"]) is True + + +def test_should_track_request_exclude_path_only() -> None: + request = _mock_request("POST", "/predict", "/predict") + assert should_track_request(request, None, ["/predict"]) is False + + +def test_should_track_request_include_allowlist() -> None: + request = _mock_request("GET", "/predict", "/predict") + other = _mock_request("GET", "/health", "/health") + include = ["GET /predict"] + assert should_track_request(request, include, []) is True + assert should_track_request(other, include, []) is False + + +def test_should_track_request_include_path_only() -> None: + get_request = _mock_request("GET", "/predict", "/predict") + post_request = _mock_request("POST", "/predict", "/predict") + include = ["/predict"] + assert should_track_request(get_request, include, []) is True + assert should_track_request(post_request, include, []) is True + + +def test_is_method_pattern_rejects_invalid_methods() -> None: + from codecarbon.integrations.fastapi._routing import is_method_pattern + + assert is_method_pattern("GET /predict") is True + assert is_method_pattern("FOO /predict") is False + assert is_method_pattern("/predict") is False + assert is_method_pattern("GET") is False + + +def test_matches_filter_pattern_non_path_literal() -> None: + from codecarbon.integrations.fastapi._routing import matches_filter_pattern + + assert ( + matches_filter_pattern( + "GET /predict", + "GET /predict", + "/predict", + "/predict", + exclude=False, + ) + is True + ) + assert ( + matches_filter_pattern( + "GET /predict", + "GET /items/1", + "/items/{item_id}", + "/items/1", + exclude=False, + ) + is False + ) + assert ( + matches_filter_pattern( + "GET /predict", + "GET /predict", + "/predict", + "/predict", + exclude=False, + ) + is True + ) + assert ( + matches_filter_pattern( + "predict", + "predict", + "/predict", + "/predict", + exclude=False, + ) + is True + ) diff --git a/tests/output_methods/test_http.py b/tests/output_methods/test_http.py index 56d909b46..7095dae26 100644 --- a/tests/output_methods/test_http.py +++ b/tests/output_methods/test_http.py @@ -147,45 +147,6 @@ def test_codecarbon_api_live_out(self): api_output.live_out(None, self.emissions_data) self.mock_add_emission.assert_called_once() - def test_codecarbon_api_live_out_creates_run_when_missing(self): - conf = { - "os": "linux", - "python_version": "3.12", - "codecarbon_version": "2.0", - "cpu_count": 4, - "cpu_model": "CPU", - "gpu_count": 0, - "gpu_model": "", - "longitude": 0.0, - "latitude": 0.0, - "region": "EU", - "provider": "AWS", - "ram_total_size": 16.0, - "tracking_mode": "machine", - } - - with patch( - "codecarbon.output_methods.http.ApiClient._create_run" - ) as mock_create_run: - api_output = CodeCarbonAPIOutput( - endpoint_url="http://test.com", - experiment_id="exp-1", - api_key=self.api_key, - conf=conf, - ) - api_output.api.run_id = None - - def create_run(experiment_id): - api_output.api.run_id = "run-created" - return "run-created" - - mock_create_run.side_effect = create_run - api_output.live_out(None, self.emissions_data) - - mock_create_run.assert_called_once_with("exp-1") - self.assertEqual(api_output.api.run_id, "run-created") - self.assertEqual(api_output.run_id, "run-created") - @patch("codecarbon.output_methods.http.logger.error") def test_codecarbon_live_out_api_call_failure(self, mock_logger): self.mock_add_emission.side_effect = Exception("Test exception") @@ -209,6 +170,52 @@ def test_codecarbon_api_out(self): api_output.out(None, self.emissions_data) self.mock_add_emission.assert_called_once() + def test_codecarbon_api_task_out(self): + from codecarbon.output_methods.emissions_data import TaskEmissionsData + + api_output = CodeCarbonAPIOutput( + endpoint_url=self.url, + experiment_id=self.experiment_id, + api_key=self.api_key, + conf=None, + ) + task_data = TaskEmissionsData( + task_name="GET /predict", + timestamp=self.emissions_data.timestamp, + project_name=self.emissions_data.project_name, + run_id=self.emissions_data.run_id, + duration=2.0, + emissions=self.emissions_data.emissions, + emissions_rate=self.emissions_data.emissions_rate, + cpu_power=self.emissions_data.cpu_power, + gpu_power=self.emissions_data.gpu_power, + ram_power=self.emissions_data.ram_power, + cpu_energy=self.emissions_data.cpu_energy, + gpu_energy=self.emissions_data.gpu_energy, + ram_energy=self.emissions_data.ram_energy, + energy_consumed=self.emissions_data.energy_consumed, + water_consumed=self.emissions_data.water_consumed, + country_name=self.emissions_data.country_name, + country_iso_code=self.emissions_data.country_iso_code, + region=self.emissions_data.region, + cloud_provider=self.emissions_data.cloud_provider, + cloud_region=self.emissions_data.cloud_region, + os=self.emissions_data.os, + python_version=self.emissions_data.python_version, + codecarbon_version=self.emissions_data.codecarbon_version, + cpu_count=self.emissions_data.cpu_count, + cpu_model=self.emissions_data.cpu_model, + gpu_count=self.emissions_data.gpu_count, + gpu_model=self.emissions_data.gpu_model, + longitude=self.emissions_data.longitude, + latitude=self.emissions_data.latitude, + ram_total_size=self.emissions_data.ram_total_size, + tracking_mode=self.emissions_data.tracking_mode, + on_cloud=self.emissions_data.on_cloud, + ) + api_output.task_out([task_data], "test_experiment") + self.mock_add_emission.assert_called_once() + @patch("codecarbon.output_methods.http.logger.error") def test_codecarbon_out_api_call_failure(self, mock_logger): self.mock_add_emission.side_effect = Exception("Test exception") diff --git a/tests/test_api_call.py b/tests/test_api_call.py index 39822ece7..e2aa28e61 100644 --- a/tests/test_api_call.py +++ b/tests/test_api_call.py @@ -199,31 +199,34 @@ def test_add_emission_returns_false_when_run_creation_fails(self): ) ) - def test_add_emission_skips_short_duration(self): - api = ApiClient( - endpoint_url="http://test.com", - experiment_id="exp-1", - conf=conf, - create_run_automatically=False, - ) - api.run_id = "run-1" + def test_add_emission_rounds_subsecond_duration_to_one_second(self): + with requests_mock.Mocker() as m: + m.post("http://test.com/emissions", json={"id": "em-1"}, status_code=201) + api = ApiClient( + endpoint_url="http://test.com", + experiment_id="exp-1", + conf=conf, + create_run_automatically=False, + ) + api.run_id = "run-1" - self.assertFalse( - api.add_emission( - { - "duration": 0.5, - "emissions": 1.0, - "emissions_rate": 1.0, - "cpu_power": 1.0, - "gpu_power": 0.0, - "ram_power": 0.5, - "cpu_energy": 0.1, - "gpu_energy": 0.0, - "ram_energy": 0.1, - "energy_consumed": 0.2, - } + self.assertTrue( + api.add_emission( + { + "duration": 0.5, + "emissions": 1.0, + "emissions_rate": 1.0, + "cpu_power": 1.0, + "gpu_power": 0.0, + "ram_power": 0.5, + "cpu_energy": 0.1, + "gpu_energy": 0.0, + "ram_energy": 0.1, + "energy_consumed": 0.2, + } + ) ) - ) + self.assertEqual(m.last_request.json()["duration"], 1) def test_add_emission_returns_false_on_unsuccessful_post(self): with requests_mock.Mocker() as m: diff --git a/tests/test_cpu_load.py b/tests/test_cpu_load.py index ecb9b2d27..f5cdf7e46 100644 --- a/tests/test_cpu_load.py +++ b/tests/test_cpu_load.py @@ -49,18 +49,13 @@ def test_cpu_total_power( self.assertEqual(power.W, 50) self.assertEqual(cpu.total_power().W, 50) - @mock.patch( - "codecarbon.core.powermetrics.is_powermetrics_available", return_value=False - ) def test_cpu_load_detection( self, - mocked_is_powermetrics_available, mocked_is_psutil_available, mocked_is_powergadget_available, mocked_is_rapl_available, ): tracker = OfflineEmissionsTracker(country_iso_code="FRA") - tracker._ensure_hardware_ready() for hardware in tracker._hardware: if ( isinstance(hardware, CPU) and hardware._mode == MODE_CPU_LOAD diff --git a/tests/test_custom_handler.py b/tests/test_custom_handler.py index 8adcf7c37..570d2df73 100644 --- a/tests/test_custom_handler.py +++ b/tests/test_custom_handler.py @@ -32,8 +32,7 @@ def test_carbon_tracker_custom_handler(self): tracker = EmissionsTracker( project_name=self.project_name, output_handlers=[handler_0, handler_1], - api_call_interval=2, - measure_power_secs=999, + api_call_interval=1, ) tracker.start() heavy_computation(run_time_secs=1) @@ -53,8 +52,7 @@ def test_decorator_flush(self): project_name=self.project_name, save_to_logger=True, output_handlers=[handler_0, handler_1], - api_call_interval=2, - measure_power_secs=999, + api_call_interval=1, ) def dummy_train_model(): heavy_computation(run_time_secs=1) diff --git a/tests/test_emissions_tracker_constant.py b/tests/test_emissions_tracker_constant.py index 724cbe9ab..65b17c666 100644 --- a/tests/test_emissions_tracker_constant.py +++ b/tests/test_emissions_tracker_constant.py @@ -5,6 +5,7 @@ from unittest import mock import pandas as pd +import psutil from codecarbon.core import cpu from codecarbon.emissions_tracker import ( @@ -88,15 +89,14 @@ def test_carbon_tracker_offline_constant_force_cpu_power( assertdf = pd.read_csv(self.emissions_file_path) self.assertEqual(USER_INPUT_CPU_POWER / 2, assertdf["cpu_power"][0]) - @mock.patch("codecarbon.external.hardware.psutil.cpu_percent", return_value=50.0) @mock.patch.object(cpu.TDP, "_get_cpu_power_from_registry") @mock.patch.object(cpu, "is_psutil_available") - def test_carbon_tracker_offline_load_force_cpu_power( - self, mock_psutil_available, mock_tdp, mock_cpu_percent - ): + def test_carbon_tracker_offline_load_force_cpu_power(self, mock_tdp, mock_psutil): + # Same as test_carbon_tracker_offline_constant test but this time forcing the default cpu power USER_INPUT_CPU_POWER = 1_000 + # Mock the output of tdp mock_tdp.return_value = 500 - mock_psutil_available.return_value = True + mock_psutil.return_value = True tracker = OfflineEmissionsTracker( country_iso_code="USA", output_dir=self.emissions_path, @@ -108,11 +108,17 @@ def test_carbon_tracker_offline_load_force_cpu_power( emissions = tracker.stop() assert isinstance(emissions, float) self.assertNotEqual(emissions, 0.0) - cpu_load = 0.5 + # Get CPU load (measured after test; may differ from load during test) + cpu_load = psutil.cpu_percent(interval=1) / 100.0 + # Assert the content stored. cpu_power should be approximately load * min(TDP, forced CPU power) assertdf = pd.read_csv(self.emissions_file_path) - load_factor = 0.1 + 0.9 * (cpu_load**3) - expected_power = USER_INPUT_CPU_POWER * load_factor - self.assertAlmostEqual(assertdf["cpu_power"][0], expected_power, delta=50) + tolerance = 350 + self.assertLess( + assertdf["cpu_power"][0], USER_INPUT_CPU_POWER * cpu_load + tolerance + ) + self.assertGreater( + assertdf["cpu_power"][0], USER_INPUT_CPU_POWER * cpu_load - tolerance + ) def test_decorator_constant(self): @track_emissions( @@ -145,7 +151,6 @@ def test_carbon_tracker_offline_region_error(self): ) tracker.start() tracker._measure_power_and_energy() - tracker._ensure_emissions_engine() cloud: CloudMetadata = tracker._get_cloud_metadata() try: diff --git a/tests/test_hardware_cache.py b/tests/test_hardware_cache.py deleted file mode 100644 index 7fa5a0f66..000000000 --- a/tests/test_hardware_cache.py +++ /dev/null @@ -1,288 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import patch - -import pytest - -from codecarbon.core import hardware_cache -from codecarbon.external.hardware import CPU -from codecarbon.external.ram import RAM - - -def make_tracker(**overrides): - defaults = { - "_tracking_mode": "machine", - "_force_cpu_power": None, - "_force_ram_power": None, - "_conf": {}, - "_gpu_ids": None, - "_rapl_include_dram": False, - "_rapl_prefer_psys": False, - "_output_dir": "out", - "_hardware": [], - } - defaults.update(overrides) - return SimpleNamespace(**defaults) - - -def test_make_key_normalizes_gpu_ids(): - tracker = make_tracker(_gpu_ids=[0, 1]) - key = hardware_cache.make_key(tracker) - assert key.gpu_ids == ("0", "1") - - -def test_make_key_treats_equivalent_gpu_id_types_as_same_key(): - key_int = hardware_cache.make_key(make_tracker(_gpu_ids=[0])) - key_str = hardware_cache.make_key(make_tracker(_gpu_ids=["0"])) - assert key_int == key_str - - -def test_spec_and_rebuild_roundtrip_for_cpu(): - cpu_hw = CPU.from_utils("out", "cpu_load", "Test CPU", 100) - spec = hardware_cache._spec_from_hardware(cpu_hw) - rebuilt = hardware_cache._hardware_from_spec(spec, "out2") - assert type(rebuilt).__name__ == "CPU" - assert rebuilt._model == "Test CPU" - assert rebuilt._mode == "cpu_load" - - -def test_spec_from_hardware_gpu_and_rapl_cpu(): - gpu_hw = type("GPU", (), {"gpu_ids": [0, 1]})() - assert hardware_cache._spec_from_hardware(gpu_hw) == { - "kind": "gpu", - "gpu_ids": ["0", "1"], - } - - gpu_hw_no_ids = type("GPU", (), {"gpu_ids": None})() - assert hardware_cache._spec_from_hardware(gpu_hw_no_ids) == { - "kind": "gpu", - "gpu_ids": None, - } - - gpu_hw_empty_ids = type("GPU", (), {"gpu_ids": []})() - assert hardware_cache._spec_from_hardware(gpu_hw_empty_ids) == { - "kind": "gpu", - "gpu_ids": None, - } - - -def test_capture_serializes_gpu_hardware(): - gpu_hw = type("GPU", (), {"gpu_ids": (0, 1)})() - tracker = make_tracker(_hardware=[gpu_hw]) - resource_tracker = SimpleNamespace( - tracker=tracker, - ram_tracker="ram", - cpu_tracker="cpu", - gpu_tracker="gpu", - ) - - plan = hardware_cache.capture(resource_tracker) - - assert plan.hardware_specs == [{"kind": "gpu", "gpu_ids": ["0", "1"]}] - - -def test_hardware_kind_apple_chip(): - apple_hw = type("AppleSiliconChip", (), {})() - assert hardware_cache._hardware_kind(apple_hw) == "apple_chip" - - -def test_spec_from_hardware_apple_chip(): - apple_hw = type( - "AppleSiliconChip", - (), - {"_model": "Apple M1", "chip_part": "CPU"}, - )() - assert hardware_cache._spec_from_hardware(apple_hw) == { - "kind": "apple_chip", - "model": "Apple M1", - "chip_part": "CPU", - } - - -def test_hardware_from_spec_rebuilds_gpu(): - fake_gpu = SimpleNamespace(gpu_ids=[0]) - with patch( - "codecarbon.external.hardware.GPU.from_utils", - return_value=fake_gpu, - ) as mock_from_utils: - rebuilt = hardware_cache._hardware_from_spec( - {"kind": "gpu", "gpu_ids": ["0"]}, - "out", - ) - mock_from_utils.assert_called_once_with(gpu_ids=["0"]) - assert rebuilt is fake_gpu - - -def test_hardware_from_spec_rejects_unknown_kind(): - with pytest.raises(ValueError, match="Unknown hardware spec kind"): - hardware_cache._hardware_from_spec({"kind": "unknown"}, "out") - - -def test_spec_from_hardware_intel_rapl_cpu(): - cpu_hw = type( - "CPU", - (), - { - "_mode": "intel_rapl", - "_model": "Intel CPU", - "_tdp": 65, - "_tracking_mode": "machine", - "_intel_interface": SimpleNamespace( - rapl_include_dram=True, - rapl_prefer_psys=True, - ), - }, - )() - spec = hardware_cache._spec_from_hardware(cpu_hw) - assert spec["rapl_include_dram"] is True - assert spec["rapl_prefer_psys"] is True - assert spec["rapl_dir"] == "/sys/class/powercap/intel-rapl/subsystem" - - -def test_spec_and_rebuild_roundtrip_for_apple_chip(): - spec = {"kind": "apple_chip", "model": "Apple M1", "chip_part": "CPU"} - fake_chip = SimpleNamespace(_model="Apple M1") - with patch( - "codecarbon.external.hardware.AppleSiliconChip", - return_value=fake_chip, - ) as mock_chip_cls: - rebuilt = hardware_cache._hardware_from_spec(spec, "out") - mock_chip_cls.assert_called_once_with( - output_dir="out", - model="Apple M1", - chip_part="CPU", - ) - assert rebuilt._model == "Apple M1" - - -def test_capture_and_apply_restore_hardware_plan(): - tracker = make_tracker( - _conf={ - "cpu_count": 8, - "cpu_physical_count": 4, - "cpu_model": "Cached CPU", - "gpu_count": 0, - "gpu_model": "", - "gpu_ids": ["0"], - }, - _gpu_ids=[0], - _hardware=[RAM(tracking_mode="machine")], - ) - resource_tracker = SimpleNamespace( - tracker=tracker, - ram_tracker="cached_ram", - cpu_tracker="cached_cpu", - gpu_tracker="cached_gpu", - ) - plan = hardware_cache.capture(resource_tracker) - - tracker2 = make_tracker() - rt2 = SimpleNamespace( - tracker=tracker2, - ram_tracker="old", - cpu_tracker="old", - gpu_tracker="old", - ) - hardware_cache.apply(rt2, plan) - - assert rt2.ram_tracker == "cached_ram" - assert rt2.cpu_tracker == "cached_cpu" - assert tracker2._conf["cpu_model"] == "Cached CPU" - assert tracker2._conf["cpu_count"] == 8 - assert tracker2._conf["cpu_physical_count"] == 4 - assert tracker2._gpu_ids == ["0"] - assert len(tracker2._hardware) == 1 - assert type(tracker2._hardware[0]).__name__ == "RAM" - - -def test_get_or_run_setup_runs_setup_once(): - tracker = make_tracker() - resource_tracker = SimpleNamespace( - tracker=tracker, - ram_tracker="Unspecified", - cpu_tracker="Unspecified", - gpu_tracker="Unspecified", - ) - calls = {"count": 0} - - def setup_fn(): - calls["count"] += 1 - resource_tracker.ram_tracker = "ran" - - hardware_cache.clear_cache() - hardware_cache.get_or_run_setup(resource_tracker, setup_fn) - hardware_cache.get_or_run_setup(resource_tracker, setup_fn) - - assert calls["count"] == 1 - assert resource_tracker.ram_tracker == "ran" - - -def test_hardware_kind_rejects_unknown_type(): - with pytest.raises(TypeError): - hardware_cache._hardware_kind(object()) - - -def test_clear_cache_resets_probe_caches(): - from codecarbon.core.cpu import clear_powergadget_cache, is_powergadget_available - from codecarbon.core.powermetrics import ( - clear_powermetrics_cache, - is_powermetrics_available, - ) - - clear_powergadget_cache() - clear_powermetrics_cache() - with patch("codecarbon.core.cpu.IntelPowerGadget", side_effect=Exception("nope")): - assert is_powergadget_available() is False - with patch( - "codecarbon.core.powermetrics.ApplePowermetrics", side_effect=Exception("nope") - ): - assert is_powermetrics_available() is False - - hardware_cache.clear_cache() - - assert is_powergadget_available.cache_info().currsize == 0 - assert is_powermetrics_available.cache_info().currsize == 0 - - -def test_get_cached_tdp_reuses_instance(): - hardware_cache.clear_cache() - fake_cpu = SimpleNamespace(TDP=lambda: SimpleNamespace(model="cached")) - first = hardware_cache.get_cached_tdp(fake_cpu) - second = hardware_cache.get_cached_tdp(fake_cpu) - assert first is second - - -def test_canonical_gpu_ids_accepts_scalar(): - assert hardware_cache._canonical_gpu_ids(0) == ("0",) - - -def test_spec_from_hardware_raises_for_unhandled_kind(): - class UnhandledKind: - value = "unhandled" - - with patch.object(hardware_cache, "_hardware_kind", return_value=UnhandledKind()): - with pytest.raises(TypeError, match="Unsupported hardware type"): - hardware_cache._spec_from_hardware(object()) - - -def test_hardware_from_spec_raises_when_kind_not_handled(): - sentinel = object() - real = hardware_cache.HardwareKind - - with patch("codecarbon.core.hardware_cache.HardwareKind") as mock_cls: - mock_cls.side_effect = lambda value: sentinel - mock_cls.RAM = real.RAM - mock_cls.CPU = real.CPU - mock_cls.APPLE_CHIP = real.APPLE_CHIP - mock_cls.GPU = real.GPU - - with pytest.raises(ValueError, match="Unknown hardware spec kind"): - hardware_cache._hardware_from_spec({"kind": "ram"}, "out") - - -def test_spec_and_rebuild_roundtrip_for_ram(): - ram_hw = RAM(tracking_mode="machine", force_ram_power=12.5) - spec = hardware_cache._spec_from_hardware(ram_hw) - rebuilt = hardware_cache._hardware_from_spec(spec, "out2") - assert type(rebuilt).__name__ == "RAM" - assert rebuilt._tracking_mode == "machine" - assert rebuilt._force_ram_power == 12.5 diff --git a/tests/test_http_request_tracking.py b/tests/test_http_request_tracking.py new file mode 100644 index 000000000..060fafbaf --- /dev/null +++ b/tests/test_http_request_tracking.py @@ -0,0 +1,301 @@ +"""Tests for per-request HTTP tracking on a shared EmissionsTracker.""" + +import sys +import threading +import time +import unittest +from unittest import mock + +from codecarbon.emissions_tracker import EmissionsTracker, HttpRequestBaseline +from codecarbon.external.geography import CloudMetadata +from tests.fake_modules import pynvml as fake_pynvml +from tests.testdata import TWO_GPU_DETAILS_RESPONSE, TWO_GPU_DETAILS_RESPONSE_HANDLES +from tests.testutils import get_custom_mock_open + +empty_conf = "[codecarbon]" + +if sys.platform == "darwin": + mock_platform_cli_setup = mock.patch( + "codecarbon.core.powermetrics.ApplePowermetrics._setup_cli" + ) +else: + mock_platform_cli_setup = mock.patch( + "codecarbon.core.cpu.IntelPowerGadget._setup_cli" + ) + + +def _build_tracker(**kwargs: object) -> EmissionsTracker: + defaults: dict[str, object] = { + "project_name": "http-request-test", + "save_to_file": False, + "save_to_api": False, + "save_to_logger": False, + "allow_multiple_runs": True, + "measure_power_secs": 10, + } + defaults.update(kwargs) + return EmissionsTracker(**defaults) + + +@mock.patch("codecarbon.core.gpu.pynvml", fake_pynvml) +@mock.patch("codecarbon.core.gpu.is_nvidia_system", return_value=True) +@mock.patch("codecarbon.core.gpu.is_gpu_details_available", return_value=True) +@mock.patch( + "codecarbon.external.hardware.AllGPUDevices.get_gpu_details", + return_value=TWO_GPU_DETAILS_RESPONSE, +) +@mock.patch( + "codecarbon.emissions_tracker.EmissionsTracker._get_cloud_metadata", + return_value=CloudMetadata(provider=None, region=None), +) +@mock.patch("codecarbon.core.cpu.IntelPowerGadget._log_values") +@mock_platform_cli_setup +class TestHttpRequestTracking(unittest.TestCase): + def setUp(self) -> None: + fake_pynvml.DETAILS = TWO_GPU_DETAILS_RESPONSE_HANDLES + patcher = mock.patch( + "builtins.open", new_callable=get_custom_mock_open(empty_conf, empty_conf) + ) + self.addCleanup(patcher.stop) + patcher.start() + + def test_mark_http_request_start_requires_started_tracker( + self, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker() + with self.assertRaises(RuntimeError): + tracker.mark_http_request_start("GET /predict") + + def test_http_request_baseline_round_trip( + self, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker() + tracker.start() + baseline = tracker.mark_http_request_start("GET /predict") + emissions_data = tracker.finish_http_request(baseline) + tracker.stop() + self.assertIsNotNone(emissions_data) + self.assertEqual(baseline.task_name.split("_")[0], "GET /predict") + + def test_http_request_task_names_are_unique_for_same_route( + self, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker() + tracker.start() + first = tracker.mark_http_request_start("GET /predict") + second = tracker.mark_http_request_start("GET /predict") + tracker.stop() + self.assertNotEqual(first.task_name, second.task_name) + self.assertTrue(second.task_name.startswith("GET /predict")) + + def test_finish_http_request_unknown_task_returns_none( + self, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker() + tracker.start() + baseline = HttpRequestBaseline( + task_name="missing-task", + started_at=0.0, + duration_at_start=0.0, + emissions=0.0, + cpu_energy=0.0, + gpu_energy=0.0, + ram_energy=0.0, + energy_consumed=0.0, + water_consumed=0.0, + ) + result = tracker.finish_http_request(baseline) + tracker.stop() + self.assertIsNone(result) + + def test_persist_completed_task_skips_when_api_disabled( + self, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker(save_to_api=False) + tracker.start() + tracker.start_task("GET /predict") + tracker.stop_task("GET /predict") + tracker.persist_completed_task("GET /predict") + tracker.stop() + + @mock.patch("codecarbon.output.CodeCarbonAPIOutput.task_out") + def test_persist_completed_task_uploads_finished_task( + self, + mock_task_out, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker( + save_to_api=True, + experiment_id="00000000-0000-0000-0000-000000000001", + api_key="test-key", + ) + tracker.start() + baseline = tracker.mark_http_request_start("GET /predict") + tracker.finish_http_request(baseline) + tracker.persist_completed_task(baseline.task_name) + tracker.stop() + mock_task_out.assert_called_once() + + def test_mark_http_request_start_empty_task_name_gets_uuid( + self, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker() + tracker.start() + baseline = tracker.mark_http_request_start("") + tracker.finish_http_request(baseline) + tracker.stop() + self.assertTrue(baseline.task_name) + self.assertNotEqual(baseline.task_name, "") + + @mock.patch("codecarbon.output.CodeCarbonAPIOutput.task_out") + def test_persist_completed_task_skips_missing_and_incomplete_tasks( + self, + mock_task_out, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker( + save_to_api=True, + experiment_id="00000000-0000-0000-0000-000000000001", + api_key="test-key", + ) + tracker.start() + tracker.persist_completed_task("does-not-exist") + baseline = tracker.mark_http_request_start("GET /predict") + tracker.persist_completed_task(baseline.task_name) + mock_task_out.assert_not_called() + tracker.finish_http_request(baseline) + tracker.persist_completed_task(baseline.task_name) + tracker.stop() + mock_task_out.assert_called_once() + + @mock.patch("codecarbon.output.CodeCarbonAPIOutput.task_out") + def test_persist_completed_task_skips_already_uploaded( + self, + mock_task_out, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker( + save_to_api=True, + experiment_id="00000000-0000-0000-0000-000000000001", + api_key="test-key", + ) + tracker.start() + baseline = tracker.mark_http_request_start("GET /predict") + tracker.finish_http_request(baseline) + tracker.persist_completed_task(baseline.task_name) + tracker.persist_completed_task(baseline.task_name) + tracker.stop() + mock_task_out.assert_called_once() + + def test_finish_http_request_reuses_cached_cloud_metadata( + self, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker() + tracker.start() + first = tracker.mark_http_request_start("GET /predict") + tracker.finish_http_request(first) + cloud_calls_after_first_finish = mock_cloud.call_count + second = tracker.mark_http_request_start("GET /predict") + tracker.finish_http_request(second) + tracker.stop() + self.assertEqual(mock_cloud.call_count, cloud_calls_after_first_finish) + + def test_mark_does_not_block_on_slow_finish( + self, + mock_cli_setup, + mock_log_values, + mock_cloud, + mock_gpu_details, + mock_gpu_available, + mock_nvidia, + ) -> None: + tracker = _build_tracker() + tracker.start() + baseline = tracker.mark_http_request_start("GET /predict") + tracker._last_measured_time = 0.0 + measure_started = threading.Event() + release_measure = threading.Event() + + original_run = tracker._run_power_measurement + + def slow_measure() -> None: + measure_started.set() + release_measure.wait(timeout=5.0) + original_run() + + with mock.patch.object( + tracker, "_run_power_measurement", side_effect=slow_measure + ): + finish_thread = threading.Thread( + target=tracker.finish_http_request, args=(baseline,) + ) + finish_thread.start() + assert measure_started.wait(timeout=2.0) + + mark_started = time.perf_counter() + second = tracker.mark_http_request_start("GET /predict") + elapsed = time.perf_counter() - mark_started + + release_measure.set() + finish_thread.join(timeout=5.0) + + tracker.stop() + self.assertLess(elapsed, 0.2) + self.assertNotEqual(baseline.task_name, second.task_name) diff --git a/tests/test_input.py b/tests/test_input.py index 875e7e99a..89739d490 100644 --- a/tests/test_input.py +++ b/tests/test_input.py @@ -12,13 +12,10 @@ class TestDataSourceCaching(unittest.TestCase): """Test that DataSource uses module-level cache for static data.""" def test_cache_populated_at_import(self): - """Verify that _CACHE is populated on first data access.""" - from codecarbon.input import _CACHE, DataSource - - ds = DataSource() - ds.get_global_energy_mix_data() + """Verify that _CACHE is populated when module is imported.""" + from codecarbon.input import _CACHE - # Static data should be loaded after first access + # All static data should be pre-loaded self.assertIn("global_energy_mix", _CACHE) self.assertIn("cloud_emissions", _CACHE) self.assertIn("carbon_intensity_per_source", _CACHE) diff --git a/tests/test_offline_emissions_tracker.py b/tests/test_offline_emissions_tracker.py index 36447409d..07adf403c 100644 --- a/tests/test_offline_emissions_tracker.py +++ b/tests/test_offline_emissions_tracker.py @@ -67,14 +67,3 @@ def test_offline_tracker_task(self): self.assertGreater(task_emission_data.emissions, 0.0) self.assertEqual(task_emission_data.country_name, None) - - def test_resolve_offline_country_name_logs_on_invalid_iso(self): - tracker = OfflineEmissionsTracker( - country_iso_code="INVALID", - save_to_file=False, - ) - with self.assertLogs("codecarbon", level="ERROR") as logs: - tracker._resolve_offline_country_name() - self.assertTrue( - any("Does not support country" in message for message in logs.output) - ) diff --git a/tests/test_powermetrics.py b/tests/test_powermetrics.py index b20f5df2c..2fbcba431 100644 --- a/tests/test_powermetrics.py +++ b/tests/test_powermetrics.py @@ -15,32 +15,6 @@ def __init__(self, stderr="", returncode=0): def communicate(self): return ("", self._stderr) - def poll(self): - return self.returncode - - def kill(self): - return None - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - -class HangingProcess: - def __init__(self): - self.killed = False - - def poll(self): - return None - - def kill(self): - self.killed = True - - def communicate(self): - return ("", "") - def __enter__(self): return self @@ -74,68 +48,11 @@ def test_get_details(self, mock_setup, mock_log_values): assert cpu_details == expected_details def test_is_powermetrics_available_returns_false_on_instantiation_error(self): - from codecarbon.core.powermetrics import clear_powermetrics_cache - - clear_powermetrics_cache() with mock.patch( "codecarbon.core.powermetrics.ApplePowermetrics", side_effect=Exception("boom"), ): assert is_powermetrics_available() is False - clear_powermetrics_cache() - - def test_is_powermetrics_available_returns_cached_value(self): - from codecarbon.core.powermetrics import clear_powermetrics_cache - - clear_powermetrics_cache() - with ( - mock.patch("codecarbon.core.powermetrics.ApplePowermetrics"), - mock.patch( - "codecarbon.core.powermetrics._has_powermetrics_sudo", - return_value=True, - ), - ): - assert is_powermetrics_available() is True - with mock.patch( - "codecarbon.core.powermetrics.ApplePowermetrics", - side_effect=Exception("should not instantiate"), - ): - assert is_powermetrics_available() is True - clear_powermetrics_cache() - - def test_is_powermetrics_available_probes_sudo_when_uncached(self): - from codecarbon.core.powermetrics import clear_powermetrics_cache - - clear_powermetrics_cache() - with ( - mock.patch("codecarbon.core.powermetrics.ApplePowermetrics"), - mock.patch( - "codecarbon.core.powermetrics._has_powermetrics_sudo", - return_value=True, - ) as mock_sudo, - ): - assert is_powermetrics_available() is True - mock_sudo.assert_called_once() - clear_powermetrics_cache() - - def test_has_powermetrics_sudo_kills_process_on_timeout(self): - hanging = HangingProcess() - with ( - mock.patch( - "codecarbon.core.powermetrics.shutil.which", - side_effect=["sudo-path", "powermetrics-path"], - ), - mock.patch( - "codecarbon.core.powermetrics.subprocess.Popen", - return_value=hanging, - ), - mock.patch( - "codecarbon.core.powermetrics.time.time", side_effect=[0, 0, 10] - ), - mock.patch("codecarbon.core.powermetrics.time.sleep"), - ): - assert powermetrics_module._has_powermetrics_sudo() is False - assert hanging.killed is True def test_has_powermetrics_sudo_returns_false_when_sudo_missing(self): with mock.patch("codecarbon.core.powermetrics.shutil.which", return_value=None): diff --git a/tests/test_resource_tracker.py b/tests/test_resource_tracker.py index c20fcf1d5..632aee464 100644 --- a/tests/test_resource_tracker.py +++ b/tests/test_resource_tracker.py @@ -204,7 +204,7 @@ def test_set_cpu_tracking_force_mode_uses_cpu_load_and_returns(): fake_tdp = SimpleNamespace(tdp=20, model="CPU") with ( - patch("codecarbon.core.resource_tracker.get_cached_tdp", return_value=fake_tdp), + patch("codecarbon.core.resource_tracker.cpu.TDP", return_value=fake_tdp), patch.object( resource_tracker, "_setup_cpu_load_mode", return_value=True ) as mock_setup, @@ -219,9 +219,6 @@ def test_set_cpu_tracking_prefers_power_gadget(): resource_tracker = ResourceTracker(tracker) with ( - patch("codecarbon.core.resource_tracker.is_mac_os", return_value=False), - patch("codecarbon.core.resource_tracker.is_linux_os", return_value=False), - patch("codecarbon.core.resource_tracker.is_windows_os", return_value=True), patch( "codecarbon.core.resource_tracker.cpu.is_powergadget_available", return_value=True, @@ -240,134 +237,11 @@ def test_set_cpu_tracking_prefers_power_gadget(): mock_power_gadget.assert_called_once_with() -def test_set_cpu_tracking_mac_arm_prefers_cpu_load_over_powermetrics(): - tracker = make_tracker() - resource_tracker = ResourceTracker(tracker) - - with ( - patch("codecarbon.core.resource_tracker.is_mac_os", return_value=True), - patch("codecarbon.core.resource_tracker.is_linux_os", return_value=False), - patch("codecarbon.core.resource_tracker.is_windows_os", return_value=False), - patch( - "codecarbon.core.resource_tracker.detect_cpu_model", - return_value="Apple M1 Max", - ), - patch( - "codecarbon.core.resource_tracker.cpu.is_powergadget_available", - return_value=True, - ), - patch( - "codecarbon.core.resource_tracker.powermetrics.is_powermetrics_available", - return_value=True, - ), - patch.object(resource_tracker, "_setup_power_gadget") as mock_power_gadget, - patch.object(resource_tracker, "_setup_powermetrics") as mock_powermetrics, - patch.object( - resource_tracker, "_setup_cpu_load_fast", return_value=True - ) as mock_cpu_load, - ): - resource_tracker.set_CPU_tracking() - - mock_power_gadget.assert_not_called() - mock_powermetrics.assert_not_called() - mock_cpu_load.assert_called_once_with("Apple M1 Max") - - -def test_setup_cpu_load_fast_returns_false_without_psutil(): - tracker = make_tracker() - resource_tracker = ResourceTracker(tracker) - - with patch( - "codecarbon.core.resource_tracker.cpu.is_psutil_available", - return_value=False, - ): - assert resource_tracker._setup_cpu_load_fast("Intel CPU") is False - - -def test_try_platform_cpu_backend_mac_intel_uses_power_gadget(): - tracker = make_tracker() - resource_tracker = ResourceTracker(tracker) - - with ( - patch("codecarbon.core.resource_tracker.is_mac_os", return_value=True), - patch("codecarbon.core.resource_tracker.is_linux_os", return_value=False), - patch("codecarbon.core.resource_tracker.is_windows_os", return_value=False), - patch( - "codecarbon.core.resource_tracker.detect_cpu_model", - return_value="Intel(R) Core(TM) i7", - ), - patch("codecarbon.core.resource_tracker.is_mac_arm", return_value=False), - patch( - "codecarbon.core.resource_tracker.cpu.is_powergadget_available", - return_value=True, - ), - patch.object(resource_tracker, "_setup_power_gadget") as mock_power_gadget, - ): - assert resource_tracker._try_platform_cpu_backend() is True - - mock_power_gadget.assert_called_once_with() - - -def test_try_platform_cpu_backend_mac_intel_falls_back_to_powermetrics(): - tracker = make_tracker() - resource_tracker = ResourceTracker(tracker) - - with ( - patch("codecarbon.core.resource_tracker.is_mac_os", return_value=True), - patch("codecarbon.core.resource_tracker.is_linux_os", return_value=False), - patch("codecarbon.core.resource_tracker.is_windows_os", return_value=False), - patch( - "codecarbon.core.resource_tracker.detect_cpu_model", - return_value="Intel(R) Core(TM) i7", - ), - patch("codecarbon.core.resource_tracker.is_mac_arm", return_value=False), - patch( - "codecarbon.core.resource_tracker.cpu.is_powergadget_available", - return_value=False, - ), - patch( - "codecarbon.core.resource_tracker.powermetrics.is_powermetrics_available", - return_value=True, - ), - patch.object(resource_tracker, "_setup_powermetrics") as mock_powermetrics, - ): - assert resource_tracker._try_platform_cpu_backend() is True - - mock_powermetrics.assert_called_once_with() - - -def test_set_cpu_tracking_mac_arm_falls_back_to_powermetrics_when_cpu_load_unavailable(): - tracker = make_tracker() - resource_tracker = ResourceTracker(tracker) - - with ( - patch("codecarbon.core.resource_tracker.is_mac_os", return_value=True), - patch("codecarbon.core.resource_tracker.is_linux_os", return_value=False), - patch("codecarbon.core.resource_tracker.is_windows_os", return_value=False), - patch( - "codecarbon.core.resource_tracker.detect_cpu_model", - return_value="Apple M4", - ), - patch( - "codecarbon.core.resource_tracker.powermetrics.is_powermetrics_available", - return_value=True, - ), - patch.object(resource_tracker, "_setup_cpu_load_fast", return_value=False), - patch.object(resource_tracker, "_setup_powermetrics") as mock_powermetrics, - ): - resource_tracker.set_CPU_tracking() - - mock_powermetrics.assert_called_once_with() - - def test_set_cpu_tracking_prefers_rapl_before_powermetrics(): tracker = make_tracker() resource_tracker = ResourceTracker(tracker) with ( - patch("codecarbon.core.resource_tracker.is_linux_os", return_value=True), - patch("codecarbon.core.resource_tracker.is_mac_os", return_value=False), - patch("codecarbon.core.resource_tracker.is_windows_os", return_value=False), patch( "codecarbon.core.resource_tracker.cpu.is_powergadget_available", return_value=False, @@ -403,7 +277,7 @@ def test_set_cpu_tracking_falls_back_when_forced_power_is_set(): "codecarbon.core.resource_tracker.powermetrics.is_powermetrics_available", return_value=True, ), - patch("codecarbon.core.resource_tracker.get_cached_tdp", return_value=fake_tdp), + patch("codecarbon.core.resource_tracker.cpu.TDP", return_value=fake_tdp), patch.object(resource_tracker, "_setup_fallback_tracking") as mock_fallback, ): resource_tracker.set_CPU_tracking() @@ -476,121 +350,3 @@ def test_set_cpu_gpu_ram_tracking_calls_all_setup_steps(): mock_ram.assert_called_once_with() mock_cpu.assert_called_once_with() mock_gpu.assert_called_once_with() - - -def test_hardware_cache_reuses_setup(): - from codecarbon.core import hardware_cache - - hardware_cache.clear_cache() - key = hardware_cache.make_key(make_tracker()) - hardware_cache._plans[key] = hardware_cache._HardwarePlan( - ram_tracker="cached_ram", - cpu_tracker="cached_cpu", - gpu_tracker="cached_gpu", - conf={"cpu_model": "Cached CPU", "gpu_count": 0, "gpu_model": ""}, - hardware_specs=[], - ) - - tracker2 = make_tracker() - rt2 = ResourceTracker(tracker2) - with ( - patch.object(rt2, "set_RAM_tracking") as mock_ram, - patch.object(rt2, "set_CPU_tracking") as mock_cpu, - patch.object(rt2, "set_GPU_tracking") as mock_gpu, - ): - rt2.set_CPU_GPU_ram_tracking() - mock_ram.assert_not_called() - mock_cpu.assert_not_called() - mock_gpu.assert_not_called() - - assert rt2.cpu_tracker == "cached_cpu" - assert tracker2._conf.get("cpu_model") == "Cached CPU" - hardware_cache.clear_cache() - - -def test_setup_power_gadget_configures_tracker(): - tracker = make_tracker() - resource_tracker = ResourceTracker(tracker) - hardware_cpu = MagicMock() - hardware_cpu.get_model.return_value = "Intel CPU" - - with patch( - "codecarbon.core.resource_tracker.CPU.from_utils", return_value=hardware_cpu - ) as mock_from_utils: - assert resource_tracker._setup_power_gadget() is True - - mock_from_utils.assert_called_once_with( - "out", - "intel_power_gadget", - tracking_mode="machine", - ) - assert resource_tracker.cpu_tracker == "Power Gadget" - assert tracker._conf["cpu_model"] == "Intel CPU" - assert tracker._hardware == [hardware_cpu] - - -def test_setup_fallback_tracking_uses_forced_cpu_power(): - tracker = make_tracker(_force_cpu_power=99) - resource_tracker = ResourceTracker(tracker) - hardware_cpu = MagicMock() - tdp = SimpleNamespace(model="Matched CPU") - - with ( - patch( - "codecarbon.core.resource_tracker.cpu.is_psutil_available", - return_value=True, - ), - patch( - "codecarbon.core.resource_tracker.CPU.from_utils", return_value=hardware_cpu - ) as mock_from_utils, - patch.object( - resource_tracker, "_get_install_instructions", return_value="instructions" - ), - ): - resource_tracker._setup_fallback_tracking(tdp, None) - - mock_from_utils.assert_called_once_with( - "out", - MODE_CPU_LOAD, - "Matched CPU", - 99, - tracking_mode="machine", - ) - assert resource_tracker.cpu_tracker == MODE_CPU_LOAD - assert tracker._conf["cpu_model"] == "Matched CPU" - - -def test_setup_fallback_tracking_cpu_load_when_tdp_falsy(): - tracker = make_tracker() - resource_tracker = ResourceTracker(tracker) - hardware_cpu = MagicMock() - - class FalseyTDP: - model = "Unknown CPU" - - def __bool__(self): - return False - - with ( - patch( - "codecarbon.core.resource_tracker.cpu.is_psutil_available", - return_value=True, - ), - patch( - "codecarbon.core.resource_tracker.CPU.from_utils", return_value=hardware_cpu - ) as mock_from_utils, - patch.object( - resource_tracker, "_get_install_instructions", return_value="instructions" - ), - ): - resource_tracker._setup_fallback_tracking(FalseyTDP(), None) - - mock_from_utils.assert_called_once_with( - "out", - MODE_CPU_LOAD, - "Unknown CPU", - None, - tracking_mode="machine", - ) - assert resource_tracker.cpu_tracker == MODE_CPU_LOAD - assert tracker._hardware == [hardware_cpu] diff --git a/uv.lock b/uv.lock index 1bc68ae4c..74d38f0e9 100644 --- a/uv.lock +++ b/uv.lock @@ -32,6 +32,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "anyio" +version = "4.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "idna" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/14/2c5dd9f512b66549ae92767a9c7b330ae88e1932ca57876909410251fe13/anyio-4.13.0.tar.gz", hash = "sha256:334b70e641fd2221c1505b3890c69882fe4a2df910cba14d97019b90b24439dc", size = 231622, upload-time = "2026-03-24T12:59:09.671Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, +] + [[package]] name = "arrow" version = "1.4.0" @@ -434,6 +448,10 @@ carbonboard = [ { name = "dash-bootstrap-components" }, { name = "fire" }, ] +fastapi = [ + { name = "fastapi" }, + { name = "httpx" }, +] viz-legacy = [ { name = "dash" }, { name = "dash-bootstrap-components" }, @@ -444,6 +462,8 @@ viz-legacy = [ dev = [ { name = "black" }, { name = "bumpver" }, + { name = "fastapi" }, + { name = "httpx" }, { name = "jsonschema" }, { name = "logfire" }, { name = "mktestdocs" }, @@ -479,8 +499,10 @@ requires-dist = [ { name = "dash", marker = "extra == 'viz-legacy'" }, { name = "dash-bootstrap-components", marker = "extra == 'carbonboard'", specifier = ">1.0.0" }, { name = "dash-bootstrap-components", marker = "extra == 'viz-legacy'", specifier = ">1.0.0" }, + { name = "fastapi", marker = "extra == 'fastapi'", specifier = ">=0.100" }, { name = "fire", marker = "extra == 'carbonboard'" }, { name = "fire", marker = "extra == 'viz-legacy'" }, + { name = "httpx", marker = "extra == 'fastapi'" }, { name = "joserfc", specifier = ">=1.0.0" }, { name = "nvidia-ml-py" }, { name = "pandas", marker = "python_full_version < '3.14'" }, @@ -496,12 +518,14 @@ requires-dist = [ { name = "rich" }, { name = "typer" }, ] -provides-extras = ["carbonboard", "viz-legacy"] +provides-extras = ["carbonboard", "viz-legacy", "fastapi"] [package.metadata.requires-dev] dev = [ { name = "black" }, { name = "bumpver" }, + { name = "fastapi", specifier = ">=0.100" }, + { name = "httpx" }, { name = "jsonschema" }, { name = "logfire", specifier = ">=1.0.1" }, { name = "mktestdocs" }, @@ -786,6 +810,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c1/ea/53f2148663b321f21b5a606bd5f191517cf40b7072c0497d3c92c4a13b1e/executing-2.2.1-py2.py3-none-any.whl", hash = "sha256:760643d3452b4d777d295bb167ccc74c64a81df23fb5e08eff250c425a4b2017", size = 28317, upload-time = "2025-09-01T09:48:08.5Z" }, ] +[[package]] +name = "fastapi" +version = "0.136.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-doc" }, + { name = "pydantic" }, + { name = "starlette" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5d/45/c130091c2dfa061bbfe3150f2a5091ef1adf149f2a8d2ae769ecaf6e99a2/fastapi-0.136.1.tar.gz", hash = "sha256:7af665ad7acfa0a3baf8983d393b6b471b9da10ede59c60045f49fbc89a0fa7f", size = 397448, upload-time = "2026-04-23T16:49:44.046Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/ff/2e4eca3ade2c22fe1dea7043b8ee9dabe47753349eb1b56a202de8af6349/fastapi-0.136.1-py3-none-any.whl", hash = "sha256:a6e9d7eeada96c93a4d69cb03836b44fa34e2854accb7244a1ece36cd4781c3f", size = 117683, upload-time = "2026-04-23T16:49:42.437Z" }, +] + [[package]] name = "filelock" version = "3.29.1" @@ -857,6 +897,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/11/8c/c9138d881c79aa0ea9ed83cbd58d5ca75624378b38cee225dcf5c42cc91f/griffelib-2.0.2-py3-none-any.whl", hash = "sha256:925c857658fb1ba40c0772c37acbc2ab650bd794d9c1b9726922e36ea4117ea1", size = 142357, upload-time = "2026-03-27T11:34:46.275Z" }, ] +[[package]] +name = "h11" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, +] + +[[package]] +name = "httpcore" +version = "1.0.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, +] + +[[package]] +name = "httpx" +version = "0.28.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "certifi" }, + { name = "httpcore" }, + { name = "idna" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, +] + [[package]] name = "identify" version = "2.6.19" @@ -3106,6 +3183,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5e/f5/0c41cb68dcae6b7de4fac4188a3a9589e21fb31df21ea3a2e888db95e6c9/soupsieve-2.8.4-py3-none-any.whl", hash = "sha256:e7e6b0769c8f51ed59acab6e994b00621096cfb1c640a7509295987388fbaf65", size = 37304, upload-time = "2026-05-24T13:55:55.406Z" }, ] +[[package]] +name = "starlette" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/81/69/17425771797c36cded50b7fe44e850315d039f28b15901ab44839e70b593/starlette-1.0.0.tar.gz", hash = "sha256:6a4beaf1f81bb472fd19ea9b918b50dc3a77a6f2e190a12954b25e6ed5eea149", size = 2655289, upload-time = "2026-03-22T18:29:46.779Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/c9/584bc9651441b4ba60cc4d557d8a547b5aff901af35bda3a4ee30c819b82/starlette-1.0.0-py3-none-any.whl", hash = "sha256:d3ec55e0bb321692d275455ddfd3df75fff145d009685eb40dc91fc66b03d38b", size = 72651, upload-time = "2026-03-22T18:29:45.111Z" }, +] + [[package]] name = "taskipy" version = "1.14.1"