feat(api): reuse connections and retry transient API failures

codecarbon/core/api_client.py

@@ -1,16 +1,16 @@
 """
 
 Based on https://kernelpanic.io/the-modern-way-to-call-apis-in-python
-
-TODO : use async call to API
 """
 
-# from httpx import AsyncClient
 import dataclasses
 import json
+import time
 from datetime import datetime, timedelta, tzinfo
 
 import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 
 from codecarbon.core.schemas import (
     EmissionCreate,
@@ -44,6 +44,51 @@ def _measurement_timestamp(carbon_emission: dict) -> str:
         return get_datetime_with_timezone()
 
 
+# Failures where the request plausibly never reached the application, so a
+# retry cannot duplicate work: no upstream was reachable (502/503) or we were
+# told to slow down before being served (429).
+_POST_SAFE_STATUSES = (429, 502, 503)
+# Reads are idempotent, so a retry costs at most a wasted round trip.
+_READ_SAFE_STATUSES = (429, 500, 502, 503, 504)
+
+
+def _build_session(
+    retries: int, backoff: float, statuses, retry_read: bool
+) -> requests.Session:
+    """
+    A Session so sockets (and the TLS handshake) are reused across calls, with
+    retry and exponential backoff on the failures that are worth retrying.
+
+    :statuses: response codes to retry.
+    :retry_read: whether to retry a read timeout or a truncated response, i.e.
+        a failure that happened *after* the request reached the server.
+    """
+    retry_kwargs = {
+        "total": retries,
+        "connect": retries,
+        # False, not 0: urllib3 then re-raises the original ReadTimeout instead
+        # of burning the budget and reporting an exhausted-retries error.
+        "read": retries if retry_read else False,
+        "status": retries,
+        "backoff_factor": backoff,
+        "status_forcelist": statuses,
+        "allowed_methods": frozenset({"GET", "POST", "PATCH", "PUT", "DELETE"}),
+        "raise_on_status": False,
+    }
+    try:
+        # Spreads the retry storm when a whole fleet reconnects after an outage.
+        retry = Retry(backoff_jitter=1.0, **retry_kwargs)
+    except TypeError:
+        # urllib3 < 2.0 has no backoff_jitter. We are a library in other
+        # people's environments, so degrade instead of pinning urllib3.
+        retry = Retry(**retry_kwargs)
+    adapter = HTTPAdapter(max_retries=retry, pool_maxsize=4)
+    session = requests.Session()
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+    return session
+
+
 class ApiClient:  # (AsyncClient)
     """
     This class call the Code Carbon API
@@ -59,6 +104,9 @@ def __init__(
         access_token=None,
         conf=None,
         create_run_automatically=True,
+        timeout=(3.05, 10),
+        retries=2,
+        backoff=0.5,
     ):
         """
         :endpoint_url: URL of the API endpoint
@@ -67,13 +115,47 @@ def __init__(
         :access_token: Code Carbon API access token
         :conf: Metadata of the experiment
         :create_run_automatically: If False, do not create a run. To use API in read only mode.
+        :timeout: requests timeout, either seconds or a (connect, read) tuple.
+        :retries: number of retries after the first attempt.
+        :backoff: backoff factor between retries, in seconds (0.5 -> 0.5s, 1s, 2s...).
+
+        Note that `timeout` and `retries` multiply, so raise one only while
+        looking at the other. Two different worst cases are worth keeping apart:
+
+        - a *hung* endpoint, which accepts the connection and never answers:
+          `read * (retries + 1)` plus backoff, since only the read times out.
+        - the *full* retry chain, where the connection also has to time out:
+          `(connect + read) * (retries + 1)` plus backoff.
+
+        POSTs do not retry read timeouts (see `_post_session`), so their hung
+        case is a single `connect + read` and the chain above is a GET bound.
         """
-        # super().__init__(base_url=endpoint_url) # (AsyncClient)
         self.url = endpoint_url
         self.experiment_id = experiment_id
         self.api_key = api_key
         self.conf = conf
         self.access_token = access_token
+        self._timeout = timeout
+        self._session = _build_session(
+            retries, backoff, _READ_SAFE_STATUSES, retry_read=True
+        )
+        # POSTs create rows. carbonserver has no idempotency key and the
+        # dashboard sums emission rows, so a POST replayed after the server
+        # already committed the insert inflates a user's reported emissions
+        # with nothing to show for it: a dropped row is visible, a duplicated
+        # one is not. This session therefore only retries POSTs that
+        # plausibly never reached the application -- connection errors,
+        # connect timeouts, 429/502/503. Read timeouts, truncated responses,
+        # 500 and 504 are *not* retried: the request landed, and the insert
+        # may well have gone through. Widen this only once the API accepts an
+        # idempotency key.
+        self._post_session = _build_session(
+            retries, backoff, _POST_SAFE_STATUSES, retry_read=False
+        )
+        # Run creation is the most expensive write path. When it fails, back off
+        # instead of re-attempting on every measurement tick.
+        self._create_run_not_before = 0.0
+        self._create_run_backoff = 0.0
         if self.experiment_id is not None and create_run_automatically:
             self._create_run(self.experiment_id)
 
@@ -96,11 +178,16 @@ def _request(self, method, url, payload=None, expected_status=200):
         :expected_status: the http code the API returns when the call succeeds
         """
         headers = self._get_headers()
-        response = method(url=url, json=payload, timeout=2, headers=headers)
+        response = method(url=url, json=payload, timeout=self._timeout, headers=headers)
         if response.status_code != expected_status:
             self._raise_api_error(url, payload or {}, response)
         return response
 
+    def close(self):
+        """Release the pooled sockets. Safe to call more than once."""
+        self._session.close()
+        self._post_session.close()
+
     def set_access_token(self, token: str):
         """This method sets the access token to be used for the API.
         Args:
@@ -113,14 +200,14 @@ def check_auth(self):
         Check API access to user account
         """
         url = self.url + "/auth/check"
-        return self._request(requests.get, url).json()
+        return self._request(self._session.get, url).json()
 
     def get_list_organizations(self):
         """
         List all organizations
         """
         url = self.url + "/organizations"
-        return self._request(requests.get, url).json()
+        return self._request(self._session.get, url).json()
 
     def check_organization_exists(self, organization_name: str):
         """
@@ -145,30 +232,30 @@ def create_organization(self, organization: OrganizationCreate):
             return organization
         else:
             return self._request(
-                requests.post, url, payload=payload, expected_status=201
+                self._post_session.post, url, payload=payload, expected_status=201
             ).json()
 
     def get_organization(self, organization_id):
         """
         Get an organization
         """
         url = self.url + "/organizations/" + organization_id
-        return self._request(requests.get, url).json()
+        return self._request(self._session.get, url).json()
 
     def update_organization(self, organization: OrganizationCreate):
         """
         Update an organization
         """
         payload = dataclasses.asdict(organization)
         url = self.url + "/organizations/" + organization.id
-        return self._request(requests.patch, url, payload=payload).json()
+        return self._request(self._session.patch, url, payload=payload).json()
 
     def list_projects_from_organization(self, organization_id):
         """
         List all projects
         """
         url = self.url + "/organizations/" + organization_id + "/projects"
-        return self._request(requests.get, url).json()
+        return self._request(self._session.get, url).json()
 
     def create_project(self, project: ProjectCreate):
         """
@@ -177,15 +264,15 @@ def create_project(self, project: ProjectCreate):
         payload = dataclasses.asdict(project)
         url = self.url + "/projects"
         return self._request(
-            requests.post, url, payload=payload, expected_status=201
+            self._post_session.post, url, payload=payload, expected_status=201
         ).json()
 
     def get_project(self, project_id):
         """
         Get a project
         """
         url = self.url + "/projects/" + project_id
-        return self._request(requests.get, url).json()
+        return self._request(self._session.get, url).json()
 
     def add_emission(self, carbon_emission: dict):
         assert self.experiment_id is not None
@@ -226,7 +313,9 @@ def add_emission(self, carbon_emission: dict):
         try:
             payload = dataclasses.asdict(emission)
             url = self.url + "/emissions"
-            self._request(requests.post, url, payload=payload, expected_status=201)
+            self._request(
+                self._post_session.post, url, payload=payload, expected_status=201
+            )
             logger.debug(f"ApiClient - Successful upload emission {payload} to {url}")
         except requests.exceptions.HTTPError:
             # Already logged by _raise_api_error, do not log it twice.
@@ -236,7 +325,39 @@ def add_emission(self, carbon_emission: dict):
             raise
         return True
 
+    # Bounds on the run-creation retry delay, in seconds.
+    _CREATE_RUN_BACKOFF_MIN = 30.0
+    _CREATE_RUN_BACKOFF_MAX = 900.0
+
     def _create_run(self, experiment_id: str):
+        """
+        Create a run, backing off after a failure.
+
+        Without the backoff every client in a fleet re-attempts run creation on
+        every measurement tick for as long as the API is down, which hammers the
+        most expensive write path exactly when it is least able to take it.
+        Returns None without calling the API while the backoff is in effect.
+        """
+        if time.monotonic() < self._create_run_not_before:
+            logger.debug(
+                "ApiClient run creation is backing off after a previous failure, "
+                "skipping this attempt."
+            )
+            return None
+        try:
+            run_id = self._create_run_once(experiment_id)
+        except Exception:
+            self._create_run_backoff = min(
+                max(self._create_run_backoff * 2, self._CREATE_RUN_BACKOFF_MIN),
+                self._CREATE_RUN_BACKOFF_MAX,
+            )
+            self._create_run_not_before = time.monotonic() + self._create_run_backoff
+            raise
+        self._create_run_backoff = 0.0
+        self._create_run_not_before = 0.0
+        return run_id
+
+    def _create_run_once(self, experiment_id: str):
         """
         Create the experiment for project_id
         """
@@ -269,7 +390,9 @@ def _create_run(self, experiment_id: str):
             )
             payload = dataclasses.asdict(run)
             url = self.url + "/runs"
-            r = self._request(requests.post, url, payload=payload, expected_status=201)
+            r = self._request(
+                self._post_session.post, url, payload=payload, expected_status=201
+            )
             self.run_id = r.json()["id"]
             logger.info(
                 "ApiClient Successfully registered your run on the API.\n\n"
@@ -295,7 +418,7 @@ def list_experiments_from_project(self, project_id: str):
         List all experiments for a project
         """
         url = self.url + "/projects/" + project_id + "/experiments"
-        return self._request(requests.get, url).json()
+        return self._request(self._session.get, url).json()
 
     def set_experiment(self, experiment_id: str):
         """
@@ -311,15 +434,15 @@ def add_experiment(self, experiment: ExperimentCreate):
         payload = dataclasses.asdict(experiment)
         url = self.url + "/experiments"
         return self._request(
-            requests.post, url, payload=payload, expected_status=201
+            self._post_session.post, url, payload=payload, expected_status=201
         ).json()
 
     def get_experiment(self, experiment_id):
         """
         Get an experiment by id
         """
         url = self.url + "/experiments/" + experiment_id
-        return self._request(requests.get, url).json()
+        return self._request(self._session.get, url).json()
 
     def _raise_api_error(self, url, payload, response):
         """

codecarbon/emissions_tracker.py

@@ -392,6 +392,8 @@ def __init__(
         api_call_interval: Optional[int] = _sentinel,
         api_endpoint: Optional[str] = _sentinel,
         api_key: Optional[str] = _sentinel,
+        api_timeout: Optional[float] = _sentinel,
+        api_retries: Optional[int] = _sentinel,
         output_dir: Optional[str] = _sentinel,
         output_file: Optional[str] = _sentinel,
         output_methods: Optional[List[OutputMethod]] = _sentinel,
@@ -437,6 +439,13 @@ def __init__(
         :param api_endpoint: Optional URL of Code Carbon API endpoint for sending
                              emissions data.
         :param api_key: API key for Code Carbon API (mandatory!).
+        :param api_timeout: Read timeout in seconds for API calls (default: 5).
+                            Connect timeout is fixed at 3.05s.
+        :param api_retries: Retries after the first failed API attempt (default: 2).
+                            An emission POST blocks a measurement for at most
+                            about 3.05 * api_retries + (3.05 + api_timeout),
+                            ~16s with the defaults: only connect failures are
+                            retried on POST, so read timeouts happen once.
         :param output_dir: Directory path to which the experiment details are logged,
                            defaults to current directory.
         :param output_file: Name of the output CSV file, defaults to `emissions.csv`.
@@ -554,6 +563,8 @@ def __init__(
         self._set_from_conf(api_call_interval, "api_call_interval", 8, int)
         self._set_from_conf(api_endpoint, "api_endpoint", "https://api.codecarbon.io")
         self._set_from_conf(api_key, "api_key", "api_key")
+        self._set_from_conf(api_timeout, "api_timeout", 5, float)
+        self._set_from_conf(api_retries, "api_retries", 2, int)
         self._configure_electricitymaps_token(
             electricitymaps_api_token, co2_signal_api_token
         )
@@ -644,6 +655,8 @@ def _init_output_methods(self, *, api_key: str = None):
                 experiment_id=self._experiment_id,
                 api_key=api_key,
                 conf=self._conf,
+                timeout=(3.05, self._api_timeout),
+                retries=self._api_retries,
             )
             self.run_id = cc_api__out.run_id
             self._output_handlers.append(cc_api__out)

codecarbon/output_methods/http.py

@@ -46,17 +46,47 @@ def __init__(
         experiment_id: str,
         api_key: str,
         conf,
+        timeout=(3.05, 5),
+        retries: int = 2,
+        backoff: float = 0.5,
     ):
+        """
+        The timeout and retry budget here is deliberately tighter than
+        ApiClient's own default: this send runs inline on the scheduler thread,
+        so a slow failure delays the next measurement.
+
+        With the defaults below, measured against a loopback stub:
+
+        - hung endpoint (connects, never answers): ~5s. POSTs do not retry
+          read timeouts, see ApiClient._post_session, so this is one read
+          timeout and not three.
+        - unreachable endpoint (every connect times out): ~11s, i.e.
+          3.05 * 3 plus jittered backoff.
+        - worst case, connect timeouts until the last attempt connects and
+          then hangs: ~16s.
+
+        All under the tracker's own 45s stale-measurement warning at
+        emissions_tracker.py, which is the number that actually matters: this
+        call blocks the scheduler thread.
+        """
         self.endpoint_url: str = endpoint_url
         self.api = ApiClient(
             endpoint_url=endpoint_url,
             experiment_id=experiment_id,
             api_key=api_key,
             conf=conf,
             create_run_automatically=False,
+            timeout=timeout,
+            retries=retries,
+            backoff=backoff,
         )
         self.run_id = self.api.run_id
 
+    def exit(self) -> None:
+        # A Session holds sockets; long-lived processes creating many trackers
+        # would otherwise leak file descriptors.
+        self.api.close()
+
     def _ensure_api_run(self) -> None:
         if self.api.run_id is None and self.api.experiment_id is not None:
             self.api._create_run(self.api.experiment_id)