diff --git a/carbonserver/carbonserver/api/infra/database/telemetry_sql_models.py b/carbonserver/carbonserver/api/infra/database/telemetry_sql_models.py index 243ae79ee..aa999d025 100644 --- a/carbonserver/carbonserver/api/infra/database/telemetry_sql_models.py +++ b/carbonserver/carbonserver/api/infra/database/telemetry_sql_models.py @@ -2,7 +2,7 @@ import uuid -from sqlalchemy import JSON, Boolean, Column, DateTime, Float, Integer, String +from sqlalchemy import Column, DateTime, Float, Integer, String from sqlalchemy.dialects.postgresql import UUID from carbonserver.database.database import Base @@ -12,7 +12,7 @@ class Telemetry(Base): __tablename__ = "telemetry" id = Column(UUID(as_uuid=True), primary_key=True, index=True, default=uuid.uuid4) - timestamp = Column(DateTime, nullable=False) + timestamp = Column(DateTime(timezone=True), nullable=False) telemetry_level = Column(String, nullable=False) os = Column(String, nullable=True) @@ -21,8 +21,6 @@ class Telemetry(Base): region = Column(String, nullable=True) cloud_provider = Column(String, nullable=True) cloud_region = Column(String, nullable=True) - longitude = Column(Float, nullable=True) - latitude = Column(Float, nullable=True) cpu_count = Column(Integer, nullable=True) cpu_physical_count = Column(Integer, nullable=True) @@ -38,64 +36,10 @@ class Telemetry(Base): python_version = Column(String, nullable=True) python_implementation = Column(String, nullable=True) - python_executable_hash = Column(String, nullable=True) python_env_type = Column(String, nullable=True) codecarbon_version = Column(String, nullable=True) codecarbon_install_method = Column(String, nullable=True) - total_emissions_kg = Column(Float, nullable=True) - emissions_rate_kg_per_sec = Column(Float, nullable=True) - energy_consumed_kwh = Column(Float, nullable=True) - cpu_energy_kwh = Column(Float, nullable=True) - gpu_energy_kwh = Column(Float, nullable=True) - ram_energy_kwh = Column(Float, nullable=True) - duration_seconds = Column(Float, nullable=True) - cpu_utilization_avg = Column(Float, nullable=True) - gpu_utilization_avg = Column(Float, nullable=True) - ram_utilization_avg = Column(Float, nullable=True) - - tracking_mode = Column(String, nullable=True) - api_mode = Column(String, nullable=True) - output_methods = Column(JSON, nullable=True) - hardware_tracked = Column(JSON, nullable=True) - task_tracking_used = Column(Boolean, nullable=True) - decorator_vs_context = Column(String, nullable=True) - measure_power_interval_secs = Column(Float, nullable=True) - - hardware_detection_success = Column(Boolean, nullable=True) - rapl_available = Column(Boolean, nullable=True) - gpu_detection_method = Column(String, nullable=True) - first_measurement_time_ms = Column(Float, nullable=True) - tracking_overhead_percent = Column(Float, nullable=True) - errors_encountered = Column(JSON, nullable=True) - warning_count = Column(Integer, nullable=True) - - ide_used = Column(String, nullable=True) - notebook_environment = Column(String, nullable=True) - ci_environment = Column(String, nullable=True) - python_package_manager = Column(String, nullable=True) - framework_detected = Column(String, nullable=True) - - has_torch = Column(Boolean, nullable=True) - torch_version = Column(String, nullable=True) - has_transformers = Column(Boolean, nullable=True) - transformers_version = Column(String, nullable=True) - has_diffusers = Column(Boolean, nullable=True) - diffusers_version = Column(String, nullable=True) - has_tensorflow = Column(Boolean, nullable=True) - tensorflow_version = Column(String, nullable=True) - has_keras = Column(Boolean, nullable=True) - keras_version = Column(String, nullable=True) - has_pytorch_lightning = Column(Boolean, nullable=True) - pytorch_lightning_version = Column(String, nullable=True) - has_fastai = Column(Boolean, nullable=True) - fastai_version = Column(String, nullable=True) - ml_framework_primary = Column(String, nullable=True) - - container_runtime = Column(String, nullable=True) - in_container = Column(Boolean, nullable=True) - host_machine_hash = Column(String, nullable=True) - def __repr__(self): return ( f'= RATE_LIMIT check and both be admitted. +_recent_requests_lock = threading.Lock() + router = APIRouter() +def _rate_limited(host: str) -> bool: + now = time.monotonic() + with _recent_requests_lock: + if host not in _recent_requests and len(_recent_requests) >= MAX_TRACKED_IPS: + _recent_requests.clear() + hits = _recent_requests[host] + while hits and now - hits[0] > RATE_WINDOW_SECONDS: + hits.popleft() + if len(hits) >= RATE_LIMIT: + return True + hits.append(now) + return False + + @router.post( "/telemetry", tags=TELEMETRY_ROUTER_TAGS, @@ -24,8 +55,15 @@ @inject def add_telemetry( telemetry: TelemetryCreate, + request: Request, telemetry_service: TelemetryService = Depends( Provide[ServerContainer.telemetry_service] ), ) -> UUID: + host = request.client.host if request.client else "unknown" + if _rate_limited(host): + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail="Too many telemetry requests", + ) return telemetry_service.add_telemetry(telemetry) diff --git a/carbonserver/carbonserver/api/schemas_telemetry.py b/carbonserver/carbonserver/api/schemas_telemetry.py index 5517c2ff4..7b321502f 100644 --- a/carbonserver/carbonserver/api/schemas_telemetry.py +++ b/carbonserver/carbonserver/api/schemas_telemetry.py @@ -2,15 +2,17 @@ from datetime import datetime from enum import Enum -from typing import List, Optional +from typing import Annotated, Optional from pydantic import BaseModel, ConfigDict, Field, model_validator +#: Every free-text field is capped so one request cannot store unbounded data. +Str = Annotated[str, Field(max_length=256)] + class TelemetryLevel(str, Enum): disabled = "disabled" minimal = "minimal" - extensive = "extensive" class TelemetryBase(BaseModel): @@ -35,141 +37,38 @@ class TelemetryBase(BaseModel): timestamp: datetime telemetry_level: TelemetryLevel - os: Optional[str] = None - country_name: Optional[str] = None + os: Optional[Str] = None + country_name: Optional[Str] = None country_iso_code: Optional[str] = Field(default=None, min_length=2, max_length=3) - region: Optional[str] = None - cloud_provider: Optional[str] = None - cloud_region: Optional[str] = None - longitude: Optional[float] = Field(default=None, ge=-180, le=180) - latitude: Optional[float] = Field(default=None, ge=-90, le=90) + region: Optional[Str] = None + cloud_provider: Optional[Str] = None + cloud_region: Optional[Str] = None cpu_count: Optional[int] = Field(default=None, ge=0) cpu_physical_count: Optional[int] = Field(default=None, ge=0) - cpu_model: Optional[str] = None - cpu_architecture: Optional[str] = None + cpu_model: Optional[Str] = None + cpu_architecture: Optional[Str] = None gpu_count: Optional[int] = Field(default=None, ge=0) - gpu_model: Optional[str] = None - gpu_driver_version: Optional[str] = None + gpu_model: Optional[Str] = None + gpu_driver_version: Optional[Str] = None gpu_memory_total_gb: Optional[float] = Field(default=None, ge=0) ram_total_size_gb: Optional[float] = Field(default=None, ge=0) - cuda_version: Optional[str] = None - cudnn_version: Optional[str] = None + cuda_version: Optional[Str] = None + cudnn_version: Optional[Str] = None - python_version: Optional[str] = None - python_implementation: Optional[str] = None - python_executable_hash: Optional[str] = Field( - default=None, min_length=64, max_length=64 - ) - python_env_type: Optional[str] = None - codecarbon_version: Optional[str] = None - codecarbon_install_method: Optional[str] = None - - total_emissions_kg: Optional[float] = Field(default=None, ge=0) - emissions_rate_kg_per_sec: Optional[float] = Field(default=None, ge=0) - energy_consumed_kwh: Optional[float] = Field(default=None, ge=0) - cpu_energy_kwh: Optional[float] = Field(default=None, ge=0) - gpu_energy_kwh: Optional[float] = Field(default=None, ge=0) - ram_energy_kwh: Optional[float] = Field(default=None, ge=0) - duration_seconds: Optional[float] = Field(default=None, ge=0) - cpu_utilization_avg: Optional[float] = Field(default=None, ge=0, le=100) - gpu_utilization_avg: Optional[float] = Field(default=None, ge=0, le=100) - ram_utilization_avg: Optional[float] = Field(default=None, ge=0, le=100) - - tracking_mode: Optional[str] = None - api_mode: Optional[str] = None - output_methods: Optional[List[str]] = None - hardware_tracked: Optional[List[str]] = None - task_tracking_used: Optional[bool] = None - decorator_vs_context: Optional[str] = None - measure_power_interval_secs: Optional[float] = Field(default=None, ge=0) - - hardware_detection_success: Optional[bool] = None - rapl_available: Optional[bool] = None - gpu_detection_method: Optional[str] = None - first_measurement_time_ms: Optional[float] = Field(default=None, ge=0) - tracking_overhead_percent: Optional[float] = Field(default=None, ge=0) - errors_encountered: Optional[List[str]] = None - warning_count: Optional[int] = Field(default=None, ge=0) - - ide_used: Optional[str] = None - notebook_environment: Optional[str] = None - ci_environment: Optional[str] = None - python_package_manager: Optional[str] = None - framework_detected: Optional[str] = None - - has_torch: Optional[bool] = None - torch_version: Optional[str] = None - has_transformers: Optional[bool] = None - transformers_version: Optional[str] = None - has_diffusers: Optional[bool] = None - diffusers_version: Optional[str] = None - has_tensorflow: Optional[bool] = None - tensorflow_version: Optional[str] = None - has_keras: Optional[bool] = None - keras_version: Optional[str] = None - has_pytorch_lightning: Optional[bool] = None - pytorch_lightning_version: Optional[str] = None - has_fastai: Optional[bool] = None - fastai_version: Optional[str] = None - ml_framework_primary: Optional[str] = None - - container_runtime: Optional[str] = None - in_container: Optional[bool] = None - host_machine_hash: Optional[str] = None + python_version: Optional[Str] = None + python_implementation: Optional[Str] = None + python_env_type: Optional[Str] = None + codecarbon_version: Optional[Str] = None + codecarbon_install_method: Optional[Str] = None @model_validator(mode="after") - def validate_telemetry_level(self): + def reject_disabled_level(self): if self.telemetry_level == TelemetryLevel.disabled: raise ValueError("Disabled telemetry must not be submitted") - - if self.telemetry_level == TelemetryLevel.minimal: - extensive_fields = set(type(self).model_fields) - MINIMAL_TELEMETRY_FIELDS - submitted_extensive_fields = [ - field - for field in extensive_fields - if getattr(self, field) not in (None, [], {}) - ] - if submitted_extensive_fields: - fields = ", ".join(sorted(submitted_extensive_fields)) - raise ValueError( - f"Minimal telemetry cannot include extensive fields: {fields}" - ) - return self -MINIMAL_TELEMETRY_FIELDS = { - "timestamp", - "telemetry_level", - "os", - "country_name", - "country_iso_code", - "region", - "cloud_provider", - "cloud_region", - "longitude", - "latitude", - "cpu_count", - "cpu_physical_count", - "cpu_model", - "cpu_architecture", - "gpu_count", - "gpu_model", - "gpu_driver_version", - "gpu_memory_total_gb", - "ram_total_size_gb", - "cuda_version", - "cudnn_version", - "python_version", - "python_implementation", - "python_executable_hash", - "python_env_type", - "codecarbon_version", - "codecarbon_install_method", -} - - class TelemetryCreate(TelemetryBase): pass diff --git a/carbonserver/carbonserver/database/alembic/versions/20260927_add_telemetry_table.py b/carbonserver/carbonserver/database/alembic/versions/20260927_add_telemetry_table.py new file mode 100644 index 000000000..1dbf71558 --- /dev/null +++ b/carbonserver/carbonserver/database/alembic/versions/20260927_add_telemetry_table.py @@ -0,0 +1,196 @@ +"""add telemetry table + +Revision ID: 20260927_add_telemetry +Revises: 20251119_add_utilization +Create Date: 2026-09-27 12:00:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "20260927_add_telemetry" +down_revision = "20251119_add_utilization" +branch_labels = None +depends_on = None + +# New columns, as created by ``create_table`` for a fresh database. A naive +# timestamp is ambiguous across servers in different zones, so it is always +# timezone-aware (``timestamptz`` on Postgres). +NEW_COLUMNS = [ + sa.Column("timestamp", sa.DateTime(timezone=True), nullable=False), + sa.Column("telemetry_level", sa.String, nullable=False), + sa.Column("os", sa.String, nullable=True), + sa.Column("country_name", sa.String, nullable=True), + sa.Column("country_iso_code", sa.String, nullable=True), + sa.Column("region", sa.String, nullable=True), + sa.Column("cloud_provider", sa.String, nullable=True), + sa.Column("cloud_region", sa.String, nullable=True), + sa.Column("cpu_count", sa.Integer, nullable=True), + sa.Column("cpu_physical_count", sa.Integer, nullable=True), + sa.Column("cpu_model", sa.String, nullable=True), + sa.Column("cpu_architecture", sa.String, nullable=True), + sa.Column("gpu_count", sa.Integer, nullable=True), + sa.Column("gpu_model", sa.String, nullable=True), + sa.Column("gpu_driver_version", sa.String, nullable=True), + sa.Column("gpu_memory_total_gb", sa.Float, nullable=True), + sa.Column("ram_total_size_gb", sa.Float, nullable=True), + sa.Column("cuda_version", sa.String, nullable=True), + sa.Column("cudnn_version", sa.String, nullable=True), + sa.Column("python_version", sa.String, nullable=True), + sa.Column("python_implementation", sa.String, nullable=True), + sa.Column("python_env_type", sa.String, nullable=True), + sa.Column("codecarbon_version", sa.String, nullable=True), + sa.Column("codecarbon_install_method", sa.String, nullable=True), +] + +# Columns that existed on the old ``create_all``-generated ``telemetry`` table +# (see ``telemetry_sql_models.py`` on master before this migration) but are not +# part of the new, minimal-level-only model. Dropped on upgrade, restored +# (nullable) on downgrade so no data is lost either way. +OLD_ONLY_COLUMNS = [ + sa.Column("longitude", sa.Float, nullable=True), + sa.Column("latitude", sa.Float, nullable=True), + sa.Column("python_executable_hash", sa.String, nullable=True), + sa.Column("total_emissions_kg", sa.Float, nullable=True), + sa.Column("emissions_rate_kg_per_sec", sa.Float, nullable=True), + sa.Column("energy_consumed_kwh", sa.Float, nullable=True), + sa.Column("cpu_energy_kwh", sa.Float, nullable=True), + sa.Column("gpu_energy_kwh", sa.Float, nullable=True), + sa.Column("ram_energy_kwh", sa.Float, nullable=True), + sa.Column("duration_seconds", sa.Float, nullable=True), + sa.Column("cpu_utilization_avg", sa.Float, nullable=True), + sa.Column("gpu_utilization_avg", sa.Float, nullable=True), + sa.Column("ram_utilization_avg", sa.Float, nullable=True), + sa.Column("tracking_mode", sa.String, nullable=True), + sa.Column("api_mode", sa.String, nullable=True), + sa.Column("output_methods", sa.JSON, nullable=True), + sa.Column("hardware_tracked", sa.JSON, nullable=True), + sa.Column("task_tracking_used", sa.Boolean, nullable=True), + sa.Column("decorator_vs_context", sa.String, nullable=True), + sa.Column("measure_power_interval_secs", sa.Float, nullable=True), + sa.Column("hardware_detection_success", sa.Boolean, nullable=True), + sa.Column("rapl_available", sa.Boolean, nullable=True), + sa.Column("gpu_detection_method", sa.String, nullable=True), + sa.Column("first_measurement_time_ms", sa.Float, nullable=True), + sa.Column("tracking_overhead_percent", sa.Float, nullable=True), + sa.Column("errors_encountered", sa.JSON, nullable=True), + sa.Column("warning_count", sa.Integer, nullable=True), + sa.Column("ide_used", sa.String, nullable=True), + sa.Column("notebook_environment", sa.String, nullable=True), + sa.Column("ci_environment", sa.String, nullable=True), + sa.Column("python_package_manager", sa.String, nullable=True), + sa.Column("framework_detected", sa.String, nullable=True), + sa.Column("has_torch", sa.Boolean, nullable=True), + sa.Column("torch_version", sa.String, nullable=True), + sa.Column("has_transformers", sa.Boolean, nullable=True), + sa.Column("transformers_version", sa.String, nullable=True), + sa.Column("has_diffusers", sa.Boolean, nullable=True), + sa.Column("diffusers_version", sa.String, nullable=True), + sa.Column("has_tensorflow", sa.Boolean, nullable=True), + sa.Column("tensorflow_version", sa.String, nullable=True), + sa.Column("has_keras", sa.Boolean, nullable=True), + sa.Column("keras_version", sa.String, nullable=True), + sa.Column("has_pytorch_lightning", sa.Boolean, nullable=True), + sa.Column("pytorch_lightning_version", sa.String, nullable=True), + sa.Column("has_fastai", sa.Boolean, nullable=True), + sa.Column("fastai_version", sa.String, nullable=True), + sa.Column("ml_framework_primary", sa.String, nullable=True), + sa.Column("container_runtime", sa.String, nullable=True), + sa.Column("in_container", sa.Boolean, nullable=True), + sa.Column("host_machine_hash", sa.String, nullable=True), +] + + +def _bind_is_postgres(bind) -> bool: + """SQLite (used by the migration test) can't ALTER a column's type.""" + return bind.dialect.name == "postgresql" + + +def _ensure_retention(bind) -> None: + """Delete telemetry rows older than 3 years, enforced in the database. + + A statement-level trigger fires after every insert and purges rows past + the retention window, so the policy holds regardless of what the + application does. The index keeps the purge's WHERE clause a cheap range + scan. Postgres only; SQLite (used by the migration test) has no trigger. + """ + if not _bind_is_postgres(bind): + return + inspector = sa.inspect(bind) + existing_indexes = {ix["name"] for ix in inspector.get_indexes("telemetry")} + if "ix_telemetry_timestamp" not in existing_indexes: + op.create_index("ix_telemetry_timestamp", "telemetry", ["timestamp"]) + op.execute(""" + CREATE OR REPLACE FUNCTION telemetry_purge() RETURNS trigger AS $$ + BEGIN + DELETE FROM telemetry WHERE timestamp < now() - interval '3 years'; + RETURN NULL; + END; + $$ LANGUAGE plpgsql; + """) + op.execute("DROP TRIGGER IF EXISTS telemetry_retention ON telemetry;") + op.execute(""" + CREATE TRIGGER telemetry_retention + AFTER INSERT ON telemetry + FOR EACH STATEMENT EXECUTE FUNCTION telemetry_purge(); + """) + + +def upgrade(): + """Anonymous SDK telemetry, one row per process (minimal level only). + + Some production databases already have a ``telemetry`` table, created by + ``Base.metadata.create_all`` before this migration existed (see #1171), + with the old wide schema (``latitude``/``longitude``/etc). In that case + only the schema is adjusted in place, so existing rows are preserved; + otherwise the table is created fresh. + """ + bind = op.get_bind() + inspector = sa.inspect(bind) + if inspector.has_table("telemetry"): + existing = {c["name"] for c in inspector.get_columns("telemetry")} + for column in OLD_ONLY_COLUMNS: + if column.name in existing: + op.drop_column("telemetry", column.name) + for column in NEW_COLUMNS: + if column.name not in existing: + op.add_column("telemetry", column.copy()) + if "timestamp" in existing and _bind_is_postgres(bind): + op.alter_column( + "telemetry", + "timestamp", + type_=sa.DateTime(timezone=True), + postgresql_using="timestamp AT TIME ZONE 'UTC'", + ) + _ensure_retention(bind) + return + op.create_table( + "telemetry", + sa.Column("id", postgresql.UUID(as_uuid=True), primary_key=True), + *(column.copy() for column in NEW_COLUMNS), + ) + op.create_index("ix_telemetry_id", "telemetry", ["id"]) + _ensure_retention(bind) + + +def downgrade(): + """Restore the old, wider schema (as nullable columns), preserving rows.""" + bind = op.get_bind() + if _bind_is_postgres(bind): + op.execute("DROP TRIGGER IF EXISTS telemetry_retention ON telemetry;") + op.execute("DROP FUNCTION IF EXISTS telemetry_purge();") + op.execute("DROP INDEX IF EXISTS ix_telemetry_timestamp;") + existing = {c["name"] for c in sa.inspect(bind).get_columns("telemetry")} + for column in OLD_ONLY_COLUMNS: + if column.name not in existing: + op.add_column("telemetry", column.copy()) + if _bind_is_postgres(bind): + op.alter_column( + "telemetry", + "timestamp", + type_=sa.DateTime(), + postgresql_using="timestamp AT TIME ZONE 'UTC'", + ) diff --git a/carbonserver/docker/entrypoint.sh b/carbonserver/docker/entrypoint.sh index dbcf3d395..c006a7f3b 100644 --- a/carbonserver/docker/entrypoint.sh +++ b/carbonserver/docker/entrypoint.sh @@ -18,4 +18,9 @@ else fi echo "Starting uvicorn server..." # uvicorn --reload main:app --host 0.0.0.0 --port 8000 +# ponytail: --forwarded-allow-ips=* trusts X-Forwarded-For from any peer, so +# request.client.host (used e.g. by the telemetry rate limiter) can be +# spoofed by the caller unless a real reverse proxy always overwrites that +# header first. Kept as-is (maintainer decision); tighten to the proxy's +# actual IP/CIDR if that stops being guaranteed. uvicorn main:app --host 0.0.0.0 --port 8000 --proxy-headers --forwarded-allow-ips=* diff --git a/carbonserver/main.py b/carbonserver/main.py index c1f5dcaa6..7fc0cfe14 100644 --- a/carbonserver/main.py +++ b/carbonserver/main.py @@ -9,7 +9,7 @@ from starlette.responses import JSONResponse from carbonserver.api.errors import DBException, UserException, get_http_exception -from carbonserver.api.infra.database import sql_models, telemetry_sql_models +from carbonserver.api.infra.database import sql_models from carbonserver.api.routers import ( authenticate, emissions, @@ -81,7 +81,6 @@ def init_db(container): db = container.db() db.create_database() sql_models.Base.metadata.create_all(bind=engine) - telemetry_sql_models.Base.metadata.create_all(bind=engine) def init_server(container): diff --git a/carbonserver/tests/api/integration/test_telemetry_e2e_db.py b/carbonserver/tests/api/integration/test_telemetry_e2e_db.py new file mode 100644 index 000000000..254274ced --- /dev/null +++ b/carbonserver/tests/api/integration/test_telemetry_e2e_db.py @@ -0,0 +1,73 @@ +"""Check the telemetry rows written by scripts/e2e_telemetry.sh. + +Reads the database directly (``DATABASE_URL``), so it checks what was stored, +not what the API answered. Skipped unless the script sets +``TELEMETRY_E2E_EXPECTED_ROWS``. +""" + +import os + +import pytest +from sqlalchemy import create_engine, inspect, text + +# The fields docs/how-to/telemetry.md lists, plus the row id. +ALLOWED_COLUMNS = { + "id", + "timestamp", + "telemetry_level", + "os", + "country_name", + "country_iso_code", + "region", + "cloud_provider", + "cloud_region", + "cpu_count", + "cpu_physical_count", + "cpu_model", + "cpu_architecture", + "gpu_count", + "gpu_model", + "gpu_memory_total_gb", + "gpu_driver_version", + "cuda_version", + "cudnn_version", + "ram_total_size_gb", + "python_version", + "python_implementation", + "python_env_type", + "codecarbon_version", + "codecarbon_install_method", +} + +pytestmark = [ + pytest.mark.integ_test, + pytest.mark.skipif( + not os.getenv("TELEMETRY_E2E_EXPECTED_ROWS"), + reason="run through scripts/e2e_telemetry.sh", + ), +] + + +@pytest.fixture(scope="module") +def engine(): + return create_engine(os.environ["DATABASE_URL"]) + + +def test_table_has_no_coordinates_and_only_allowed_columns(engine): + columns = {column["name"] for column in inspect(engine).get_columns("telemetry")} + assert not {"latitude", "longitude"} & columns + assert columns <= ALLOWED_COLUMNS + + +def test_one_minimal_row_per_process(engine): + with engine.connect() as conn: + rows = conn.execute(text("SELECT * FROM telemetry")).mappings().all() + + for row in rows: + print({key: value for key, value in row.items() if key != "id"}) + assert len(rows) == int(os.environ["TELEMETRY_E2E_EXPECTED_ROWS"]) + for row in rows: + assert row["telemetry_level"] == "minimal" + assert {key for key, value in row.items() if value is not None} <= ( + ALLOWED_COLUMNS + ) diff --git a/carbonserver/tests/api/integration/test_telemetry_local_api.py b/carbonserver/tests/api/integration/test_telemetry_local_api.py new file mode 100644 index 000000000..d078f7d8d --- /dev/null +++ b/carbonserver/tests/api/integration/test_telemetry_local_api.py @@ -0,0 +1,88 @@ +"""Integration tests for telemetry against a running local carbonserver API.""" + +import os +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import requests + +REPO_ROOT = Path(__file__).resolve().parents[4] +sys.path.insert(0, str(REPO_ROOT)) + +from codecarbon.core.telemetry.client import post_private # noqa: E402 +from codecarbon.core.telemetry.collect import build_payload # noqa: E402 +from codecarbon.core.telemetry.schemas import TelemetryLevel # noqa: E402 +from codecarbon.core.telemetry.settings import TelemetrySettings # noqa: E402 +from codecarbon.output_methods.emissions_data import EmissionsData # noqa: E402 + +URL = os.getenv("CODECARBON_API_URL") +if URL is None: + pytest.exit("CODECARBON_API_URL is not defined (e.g. http://localhost:8008)") + + +def _local_settings() -> TelemetrySettings: + return TelemetrySettings.resolve( + external_conf={ + "telemetry_level": "minimal", + "telemetry_api_url": URL.rstrip("/"), + } + ) + + +def _sample_emissions() -> EmissionsData: + return EmissionsData( + timestamp="2026-01-01T00:00:00", + project_name="telemetry-local", + run_id="local-run", + experiment_id="e", + duration=10.0, + emissions=0.001, + emissions_rate=0.0001, + cpu_power=0.0, + gpu_power=0.0, + ram_power=0.0, + cpu_energy=0.0, + gpu_energy=0.0, + ram_energy=0.0, + energy_consumed=0.01, + water_consumed=0.0, + country_name="France", + country_iso_code="FRA", + region="idf", + cloud_provider="", + cloud_region="", + os="Linux", + python_version="3.12", + codecarbon_version="3.2.8", + cpu_count=4, + cpu_model="test-cpu", + gpu_count=0, + gpu_model="", + longitude=0.0, + latitude=0.0, + ram_total_size=16.0, + tracking_mode="process", + ) + + +def test_local_api_is_up(): + response = requests.get(URL.rstrip("/") + "/", timeout=5) + assert response.status_code == 200 + assert response.json()["status"] == "OK" + + +def test_local_telemetry_post_accepts_sdk_payload(): + settings = _local_settings() + tracker = SimpleNamespace( + _conf={ + "os": "Linux-5.10.0-x86_64", + "codecarbon_version": "3.2.8", + "cpu_count": 4, + "python_version": "3.12", + "tracking_mode": "process", + }, + ) + payload = build_payload(tracker, _sample_emissions(), level=TelemetryLevel.minimal) + assert post_private(settings, payload) is True diff --git a/carbonserver/tests/api/routers/test_telemetry.py b/carbonserver/tests/api/routers/test_telemetry.py index a405b59e8..e93e4e9b9 100644 --- a/carbonserver/tests/api/routers/test_telemetry.py +++ b/carbonserver/tests/api/routers/test_telemetry.py @@ -41,40 +41,126 @@ def client(custom_test_server): yield TestClient(custom_test_server) -def test_add_telemetry(client, custom_test_server): +@pytest.fixture(autouse=True) +def _reset_rate_limit(): + telemetry._recent_requests.clear() + yield + telemetry._recent_requests.clear() + + +@pytest.fixture +def repository_mock(custom_test_server): repository_mock = mock.Mock(spec=TelemetryRepository) repository_mock.add_telemetry.return_value = UUID(TELEMETRY_ID) - with custom_test_server.container.telemetry_repository.override(repository_mock): - response = client.post("/telemetry", json=MINIMAL_TELEMETRY_TO_CREATE) + yield repository_mock + + +def test_add_telemetry_needs_no_token(client, repository_mock): + response = client.post("/telemetry", json=MINIMAL_TELEMETRY_TO_CREATE) assert response.status_code == status.HTTP_201_CREATED assert response.json() == TELEMETRY_ID + repository_mock.add_telemetry.assert_called_once() -def test_minimal_telemetry_rejects_extensive_fields(client, custom_test_server): - repository_mock = mock.Mock(spec=TelemetryRepository) - telemetry_with_extensive_field = { - **MINIMAL_TELEMETRY_TO_CREATE, - "total_emissions_kg": 0.42, - } +def test_unknown_fields_are_rejected(client, repository_mock): + response = client.post( + "/telemetry", + json={**MINIMAL_TELEMETRY_TO_CREATE, "total_emissions_kg": 0.42}, + ) - with custom_test_server.container.telemetry_repository.override(repository_mock): - response = client.post("/telemetry", json=telemetry_with_extensive_field) + assert response.status_code == 422 + repository_mock.add_telemetry.assert_not_called() + + +def test_disabled_telemetry_is_rejected(client, repository_mock): + response = client.post( + "/telemetry", + json={**MINIMAL_TELEMETRY_TO_CREATE, "telemetry_level": "disabled"}, + ) assert response.status_code == 422 repository_mock.add_telemetry.assert_not_called() -def test_disabled_telemetry_is_rejected(client, custom_test_server): - repository_mock = mock.Mock(spec=TelemetryRepository) - disabled_telemetry = { - **MINIMAL_TELEMETRY_TO_CREATE, - "telemetry_level": "disabled", - } +def test_coordinates_are_rejected(client, repository_mock): + response = client.post( + "/telemetry", + json={**MINIMAL_TELEMETRY_TO_CREATE, "latitude": 48.8, "longitude": 2.3}, + ) - with custom_test_server.container.telemetry_repository.override(repository_mock): - response = client.post("/telemetry", json=disabled_telemetry) + assert response.status_code == 422 + repository_mock.add_telemetry.assert_not_called() + + +def test_oversized_string_is_rejected(client, repository_mock): + response = client.post( + "/telemetry", + json={**MINIMAL_TELEMETRY_TO_CREATE, "cpu_model": "x" * 257}, + ) assert response.status_code == 422 repository_mock.add_telemetry.assert_not_called() + + +def test_rate_limit_returns_429_per_ip(client, repository_mock, monkeypatch): + monkeypatch.setattr(telemetry, "RATE_LIMIT", 2) + + codes = [ + client.post("/telemetry", json=MINIMAL_TELEMETRY_TO_CREATE).status_code + for _ in range(3) + ] + + assert codes == [201, 201, 429] + assert repository_mock.add_telemetry.call_count == 2 + # Another IP has its own budget. + other = TestClient(client.app, client=("10.0.0.2", 50000)) + assert other.post("/telemetry", json=MINIMAL_TELEMETRY_TO_CREATE).status_code == 201 + + +def test_rate_limit_window_expires(monkeypatch): + monkeypatch.setattr(telemetry, "RATE_LIMIT", 1) + now = [1000.0] + monkeypatch.setattr(telemetry.time, "monotonic", lambda: now[0]) + + assert telemetry._rate_limited("1.2.3.4") is False + assert telemetry._rate_limited("1.2.3.4") is True + now[0] += telemetry.RATE_WINDOW_SECONDS + 1 + assert telemetry._rate_limited("1.2.3.4") is False + + +def test_rate_limited_is_atomic_under_concurrency(monkeypatch): + """FastAPI runs sync endpoints in a threadpool; _rate_limited must not + let concurrent callers both slip past the limit check (see the lock in + carbonserver/carbonserver/api/routers/telemetry.py).""" + import threading + + monkeypatch.setattr(telemetry, "RATE_LIMIT", 5) + admitted = [] + admitted_lock = threading.Lock() + barrier = threading.Barrier(20) + + def call(): + barrier.wait() + limited = telemetry._rate_limited("9.9.9.9") + if not limited: + with admitted_lock: + admitted.append(1) + + threads = [threading.Thread(target=call) for _ in range(20)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert len(admitted) == telemetry.RATE_LIMIT + + +def test_rate_limit_memory_is_bounded(monkeypatch): + monkeypatch.setattr(telemetry, "MAX_TRACKED_IPS", 2) + + for host in ("1.1.1.1", "2.2.2.2", "3.3.3.3"): + telemetry._rate_limited(host) + + assert list(telemetry._recent_requests) == ["3.3.3.3"] diff --git a/carbonserver/tests/api/test_telemetry_schema_drift.py b/carbonserver/tests/api/test_telemetry_schema_drift.py index abc20d764..388e39963 100644 --- a/carbonserver/tests/api/test_telemetry_schema_drift.py +++ b/carbonserver/tests/api/test_telemetry_schema_drift.py @@ -7,7 +7,9 @@ from carbonserver.api.schemas_telemetry import TelemetryCreate as ServerTelemetryCreate REPO_ROOT = Path(__file__).resolve().parents[3] -CORE_TELEMETRY_SCHEMA_PATH = REPO_ROOT / "codecarbon" / "core" / "telemetry_schemas.py" +CORE_TELEMETRY_SCHEMA_PATH = ( + REPO_ROOT / "codecarbon" / "core" / "telemetry" / "schemas.py" +) def _load_core_telemetry_create(): diff --git a/carbonserver/tests/database/test_telemetry_migration.py b/carbonserver/tests/database/test_telemetry_migration.py new file mode 100644 index 000000000..f0b75d9b6 --- /dev/null +++ b/carbonserver/tests/database/test_telemetry_migration.py @@ -0,0 +1,262 @@ +"""Cover the ``add telemetry table`` migration against both a fresh database +and one carrying the pre-existing, ``create_all``-generated ``telemetry`` +table (old, wide schema, see #1171).""" + +import importlib.util +import sys +from pathlib import Path +from unittest.mock import patch + +import sqlalchemy as sa +from alembic.operations import Operations +from alembic.runtime.migration import MigrationContext +from sqlalchemy.dialects.postgresql import UUID as PGUUID +from sqlalchemy.ext.compiler import compiles + + +# Production is Postgres-only; sqlite has no native UUID type, so teach it to +# render one as CHAR(32) purely so the "fresh table" path is exercisable here. +@compiles(PGUUID, "sqlite") +def _compile_uuid_sqlite(type_, compiler, **kw): + return "CHAR(32)" + + +MIGRATION_PATH = ( + Path(__file__).parents[2] + / "carbonserver" + / "database" + / "alembic" + / "versions" + / "20260927_add_telemetry_table.py" +) + + +def _load_migration(): + spec = importlib.util.spec_from_file_location( + "telemetry_migration_20260927", MIGRATION_PATH + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +OLD_TABLE_EXTRA_COLUMNS = { + "longitude", + "latitude", + "python_executable_hash", + "total_emissions_kg", + "host_machine_hash", +} +NEW_ONLY_COLUMNS = { + "cpu_architecture", + "python_env_type", + "codecarbon_install_method", +} + + +def _create_old_table(conn, module, skip_columns=()): + """Build the full pre-existing, ``create_all``-generated old schema.""" + all_columns = ["id VARCHAR PRIMARY KEY"] + for column in [*module.NEW_COLUMNS, *module.OLD_ONLY_COLUMNS]: + if column.name in skip_columns: + continue + sql_type = "VARCHAR" + if isinstance(column.type, (sa.Integer, sa.Float)): + sql_type = "FLOAT" if isinstance(column.type, sa.Float) else "INTEGER" + elif isinstance(column.type, sa.DateTime): + sql_type = "DATETIME" + elif isinstance(column.type, sa.Boolean): + sql_type = "BOOLEAN" + all_columns.append(f"{column.name} {sql_type}") + conn.execute(sa.text(f"CREATE TABLE telemetry ({', '.join(all_columns)})")) + conn.execute( + sa.text( + "INSERT INTO telemetry (id, timestamp, telemetry_level, os, " + "longitude, latitude) VALUES " + "('row-1', '2026-01-01 00:00:00', 'extensive', 'Linux', 1.0, 2.0)" + ) + ) + + +def test_timestamp_column_is_timezone_aware_on_create(): + """Naive timestamps are ambiguous across servers in different zones; + the column must be created as ``timestamptz``.""" + module = _load_migration() + (timestamp_column,) = [c for c in module.NEW_COLUMNS if c.name == "timestamp"] + assert isinstance(timestamp_column.type, sa.DateTime) + assert timestamp_column.type.timezone is True + + +def test_bind_is_postgres_helper(): + module = _load_migration() + + class _FakeDialect: + def __init__(self, name): + self.name = name + + class _FakeBind: + def __init__(self, name): + self.dialect = _FakeDialect(name) + + assert module._bind_is_postgres(_FakeBind("postgresql")) is True + assert module._bind_is_postgres(_FakeBind("sqlite")) is False + + +def test_upgrade_alters_timestamp_to_timezone_aware_on_existing_postgres_table(): + """When the table pre-exists on Postgres, the migration must alter the + ``timestamp`` column type to timestamptz rather than leaving it naive.""" + module = _load_migration() + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + _create_old_table(conn, module) + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + with ( + patch.object(module, "_bind_is_postgres", return_value=True), + patch("alembic.op.alter_column") as mock_alter, + patch("alembic.op.create_index"), + patch("alembic.op.execute"), + ): + module.upgrade() + (call,) = mock_alter.call_args_list + args, kwargs = call + assert args == ("telemetry", "timestamp") + assert kwargs["type_"].timezone is True + assert kwargs["postgresql_using"] == "timestamp AT TIME ZONE 'UTC'" + + +def test_upgrade_creates_table_when_absent(): + module = _load_migration() + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + module.upgrade() + columns = {c["name"] for c in sa.inspect(conn).get_columns("telemetry")} + assert "cpu_architecture" in columns + assert not (columns & OLD_TABLE_EXTRA_COLUMNS) + + +def test_upgrade_migrates_pre_existing_old_schema_table_preserving_rows(): + module = _load_migration() + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + _create_old_table(conn, module) + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + module.upgrade() + columns = {c["name"] for c in sa.inspect(conn).get_columns("telemetry")} + rows = conn.execute(sa.text("SELECT id, telemetry_level FROM telemetry")).all() + assert NEW_ONLY_COLUMNS <= columns + assert not (columns & OLD_TABLE_EXTRA_COLUMNS) + # The pre-existing row survives the schema change untouched. + assert rows == [("row-1", "extensive")] + + +def test_upgrade_adds_a_genuinely_missing_new_column(): + """Covers the "add any missing column" branch: an old table that lacks + one of the new model's columns gets it added (nullable).""" + module = _load_migration() + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + _create_old_table(conn, module, skip_columns={"cpu_architecture"}) + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + module.upgrade() + columns = {c["name"] for c in sa.inspect(conn).get_columns("telemetry")} + assert "cpu_architecture" in columns + + +def test_downgrade_alters_timestamp_back_to_naive_on_postgres(): + module = _load_migration() + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + _create_old_table(conn, module) + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + module.upgrade() + with ( + patch.object(module, "_bind_is_postgres", return_value=True), + patch("alembic.op.alter_column") as mock_alter, + patch("alembic.op.execute"), + ): + module.downgrade() + (call,) = mock_alter.call_args_list + args, kwargs = call + assert args == ("telemetry", "timestamp") + assert kwargs["type_"].timezone is False + assert kwargs["postgresql_using"] == "timestamp AT TIME ZONE 'UTC'" + + +def test_ensure_retention_is_a_noop_on_sqlite(): + """SQLite has no plpgsql triggers; ``_ensure_retention`` must not touch + the connection at all, and the migration must still succeed.""" + module = _load_migration() + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + module.upgrade() # must not raise + indexes = {ix["name"] for ix in sa.inspect(conn).get_indexes("telemetry")} + assert "ix_telemetry_timestamp" not in indexes + + +def test_upgrade_creates_retention_index_trigger_and_function_on_postgres(): + """On Postgres, upgrade must create the timestamp index, the purge + function and the statement-level AFTER INSERT trigger.""" + module = _load_migration() + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + with ( + patch.object(module, "_bind_is_postgres", return_value=True), + patch("alembic.op.create_index") as mock_create_index, + patch("alembic.op.execute") as mock_execute, + ): + module.upgrade() + index_calls = [c.args for c in mock_create_index.call_args_list] + assert ("ix_telemetry_timestamp", "telemetry", ["timestamp"]) in index_calls + executed_sql = "\n".join(c.args[0] for c in mock_execute.call_args_list) + assert "CREATE OR REPLACE FUNCTION telemetry_purge()" in executed_sql + assert "DELETE FROM telemetry WHERE timestamp < now() - interval '3 years'" in ( + executed_sql + ) + assert "CREATE TRIGGER telemetry_retention" in executed_sql + assert "FOR EACH STATEMENT" in executed_sql + + +def test_downgrade_drops_retention_trigger_function_and_index_on_postgres(): + module = _load_migration() + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + _create_old_table(conn, module) + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + module.upgrade() + with ( + patch.object(module, "_bind_is_postgres", return_value=True), + patch("alembic.op.alter_column"), + patch("alembic.op.execute") as mock_execute, + ): + module.downgrade() + executed_sql = "\n".join(c.args[0] for c in mock_execute.call_args_list) + assert "DROP TRIGGER IF EXISTS telemetry_retention" in executed_sql + assert "DROP FUNCTION IF EXISTS telemetry_purge()" in executed_sql + assert "DROP INDEX IF EXISTS ix_telemetry_timestamp" in executed_sql + + +def test_downgrade_restores_old_columns_as_nullable_preserving_rows(): + module = _load_migration() + engine = sa.create_engine("sqlite:///:memory:") + with engine.connect() as conn: + _create_old_table(conn, module) + ctx = MigrationContext.configure(conn) + with Operations.context(ctx): + module.upgrade() + module.downgrade() + columns = {c["name"] for c in sa.inspect(conn).get_columns("telemetry")} + rows = conn.execute(sa.text("SELECT id FROM telemetry")).all() + assert OLD_TABLE_EXTRA_COLUMNS <= columns + assert rows == [("row-1",)] diff --git a/carbonserver/tests/database/test_telemetry_retention_db.py b/carbonserver/tests/database/test_telemetry_retention_db.py new file mode 100644 index 000000000..ae715cecd --- /dev/null +++ b/carbonserver/tests/database/test_telemetry_retention_db.py @@ -0,0 +1,56 @@ +"""Real-Postgres check for the telemetry retention trigger (see +``20260927_add_telemetry_table.py``): rows older than 3 years must be purged +on the next insert, and recent rows must survive. + +Requires a live Postgres reachable via ``DATABASE_URL`` with the telemetry +migration already applied. Skipped unless ``DATABASE_URL`` is set. +""" + +import os +import uuid +from datetime import datetime, timedelta, timezone + +import pytest +from sqlalchemy import create_engine, text + +pytestmark = [ + pytest.mark.integ_test, + pytest.mark.skipif( + not os.getenv("DATABASE_URL"), reason="requires a live Postgres DATABASE_URL" + ), +] + + +@pytest.fixture(scope="module") +def engine(): + return create_engine(os.environ["DATABASE_URL"]) + + +def _insert(conn, timestamp): + row_id = uuid.uuid4() + conn.execute( + text( + "INSERT INTO telemetry (id, timestamp, telemetry_level) " + "VALUES (:id, :ts, 'minimal')" + ), + {"id": row_id, "ts": timestamp}, + ) + conn.commit() + return row_id + + +def test_old_rows_are_purged_on_next_insert_and_recent_rows_survive(engine): + now = datetime.now(timezone.utc) + with engine.connect() as conn: + old_row_id = _insert(conn, now - timedelta(days=4 * 365)) + recent_row_id = _insert(conn, now) + # A further insert fires the AFTER INSERT statement trigger that + # purges anything older than 3 years. + trigger_row_id = _insert(conn, now) + + ids = { + row[0] for row in conn.execute(text("SELECT id FROM telemetry")).fetchall() + } + assert old_row_id not in ids + assert recent_row_id in ids + assert trigger_row_id in ids diff --git a/codecarbon/cli/main.py b/codecarbon/cli/main.py index f81fb302a..8595f0520 100644 --- a/codecarbon/cli/main.py +++ b/codecarbon/cli/main.py @@ -18,6 +18,7 @@ get_existing_exp_id, overwrite_local_config, ) +from codecarbon.cli.telemetry_cli import ask_telemetry_level_once, telemetry_app API_URL = os.environ.get("API_URL", "https://dashboard.codecarbon.io/api") @@ -25,6 +26,7 @@ DEFAULT_ORGANIzATION_ID = "e60afa92-17b7-4720-91a0-1ae91e409ba1" codecarbon = typer.Typer(no_args_is_help=True) +codecarbon.add_typer(telemetry_app, name="telemetry") def main(): @@ -191,6 +193,7 @@ def config(): ) print("Welcome to CodeCarbon configuration wizard") + ask_telemetry_level_once() home = Path.home() global_path = (home / ".codecarbon.config").expanduser().resolve() @@ -418,8 +421,16 @@ def monitor( str, typer.Option(help="Log level (critical, error, warning, info, debug)"), ] = "error", + telemetry_level: Annotated[ + Optional[str], + typer.Option( + help="Override telemetry level for this run only (disabled or minimal).", + ), + ] = None, ): """Monitor your machine's carbon emissions.""" + if telemetry_level is None and not offline: + ask_telemetry_level_once() external_conf = _external_config() @@ -439,6 +450,10 @@ def monitor( # Nothing configures it: keep the defaults advertised by `--help` # (and an unattended monitor quiet) instead of the tracker's own. tracker_args[name] = value + if telemetry_level is not None: + from codecarbon.cli.telemetry_cli import normalize_telemetry_level + + tracker_args["telemetry_level"] = normalize_telemetry_level(telemetry_level) # Set up the tracker arguments based on mode (offline vs online) and validate required args for each mode if offline: diff --git a/codecarbon/cli/telemetry_cli.py b/codecarbon/cli/telemetry_cli.py new file mode 100644 index 000000000..3dcb22d7d --- /dev/null +++ b/codecarbon/cli/telemetry_cli.py @@ -0,0 +1,284 @@ +"""CLI commands to configure CodeCarbon product telemetry tiers.""" + +import sys +from pathlib import Path +from typing import Optional + +import questionary +import typer +from rich import print +from typing_extensions import Annotated + +from codecarbon.cli.cli_utils import ( + create_new_config_file, + get_config, + overwrite_local_config, +) +from codecarbon.core.config import ( + _config_file_paths, + get_config_file_settings, + get_hierarchical_config, +) +from codecarbon.core.telemetry.schemas import TelemetryLevel +from codecarbon.core.telemetry.settings import ( + DEFAULT_TELEMETRY_LEVEL, + TelemetrySettings, + parse_telemetry_level, +) + +telemetry_app = typer.Typer( + help="Configure product telemetry (disabled or minimal).", + no_args_is_help=False, +) + +TIER_DESCRIPTIONS = { + "disabled": "Nothing is sent.", + "minimal": ( + "OS, Python/CodeCarbon versions, CPU/GPU model and count, RAM size, " + "country/region, cloud provider. Sent once per process." + ), +} + + +def normalize_telemetry_level(level: str) -> str: + """Validate and normalize a telemetry tier string for CLI use. + + Args: + level: User-provided tier name. + + Returns: + Canonical tier value. + + Raises: + typer.BadParameter: If the level is not a valid ``TelemetryLevel``. + """ + try: + return parse_telemetry_level(level).value + except ValueError as error: + raise typer.BadParameter(str(error)) from error + + +def resolve_config_path(config: Optional[Path], *, create: bool = False) -> Path: + """Resolve which config file to read or write. + + Args: + config: Explicit path from ``--config``, if any. + create: When True and no file exists, create ``./.codecarbon.config``. + + Returns: + Resolved config file path. + """ + if config is not None: + path = config.expanduser().resolve() + if create and not path.exists(): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("[codecarbon]\n", encoding="utf-8") + return path + global_path, local_path = (Path(p) for p in _config_file_paths()) + if local_path.exists(): + return local_path + if global_path.exists(): + return global_path + if create: + local_path.write_text("[codecarbon]\n", encoding="utf-8") + return local_path + return local_path + + +def pick_config_path_interactive() -> Path: + """Prompt for which config file to update. + + Returns: + Path chosen by the user. + """ + global_path, local_path = (Path(p) for p in _config_file_paths()) + options = [] + if global_path.exists(): + options.append(str(global_path)) + if local_path.exists() and str(local_path) not in options: + options.append(str(local_path)) + options.append("Create new config file") + choice = questionary.select( + "Which configuration file should store telemetry_level?", + choices=options, + ).ask() + if choice is None: + raise typer.Exit(0) + if choice == "Create new config file": + return create_new_config_file() + return Path(choice).expanduser().resolve() + + +def write_telemetry_level(path: Path, level: str) -> None: + """Persist ``telemetry_level`` to a config file. + + Args: + path: Target ``.codecarbon.config`` path. + level: Validated tier value. + """ + if not path.exists(): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("[codecarbon]\n", encoding="utf-8") + overwrite_local_config("telemetry_level", level, path=path) + + +def print_telemetry_status(config_path: Optional[Path] = None) -> None: + """Print resolved telemetry settings. + + Without ``config_path``, uses the same merged file settings and env overlay + as ``EmissionsTracker``. With ``config_path``, inspects that file only. + + Args: + config_path: Optional single config file to inspect. + """ + if config_path is not None: + path = config_path.expanduser().resolve() + if not path.exists(): + print(f"[yellow]Config file not found:[/yellow] {path}") + print(f"Default tier: {DEFAULT_TELEMETRY_LEVEL.value} (not explicit)") + return + file_settings = get_config(path) + external_conf = dict(file_settings) + source_label = str(path) + else: + file_settings = get_config_file_settings() + external_conf = get_hierarchical_config() + source_label = "merged ~/.codecarbon.config + ./.codecarbon.config" + + settings = TelemetrySettings.resolve(external_conf=external_conf) + level = settings.level + explicit = settings.is_explicit + stored = file_settings.get("telemetry_level") + print(f"Config source: {source_label}") + print(f"telemetry_level in file(s): {stored!r}") + print(f"Resolved tier: {level.value}") + print(f"Explicitly configured: {explicit}") + if level != TelemetryLevel.disabled: + print(f"Sent once per process to {settings.api_url}/telemetry.") + if not explicit: + print("Opt out with: codecarbon telemetry set disabled") + + +def ask_telemetry_level_once() -> None: + """Ask for a telemetry tier on first interactive CLI use, and remember it. + + Does nothing without a TTY (CI, SLURM, pipes) or once a tier is configured, + so it never blocks an unattended run: the tracker's one-time notice covers + those. + """ + if not (sys.stdin.isatty() and sys.stdout.isatty()): + return + if TelemetrySettings.resolve(external_conf=get_hierarchical_config()).is_explicit: + return + print("CodeCarbon product telemetry: which level do you want?") + level = questionary.select( + "telemetry_level:", + choices=[ + questionary.Choice(f"{name} — {text}", name) + for name, text in TIER_DESCRIPTIONS.items() + ], + default=DEFAULT_TELEMETRY_LEVEL.value, + ).ask() + if level is None: # Ctrl-C: keep the default, ask again next time + return + path = Path(_config_file_paths()[0]) + write_telemetry_level(path, level) + print(f"Saved telemetry_level = {level} in {path}") + + +def run_telemetry_interactive(config: Optional[Path] = None) -> None: + """Run the interactive telemetry configuration wizard. + + Args: + config: Optional fixed config path; otherwise prompt for file choice. + """ + print("CodeCarbon product telemetry") + print( + "Separate from your dashboard experiment (codecarbon config). " + "Controls optional, anonymous usage analytics.\n" + ) + path = resolve_config_path(config) if config else None + if path is None or (config is None and not path.exists()): + path = pick_config_path_interactive() + else: + path = resolve_config_path(config, create=True) + + choices = [ + questionary.Choice("disabled — " + TIER_DESCRIPTIONS["disabled"], "disabled"), + questionary.Choice("minimal — " + TIER_DESCRIPTIONS["minimal"], "minimal"), + ] + try: + current = get_config(path).get("telemetry_level") + except FileNotFoundError: + current = None + valid_levels = {member.value for member in TelemetryLevel} + default = current if current in valid_levels else "minimal" + level = questionary.select( + "Select telemetry_level:", + choices=choices, + default=default, + ).ask() + if level is None: + raise typer.Exit(0) + level = normalize_telemetry_level(level) + write_telemetry_level(path, level) + print(f"[green]Saved[/green] telemetry_level = {level} in {path}") + + +@telemetry_app.callback(invoke_without_command=True) +def telemetry_entry( + ctx: typer.Context, + config: Annotated[ + Optional[Path], + typer.Option( + "--config", + help="Path to .codecarbon.config (default: local then global).", + ), + ] = None, +) -> None: + """Configure telemetry interactively when no subcommand is given.""" + if ctx.invoked_subcommand is None: + run_telemetry_interactive(config=config) + + +@telemetry_app.command("status") +def status( + config: Annotated[ + Optional[Path], + typer.Option( + "--config", + help="Inspect one file only; default matches EmissionsTracker merge.", + ), + ] = None, +) -> None: + """Print resolved telemetry tier and configuration source.""" + print_telemetry_status(config_path=config) + + +@telemetry_app.command("set") +def set_level( + level: Annotated[ + str, + typer.Argument(help="Telemetry tier: disabled or minimal."), + ], + config: Annotated[ + Optional[Path], + typer.Option( + "--config", + help="Path to .codecarbon.config (creates ./.codecarbon.config if missing).", + ), + ] = None, +) -> None: + """Write telemetry_level to a config file. + + Without ``--config``, this always writes the GLOBAL config file (creating + it if missing), the same file ``ask_telemetry_level_once`` writes to: an + opt-out must not depend on which local file happens to exist yet. + """ + if config is not None: + path = resolve_config_path(config, create=True) + else: + path = Path(_config_file_paths()[0]) + normalized = normalize_telemetry_level(level) + write_telemetry_level(path, normalized) + print(f"[green]Saved[/green] telemetry_level = {normalized} in {path}") diff --git a/codecarbon/core/config.py b/codecarbon/core/config.py index c94835ad8..434cd1337 100644 --- a/codecarbon/core/config.py +++ b/codecarbon/core/config.py @@ -113,6 +113,27 @@ def normalize_gpu_ids( return None +def _config_file_paths() -> tuple[str, str]: + """Return resolved paths for global and local CodeCarbon config files.""" + global_path = str((Path.home() / CONFIG_FILE_NAME).expanduser().resolve()) + local_path = str((Path.cwd() / CONFIG_FILE_NAME).expanduser().resolve()) + return global_path, local_path + + +def get_config_file_settings() -> dict[str, str]: + """Return the ``[codecarbon]`` section from config files without environment overlay. + + Reads ``~/.codecarbon.config`` then ``./.codecarbon.config`` (local overrides global). + + Returns: + Configuration dict from files only. Empty when no file or section exists. + """ + conf: dict[str, str] = {} + for path in _config_file_paths(): + conf.update(read_config_file(path)) + return conf + + def read_config_file(path: str) -> dict: """ Read the `[codecarbon]` section of a single configuration file. @@ -167,13 +188,10 @@ def get_hierarchical_config_with_sources() -> Tuple[dict, Dict[str, str]]: maps a key to its value (**all values are strings**) and `sources` maps a key to the layer it was read from. """ - global_path = str((Path.home() / CONFIG_FILE_NAME).expanduser().resolve()) - local_path = str((Path.cwd() / CONFIG_FILE_NAME).expanduser().resolve()) - config = {} sources = {} - for path in (global_path, local_path): + for path in _config_file_paths(): for key, value in read_config_file(path).items(): config[key] = value sources[key] = path diff --git a/codecarbon/core/telemetry/__init__.py b/codecarbon/core/telemetry/__init__.py new file mode 100644 index 000000000..9c738afa6 --- /dev/null +++ b/codecarbon/core/telemetry/__init__.py @@ -0,0 +1,5 @@ +"""Product telemetry sent at tracker stop (minimal tier).""" + +from codecarbon.core.telemetry.dispatcher import Telemetry + +__all__ = ["Telemetry"] diff --git a/codecarbon/core/telemetry/client.py b/codecarbon/core/telemetry/client.py new file mode 100644 index 000000000..17ad02aef --- /dev/null +++ b/codecarbon/core/telemetry/client.py @@ -0,0 +1,47 @@ +"""HTTP client for product telemetry.""" + +from __future__ import annotations + +import time +from typing import Optional + +import requests + +from codecarbon.core.telemetry.schemas import TelemetryCreate +from codecarbon.core.telemetry.settings import TelemetrySettings +from codecarbon.external.logger import logger + +DEFAULT_TIMEOUT = 2.0 + + +def remaining_time(deadline: Optional[float]) -> float: + """Seconds left before ``deadline`` (a ``time.monotonic()`` value).""" + if deadline is None: + return DEFAULT_TIMEOUT + return max(0.0, deadline - time.monotonic()) + + +def post_private( + settings: TelemetrySettings, payload: dict, deadline: Optional[float] = None +) -> bool: + timeout = remaining_time(deadline) + if timeout <= 0: + logger.debug("Telemetry not sent: time budget exhausted.") + return False + body = TelemetryCreate(**payload).model_dump(mode="json", exclude_none=True) + telemetry_url = f"{settings.api_url.rstrip('/')}/telemetry" + try: + response = requests.post(url=telemetry_url, json=body, timeout=timeout) + except Exception: + logger.debug("Telemetry request failed.", exc_info=True) + return False + if response.status_code == 201: + return True + # Best effort: never bother the user about a telemetry endpoint. + logger.debug( + "Telemetry API %s returned %s: %s", + telemetry_url, + response.status_code, + response.text, + ) + return False diff --git a/codecarbon/core/telemetry/collect.py b/codecarbon/core/telemetry/collect.py new file mode 100644 index 000000000..d11bc7f56 --- /dev/null +++ b/codecarbon/core/telemetry/collect.py @@ -0,0 +1,151 @@ +"""Collect private product telemetry from tracker state.""" + +from __future__ import annotations + +import json +import os +import platform +import sys +from datetime import datetime, timezone +from typing import Any + +from codecarbon.core.gpu import is_nvidia_system +from codecarbon.core.telemetry.schemas import TelemetryLevel +from codecarbon.output_methods.emissions_data import EmissionsData + + +def _strip_empty(data: dict[str, Any]) -> dict[str, Any]: + return { + key: value for key, value in data.items() if value not in (None, "", [], {}) + } + + +def _cloud_region( + emissions: EmissionsData, +) -> tuple[str | None, str | None, str | None]: + # Reuse what the tracker already detected: probing the cloud metadata + # endpoints again costs up to a few seconds off-cloud. + cloud_provider = emissions.cloud_provider or None + cloud_region = emissions.cloud_region or None + region = emissions.region + if emissions.on_cloud == "Y" and cloud_region: + region = region or cloud_region + return cloud_provider, cloud_region, region + + +def _detect_python_env_type() -> str | None: + if os.environ.get("CONDA_DEFAULT_ENV"): + return "conda" + if os.environ.get("VIRTUAL_ENV"): + return "venv" + if sys.prefix != getattr(sys, "base_prefix", sys.prefix): + return "venv" + return "system" + + +def _detect_codecarbon_install_method() -> str | None: + try: + from importlib.metadata import distribution + + dist = distribution("codecarbon") + # PEP 610: editable installs record ``dir_info.editable`` in direct_url.json. + direct_url = dist.read_text("direct_url.json") + if direct_url and json.loads(direct_url).get("dir_info", {}).get("editable"): + return "editable" + installer = (dist.metadata.get("Installer") or "").lower() + if "uv" in installer: + return "uv" + if "pip" in installer: + return "pip" + except Exception: + pass + return None + + +def _cudnn_version() -> str | None: + # Only report torch if the user's process already imported it: importing it + # here can take seconds and load CUDA libraries, well past the time budget. + torch = sys.modules.get("torch") + if torch is None: + return None + try: + version = torch.backends.cudnn.version() + return str(version) if version is not None else None + except Exception: + return None + + +def _gpu_static_fields() -> dict[str, Any]: + if not is_nvidia_system(): + return {} + try: + import pynvml + + pynvml.nvmlInit() + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + mem = pynvml.nvmlDeviceGetMemoryInfo(handle) + cuda_version = pynvml.nvmlSystemGetCudaDriverVersion_v2() + if isinstance(cuda_version, int): + cuda_version = f"{cuda_version // 1000}.{(cuda_version % 1000) // 10}" + return { + "gpu_memory_total_gb": mem.total / (1024**3), + "gpu_driver_version": pynvml.nvmlSystemGetDriverVersion(), + "cuda_version": cuda_version, + } + except Exception: + return {} + + +def _minimal_payload( + tracker: Any, emissions: EmissionsData, level: TelemetryLevel +) -> dict[str, Any]: + """Every field of ``TelemetryCreate``; empty values are dropped.""" + conf = getattr(tracker, "_conf", {}) + cloud_provider, cloud_region, region = _cloud_region(emissions) + region = region or conf.get("region") + country_name, country_iso_code = emissions.country_name, emissions.country_iso_code + # A failed geolocation falls back to Canada; that is a guess, not a location. + # On cloud the location comes from the cloud region, so it is still valid. + if emissions.on_cloud != "Y" and getattr( + getattr(tracker, "_geo", None), "is_default", False + ): + country_name = country_iso_code = region = None + + payload = { + # Truncated to the hour: enough to see usage trends over time without + # pinning a row to the exact second a process ran. + "timestamp": datetime.now(timezone.utc).replace( + minute=0, second=0, microsecond=0 + ), + "telemetry_level": level.value, + "os": conf.get("os") or platform.platform(), + "country_name": country_name, + "country_iso_code": country_iso_code, + "region": region, + "cloud_provider": cloud_provider, + "cloud_region": cloud_region, + "cpu_count": conf.get("cpu_count"), + "cpu_physical_count": conf.get("cpu_physical_count"), + "cpu_model": conf.get("cpu_model"), + "cpu_architecture": platform.machine(), + "gpu_count": conf.get("gpu_count"), + "gpu_model": conf.get("gpu_model"), + "ram_total_size_gb": conf.get("ram_total_size"), + "python_version": conf.get("python_version") or platform.python_version(), + "python_implementation": platform.python_implementation(), + "python_env_type": _detect_python_env_type(), + "codecarbon_version": conf.get("codecarbon_version"), + "codecarbon_install_method": _detect_codecarbon_install_method(), + "cudnn_version": _cudnn_version(), + **_gpu_static_fields(), + } + return _strip_empty(payload) + + +def build_payload( + tracker: Any, + emissions: EmissionsData, + level: TelemetryLevel = TelemetryLevel.minimal, +) -> dict[str, Any]: + """Build a validated telemetry payload dict for ``POST /telemetry``.""" + return _minimal_payload(tracker, emissions, level) diff --git a/codecarbon/core/telemetry/dispatcher.py b/codecarbon/core/telemetry/dispatcher.py new file mode 100644 index 000000000..80b2773af --- /dev/null +++ b/codecarbon/core/telemetry/dispatcher.py @@ -0,0 +1,140 @@ +"""Per-tracker telemetry dispatcher.""" + +from __future__ import annotations + +import atexit +import os +import sys +import threading +import time +from typing import Any + +from codecarbon.core.telemetry.client import post_private +from codecarbon.core.telemetry.collect import build_payload +from codecarbon.core.telemetry.schemas import TelemetryLevel +from codecarbon.core.telemetry.settings import TelemetrySettings +from codecarbon.external.logger import logger +from codecarbon.output_methods.emissions_data import EmissionsData + +#: Total wall-clock budget for one send. +TELEMETRY_TIMEOUT_SECONDS = 2.0 + +#: How long interpreter exit waits for pending sends. +EXIT_JOIN_SECONDS = 1.0 + +THREAD_NAME = "codecarbon-telemetry" + +TELEMETRY_NOTICE = ( + "CodeCarbon telemetry is on by default at level %r: once per process, " + "stop() sends environment and hardware info (OS, Python and CodeCarbon " + "versions, CPU/GPU model and count, RAM size, country/region, cloud " + "provider); no code, data, file paths or coordinates. Choose a level " + "with `codecarbon telemetry set minimal|disabled` or " + "CODECARBON_TELEMETRY_LEVEL to silence this notice. Details: " + "https://docs.codecarbon.io/latest/how-to/telemetry/" +) + +_pending: set[threading.Thread] = set() + +#: Telemetry describes the environment, which does not change within a +#: process, so only the first qualifying stop() sends it. +_sent_lock = threading.Lock() +_sent = False + +#: The implicit-level notice is shown at most once per process. +_notice_lock = threading.Lock() +_notice_shown = False + + +def _reset_after_fork() -> None: + """A forked child is its own process: let it send its own telemetry too. + + Also rebuilds the locks, since a lock inherited across ``fork()`` may be + left held if another thread owned it at fork time. + """ + global _sent, _sent_lock, _notice_shown, _notice_lock + _sent = False + _sent_lock = threading.Lock() + _notice_shown = False + _notice_lock = threading.Lock() + _pending.clear() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_after_fork) + + +@atexit.register +def _join_pending() -> None: + """Give in-flight sends a short, bounded chance to finish at exit.""" + deadline = time.monotonic() + EXIT_JOIN_SECONDS + for thread in [t for t in list(_pending) if t.is_alive()]: + thread.join(max(0.0, deadline - time.monotonic())) + + +class Telemetry: + """Per-tracker telemetry dispatcher.""" + + def __init__(self, settings: TelemetrySettings) -> None: + self.settings = settings + #: The last send thread, kept only so tests can join it. + self._thread: threading.Thread | None = None + + def notice_once_if_implicit(self) -> None: + """Explain the default tier once per process while none was chosen.""" + global _notice_shown + if self.settings.is_explicit or self.settings.level == TelemetryLevel.disabled: + return + with _notice_lock: + if _notice_shown: + return + _notice_shown = True + # `codecarbon monitor` defaults to log_level=error, which would hide a + # WARNING-only notice while still marking it as shown. Print to stderr + # so the one-time notice is seen regardless of the configured log level. + print(TELEMETRY_NOTICE % self.settings.level.value, file=sys.stderr) + logger.debug(TELEMETRY_NOTICE, self.settings.level.value) + + def send_at_stop(self, tracker: Any, emissions: EmissionsData) -> None: + """Send product telemetry on the first qualifying ``stop()`` of the process.""" + global _sent + if self.settings.level == TelemetryLevel.disabled: + return + if emissions.duration is not None and emissions.duration < 1: + logger.debug("Telemetry not sent: run shorter than 1 second.") + return + # Payload building (NVML, package lookups) and the network both happen + # on this thread: stop() never waits on either. At exit it is joined for + # at most EXIT_JOIN_SECONDS, then dropped. The whole claim-and-start + # happens under the lock so two concurrent stop()s can't both start a + # thread, and _sent is only set once a thread has actually started, so + # stop() can never crash from telemetry and a failed start doesn't + # permanently suppress every later send. + with _sent_lock: + if _sent: + return + try: + thread = threading.Thread( + target=self._send, + args=(tracker, emissions), + name=THREAD_NAME, + daemon=True, + ) + thread.start() + except Exception: + logger.debug("Telemetry thread failed to start.", exc_info=True) + return + self._thread = thread + _pending.add(thread) + _sent = True + + def _send(self, tracker: Any, emissions: EmissionsData) -> None: + """Build and post the payload under one time budget.""" + deadline = time.monotonic() + TELEMETRY_TIMEOUT_SECONDS + try: + payload = build_payload(tracker, emissions, level=self.settings.level) + post_private(self.settings, payload, deadline=deadline) + except Exception: + logger.debug("Telemetry send failed.", exc_info=True) + finally: + _pending.discard(threading.current_thread()) diff --git a/codecarbon/core/telemetry/schemas.py b/codecarbon/core/telemetry/schemas.py new file mode 100644 index 000000000..93ad6ebf8 --- /dev/null +++ b/codecarbon/core/telemetry/schemas.py @@ -0,0 +1,57 @@ +"""Telemetry payload schemas aligned with carbonserver ``telemetry_sql_models.Telemetry``.""" + +from datetime import datetime +from enum import Enum +from typing import Annotated, Optional + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +#: Every free-text field is capped so one request cannot store unbounded data. +Str = Annotated[str, Field(max_length=256)] + + +class TelemetryLevel(str, Enum): + disabled = "disabled" + minimal = "minimal" + + +class TelemetryBase(BaseModel): + model_config = ConfigDict(extra="forbid", use_enum_values=True) + + timestamp: datetime + telemetry_level: TelemetryLevel + + os: Optional[Str] = None + country_name: Optional[Str] = None + country_iso_code: Optional[str] = Field(default=None, min_length=2, max_length=3) + region: Optional[Str] = None + cloud_provider: Optional[Str] = None + cloud_region: Optional[Str] = None + + cpu_count: Optional[int] = Field(default=None, ge=0) + cpu_physical_count: Optional[int] = Field(default=None, ge=0) + cpu_model: Optional[Str] = None + cpu_architecture: Optional[Str] = None + gpu_count: Optional[int] = Field(default=None, ge=0) + gpu_model: Optional[Str] = None + gpu_driver_version: Optional[Str] = None + gpu_memory_total_gb: Optional[float] = Field(default=None, ge=0) + ram_total_size_gb: Optional[float] = Field(default=None, ge=0) + cuda_version: Optional[Str] = None + cudnn_version: Optional[Str] = None + + python_version: Optional[Str] = None + python_implementation: Optional[Str] = None + python_env_type: Optional[Str] = None + codecarbon_version: Optional[Str] = None + codecarbon_install_method: Optional[Str] = None + + @model_validator(mode="after") + def reject_disabled_level(self): + if self.telemetry_level == TelemetryLevel.disabled: + raise ValueError("Disabled telemetry must not be submitted") + return self + + +class TelemetryCreate(TelemetryBase): + pass diff --git a/codecarbon/core/telemetry/settings.py b/codecarbon/core/telemetry/settings.py new file mode 100644 index 000000000..46e827761 --- /dev/null +++ b/codecarbon/core/telemetry/settings.py @@ -0,0 +1,78 @@ +"""Resolve telemetry tier and API URL from config and environment.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any + +from codecarbon.core.telemetry.schemas import TelemetryLevel +from codecarbon.external.logger import logger + +DEFAULT_TELEMETRY_API_URL = "https://api.codecarbon.io" +DEFAULT_TELEMETRY_LEVEL = TelemetryLevel.minimal + +TELEMETRY_LEVEL_CONFIG_KEY = "telemetry_level" + +API_URL_CONFIG_KEYS = ("telemetry_api_url",) +API_URL_ENV_VAR = "CODECARBON_TELEMETRY_API_URL" + + +#: Older releases used a wider ``extensive`` tier; map it onto the tier that +#: replaced it rather than treating it as unparseable. +LEGACY_LEVEL_ALIASES = {"extensive": TelemetryLevel.minimal} + + +def parse_telemetry_level(raw: str | TelemetryLevel) -> TelemetryLevel: + """Parse a telemetry tier name or enum member.""" + if isinstance(raw, TelemetryLevel): + return raw + normalized = str(raw).lower() + if normalized in LEGACY_LEVEL_ALIASES: + return LEGACY_LEVEL_ALIASES[normalized] + try: + return TelemetryLevel(normalized) + except ValueError as error: + raise ValueError( + f"Invalid telemetry_level {raw!r}. Choose: disabled or minimal." + ) from error + + +@dataclass(frozen=True) +class TelemetrySettings: + """Resolved telemetry tier and API URL.""" + + level: TelemetryLevel + is_explicit: bool + api_url: str + + @classmethod + def resolve( + cls, + *, + external_conf: dict[str, Any] | None = None, + override: str | TelemetryLevel | None = None, + ) -> TelemetrySettings: + """Resolve tier (override > config/env > default minimal) and API URL.""" + conf = external_conf or {} + raw = override if override is not None else conf.get(TELEMETRY_LEVEL_CONFIG_KEY) + level = DEFAULT_TELEMETRY_LEVEL + if raw is not None: + try: + level = parse_telemetry_level(raw) + except ValueError: + # An unrecognized value (typo, or a privacy-intent string like + # "off"/"false"/"none"/"0") must never be treated as consent + # to send telemetry: fail closed to disabled, not minimal. + level = TelemetryLevel.disabled + logger.error( + "Invalid telemetry_level provided; falling back to %r", + TelemetryLevel.disabled.value, + ) + api_url = next( + (conf[key] for key in API_URL_CONFIG_KEYS if conf.get(key)), + os.environ.get(API_URL_ENV_VAR) or DEFAULT_TELEMETRY_API_URL, + ) + return cls( + level=level, is_explicit=raw is not None, api_url=api_url.rstrip("/") + ) diff --git a/codecarbon/core/telemetry_client.py b/codecarbon/core/telemetry_client.py deleted file mode 100644 index 8dfdad05a..000000000 --- a/codecarbon/core/telemetry_client.py +++ /dev/null @@ -1,61 +0,0 @@ -import json -from typing import Optional, Union - -import requests - -from codecarbon.core.telemetry_schemas import TelemetryCreate -from codecarbon.external.logger import logger - - -class TelemetryClient: - """ - Client dedicated to sending CodeCarbon telemetry payloads. - """ - - def __init__( - self, - endpoint_url="https://api.codecarbon.io", - telemetry: Optional[Union[TelemetryCreate, dict]] = None, - ): - self.endpoint_url = endpoint_url.rstrip("/") - self.telemetry_url = self.endpoint_url + "/telemetry" - self.headers = {"Content-Type": "application/json"} - self.telemetry = self._validate_telemetry(telemetry) if telemetry else None - - def add_telemetry(self, telemetry: Optional[Union[TelemetryCreate, dict]] = None): - telemetry_payload = ( - self._validate_telemetry(telemetry) if telemetry else self.telemetry - ) - if telemetry_payload is None: - logger.error("TelemetryClient.add_telemetry() needs a telemetry payload") - return None - payload = telemetry_payload.model_dump(mode="json", exclude_none=True) - - try: - response = requests.post( - url=self.telemetry_url, - json=payload, - timeout=2, - headers=self.headers, - ) - if response.status_code != 201: - self._log_error(payload, response) - return None - return response.json() - except Exception as e: - logger.error(e, exc_info=True) - return None - - @staticmethod - def _validate_telemetry(telemetry: Union[TelemetryCreate, dict]) -> TelemetryCreate: - if isinstance(telemetry, TelemetryCreate): - return telemetry - return TelemetryCreate(**telemetry) - - def _log_error(self, payload, response): - logger.error( - f"TelemetryClient Error when calling the API on {self.telemetry_url} with : {json.dumps(payload)}" - ) - logger.error( - f"TelemetryClient API return http code {response.status_code} and answer : {response.text}" - ) diff --git a/codecarbon/core/telemetry_schemas.py b/codecarbon/core/telemetry_schemas.py deleted file mode 100644 index ea6249b65..000000000 --- a/codecarbon/core/telemetry_schemas.py +++ /dev/null @@ -1,156 +0,0 @@ -from datetime import datetime -from enum import Enum -from typing import List, Optional - -from pydantic import BaseModel, ConfigDict, Field, model_validator - - -class TelemetryLevel(str, Enum): - disabled = "disabled" - minimal = "minimal" - extensive = "extensive" - - -class TelemetryBase(BaseModel): - model_config = ConfigDict(extra="forbid", use_enum_values=True) - - timestamp: datetime - telemetry_level: TelemetryLevel - - os: Optional[str] = None - country_name: Optional[str] = None - country_iso_code: Optional[str] = Field(default=None, min_length=2, max_length=3) - region: Optional[str] = None - cloud_provider: Optional[str] = None - cloud_region: Optional[str] = None - longitude: Optional[float] = Field(default=None, ge=-180, le=180) - latitude: Optional[float] = Field(default=None, ge=-90, le=90) - - cpu_count: Optional[int] = Field(default=None, ge=0) - cpu_physical_count: Optional[int] = Field(default=None, ge=0) - cpu_model: Optional[str] = None - cpu_architecture: Optional[str] = None - gpu_count: Optional[int] = Field(default=None, ge=0) - gpu_model: Optional[str] = None - gpu_driver_version: Optional[str] = None - gpu_memory_total_gb: Optional[float] = Field(default=None, ge=0) - ram_total_size_gb: Optional[float] = Field(default=None, ge=0) - cuda_version: Optional[str] = None - cudnn_version: Optional[str] = None - - python_version: Optional[str] = None - python_implementation: Optional[str] = None - python_executable_hash: Optional[str] = Field( - default=None, min_length=64, max_length=64 - ) - python_env_type: Optional[str] = None - codecarbon_version: Optional[str] = None - codecarbon_install_method: Optional[str] = None - - total_emissions_kg: Optional[float] = Field(default=None, ge=0) - emissions_rate_kg_per_sec: Optional[float] = Field(default=None, ge=0) - energy_consumed_kwh: Optional[float] = Field(default=None, ge=0) - cpu_energy_kwh: Optional[float] = Field(default=None, ge=0) - gpu_energy_kwh: Optional[float] = Field(default=None, ge=0) - ram_energy_kwh: Optional[float] = Field(default=None, ge=0) - duration_seconds: Optional[float] = Field(default=None, ge=0) - cpu_utilization_avg: Optional[float] = Field(default=None, ge=0, le=100) - gpu_utilization_avg: Optional[float] = Field(default=None, ge=0, le=100) - ram_utilization_avg: Optional[float] = Field(default=None, ge=0, le=100) - - tracking_mode: Optional[str] = None - api_mode: Optional[str] = None - output_methods: Optional[List[str]] = None - hardware_tracked: Optional[List[str]] = None - task_tracking_used: Optional[bool] = None - decorator_vs_context: Optional[str] = None - measure_power_interval_secs: Optional[float] = Field(default=None, ge=0) - - hardware_detection_success: Optional[bool] = None - rapl_available: Optional[bool] = None - gpu_detection_method: Optional[str] = None - first_measurement_time_ms: Optional[float] = Field(default=None, ge=0) - tracking_overhead_percent: Optional[float] = Field(default=None, ge=0) - errors_encountered: Optional[List[str]] = None - warning_count: Optional[int] = Field(default=None, ge=0) - - ide_used: Optional[str] = None - notebook_environment: Optional[str] = None - ci_environment: Optional[str] = None - python_package_manager: Optional[str] = None - framework_detected: Optional[str] = None - - has_torch: Optional[bool] = None - torch_version: Optional[str] = None - has_transformers: Optional[bool] = None - transformers_version: Optional[str] = None - has_diffusers: Optional[bool] = None - diffusers_version: Optional[str] = None - has_tensorflow: Optional[bool] = None - tensorflow_version: Optional[str] = None - has_keras: Optional[bool] = None - keras_version: Optional[str] = None - has_pytorch_lightning: Optional[bool] = None - pytorch_lightning_version: Optional[str] = None - has_fastai: Optional[bool] = None - fastai_version: Optional[str] = None - ml_framework_primary: Optional[str] = None - - container_runtime: Optional[str] = None - in_container: Optional[bool] = None - host_machine_hash: Optional[str] = None - - @model_validator(mode="after") - def validate_telemetry_level(self): - if self.telemetry_level == TelemetryLevel.disabled: - raise ValueError("Disabled telemetry must not be submitted") - - if self.telemetry_level == TelemetryLevel.minimal: - extensive_fields = set(type(self).model_fields) - MINIMAL_TELEMETRY_FIELDS - submitted_extensive_fields = [ - field - for field in extensive_fields - if getattr(self, field) not in (None, [], {}) - ] - if submitted_extensive_fields: - fields = ", ".join(sorted(submitted_extensive_fields)) - raise ValueError( - f"Minimal telemetry cannot include extensive fields: {fields}" - ) - - return self - - -MINIMAL_TELEMETRY_FIELDS = { - "timestamp", - "telemetry_level", - "os", - "country_name", - "country_iso_code", - "region", - "cloud_provider", - "cloud_region", - "longitude", - "latitude", - "cpu_count", - "cpu_physical_count", - "cpu_model", - "cpu_architecture", - "gpu_count", - "gpu_model", - "gpu_driver_version", - "gpu_memory_total_gb", - "ram_total_size_gb", - "cuda_version", - "cudnn_version", - "python_version", - "python_implementation", - "python_executable_hash", - "python_env_type", - "codecarbon_version", - "codecarbon_install_method", -} - - -class TelemetryCreate(TelemetryBase): - pass diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 4bd0ea99d..5f1fc46b3 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -21,6 +21,9 @@ from codecarbon._version import __version__ from codecarbon.core.config import get_hierarchical_config, normalize_gpu_ids +from codecarbon.core.telemetry import Telemetry +from codecarbon.core.telemetry.schemas import TelemetryLevel +from codecarbon.core.telemetry.settings import TelemetrySettings from codecarbon.core.units import Energy, Power, Time, Water from codecarbon.core.util import count_cpus, count_physical_cpus, suppress from codecarbon.external.hardware import CPU, GPU, AppleSiliconChip @@ -425,6 +428,7 @@ def __init__( allow_multiple_runs: Optional[bool] = _sentinel, rapl_include_dram: Optional[bool] = _sentinel, rapl_prefer_psys: Optional[bool] = _sentinel, + telemetry_level: Optional[str] = _sentinel, ): """ :param project_name: Project name for current experiment run, default name @@ -522,10 +526,26 @@ def __init__( (CPU + chipset + PCIe). When False, uses package domains which are more reliable. Note: psys can report higher values than CPU TDP and may be unreliable on older systems. + :param telemetry_level: Telemetry tier (``disabled`` or ``minimal``). + Overrides config file and ``CODECARBON_TELEMETRY_LEVEL`` when set. + Defaults to ``minimal``. """ - # logger.info("base tracker init") self._external_conf = get_hierarchical_config() + # Resolve the tier from the constructor kwarg first so that + # ``EmissionsTracker(telemetry_level="disabled")`` really wins. + # Offline mode is chosen for no-network runs, so it must never call + # out, including for telemetry, regardless of config/env/kwarg. + if isinstance(self, OfflineEmissionsTracker): + telemetry_settings = TelemetrySettings.resolve( + external_conf={}, override=TelemetryLevel.disabled + ) + else: + telemetry_settings = TelemetrySettings.resolve( + external_conf=self._external_conf, + override=None if telemetry_level is _sentinel else telemetry_level, + ) + self._telemetry = Telemetry(telemetry_settings) self._set_from_conf( force_carbon_intensity_g_co2e_kwh, "force_carbon_intensity_g_co2e_kwh", @@ -606,8 +626,13 @@ def __init__( self._initialize_runtime_state() self._initialize_scheduler_state() self._initialize_emissions_context() + self._telemetry.notice_once_if_implicit() self._init_output_methods(api_key=self._api_key) + @suppress(Exception) + def _send_telemetry_at_stop(self, emissions_data: EmissionsData) -> None: + self._telemetry.send_at_stop(self, emissions_data) + def _init_output_methods(self, *, api_key: str = None): """ Prepare the different output methods based on ``self._output_methods``. @@ -951,6 +976,7 @@ def stop(self) -> Optional[float]: emissions_data = self._prepare_emissions_data() emissions_data_delta = self._compute_emissions_delta(emissions_data) + self._send_telemetry_at_stop(emissions_data) self._persist_data( total_emissions=emissions_data, @@ -1542,6 +1568,7 @@ def track_emissions( allow_multiple_runs: Optional[bool] = _sentinel, rapl_include_dram: Optional[bool] = _sentinel, rapl_prefer_psys: Optional[bool] = _sentinel, + telemetry_level: Optional[str] = _sentinel, ): """ Decorator that supports both `EmissionsTracker` and `OfflineEmissionsTracker` @@ -1626,6 +1653,7 @@ def track_emissions( When True, measures CPU package + DRAM. :param rapl_prefer_psys: Prefer psys over package domains for RAPL on Linux (default: False). When True, uses total platform power. + :param telemetry_level: Telemetry tier (``disabled`` or ``minimal``). :return: The decorated function """ @@ -1682,6 +1710,7 @@ def wrapped_fn(*args, **kwargs): allow_multiple_runs=allow_multiple_runs, rapl_include_dram=rapl_include_dram, rapl_prefer_psys=rapl_prefer_psys, + telemetry_level=telemetry_level, ) else: tracker = EmissionsTracker( @@ -1718,6 +1747,7 @@ def wrapped_fn(*args, **kwargs): allow_multiple_runs=allow_multiple_runs, rapl_include_dram=rapl_include_dram, rapl_prefer_psys=rapl_prefer_psys, + telemetry_level=telemetry_level, ) tracker.start() try: diff --git a/codecarbon/external/geography.py b/codecarbon/external/geography.py index 075824959..4ed8b42ba 100644 --- a/codecarbon/external/geography.py +++ b/codecarbon/external/geography.py @@ -66,7 +66,10 @@ def __init__( latitude: Optional[float] = None, longitude: Optional[float] = None, country_2letter_iso_code: Optional[str] = None, + is_default: bool = False, ): + #: True when both geolocation APIs failed and this is the Canada fallback. + self.is_default = is_default self.country_iso_code = ( None if country_iso_code is None else country_iso_code.upper() ) @@ -150,4 +153,5 @@ def from_geo_js(cls, url: str) -> "GeoMetadata": latitude=46.8, longitude=-71.2, country_2letter_iso_code="CA", + is_default=True, ) diff --git a/docs/how-to/configuration.md b/docs/how-to/configuration.md index 9075470e9..830b4c1fd 100644 --- a/docs/how-to/configuration.md +++ b/docs/how-to/configuration.md @@ -214,3 +214,9 @@ os.environ["HTTPS_PROXY"] = "http://0.0.0.0:0000" For more information, please read the [requests library proxy documentation](https://requests.readthedocs.io/en/latest/user/advanced/#proxies) + +## Product telemetry + +Optional library telemetry (`telemetry_level`: `disabled` or `minimal`) is configured separately from dashboard API settings. Set it in `.codecarbon.config`, via `CODECARBON_TELEMETRY_LEVEL`, or with `EmissionsTracker(telemetry_level=...)` (argument wins). `minimal` sends anonymous product telemetry once per process; see [Product telemetry](telemetry.md). + +See [Product telemetry](telemetry.md) for what is collected and how to opt out. diff --git a/docs/how-to/telemetry.md b/docs/how-to/telemetry.md new file mode 100644 index 000000000..37067ab04 --- /dev/null +++ b/docs/how-to/telemetry.md @@ -0,0 +1,112 @@ +# Product telemetry + +CodeCarbon sends **anonymous product telemetry** to help improve the library: the hardware and software environment it runs on. It is **on by default (opt-out)**. This is separate from sending **your** emissions to the [dashboard](cloud-api.md) with `save_to_api=True`. + +## Telemetry vs your dashboard data + +| | Product telemetry | Your emissions (`save_to_api`) | +|--|-------------------|--------------------------------| +| Purpose | Improve CodeCarbon (aggregate usage) | Your projects and experiments | +| Config | `telemetry_level`, `codecarbon telemetry` | `codecarbon config`, `experiment_id` | +| Endpoint | `POST /telemetry`, no API key | Your account / experiment | + +You can use one without the other. + +## Levels + +| `telemetry_level` | What is sent | +|-------------------|--------------| +| `disabled` | Nothing | +| `minimal` (default) | The fields listed below, once per process | + +The level is resolved in this order: + +1. **Tracker or CLI argument**: `EmissionsTracker(telemetry_level=...)` or `codecarbon monitor --telemetry-level ...` +2. **`CODECARBON_TELEMETRY_LEVEL`**: overrides the config file when both are set +3. **Config file**: `telemetry_level` in `.codecarbon.config` +4. **Default:** `minimal` + +Telemetry is sent **once per process**, at the first `stop()` of a run that lasted at least one second; later `stop()` calls in the same process send nothing. The payload is built and sent on a background thread, so `stop()` never waits on it. At interpreter exit, a pending send gets at most one second before it is dropped. + +## Every field sent + +Empty or unknown values are left out rather than sent as zeros. + +| Field | Content | +|-------|---------| +| `timestamp` | Hour of the send (UTC), truncated to the hour | +| `telemetry_level` | Always `minimal` | +| `os` | Platform string, for example `Linux-5.10.0-x86_64` | +| `country_name`, `country_iso_code`, `region` | Location as already detected by the tracker | +| `cloud_provider`, `cloud_region` | Cloud provider and region, when on a cloud | +| `cpu_count`, `cpu_physical_count`, `cpu_model`, `cpu_architecture` | CPU | +| `gpu_count`, `gpu_model`, `gpu_memory_total_gb`, `gpu_driver_version`, `cuda_version`, `cudnn_version` | GPU, when present | +| `ram_total_size_gb` | Total RAM | +| `python_version`, `python_implementation`, `python_env_type` | Python interpreter and environment type (`conda`, `venv` or `system`) | +| `codecarbon_version`, `codecarbon_install_method` | CodeCarbon version and how it was installed (`pip`, `uv`, `editable`) | + +## Never collected + +- Project name, experiment id, run id, API keys +- Emissions, energy, duration or any other run measurement +- Source code, file paths, hostnames +- Coordinates (latitude and longitude) +- Voluntary [user survey](https://docs.google.com/forms/d/e/1FAIpQLSeQ5Tu_rdrpDhBJvh5R1-_iB4Ld-kgh6iNMjgaMXa8AEVPxqA/viewform) answers + +## Retention + +Telemetry rows are deleted automatically after 3 years. They do not include your IP address. + +## Configure telemetry + +### Config file + +```ini +[codecarbon] +telemetry_level = minimal +``` + +### CLI + +```bash +codecarbon telemetry set minimal +codecarbon telemetry status +codecarbon monitor --telemetry-level disabled -- python train.py +``` + +### Python + +```python +from codecarbon import EmissionsTracker + +tracker = EmissionsTracker(telemetry_level="minimal") +tracker.start() +# ... +tracker.stop() +``` + +## Offline mode + +`OfflineEmissionsTracker` and `codecarbon monitor --offline` never send telemetry, regardless of `telemetry_level` in config, environment or argument: offline mode is chosen for runs with no network access. + +## Opt out + +```ini +[codecarbon] +telemetry_level = disabled +``` + +Or set `CODECARBON_TELEMETRY_LEVEL=disabled`, or run `codecarbon telemetry set disabled` (writes to the global `~/.codecarbon.config`, creating it if missing; pass `--config` to target a specific file instead). + +## First run without explicit configuration + +Telemetry is on by default (`minimal`), and you are told about it once: + +- **Interactive CLI** (`codecarbon config` or `codecarbon monitor` in a terminal): you are asked once which level you want. The answer is saved as `telemetry_level` in `~/.codecarbon.config`. +- **Everything else** (library use, CI, SLURM, pipes): nothing ever blocks on a prompt. CodeCarbon prints a notice to stderr saying what is sent and how to choose a level. It is shown **once per process**, on every run, for as long as the level stays implicit (not set by argument, environment variable or config file) — set `telemetry_level` explicitly (via `codecarbon telemetry set minimal|disabled`, `CODECARBON_TELEMETRY_LEVEL`, or the config file) to silence it for good. + +## Related + +- [Configure CodeCarbon](configuration.md) +- [CLI reference](../reference/cli.md#codecarbon-telemetry) +- [Cloud API & dashboard](cloud-api.md) diff --git a/docs/reference/api.md b/docs/reference/api.md index 15117b123..cfe50b826 100644 --- a/docs/reference/api.md +++ b/docs/reference/api.md @@ -12,6 +12,19 @@ Parameters can be set via `EmissionsTracker()`, `OfflineEmissionsTracker()`, the If you use `CUDA_VISIBLE_DEVICES` or `ROCR_VISIBLE_DEVICES` to set GPUs, CodeCarbon will automatically populate `gpu_ids`. Manual `gpu_ids` overrides this. +## Product telemetry + +Optional library telemetry is controlled by **`telemetry_level`** on the tracker (same parameter on `OfflineEmissionsTracker` and `@track_emissions`): + +| Value | Behavior | +|-------|----------| +| `disabled` | No product telemetry | +| `minimal` (default) | Anonymous environment and hardware telemetry, once per process | + +**Resolution order:** tracker argument → `CODECARBON_TELEMETRY_LEVEL` → `.codecarbon.config` → default `minimal`. The tracker argument overrides the environment variable and config file. + +This is separate from `save_to_api` (your dashboard experiment). See [Product telemetry](../how-to/telemetry.md). + ## EmissionsTracker / BaseEmissionsTracker `EmissionsTracker` and `OfflineEmissionsTracker` inherit from `BaseEmissionsTracker`. diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 71ba05a85..ece66b1b1 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -45,6 +45,7 @@ Displays real-time emissions data for all processes on your machine. Press `Ctrl | `--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 | +| `--telemetry-level` | string | - | One-run level: `disabled` or `minimal` | **Examples:** ```bash @@ -88,6 +89,34 @@ codecarbon monitor -- node app.js --port 8080 Same options as `codecarbon monitor` apply (see above). +### `codecarbon telemetry` + +Configure **product telemetry** (library usage metadata), separate from `codecarbon config` (dashboard org/project/experiment). + +**Usage:** + +```bash +codecarbon telemetry # interactive wizard +codecarbon telemetry set # disabled | minimal +codecarbon telemetry status # resolved tier and whether it was set explicitly +``` + +**Options:** + +| Option | Description | +|--------|-------------| +| `--config PATH` | Use a specific `.codecarbon.config` (default: local then global) | + +See [Product telemetry](../how-to/telemetry.md) for what `minimal` collects and how to opt out. + +### `codecarbon monitor --telemetry-level` + +Override the telemetry tier for a single monitor run (does not update config). + +```bash +codecarbon monitor --telemetry-level disabled -- python train.py +``` + ### `codecarbon detect` Detect and print hardware information. diff --git a/mkdocs.yml b/mkdocs.yml index 77a1f7b6e..234b0a77c 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -148,6 +148,7 @@ nav: - How-to Guides: - Examples: how-to/examples.md - Configure CodeCarbon: how-to/configuration.md + - Product telemetry: how-to/telemetry.md - Compare Model Efficiency: tutorials/comparing-model-efficiency.md - Dashboard & Visualization: - Use the Cloud API & Dashboard: how-to/cloud-api.md diff --git a/scripts/e2e_telemetry.sh b/scripts/e2e_telemetry.sh new file mode 100755 index 000000000..ba0467978 --- /dev/null +++ b/scripts/e2e_telemetry.sh @@ -0,0 +1,127 @@ +#!/usr/bin/env bash +# End-to-end telemetry check against a throwaway Postgres and a local API. +# +# 1. Starts the docker-compose Postgres under its own project name, so the +# volume is separate from your dev database, and migrates it. +# 2. Runs the API with uvicorn on http://localhost:8008. +# 3. Runs one library process (two tracker stops) and one `codecarbon monitor` +# process with telemetry at `minimal`, both pointed at that API. +# 4. Checks the database: one row per process, only allowed columns. +# Everything is torn down on exit. +# +# Set E2E_ALLOW_PULL=1 to allow pulling the postgres image when it is missing. +# Set E2E_PG_PORT to change the host port for Postgres (default 55432, so a dev +# database already on 5432 does not clash). +set -euo pipefail + +ROOT="$(cd "$(dirname "$0")/.." && pwd)" +PROJECT=codecarbon-e2e +API_URL=http://localhost:8008 +PG_PORT="${E2E_PG_PORT:-55432}" +export DATABASE_URL="postgresql://codecarbon-user:supersecret@localhost:${PG_PORT}/codecarbon_db" +WORK="$(mktemp -d)" +cat >"$WORK/compose.override.yml" </dev/null || echo "WARN: API process already gone" >&2 + fi + if ! "${COMPOSE[@]}" down -v; then + echo "WARN: docker compose down failed; remove project $PROJECT by hand" >&2 + fi + rm -rf "$WORK" +} + +if ! docker info >/dev/null 2>&1; then + echo "ERROR: Docker daemon is not reachable." >&2 + exit 1 +fi +if [ -z "$(docker image ls -q postgres)" ] && [ "${E2E_ALLOW_PULL:-0}" != 1 ]; then + echo "ERROR: no postgres image present; pulling one downloads about 150 MB." >&2 + echo "Re-run with E2E_ALLOW_PULL=1 to allow it." >&2 + exit 2 +fi +trap cleanup EXIT + +echo "== Starting Postgres" +"${COMPOSE[@]}" up -d postgres +pg_ready=0 +for _ in $(seq 1 30); do + if "${COMPOSE[@]}" exec -T postgres \ + pg_isready -U codecarbon-user -d codecarbon_db >/dev/null 2>&1; then + pg_ready=1 + break + fi + sleep 1 +done +if [ "$pg_ready" -ne 1 ]; then + echo "ERROR: Postgres never became ready (pg_isready kept failing)." >&2 + exit 1 +fi + +echo "== Migrating" +(cd "$ROOT/carbonserver" && uv run --project . python -m alembic \ + -c carbonserver/database/alembic.ini upgrade head) + +if curl --noproxy '*' -fsS "$API_URL/" >/dev/null 2>&1; then + echo "ERROR: something is already listening on $API_URL before the API starts." >&2 + exit 1 +fi + +echo "== Starting API" +(cd "$ROOT/carbonserver" && exec env AUTH_PROVIDER=none ENVIRONMENT=local \ + uv run --project . uvicorn main:app --port 8008 >"$WORK/api.log" 2>&1) & +API_PID=$! +api_up=0 +for _ in $(seq 1 60); do + if ! kill -0 "$API_PID" 2>/dev/null; then + echo "ERROR: API process exited before coming up; log follows" >&2 + cat "$WORK/api.log" >&2 + exit 1 + fi + if curl --noproxy '*' -fsS "$API_URL/" >/dev/null 2>&1; then + api_up=1 + break + fi + sleep 1 +done +if [ "$api_up" -ne 1 ] || ! kill -0 "$API_PID" 2>/dev/null; then + echo "ERROR: API did not come up; log follows" >&2 + cat "$WORK/api.log" >&2 + exit 1 +fi + +export CODECARBON_TELEMETRY_API_URL="$API_URL" +export CODECARBON_TELEMETRY_LEVEL=minimal +export CODECARBON_ALLOW_MULTIPLE_RUNS=True +export NO_PROXY="localhost,127.0.0.1" +cd "$WORK" + +echo "== Library process: two stops, expect one row" +uv run --project "$ROOT" python - <<'PY' +import time + +from codecarbon import EmissionsTracker + +for _ in range(2): + tracker = EmissionsTracker(save_to_file=False, save_to_api=False) + tracker.start() + time.sleep(2) + tracker.stop() +PY + +echo "== CLI monitor process: expect one row" +uv run --project "$ROOT" codecarbon monitor --no-api --telemetry-level minimal \ + -- python -c "import time; time.sleep(2)" + +echo "== Checking the database" +(cd "$ROOT/carbonserver" && TELEMETRY_E2E_EXPECTED_ROWS=2 \ + uv run --project . --extra dev python -m pytest -q -s -p no:cacheprovider \ + tests/api/integration/test_telemetry_e2e_db.py) diff --git a/tests/cli/test_cli_main.py b/tests/cli/test_cli_main.py index 09585126f..48e7be489 100644 --- a/tests/cli/test_cli_main.py +++ b/tests/cli/test_cli_main.py @@ -2,6 +2,7 @@ import os from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import requests @@ -387,6 +388,36 @@ def stop(self): assert calls["kwargs"]["region"] == "IDF" +def test_monitor_offline_skips_telemetry_prompt(monkeypatch): + """``codecarbon monitor --offline`` never sends telemetry, so it must not + ask the interactive telemetry-level question either.""" + + class FakeOfflineTracker: + def __init__(self, **kwargs): + self._another_instance_already_running = True + + def start(self): + pass + + def stop(self): + return None + + monkeypatch.setattr( + "codecarbon.emissions_tracker.OfflineEmissionsTracker", FakeOfflineTracker + ) + monkeypatch.setattr(cli_main.signal, "signal", lambda *args, **kwargs: None) + ask_mock = MagicMock() + monkeypatch.setattr(cli_main, "ask_telemetry_level_once", ask_mock) + + runner = CliRunner() + result = runner.invoke( + cli_main.codecarbon, + ["monitor", "--offline", "--country-iso-code", "FRA"], + ) + assert result.exit_code == 0 + ask_mock.assert_not_called() + + def _fake_offline_monitor(monkeypatch, tmp_path): """Patch the offline tracker and run the monitor loop in `tmp_path`.""" calls = {} diff --git a/tests/cli/test_monitor.py b/tests/cli/test_monitor.py index f3fa4852c..0135e9581 100644 --- a/tests/cli/test_monitor.py +++ b/tests/cli/test_monitor.py @@ -2,7 +2,9 @@ import pytest import typer +from typer.testing import CliRunner +from codecarbon.cli import main as cli_main from codecarbon.cli import monitor as monitor_module @@ -226,3 +228,42 @@ def kill(self): assert exc_info.value.exit_code == 130 assert process_info["terminated"] == 1 assert process_info["killed"] == 1 + + +def test_run_and_monitor_passes_telemetry_level_override(monkeypatch): + captured = {} + + class CapturingTracker(FakeTracker): + def __init__(self, **kwargs): + captured.update(kwargs) + super().__init__() + + class FakePopen: + def __init__(self, command, text=True): + pass + + def wait(self): + return 0 + + monkeypatch.setattr( + "codecarbon.emissions_tracker.EmissionsTracker", CapturingTracker + ) + monkeypatch.setattr(monitor_module.subprocess, "Popen", FakePopen) + monkeypatch.setattr(monitor_module, "print", lambda *args, **kwargs: None) + + with pytest.raises(typer.Exit) as exc_info: + monitor_module.run_and_monitor( + SimpleNamespace(args=["echo", "ok"]), + telemetry_level="disabled", + ) + + assert exc_info.value.exit_code == 0 + assert captured.get("telemetry_level") == "disabled" + + +def test_monitor_cli_rejects_invalid_telemetry_level(): + result = CliRunner().invoke( + cli_main.codecarbon, + ["monitor", "--telemetry-level", "invalid"], + ) + assert result.exit_code != 0 diff --git a/tests/cli/test_telemetry_cli.py b/tests/cli/test_telemetry_cli.py new file mode 100644 index 000000000..38e32250c --- /dev/null +++ b/tests/cli/test_telemetry_cli.py @@ -0,0 +1,374 @@ +"""Tests for codecarbon telemetry CLI commands.""" + +import tempfile +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import typer +from typer.testing import CliRunner + +from codecarbon.cli import main as cli_main +from codecarbon.cli.telemetry_cli import ( + normalize_telemetry_level, + pick_config_path_interactive, + resolve_config_path, + telemetry_app, + write_telemetry_level, +) + + +def test_normalize_telemetry_level_accepts_valid_values(): + assert normalize_telemetry_level("MINIMAL") == "minimal" + assert normalize_telemetry_level("disabled") == "disabled" + + +def test_normalize_telemetry_level_rejects_invalid(): + with pytest.raises(typer.BadParameter): + normalize_telemetry_level("bogus") + + +def test_telemetry_set_writes_config(): + runner = CliRunner() + with tempfile.TemporaryDirectory() as tmp: + config_path = Path(tmp) / ".codecarbon.config" + result = runner.invoke( + telemetry_app, + ["set", "disabled", "--config", str(config_path)], + ) + assert result.exit_code == 0 + assert "telemetry_level = disabled" in result.output + content = config_path.read_text() + assert "telemetry_level = disabled" in content + + +def test_telemetry_set_without_config_writes_global_not_local(tmp_path, monkeypatch): + """Privacy: an opt-out without --config must land in the GLOBAL config, + like ``ask_telemetry_level_once`` does, not in a newly created local one. + """ + home = tmp_path / "home" + home.mkdir() + work = tmp_path / "work" + work.mkdir() + monkeypatch.chdir(work) + monkeypatch.setattr(Path, "home", lambda: home) + + runner = CliRunner() + result = runner.invoke(telemetry_app, ["set", "disabled"]) + assert result.exit_code == 0 + + global_path = home / ".codecarbon.config" + local_path = work / ".codecarbon.config" + assert global_path.exists() + assert "telemetry_level = disabled" in global_path.read_text() + assert not local_path.exists() + + +def test_telemetry_status_reports_stored_level(): + runner = CliRunner() + with tempfile.TemporaryDirectory() as tmp: + config_path = Path(tmp) / ".codecarbon.config" + config_path.write_text("[codecarbon]\ntelemetry_level = minimal\n") + result = runner.invoke( + telemetry_app, + ["status", "--config", str(config_path)], + ) + assert result.exit_code == 0 + assert "Resolved tier: minimal" in result.output + assert "Explicitly configured: True" in result.output + + +def test_telemetry_status_merged_matches_tracker_precedence(): + runner = CliRunner() + with tempfile.TemporaryDirectory() as tmp: + global_path = Path(tmp) / "global.config" + local_path = Path(tmp) / "local.config" + global_path.write_text("[codecarbon]\ntelemetry_level = minimal\n") + local_path.write_text("[codecarbon]\ntelemetry_level = disabled\n") + with patch( + "codecarbon.core.config._config_file_paths", + return_value=(str(global_path), str(local_path)), + ): + result = runner.invoke(telemetry_app, ["status"]) + assert result.exit_code == 0 + assert "Resolved tier: disabled" in result.output + assert "merged" in result.output + + +def test_monitor_passes_telemetry_level_override(monkeypatch): + captured = {} + + class FakeTracker: + def __init__(self, **kwargs): + captured.update(kwargs) + self._conf = {"output_file": "emissions.csv"} + + def start(self): + return None + + def stop(self): + return 0.0 + + monkeypatch.setattr("codecarbon.emissions_tracker.EmissionsTracker", FakeTracker) + + runner = CliRunner() + result = runner.invoke( + cli_main.codecarbon, + [ + "monitor", + "--no-api", + "--telemetry-level", + "disabled", + "--", + "echo", + "ok", + ], + ) + assert result.exit_code == 0 + assert captured.get("telemetry_level") == "disabled" + + +def test_monitor_rejects_invalid_telemetry_level(): + runner = CliRunner() + result = runner.invoke( + cli_main.codecarbon, + ["monitor", "--telemetry-level", "invalid"], + ) + assert result.exit_code != 0 + + +def test_telemetry_status_missing_config_file(): + runner = CliRunner() + with tempfile.TemporaryDirectory() as tmp: + missing = Path(tmp) / "missing.config" + result = runner.invoke( + telemetry_app, + ["status", "--config", str(missing)], + ) + assert result.exit_code == 0 + assert "Config file not found" in result.output + assert "not explicit" in result.output + + +def test_telemetry_status_shows_implicit_warning(): + runner = CliRunner() + with patch( + "codecarbon.cli.telemetry_cli.get_config_file_settings", + return_value={}, + ): + with patch( + "codecarbon.cli.telemetry_cli.get_hierarchical_config", + return_value={}, + ): + result = runner.invoke(telemetry_app, ["status"]) + assert "Explicitly configured: False" in result.output + assert "Sent once per process to" in result.output + + +def test_resolve_config_path_creates_explicit_file(): + with tempfile.TemporaryDirectory() as tmp: + config_path = Path(tmp) / "custom.config" + resolved = resolve_config_path(config_path, create=True) + assert resolved == config_path.resolve() + assert config_path.exists() + + +def test_resolve_config_path_prefers_local_config(tmp_path, monkeypatch): + local_path = tmp_path / ".codecarbon.config" + local_path.write_text("[codecarbon]\n", encoding="utf-8") + monkeypatch.chdir(tmp_path) + assert resolve_config_path(None) == local_path.resolve() + + +def test_resolve_config_path_uses_global_when_local_missing(tmp_path, monkeypatch): + home = tmp_path / "home" + home.mkdir() + global_path = home / ".codecarbon.config" + global_path.write_text("[codecarbon]\n", encoding="utf-8") + workdir = tmp_path / "work" + workdir.mkdir() + monkeypatch.chdir(workdir) + monkeypatch.setattr(Path, "home", lambda: home) + assert resolve_config_path(None) == global_path.resolve() + + +def test_write_telemetry_level_creates_missing_file(tmp_path): + config_path = tmp_path / "nested" / ".codecarbon.config" + write_telemetry_level(config_path, "minimal") + assert config_path.exists() + assert "telemetry_level = minimal" in config_path.read_text() + + +def test_pick_config_path_interactive_create_new(): + with patch("codecarbon.cli.telemetry_cli.questionary.select") as mock_select: + mock_select.return_value.ask.return_value = "Create new config file" + with patch( + "codecarbon.cli.telemetry_cli.create_new_config_file", + return_value=Path("/tmp/new.config"), + ) as mock_create: + path = pick_config_path_interactive() + mock_create.assert_called_once() + assert path == Path("/tmp/new.config") + + +def test_resolve_config_path_creates_local_in_cwd(tmp_path, monkeypatch): + home = tmp_path / "home" + home.mkdir() + work = tmp_path / "work" + work.mkdir() + monkeypatch.chdir(work) + monkeypatch.setattr(Path, "home", lambda: home) + resolved = resolve_config_path(None, create=True) + assert resolved == (work / ".codecarbon.config").resolve() + assert resolved.exists() + + +def test_pick_config_path_lists_existing_configs(tmp_path, monkeypatch): + home = tmp_path / "home" + home.mkdir() + global_path = home / ".codecarbon.config" + global_path.write_text("[codecarbon]\n", encoding="utf-8") + work = tmp_path / "work" + work.mkdir() + local_path = work / ".codecarbon.config" + local_path.write_text("[codecarbon]\n", encoding="utf-8") + monkeypatch.chdir(work) + monkeypatch.setattr(Path, "home", lambda: home) + with patch("codecarbon.cli.telemetry_cli.questionary.select") as mock_select: + mock_select.return_value.ask.return_value = str(local_path) + path = pick_config_path_interactive() + assert path == local_path.resolve() + + +def test_pick_config_path_lists_same_file_once(tmp_path, monkeypatch): + """Run from $HOME, the global and local config are the same file.""" + home = tmp_path / "home" + home.mkdir() + (home / ".codecarbon.config").write_text("[codecarbon]\n", encoding="utf-8") + monkeypatch.chdir(home) + monkeypatch.setattr(Path, "home", lambda: home) + with patch("codecarbon.cli.telemetry_cli.questionary.select") as mock_select: + mock_select.return_value.ask.return_value = "Create new config file" + with patch("codecarbon.cli.telemetry_cli.create_new_config_file"): + pick_config_path_interactive() + choices = mock_select.call_args.kwargs["choices"] + assert choices == [ + str((home / ".codecarbon.config").resolve()), + "Create new config file", + ] + + +def test_pick_config_path_interactive_cancel_exits(): + """Ctrl-C at the prompt makes questionary return None.""" + with patch("codecarbon.cli.telemetry_cli.questionary.select") as mock_select: + mock_select.return_value.ask.return_value = None + with pytest.raises(typer.Exit): + pick_config_path_interactive() + + +def test_telemetry_default_command_runs_interactive_wizard(tmp_path): + config_path = tmp_path / ".codecarbon.config" + config_path.write_text("[codecarbon]\n", encoding="utf-8") + runner = CliRunner() + with patch("codecarbon.cli.telemetry_cli.questionary.select") as mock_select: + mock_select.return_value.ask.return_value = "disabled" + result = runner.invoke(telemetry_app, ["--config", str(config_path)]) + assert result.exit_code == 0 + assert "telemetry_level = disabled" in config_path.read_text() + + +def test_telemetry_interactive_exits_when_selection_cancelled(tmp_path): + config_path = tmp_path / ".codecarbon.config" + config_path.write_text("[codecarbon]\n", encoding="utf-8") + runner = CliRunner() + with patch("codecarbon.cli.telemetry_cli.questionary.select") as mock_select: + mock_select.return_value.ask.return_value = None + result = runner.invoke(telemetry_app, ["--config", str(config_path)]) + assert result.exit_code == 0 + assert "telemetry_level" not in config_path.read_text() + + +def test_telemetry_interactive_prompts_for_config_path(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + runner = CliRunner() + select = MagicMock() + select.ask.return_value = "minimal" + with patch( + "codecarbon.cli.telemetry_cli.pick_config_path_interactive", + return_value=tmp_path / ".codecarbon.config", + ): + with patch( + "codecarbon.cli.telemetry_cli.questionary.select", return_value=select + ): + result = runner.invoke(telemetry_app, []) + config_path = tmp_path / ".codecarbon.config" + assert result.exit_code == 0 + assert config_path.exists() + assert "telemetry_level = minimal" in config_path.read_text() + + +def test_ask_telemetry_level_once_never_prompts_without_tty(): + from codecarbon.cli.telemetry_cli import ask_telemetry_level_once + + with ( + patch("codecarbon.cli.telemetry_cli.sys") as mock_sys, + patch("codecarbon.cli.telemetry_cli.questionary.select") as mock_select, + ): + mock_sys.stdin.isatty.return_value = False + ask_telemetry_level_once() + mock_select.assert_not_called() + + +def test_ask_telemetry_level_once_stores_answer_in_global_config(tmp_path): + from codecarbon.cli.telemetry_cli import ask_telemetry_level_once + + global_path = tmp_path / ".codecarbon.config" + with ( + patch("codecarbon.cli.telemetry_cli.sys") as mock_sys, + patch("codecarbon.cli.telemetry_cli.get_hierarchical_config", return_value={}), + patch( + "codecarbon.cli.telemetry_cli._config_file_paths", + return_value=(str(global_path), str(tmp_path / "local")), + ), + patch("codecarbon.cli.telemetry_cli.questionary.select") as mock_select, + ): + mock_sys.stdin.isatty.return_value = True + mock_sys.stdout.isatty.return_value = True + mock_select.return_value.ask.return_value = "disabled" + ask_telemetry_level_once() + assert "telemetry_level = disabled" in global_path.read_text() + + +def test_resolve_config_path_without_files_or_create(tmp_path, monkeypatch): + monkeypatch.setattr(Path, "home", lambda: tmp_path / "home") + monkeypatch.chdir(tmp_path) + resolved = resolve_config_path(None) + assert resolved == (tmp_path / ".codecarbon.config").resolve() + assert not resolved.exists() + + +@pytest.mark.parametrize( + ("conf", "answer"), + [({"telemetry_level": "minimal"}, "disabled"), ({}, None)], + ids=["already-explicit", "ctrl-c"], +) +def test_ask_telemetry_level_once_writes_nothing(tmp_path, conf, answer): + from codecarbon.cli.telemetry_cli import ask_telemetry_level_once + + global_path = tmp_path / ".codecarbon.config" + with ( + patch("codecarbon.cli.telemetry_cli.sys") as mock_sys, + patch( + "codecarbon.cli.telemetry_cli.get_hierarchical_config", return_value=conf + ), + patch( + "codecarbon.cli.telemetry_cli._config_file_paths", + return_value=(str(global_path), str(tmp_path / "local")), + ), + patch("codecarbon.cli.telemetry_cli.questionary.select") as mock_select, + ): + mock_sys.stdin.isatty.return_value = True + mock_sys.stdout.isatty.return_value = True + mock_select.return_value.ask.return_value = answer + ask_telemetry_level_once() + assert not global_path.exists() diff --git a/tests/conftest.py b/tests/conftest.py index 4d209b836..dfbb2bc17 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -22,6 +22,18 @@ def _reset_process_hardware_cache(): detect_cpu_model.cache_clear() +@pytest.fixture(autouse=True) +def _isolate_telemetry(monkeypatch): + """Keep tests off the real telemetry endpoint. + + Telemetry needs no key, so any tracker a test stops would post to + api.codecarbon.io. Tests that exercise telemetry clear this variable. + """ + monkeypatch.setenv("CODECARBON_TELEMETRY_LEVEL", "disabled") + monkeypatch.setattr("codecarbon.core.telemetry.dispatcher._sent", False) + monkeypatch.setattr("codecarbon.core.telemetry.dispatcher._notice_shown", False) + + @pytest.fixture(autouse=True) def _no_real_browser(monkeypatch): """Keep tests headless: never let code under test open a browser tab.""" diff --git a/tests/test_config.py b/tests/test_config.py index e913d9f36..f295cdea6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -30,12 +30,12 @@ def setUp(self): "CODECARBON_EXPERIMENT_ID", "CODECARBON_API_ENDPOINT", "CODECARBON_TELEMETRY", - "CODECARBON_TELEMETRY_PROJECT_TOKEN", + "CODECARBON_TELEMETRY_LEVEL", + "CODECARBON_TELEMETRY_API_URL", "codecarbon_api_key", "codecarbon_experiment_id", "codecarbon_api_endpoint", "codecarbon_telemetry", - "codecarbon_telemetry_project_token", ]: os.environ.pop(key, None) os.environ.setdefault("CODECARBON_ALLOW_MULTIPLE_RUNS", "True") diff --git a/tests/test_config_file_settings.py b/tests/test_config_file_settings.py new file mode 100644 index 000000000..6a155ef7e --- /dev/null +++ b/tests/test_config_file_settings.py @@ -0,0 +1,51 @@ +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from codecarbon.core.config import get_config_file_settings, get_hierarchical_config + + +class TestGetConfigFileSettings(unittest.TestCase): + def test_returns_empty_when_no_config_files(self): + with patch("codecarbon.core.config._config_file_paths") as mock_paths: + mock_paths.return_value = ("/nonexistent/global", "/nonexistent/local") + settings = get_config_file_settings() + self.assertEqual(settings, {}) + + def test_local_overrides_global_telemetry_level(self): + with tempfile.TemporaryDirectory() as tmp: + global_path = Path(tmp) / "global.config" + local_path = Path(tmp) / "local.config" + global_path.write_text("[codecarbon]\ntelemetry_level = minimal\n") + local_path.write_text("[codecarbon]\ntelemetry_level = disabled\n") + with patch( + "codecarbon.core.config._config_file_paths", + return_value=(str(global_path), str(local_path)), + ): + settings = get_config_file_settings() + self.assertEqual(settings["telemetry_level"], "disabled") + + def test_hierarchical_config_includes_env_but_file_settings_do_not(self): + with tempfile.TemporaryDirectory() as tmp: + local_path = Path(tmp) / ".codecarbon.config" + local_path.write_text("[codecarbon]\ntelemetry_level = minimal\n") + with patch( + "codecarbon.core.config._config_file_paths", + return_value=("/nonexistent/global", str(local_path)), + ): + with patch.dict( + os.environ, + {"CODECARBON_TELEMETRY": "disabled"}, + clear=False, + ): + file_settings = get_config_file_settings() + hierarchical = get_hierarchical_config() + self.assertEqual(file_settings.get("telemetry_level"), "minimal") + self.assertNotIn("telemetry", file_settings) + self.assertEqual(hierarchical.get("telemetry"), "disabled") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_emissions_tracker.py b/tests/test_emissions_tracker.py index a2d6b5966..531733ce2 100644 --- a/tests/test_emissions_tracker.py +++ b/tests/test_emissions_tracker.py @@ -37,7 +37,7 @@ def heavy_computation(run_time_secs: float = 3): pass -empty_conf = "[codecarbon]" +disabled_conf = "[codecarbon]\ntelemetry_level = disabled\n" if sys.platform == "darwin": @@ -79,7 +79,8 @@ def setUp(self) -> None: # ./.codecarbon.config so that the user's local configuration does not # alter tests patcher = mock.patch( - "builtins.open", new_callable=get_custom_mock_open(empty_conf, empty_conf) + "builtins.open", + new_callable=get_custom_mock_open(disabled_conf, disabled_conf), ) self.addCleanup(patcher.stop) patcher.start() diff --git a/tests/test_geography.py b/tests/test_geography.py index 8f95f7f43..a0e2b6942 100644 --- a/tests/test_geography.py +++ b/tests/test_geography.py @@ -127,3 +127,14 @@ def test_geo_metadata_CANADA(self): self.assertEqual("CAN", geo.country_iso_code) self.assertEqual("Canada", geo.country_name) self.assertEqual("ontario", geo.region) + self.assertFalse(geo.is_default) + + @responses.activate + def test_geo_metadata_marks_the_fallback_as_default(self): + responses.add(responses.GET, self.geo_js_url, body=ConnectionError("down")) + responses.add( + responses.GET, "https://ipinfo.io/json", body=ConnectionError("down") + ) + geo = GeoMetadata.from_geo_js(self.geo_js_url) + self.assertEqual("Canada", geo.country_name) + self.assertTrue(geo.is_default) diff --git a/tests/test_offline_emissions_tracker.py b/tests/test_offline_emissions_tracker.py index 36447409d..1e0b7d4f0 100644 --- a/tests/test_offline_emissions_tracker.py +++ b/tests/test_offline_emissions_tracker.py @@ -18,7 +18,7 @@ def heavy_computation(run_time_secs: float = 3): pass -empty_conf = "[codecarbon]" +disabled_conf = "[codecarbon]\ntelemetry_level = disabled\n" class TestOfflineEmissionsTracker(unittest.TestCase): @@ -32,7 +32,8 @@ def setUp(self) -> None: # ./.codecarbon.config so that the user's local configuration does not # alter tests patcher = mock.patch( - "builtins.open", new_callable=get_custom_mock_open(empty_conf, empty_conf) + "builtins.open", + new_callable=get_custom_mock_open(disabled_conf, disabled_conf), ) self.addCleanup(patcher.stop) patcher.start() diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 000000000..796679b4a --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,485 @@ +"""Telemetry behaviour as seen through the public ``EmissionsTracker`` path.""" + +import os +import sys +import time +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from codecarbon.core.telemetry.dispatcher import Telemetry +from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker +from tests.testutils import ( + ensure_telemetry_run_duration, + get_custom_mock_open, + hanging_endpoint, + join_telemetry, +) + +if sys.platform == "darwin": + mock_platform_cli_setup = patch( + "codecarbon.core.powermetrics.ApplePowermetrics._setup_cli" + ) +else: + mock_platform_cli_setup = patch("codecarbon.core.cpu.IntelPowerGadget._setup_cli") + +TELEMETRY_ENV_VARS = ( + "CODECARBON_TELEMETRY", + "CODECARBON_TELEMETRY_LEVEL", + "CODECARBON_TELEMETRY_API_URL", +) + + +def _conf(level: str = None, *, extra: str = "") -> str: + conf = "[codecarbon]\n" + if level is not None: + conf += f"telemetry_level = {level}\n" + return conf + extra + + +@mock_platform_cli_setup +class TestTrackerTelemetry(unittest.TestCase): + """Every case goes through the real constructor, not ``resolve()``.""" + + def setUp(self) -> None: + self._patchers = [] + clean_env = { + key: value + for key, value in os.environ.items() + if key.upper() not in TELEMETRY_ENV_VARS + } + self._enter(patch.dict(os.environ, clean_env, clear=True)) + self._enter(patch("codecarbon.external.geography.GeoMetadata.from_geo_js")) + + def tearDown(self) -> None: + for patcher in reversed(self._patchers): + patcher.stop() + + def _enter(self, patcher): + patcher.start() + self._patchers.append(patcher) + return patcher + + def _mock_config(self, conf: str) -> None: + self._enter( + patch("builtins.open", new_callable=get_custom_mock_open(conf, conf)) + ) + + def _mock_post(self): + return patch( + "codecarbon.core.telemetry.dispatcher.post_private", return_value=True + ) + + def _run_tracker(self, **kwargs): + with ensure_telemetry_run_duration(): + tracker = EmissionsTracker( + measure_power_secs=1, save_to_api=False, save_to_file=False, **kwargs + ) + tracker.start() + tracker.stop() + join_telemetry(tracker) + return tracker + + def test_no_telemetry_on_init(self, mock_cli_setup): + self._mock_config(_conf("minimal")) + with self._mock_post() as mock_post: + EmissionsTracker(save_to_api=False, save_to_file=False) + mock_post.assert_not_called() + + def test_minimal_sends_minimal_payload_on_stop(self, mock_cli_setup): + self._mock_config(_conf("minimal")) + with self._mock_post() as mock_post: + self._run_tracker() + mock_post.assert_called_once() + payload = mock_post.call_args[0][1] + self.assertEqual(payload["telemetry_level"], "minimal") + self.assertNotIn("total_emissions_kg", payload) + + def test_sends_once_per_process(self, mock_cli_setup): + self._mock_config(_conf("minimal")) + with self._mock_post() as mock_post: + self._run_tracker() + self._run_tracker() + mock_post.assert_called_once() + + def test_short_run_does_not_use_up_the_single_send(self, mock_cli_setup): + self._mock_config(_conf("minimal")) + with self._mock_post() as mock_post: + tracker = EmissionsTracker(save_to_api=False, save_to_file=False) + tracker._telemetry.send_at_stop(tracker, SimpleNamespace(duration=0.5)) + join_telemetry(tracker) + mock_post.assert_not_called() + self._run_tracker() + mock_post.assert_called_once() + + def test_config_disabled_sends_nothing(self, mock_cli_setup): + self._mock_config(_conf("disabled")) + with self._mock_post() as mock_post: + self._run_tracker() + mock_post.assert_not_called() + + def test_kwarg_disabled_overrides_config(self, mock_cli_setup): + """Regression: ``telemetry_level="disabled"`` used to be ignored entirely.""" + self._mock_config(_conf("minimal")) + with self._mock_post() as mock_post: + tracker = self._run_tracker(telemetry_level="disabled") + mock_post.assert_not_called() + self.assertEqual(tracker._telemetry.settings.level.value, "disabled") + self.assertNotIn("telemetry_level", tracker._conf) + + def test_kwarg_disabled_overrides_env(self, mock_cli_setup): + self._mock_config(_conf()) + with patch.dict(os.environ, {"CODECARBON_TELEMETRY_LEVEL": "minimal"}): + with self._mock_post() as mock_post: + self._run_tracker(telemetry_level="disabled") + mock_post.assert_not_called() + + def test_offline_tracker_never_sends_telemetry(self, mock_cli_setup): + """Offline mode is chosen for no-network runs: it must never call out, + including for telemetry, even when the config/env would otherwise + enable it. + """ + self._mock_config(_conf("minimal")) + with self._mock_post() as mock_post: + with ensure_telemetry_run_duration(): + tracker = OfflineEmissionsTracker( + country_iso_code="CAN", save_to_api=False, save_to_file=False + ) + self.assertEqual(tracker._telemetry.settings.level.value, "disabled") + tracker.start() + tracker.stop() + join_telemetry(tracker) + mock_post.assert_not_called() + + def test_env_level_overrides_config_file(self, mock_cli_setup): + self._mock_config(_conf("minimal")) + with patch.dict(os.environ, {"CODECARBON_TELEMETRY_LEVEL": "disabled"}): + with self._mock_post() as mock_post: + self._run_tracker() + mock_post.assert_not_called() + + def test_legacy_env_var_is_ignored(self, mock_cli_setup): + self._mock_config(_conf("minimal")) + with patch.dict(os.environ, {"CODECARBON_TELEMETRY": "disabled"}): + with self._mock_post() as mock_post: + self._run_tracker() + mock_post.assert_called_once() + + def test_api_url_comes_from_config(self, mock_cli_setup): + self._mock_config( + _conf("minimal", extra="telemetry_api_url = http://example.test/\n") + ) + tracker = EmissionsTracker(save_to_api=False, save_to_file=False) + self.assertEqual(tracker._telemetry.settings.api_url, "http://example.test") + + def test_api_url_comes_from_env(self, mock_cli_setup): + self._mock_config(_conf("minimal")) + with patch.dict( + os.environ, {"CODECARBON_TELEMETRY_API_URL": "http://env.test"} + ): + tracker = EmissionsTracker(save_to_api=False, save_to_file=False) + self.assertEqual(tracker._telemetry.settings.api_url, "http://env.test") + + def _notices(self, mock_warning): + return [ + call + for call in mock_warning.call_args_list + if call[0] and "telemetry is on by default" in str(call[0][0]) + ] + + def test_notice_printed_to_stderr_even_at_default_log_level(self, mock_cli_setup): + """``codecarbon monitor`` defaults to log_level=error, which hides a + WARNING-only notice; the notice must be visible regardless of log level. + """ + self._mock_config(_conf()) + with patch("sys.stderr") as mock_stderr: + EmissionsTracker(save_to_api=False, save_to_file=False, log_level="error") + printed = "".join( + call.args[0] for call in mock_stderr.write.call_args_list if call.args + ) + self.assertIn("telemetry is on by default", printed) + + def test_notice_once_per_process_when_level_not_explicit(self, mock_cli_setup): + self._mock_config(_conf()) + with patch("codecarbon.core.telemetry.dispatcher.logger.debug") as mock_debug: + EmissionsTracker(save_to_api=False, save_to_file=False) + EmissionsTracker(save_to_api=False, save_to_file=False) + notices = self._notices(mock_debug) + self.assertEqual(len(notices), 1) + + def test_notice_logged_at_debug_not_warning(self, mock_cli_setup): + """The notice is printed to stderr once; it must not also be logged + at WARNING, or it would show twice at warning-or-lower log levels.""" + self._mock_config(_conf()) + with ( + patch("sys.stderr"), + patch( + "codecarbon.core.telemetry.dispatcher.logger.warning" + ) as mock_warning, + ): + EmissionsTracker(save_to_api=False, save_to_file=False) + self.assertEqual(self._notices(mock_warning), []) + + def test_payload_built_off_the_stop_thread(self, mock_cli_setup): + self._mock_config(_conf("minimal")) + with ( + self._mock_post(), + patch( + "codecarbon.core.telemetry.dispatcher.build_payload", return_value={} + ) as mock_build, + ): + with ensure_telemetry_run_duration(): + tracker = EmissionsTracker( + measure_power_secs=1, save_to_api=False, save_to_file=False + ) + tracker.start() + with patch( + "codecarbon.core.telemetry.dispatcher.threading.Thread.start" + ): + tracker.stop() + mock_build.assert_not_called() + + def _timed_run(self, level: str, url: str): + """Run a tracker against ``url`` and time ``stop()`` alone.""" + env = { + "CODECARBON_TELEMETRY_LEVEL": level, + "CODECARBON_TELEMETRY_API_URL": url, + } + with patch.dict(os.environ, env), ensure_telemetry_run_duration(): + tracker = OfflineEmissionsTracker( + country_iso_code="CAN", save_to_api=False, save_to_file=False + ) + # Offline mode never sends telemetry (see + # test_offline_tracker_never_sends_telemetry); this test is only + # about the generic non-blocking behaviour of a send, so it + # re-enables telemetry on this instance directly. + from codecarbon.core.telemetry.schemas import TelemetryLevel + from codecarbon.core.telemetry.settings import TelemetrySettings + + tracker._telemetry = Telemetry( + TelemetrySettings( + level=TelemetryLevel(level), + is_explicit=True, + api_url=url, + ) + ) + tracker.start() + start = time.monotonic() + tracker.stop() + return tracker, start, time.monotonic() - start + + def test_stop_does_not_block_on_hanging_endpoint(self, mock_cli_setup): + """``stop()`` hands the send to a daemon thread, so it never waits on IO.""" + with hanging_endpoint() as url: + # The first stop() in a process pays one-off setup costs unrelated to + # telemetry, so warm those up before timing anything. + self._timed_run("disabled", url) + tracker, start, elapsed = self._timed_run("minimal", url) + thread = tracker._telemetry._thread + self.assertIsNotNone(thread) + # Synchronously this cost the full 2s request timeout. + self.assertLess(elapsed, 0.5) + # The send itself gives up once its time budget is spent. + thread.join(20) + self.assertFalse(thread.is_alive()) + self.assertLess(time.monotonic() - start, 4.0) + + def test_no_warning_when_level_set_by_kwarg(self, mock_cli_setup): + self._mock_config(_conf()) + with patch( + "codecarbon.core.telemetry.dispatcher.logger.warning" + ) as mock_warning: + EmissionsTracker( + telemetry_level="disabled", save_to_api=False, save_to_file=False + ) + self.assertEqual(self._notices(mock_warning), []) + + +class TestDispatcherEdges(unittest.TestCase): + def _telemetry(self): + from codecarbon.core.telemetry.dispatcher import Telemetry + from codecarbon.core.telemetry.settings import TelemetrySettings + + return Telemetry(TelemetrySettings.resolve(external_conf={})) + + def test_send_failure_is_swallowed(self): + from codecarbon.core.telemetry import dispatcher + + with ( + patch.object(dispatcher, "build_payload", side_effect=RuntimeError), + patch.object(dispatcher, "post_private") as mock_post, + ): + self._telemetry()._send(SimpleNamespace(), SimpleNamespace()) + mock_post.assert_not_called() + + def test_notice_shown_once_per_process_even_across_trackers(self): + """Many trackers in one process must only print the notice once.""" + from codecarbon.core.telemetry import dispatcher + + with patch.object(dispatcher.logger, "debug") as mock_debug: + self._telemetry().notice_once_if_implicit() + self._telemetry().notice_once_if_implicit() + mock_debug.assert_called_once() + + def test_notice_shown_again_in_a_fresh_process(self): + """A new process (simulated by ``_reset_after_fork``) sees it again.""" + from codecarbon.core.telemetry import dispatcher + + with patch.object(dispatcher.logger, "debug") as mock_debug: + self._telemetry().notice_once_if_implicit() + dispatcher._reset_after_fork() + self._telemetry().notice_once_if_implicit() + self.assertEqual(mock_debug.call_count, 2) + + def test_notice_never_shown_when_level_explicit(self): + from codecarbon.core.telemetry import dispatcher + from codecarbon.core.telemetry.settings import TelemetrySettings + + telemetry = dispatcher.Telemetry( + TelemetrySettings.resolve(external_conf={}, override="minimal") + ) + with patch.object(dispatcher.logger, "warning") as mock_warning: + telemetry.notice_once_if_implicit() + mock_warning.assert_not_called() + + def test_notice_never_shown_when_disabled(self): + from codecarbon.core.telemetry import dispatcher + from codecarbon.core.telemetry.settings import TelemetrySettings + + telemetry = dispatcher.Telemetry( + TelemetrySettings.resolve(external_conf={"telemetry_level": "disabled"}) + ) + with patch.object(dispatcher.logger, "warning") as mock_warning: + telemetry.notice_once_if_implicit() + mock_warning.assert_not_called() + + def test_exit_joins_pending_sends(self): + from codecarbon.core.telemetry import dispatcher + + thread = MagicMock() + thread.is_alive.return_value = True + with patch.object(dispatcher, "_pending", {thread}): + dispatcher._join_pending() + thread.join.assert_called_once() + + def test_thread_start_failure_does_not_crash_stop(self): + """A broken thread creation/start must never propagate out of stop().""" + from codecarbon.core.telemetry import dispatcher + + with ( + patch.object(dispatcher, "_sent", False), + patch.object( + dispatcher.threading, "Thread", side_effect=RuntimeError("boom") + ), + ): + telemetry = self._telemetry() + telemetry.send_at_stop( + SimpleNamespace(), SimpleNamespace(duration=5) + ) # must not raise + self.assertFalse(dispatcher._sent) + + def test_sent_marked_only_after_thread_actually_started(self): + from codecarbon.core.telemetry import dispatcher + + started = [] + real_thread_cls = dispatcher.threading.Thread + + class TrackingThread(real_thread_cls): + def start(self): + started.append(dispatcher._sent) + super().start() + + with ( + patch.object(dispatcher, "_sent", False), + patch.object(dispatcher.threading, "Thread", TrackingThread), + ): + telemetry = self._telemetry() + telemetry.send_at_stop(SimpleNamespace(), SimpleNamespace(duration=5)) + telemetry._thread.join(5) + # _sent was still False at the moment start() was called. + self.assertEqual(started, [False]) + self.assertTrue(dispatcher._sent) + + def test_fork_reset_clears_sent_flag_and_lock(self): + from codecarbon.core.telemetry import dispatcher + + dispatcher._sent = True + old_lock = dispatcher._sent_lock + try: + dispatcher._reset_after_fork() + self.assertFalse(dispatcher._sent) + self.assertIsNot(dispatcher._sent_lock, old_lock) + finally: + dispatcher._sent = False + + +class TestTelemetrySettings(unittest.TestCase): + def test_enum_level_passes_through(self): + from codecarbon.core.telemetry.schemas import TelemetryLevel + from codecarbon.core.telemetry.settings import parse_telemetry_level + + self.assertIs( + parse_telemetry_level(TelemetryLevel.disabled), TelemetryLevel.disabled + ) + + def test_legacy_extensive_level_maps_to_minimal(self): + from codecarbon.core.telemetry.settings import TelemetrySettings + + settings = TelemetrySettings.resolve( + external_conf={"telemetry_level": "extensive"} + ) + self.assertEqual(settings.level.value, "minimal") + + def test_unparseable_level_falls_back_to_disabled(self): + from codecarbon.core.telemetry.schemas import TelemetryLevel + from codecarbon.core.telemetry.settings import TelemetrySettings + + with patch("codecarbon.core.telemetry.settings.logger") as mock_logger: + settings = TelemetrySettings.resolve( + external_conf={"telemetry_level": "bogus"} + ) + self.assertIs(settings.level, TelemetryLevel.disabled) + mock_logger.error.assert_called_once() + + def test_api_endpoint_config_key_is_not_used_for_telemetry_url(self): + """``api_endpoint`` configures the dashboard experiment API + (``save_to_api``), a separate concern from telemetry; it must never + redirect telemetry to that host. + """ + from codecarbon.core.telemetry.settings import ( + DEFAULT_TELEMETRY_API_URL, + TelemetrySettings, + ) + + settings = TelemetrySettings.resolve( + external_conf={"api_endpoint": "https://dashboard.example.test"} + ) + self.assertEqual(settings.api_url, DEFAULT_TELEMETRY_API_URL) + + def test_empty_config_value_falls_back_to_disabled_not_default(self): + """An empty ``telemetry_level =`` (or ``CODECARBON_TELEMETRY_LEVEL=``) + must go through the invalid branch to disabled, not be treated as + unset and silently resolve to the default ``minimal``. + """ + from codecarbon.core.telemetry.schemas import TelemetryLevel + from codecarbon.core.telemetry.settings import TelemetrySettings + + settings = TelemetrySettings.resolve(external_conf={"telemetry_level": ""}) + self.assertIs(settings.level, TelemetryLevel.disabled) + self.assertTrue(settings.is_explicit) + + def test_privacy_intent_strings_fall_back_to_disabled_not_minimal(self): + """off/false/none/0 must never resolve to a level that sends data.""" + from codecarbon.core.telemetry.schemas import TelemetryLevel + from codecarbon.core.telemetry.settings import TelemetrySettings + + for value in ("off", "false", "none", "0"): + with self.subTest(value=value): + settings = TelemetrySettings.resolve( + external_conf={"telemetry_level": value} + ) + self.assertIs(settings.level, TelemetryLevel.disabled) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_telemetry_client.py b/tests/test_telemetry_client.py index 14ade8b7f..3a850da65 100644 --- a/tests/test_telemetry_client.py +++ b/tests/test_telemetry_client.py @@ -1,104 +1,79 @@ +import time import unittest +from unittest.mock import patch import requests_mock from pydantic import ValidationError -from codecarbon.core.telemetry_client import TelemetryClient -from codecarbon.core.telemetry_schemas import TelemetryCreate +from codecarbon.core.telemetry.client import post_private +from codecarbon.core.telemetry.schemas import TelemetryLevel +from codecarbon.core.telemetry.settings import TelemetrySettings +TELEMETRY = { + "timestamp": "2026-05-03T12:00:00+00:00", + "telemetry_level": "minimal", + "os": "Linux-5.10.0-x86_64", +} -class TestTelemetryClient(unittest.TestCase): - def test_init_sets_up_client_without_calling_api(self): - with requests_mock.Mocker() as m: - client = TelemetryClient( - endpoint_url="http://test.com/", - telemetry={ - "timestamp": "2026-05-03T12:00:00+00:00", - "telemetry_level": "minimal", - }, - ) - self.assertEqual(client.endpoint_url, "http://test.com") - self.assertEqual(client.telemetry_url, "http://test.com/telemetry") - self.assertIsInstance(client.telemetry, TelemetryCreate) - self.assertEqual(m.call_count, 0) +def _settings(api_url: str = "http://test.com") -> TelemetrySettings: + return TelemetrySettings( + level=TelemetryLevel.minimal, is_explicit=False, api_url=api_url + ) - def test_add_telemetry_posts_configured_payload(self): - telemetry = { - "timestamp": "2026-05-03T12:00:00+00:00", - "telemetry_level": "minimal", - "os": "Linux-5.10.0-x86_64", - } +class TestPostPrivate(unittest.TestCase): + def test_post_private_sends_validated_payload_without_token(self): with requests_mock.Mocker() as m: m.post( "http://test.com/telemetry", json="f52fe339-164d-4c2b-a8c0-f562dfce066d", status_code=201, ) - client = TelemetryClient( - endpoint_url="http://test.com", telemetry=telemetry - ) - - actual_telemetry_id = client.add_telemetry() + result = post_private(_settings(), TELEMETRY) - self.assertEqual( - actual_telemetry_id, "f52fe339-164d-4c2b-a8c0-f562dfce066d" - ) + self.assertTrue(result) self.assertEqual(m.call_count, 1) self.assertEqual( m.last_request.json(), - { - **telemetry, - "timestamp": "2026-05-03T12:00:00Z", - }, - ) - - def test_add_telemetry_posts_call_payload(self): - telemetry = TelemetryCreate( - timestamp="2026-05-03T12:00:00+00:00", - telemetry_level="minimal", - os="Linux-5.10.0-x86_64", - ) - - with requests_mock.Mocker() as m: - m.post( - "http://test.com/telemetry", - json="f52fe339-164d-4c2b-a8c0-f562dfce066d", - status_code=201, + {**TELEMETRY, "timestamp": "2026-05-03T12:00:00Z"}, ) - client = TelemetryClient(endpoint_url="http://test.com") - - actual_telemetry_id = client.add_telemetry(telemetry) - - self.assertEqual( - actual_telemetry_id, "f52fe339-164d-4c2b-a8c0-f562dfce066d" - ) - self.assertEqual(m.call_count, 1) - self.assertEqual( - m.last_request.json(), - { - "timestamp": "2026-05-03T12:00:00Z", - "telemetry_level": "minimal", - "os": "Linux-5.10.0-x86_64", - }, - ) - - def test_init_rejects_invalid_telemetry_without_calling_api(self): - with requests_mock.Mocker() as m: - with self.assertRaises(ValidationError): - TelemetryClient( - endpoint_url="http://test.com", - telemetry={ - "timestamp": "2026-05-03T12:00:00+00:00", - "telemetry_level": "minimal", - "total_emissions_kg": 0.42, - }, - ) - - self.assertEqual(m.call_count, 0) - - def test_add_telemetry_returns_none_without_payload(self): - client = TelemetryClient(endpoint_url="http://test.com") - - self.assertIsNone(client.add_telemetry()) + self.assertNotIn("x-api-token", m.last_request.headers) + + def test_post_private_rejects_invalid_payload(self): + with self.assertRaises(ValidationError): + post_private(_settings(), {**TELEMETRY, "unknown_field": "value"}) + + def test_post_private_logs_debug_on_non_201(self): + for status_code in (404, 429, 500): + with self.subTest(status_code=status_code): + with requests_mock.Mocker() as m: + m.post( + "http://test.com/telemetry", + text="nope", + status_code=status_code, + ) + with patch( + "codecarbon.core.telemetry.client.logger" + ) as mock_logger: + result = post_private(_settings(), TELEMETRY) + self.assertFalse(result) + mock_logger.debug.assert_called_once() + mock_logger.warning.assert_not_called() + + def test_post_private_returns_false_on_request_error(self): + with patch("codecarbon.core.telemetry.client.requests.post") as mock_post: + mock_post.side_effect = ConnectionError("network down") + with patch("codecarbon.core.telemetry.client.logger"): + result = post_private(_settings(), TELEMETRY) + self.assertFalse(result) + + def test_post_private_skips_when_deadline_passed(self): + with patch("codecarbon.core.telemetry.client.requests.post") as mock_post: + result = post_private(_settings(), TELEMETRY, deadline=time.monotonic() - 1) + self.assertFalse(result) + mock_post.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_telemetry_collect.py b/tests/test_telemetry_collect.py new file mode 100644 index 000000000..03d7fbe7e --- /dev/null +++ b/tests/test_telemetry_collect.py @@ -0,0 +1,284 @@ +import os +import sys +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from codecarbon.core.telemetry.collect import build_payload +from codecarbon.core.telemetry.schemas import ( + TelemetryBase, + TelemetryCreate, + TelemetryLevel, +) +from codecarbon.output_methods.emissions_data import EmissionsData + + +def _sample_emissions(**overrides): + base = dict( + timestamp="2026-01-01T00:00:00", + project_name="p", + run_id="r", + experiment_id="e", + duration=10.0, + emissions=0.5, + emissions_rate=0.05, + cpu_power=1.0, + gpu_power=2.0, + ram_power=0.5, + cpu_energy=0.01, + gpu_energy=0.02, + ram_energy=0.001, + energy_consumed=0.031, + water_consumed=0.0, + country_name="France", + country_iso_code="FRA", + region="idf", + cloud_provider="", + cloud_region="", + os="Linux", + python_version="3.11", + codecarbon_version="3.0", + cpu_count=4, + cpu_model="cpu", + gpu_count=1, + gpu_model="gpu", + longitude=0.0, + latitude=0.0, + ram_total_size=16.0, + tracking_mode="machine", + ) + base.update(overrides) + return EmissionsData(**base) + + +def _tracker_context(**overrides): + """Return a (tracker, emissions) pair standing in for a live tracker.""" + tracker = SimpleNamespace( + _conf=overrides.pop("conf", {"codecarbon_version": "3.0"}) + ) + return tracker, overrides.pop("emissions", _sample_emissions()) + + +def _build(ctx, level=TelemetryLevel.minimal): + tracker, emissions = ctx + return build_payload(tracker, emissions, level=level) + + +class TestTelemetryCollect(unittest.TestCase): + def test_build_payload_minimal_omits_run_metrics(self): + ctx = _tracker_context( + conf={ + "os": "Linux", + "codecarbon_version": "3.0", + "cpu_count": 4, + "tracking_mode": "machine", + }, + ) + payload = _build(ctx, level=TelemetryLevel.minimal) + + self.assertEqual(payload["telemetry_level"], "minimal") + self.assertNotIn("total_emissions_kg", payload) + + def test_timestamp_is_truncated_to_the_hour(self): + from datetime import datetime, timezone + + ctx = _tracker_context(conf={"codecarbon_version": "3.0"}) + with patch("codecarbon.core.telemetry.collect.datetime") as mock_datetime: + mock_datetime.now.return_value = datetime( + 2026, 1, 1, 12, 34, 56, 789, tzinfo=timezone.utc + ) + payload = _build(ctx, level=TelemetryLevel.minimal) + self.assertEqual( + payload["timestamp"], + datetime(2026, 1, 1, 12, 0, 0, 0, tzinfo=timezone.utc), + ) + + def test_minimal_payload_passes_schema_validation(self): + ctx = _tracker_context(conf={"os": "Linux", "codecarbon_version": "3.0"}) + payload = _build(ctx, level=TelemetryLevel.minimal) + TelemetryCreate(**payload) + + def test_payload_keys_are_schema_fields(self): + ctx = _tracker_context(conf={"codecarbon_version": "3.0"}) + payload = _build(ctx) + self.assertTrue(set(payload).issubset(TelemetryBase.model_fields)) + + def test_cloud_fields_come_from_tracker_detection(self): + """No second metadata probe: stop() must not pay for it off-cloud.""" + emissions = _sample_emissions( + on_cloud="Y", cloud_provider="aws", cloud_region="eu-west-1", region="" + ) + payload = _build( + _tracker_context(emissions=emissions), level=TelemetryLevel.minimal + ) + self.assertEqual(payload["cloud_provider"], "aws") + self.assertEqual(payload["cloud_region"], "eu-west-1") + self.assertEqual(payload["region"], "eu-west-1") + + def test_default_geo_fallback_is_not_reported_as_a_location(self): + """When both geo APIs fail the tracker assumes Canada; telemetry must + not record that guess as where the user is.""" + tracker, emissions = _tracker_context( + emissions=_sample_emissions( + country_name="Canada", country_iso_code="CAN", region="quebec" + ) + ) + tracker._geo = SimpleNamespace(is_default=True) + payload = build_payload(tracker, emissions, level=TelemetryLevel.minimal) + self.assertNotIn("country_name", payload) + self.assertNotIn("country_iso_code", payload) + self.assertNotIn("region", payload) + + def test_default_geo_fallback_keeps_cloud_location(self): + """On cloud the location comes from the cloud region, not from IP + geolocation, so a failed geo lookup must not drop it.""" + tracker, emissions = _tracker_context( + emissions=_sample_emissions( + on_cloud="Y", + cloud_provider="aws", + cloud_region="eu-west-3", + country_name="France", + country_iso_code="FRA", + region="eu-west-3", + ) + ) + tracker._geo = SimpleNamespace(is_default=True) + payload = build_payload(tracker, emissions, level=TelemetryLevel.minimal) + self.assertEqual(payload["country_name"], "France") + self.assertEqual(payload["country_iso_code"], "FRA") + self.assertEqual(payload["region"], "eu-west-3") + + def test_coordinates_are_never_sent(self): + emissions = _sample_emissions(longitude=-7.61743, latitude=33.58229) + payload = _build( + _tracker_context( + conf={"longitude": 2.35, "latitude": 48.85}, emissions=emissions + ) + ) + self.assertNotIn("longitude", payload) + self.assertNotIn("latitude", payload) + + def test_gpu_static_fields_when_nvidia_available(self): + ctx = _tracker_context() + mock_mem = MagicMock(total=8 * 1024**3) + mock_pynvml = MagicMock() + mock_pynvml.nvmlDeviceGetMemoryInfo.return_value = mock_mem + mock_pynvml.nvmlSystemGetCudaDriverVersion_v2.return_value = 12040 + mock_pynvml.nvmlSystemGetDriverVersion.return_value = "535.0" + with patch( + "codecarbon.core.telemetry.collect.is_nvidia_system", + return_value=True, + ): + with patch.dict(sys.modules, {"pynvml": mock_pynvml}): + payload = _build(ctx, level=TelemetryLevel.minimal) + self.assertEqual(payload["gpu_memory_total_gb"], 8.0) + self.assertEqual(payload["cuda_version"], "12.4") + self.assertEqual(payload["gpu_driver_version"], "535.0") + + def test_minimal_payload_detects_virtualenv(self): + ctx = _tracker_context() + with patch.dict(os.environ, {"VIRTUAL_ENV": "/venv"}, clear=False): + with patch( + "codecarbon.core.telemetry.collect.sys.prefix", + "/venv", + create=True, + ): + with patch( + "codecarbon.core.telemetry.collect.sys.base_prefix", + "/usr", + create=True, + ): + payload = _build(ctx, level=TelemetryLevel.minimal) + self.assertEqual(payload["python_env_type"], "venv") + + def test_python_env_type_variants(self): + from codecarbon.core.telemetry import collect + + cases = [ + ({"CONDA_DEFAULT_ENV": "base"}, "/usr", "conda"), + ({}, "/venv", "venv"), + ({}, "/usr", "system"), + ] + for env, prefix, expected in cases: + with self.subTest(expected=expected): + with ( + patch.dict(os.environ, env, clear=True), + patch.object(collect.sys, "prefix", prefix), + patch.object(collect.sys, "base_prefix", "/usr"), + ): + self.assertEqual(collect._detect_python_env_type(), expected) + + def test_codecarbon_install_method(self): + from codecarbon.core.telemetry import collect + + def dist(direct_url=None, installer=None): + return SimpleNamespace( + read_text=lambda name: ( + direct_url if name == "direct_url.json" else None + ), + metadata={"Installer": installer} if installer else {}, + ) + + editable_url = '{"url": "file:///src", "dir_info": {"editable": true}}' + local_url = '{"url": "file:///src", "dir_info": {}}' + cases = [ + (dist(direct_url=editable_url, installer="uv"), "editable"), + (dist(direct_url=local_url, installer="pip"), "pip"), + (dist(installer="uv"), "uv"), + (dist(installer="pip"), "pip"), + (dist(installer="poetry"), None), + (RuntimeError("no metadata"), None), + ] + for result, expected in cases: + with self.subTest(expected=expected, result=result): + with patch( + "importlib.metadata.distribution", + side_effect=[result], + ): + self.assertEqual( + collect._detect_codecarbon_install_method(), expected + ) + + def test_cudnn_version(self): + from codecarbon.core.telemetry import collect + + torch_ok = SimpleNamespace( + backends=SimpleNamespace(cudnn=SimpleNamespace(version=lambda: 8902)) + ) + torch_none = SimpleNamespace( + backends=SimpleNamespace(cudnn=SimpleNamespace(version=lambda: None)) + ) + torch_broken = SimpleNamespace(backends=None) + cases = [(torch_ok, "8902"), (torch_none, None), (torch_broken, None)] + for torch, expected in cases: + with self.subTest(expected=expected): + with patch.dict(sys.modules, {"torch": torch}): + self.assertEqual(collect._cudnn_version(), expected) + + def test_cudnn_version_does_not_import_torch(self): + """Importing torch is slow; only read it if the user already did.""" + from codecarbon.core.telemetry import collect + + with patch.dict(sys.modules): + sys.modules.pop("torch", None) + with patch("builtins.__import__") as mock_import: + self.assertIsNone(collect._cudnn_version()) + mock_import.assert_not_called() + self.assertNotIn("torch", sys.modules) + + def test_gpu_static_fields_empty_when_nvml_fails(self): + mock_pynvml = MagicMock() + mock_pynvml.nvmlInit.side_effect = RuntimeError("no driver") + with ( + patch( + "codecarbon.core.telemetry.collect.is_nvidia_system", + return_value=True, + ), + patch.dict(sys.modules, {"pynvml": mock_pynvml}), + ): + payload = _build(_tracker_context()) + self.assertNotIn("gpu_driver_version", payload) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_telemetry_payload_roundtrip.py b/tests/test_telemetry_payload_roundtrip.py new file mode 100644 index 000000000..9de8fac7e --- /dev/null +++ b/tests/test_telemetry_payload_roundtrip.py @@ -0,0 +1,129 @@ +"""Verify SDK telemetry payloads match the server schema and shared defaults.""" + +import importlib.util +import unittest +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +from codecarbon.core.telemetry.client import post_private +from codecarbon.core.telemetry.collect import build_payload +from codecarbon.core.telemetry.schemas import TelemetryCreate as CoreTelemetryCreate +from codecarbon.core.telemetry.schemas import TelemetryLevel +from codecarbon.core.telemetry.settings import TelemetrySettings +from codecarbon.output_methods.emissions_data import EmissionsData + +REPO_ROOT = Path(__file__).resolve().parents[1] +SERVER_TELEMETRY_SCHEMA_PATH = ( + REPO_ROOT / "carbonserver" / "carbonserver" / "api" / "schemas_telemetry.py" +) + + +requires_server_tree = unittest.skipUnless( + SERVER_TELEMETRY_SCHEMA_PATH.exists(), + "carbonserver source tree is not available (e.g. install-from-wheel test run)", +) + + +def _load_server_telemetry_create(): + spec = importlib.util.spec_from_file_location( + "server_telemetry_schemas", + SERVER_TELEMETRY_SCHEMA_PATH, + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.TelemetryCreate + + +def _sample_emissions() -> EmissionsData: + return EmissionsData( + timestamp="2026-01-01T00:00:00", + project_name="p", + run_id="r", + experiment_id="e", + duration=10.0, + emissions=0.5, + emissions_rate=0.05, + cpu_power=1.0, + gpu_power=2.0, + ram_power=0.5, + cpu_energy=0.01, + gpu_energy=0.02, + ram_energy=0.001, + energy_consumed=0.031, + water_consumed=0.0, + country_name="France", + country_iso_code="FRA", + region="idf", + cloud_provider="", + cloud_region="", + os="Linux-5.10.0-x86_64", + python_version="3.11.5", + codecarbon_version="3.2.6", + cpu_count=12, + cpu_model="Intel(R) Core(TM) i7-8850H CPU @ 2.60GHz", + gpu_count=0, + gpu_model="", + longitude=2.3, + latitude=48.8, + ram_total_size=16.0, + tracking_mode="machine", + ) + + +def _tracker_context(**overrides): + conf = overrides.pop( + "conf", + { + "os": "Linux-5.10.0-x86_64", + "codecarbon_version": "3.2.6", + "cpu_count": 12, + "cpu_model": "Intel(R) Core(TM) i7-8850H CPU @ 2.60GHz", + "python_version": "3.11.5", + "tracking_mode": "machine", + }, + ) + tracker = SimpleNamespace(_conf=conf) + return tracker, overrides.pop("emissions", _sample_emissions()) + + +def _sdk_request_body() -> dict: + tracker, emissions = _tracker_context() + payload = build_payload(tracker, emissions, level=TelemetryLevel.minimal) + return CoreTelemetryCreate(**payload).model_dump(mode="json", exclude_none=True) + + +@requires_server_tree +class TestTelemetryPayloadContract(unittest.TestCase): + def test_minimal_sdk_payload_is_accepted_by_server_schema(self): + ServerTelemetryCreate = _load_server_telemetry_create() + request_body = _sdk_request_body() + + parsed = ServerTelemetryCreate.model_validate(request_body) + + self.assertEqual(parsed.telemetry_level, "minimal") + self.assertEqual(parsed.os, request_body["os"]) + self.assertEqual(parsed.country_name, request_body["country_name"]) + self.assertNotIn("latitude", request_body) + + def test_post_private_sends_round_tripped_body_without_token(self): + ServerTelemetryCreate = _load_server_telemetry_create() + request_body = _sdk_request_body() + settings = TelemetrySettings.resolve( + external_conf={"telemetry_api_url": "http://test.example"} + ) + + with patch("codecarbon.core.telemetry.client.requests.post") as mock_post: + mock_post.return_value.status_code = 201 + self.assertTrue(post_private(settings, deepcopy(request_body))) + + self.assertNotIn("headers", mock_post.call_args.kwargs) + sent_body = mock_post.call_args.kwargs["json"] + ServerTelemetryCreate.model_validate(sent_body) + self.assertEqual(sent_body["telemetry_level"], "minimal") + self.assertEqual(sent_body["os"], request_body["os"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_telemetry_schemas.py b/tests/test_telemetry_schemas.py new file mode 100644 index 000000000..18dc04ead --- /dev/null +++ b/tests/test_telemetry_schemas.py @@ -0,0 +1,45 @@ +import unittest +from datetime import datetime, timezone + +import pytest +from pydantic import ValidationError + +from codecarbon.core.telemetry.schemas import TelemetryCreate, TelemetryLevel + + +class TestTelemetrySchemaValidation(unittest.TestCase): + def _minimal_payload(self) -> dict: + return { + "timestamp": datetime(2026, 1, 1, tzinfo=timezone.utc), + "telemetry_level": TelemetryLevel.minimal, + "os": "Linux", + } + + def test_rejects_disabled_telemetry_level(self): + with pytest.raises( + ValidationError, match="Disabled telemetry must not be submitted" + ): + TelemetryCreate( + **{**self._minimal_payload(), "telemetry_level": "disabled"} + ) + + def test_rejects_unknown_field(self): + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + TelemetryCreate( + **{ + **self._minimal_payload(), + "total_emissions_kg": 0.5, + } + ) + + def test_rejects_oversized_string(self): + with pytest.raises(ValidationError, match="at most 256 characters"): + TelemetryCreate(**{**self._minimal_payload(), "cpu_model": "x" * 257}) + + def test_rejects_coordinates(self): + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + TelemetryCreate(**{**self._minimal_payload(), "latitude": 48.8}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/testutils.py b/tests/testutils.py index e3d1dc2d1..fa14c1128 100644 --- a/tests/testutils.py +++ b/tests/testutils.py @@ -1,6 +1,8 @@ import builtins import unittest +from contextlib import contextmanager from pathlib import Path +from unittest.mock import patch from codecarbon.input import DataSource @@ -33,3 +35,43 @@ def conditional_open_func(path, *args, **kwargs): return conditional_open_func return mocked_open + + +def join_telemetry(tracker, timeout: float = 10.0) -> None: + """Wait for the background telemetry send that ``stop()`` fired, if any.""" + thread = getattr(getattr(tracker, "_telemetry", None), "_thread", None) + if thread is not None: + thread.join(timeout) + + +@contextmanager +def hanging_endpoint(): + """Yield a URL whose TCP connections are accepted but never answered.""" + import socket + + server = socket.socket() + server.bind(("127.0.0.1", 0)) + server.listen(8) # backlog completes the handshake; we never accept() + try: + yield "http://{}:{}".format(*server.getsockname()) + finally: + server.close() + + +@contextmanager +def ensure_telemetry_run_duration(min_seconds: float = 10.0): + """Force tracker stop emissions duration above telemetry's 1s minimum.""" + from codecarbon.emissions_tracker import BaseEmissionsTracker + + original_prepare = BaseEmissionsTracker._prepare_emissions_data + + def prepare_with_min_duration(self): + data = original_prepare(self) + if data is not None and (data.duration is None or data.duration < min_seconds): + data.duration = min_seconds + return data + + with patch.object( + BaseEmissionsTracker, "_prepare_emissions_data", prepare_with_min_duration + ): + yield