feat(api): reuse connections, add timeouts, and stamp emissions with measurement time

codecarbon/core/api_client.py

@@ -8,7 +8,8 @@
 # from httpx import AsyncClient
 import dataclasses
 import json
-from datetime import timedelta, tzinfo
+import time
+from datetime import datetime, timedelta, tzinfo
 
 import requests
 
@@ -33,6 +34,29 @@ def get_datetime_with_timezone():
     return str(arrow.now().isoformat())
 
 
+# (connect, read) seconds, replacing a flat 2s that timed out on a loaded API.
+_TIMEOUT = (3.05, 10)
+# Seconds to wait after a failed run creation before trying again, so a down
+# API costs one blocking call per minute instead of one per measurement.
+_RUN_CREATE_COOLDOWN = 60
+
+
+def _measurement_timestamp(carbon_emission: dict) -> str:
+    """
+    Offset-aware ISO timestamp of *when the measurement was taken*, taken from
+    EmissionsData.timestamp. Falls back to now for hand-built payloads that
+    carry no usable timestamp.
+    """
+    try:
+        return (
+            datetime.fromisoformat(carbon_emission["timestamp"])
+            .astimezone()
+            .isoformat()
+        )
+    except (KeyError, TypeError, ValueError):
+        return get_datetime_with_timezone()
+
+
 class ApiClient:  # (AsyncClient)
     """
     This class call the Code Carbon API
@@ -58,11 +82,14 @@ def __init__(
         :create_run_automatically: If False, do not create a run. To use API in read only mode.
         """
         # super().__init__(base_url=endpoint_url) # (AsyncClient)
+        # A Session so the socket and TLS handshake are reused across calls.
+        self._session = requests.Session()
         self.url = endpoint_url
         self.experiment_id = experiment_id
         self.api_key = api_key
         self.conf = conf
         self.access_token = access_token
+        self._run_create_failed_at = None
         if self.experiment_id is not None and create_run_automatically:
             self._create_run(self.experiment_id)
 
@@ -80,16 +107,20 @@ def _request(self, method, url, payload=None, expected_status=200):
         Call the API and return the response, raising on anything that is not
         the status code the API answers on success.
 
-        :method: the requests function to call, for example requests.get
+        :method: the session function to call, for example self._session.get
         :payload: the JSON body to send, if any
         :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=_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()
+
     def set_access_token(self, token: str):
         """This method sets the access token to be used for the API.
         Args:
@@ -102,14 +133,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):
         """
@@ -134,30 +165,30 @@ def create_organization(self, organization: OrganizationCreate):
             return organization
         else:
             return self._request(
-                requests.post, url, payload=payload, expected_status=201
+                self._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):
         """
@@ -166,36 +197,52 @@ 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._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):
+    def add_emission(self, carbon_emission: dict, final: bool = False):
         assert self.experiment_id is not None
         if self.run_id is None:
+            # Captured before the call: tells us whether _create_run is about
+            # to skip its attempt because of the cooldown, so we can log the
+            # expected per-tick skip at debug instead of as an error.
+            in_cooldown = (
+                not final
+                and self._run_create_failed_at is not None
+                and time.monotonic() - self._run_create_failed_at < _RUN_CREATE_COOLDOWN
+            )
             logger.warning(
                 "ApiClient.add_emission() need a run_id : the initial call may "
                 + "have failed. Retrying..."
             )
-            self._create_run(self.experiment_id)
+            self._create_run(self.experiment_id, bypass_cooldown=final)
             if self.run_id is None:
-                logger.error(
-                    "ApiClient.add_emission still no run_id, aborting for this time !"
-                )
-            return False
+                if in_cooldown:
+                    logger.debug(
+                        "ApiClient.add_emission still no run_id, run creation"
+                        " is in its cooldown, will retry later."
+                    )
+                else:
+                    logger.error(
+                        "ApiClient.add_emission still no run_id, aborting for this time !"
+                    )
+                return False
+            # Run creation just succeeded (e.g. a final-flush bypass): fall
+            # through and send this emission instead of dropping it.
         if carbon_emission["duration"] < 1:
             logger.warning(
                 "ApiClient : emissions not sent because of a duration smaller than 1."
             )
             return False
         emission = EmissionCreate(
-            timestamp=get_datetime_with_timezone(),
+            timestamp=_measurement_timestamp(carbon_emission),
             run_id=self.run_id,
             duration=int(carbon_emission["duration"]),
             emissions_sum=carbon_emission["emissions"],
@@ -215,7 +262,7 @@ 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._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.
@@ -225,16 +272,29 @@ def add_emission(self, carbon_emission: dict):
             raise
         return True
 
-    def _create_run(self, experiment_id: str):
+    def _create_run(self, experiment_id: str, bypass_cooldown: bool = False):
         """
         Create the experiment for project_id
+
+        :bypass_cooldown: skip the cooldown check and retry immediately, used
+            by the final flush on tracker stop/exit so the last emission is
+            not silently dropped because of a recent failure.
         """
         if self.experiment_id is None:
             # TODO : raise an Exception ?
             logger.error(
                 "ApiClient FATAL The ApiClient._create_run() needs an experiment_id !"
             )
             return None
+        if (
+            not bypass_cooldown
+            and self._run_create_failed_at is not None
+            and time.monotonic() - self._run_create_failed_at < _RUN_CREATE_COOLDOWN
+        ):
+            logger.debug("ApiClient - run creation failed recently, not retrying yet")
+            return None
+        # Cleared on success below; set now so every failure path is covered.
+        self._run_create_failed_at = time.monotonic()
         try:
             run = RunCreate(
                 timestamp=get_datetime_with_timezone(),
@@ -256,8 +316,11 @@ 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._session.post, url, payload=payload, expected_status=201
+            )
             self.run_id = r.json()["id"]
+            self._run_create_failed_at = None
             logger.info(
                 "ApiClient Successfully registered your run on the API.\n\n"
                 + f"Run ID: {self.run_id}\n"
@@ -282,7 +345,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):
         """
@@ -298,15 +361,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._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

@@ -32,6 +32,7 @@
 from codecarbon.lock import Lock
 from codecarbon.output_methods.base_output import BaseOutput, OutputMethod
 from codecarbon.output_methods.emissions_data import EmissionsData
+from codecarbon.output_methods.http import CodeCarbonAPIOutput
 
 if TYPE_CHECKING:
     from codecarbon.external.geography import CloudMetadata, GeoMetadata
@@ -925,12 +926,28 @@ def stop(self) -> Optional[float]:
         emissions_data = self._prepare_emissions_data()
         emissions_data_delta = self._compute_emissions_delta(emissions_data)
 
+        persist_started_at = time.monotonic()
         self._persist_data(
             total_emissions=emissions_data,
             delta_emissions=emissions_data_delta,
             experiment_name=self._experiment_name,
         )
 
+        # If run creation was still in its cooldown, the emission above was
+        # dropped without even trying the API. This is the last chance to
+        # send it, so bypass the cooldown and retry once instead of losing
+        # the row. But if the persist call above already made (and failed) a
+        # run-creation attempt, retrying again here would just hit the
+        # server a second time for nothing.
+        for handler in self._output_handlers:
+            if isinstance(handler, CodeCarbonAPIOutput) and handler.run_id is None:
+                already_attempted = (
+                    handler.api._run_create_failed_at is not None
+                    and handler.api._run_create_failed_at >= persist_started_at
+                )
+                if not already_attempted:
+                    handler.out(emissions_data, emissions_data_delta, final=True)
+
         self.final_emissions_data = emissions_data
         self.final_emissions = emissions_data.emissions
 

codecarbon/output_methods/http.py

@@ -57,20 +57,23 @@ def __init__(
         )
         self.run_id = self.api.run_id
 
-    def _ensure_api_run(self) -> None:
+    def exit(self) -> None:
+        self.api.close()
+
+    def _ensure_api_run(self, final: bool = False) -> None:
         if self.api.run_id is None and self.api.experiment_id is not None:
-            self.api._create_run(self.api.experiment_id)
+            self.api._create_run(self.api.experiment_id, bypass_cooldown=final)
             self.run_id = self.api.run_id
 
-    def _emit(self, delta: EmissionsData) -> None:
+    def _emit(self, delta: EmissionsData, final: bool = False) -> None:
         try:
-            self._ensure_api_run()
-            self.api.add_emission(dataclasses.asdict(delta))
+            self._ensure_api_run(final=final)
+            self.api.add_emission(dataclasses.asdict(delta), final=final)
         except Exception as e:
             logger.error(e, exc_info=True)
 
     def live_out(self, _, delta: EmissionsData):
         self._emit(delta)
 
-    def out(self, _, delta: EmissionsData):
-        self._emit(delta)
+    def out(self, _, delta: EmissionsData, final: bool = False):
+        self._emit(delta, final=final)

tests/output_methods/test_http.py

@@ -175,14 +175,14 @@ def test_codecarbon_api_live_out_creates_run_when_missing(self):
             )
             api_output.api.run_id = None
 
-            def create_run(experiment_id):
+            def create_run(experiment_id, bypass_cooldown=False):
                 api_output.api.run_id = "run-created"
                 return "run-created"
 
             mock_create_run.side_effect = create_run
             api_output.live_out(None, self.emissions_data)
 
-        mock_create_run.assert_called_once_with("exp-1")
+        mock_create_run.assert_called_once_with("exp-1", bypass_cooldown=False)
         self.assertEqual(api_output.api.run_id, "run-created")
         self.assertEqual(api_output.run_id, "run-created")
 
@@ -209,6 +209,27 @@ def test_codecarbon_api_out(self):
         api_output.out(None, self.emissions_data)
         self.mock_add_emission.assert_called_once()
 
+    def test_codecarbon_api_out_final_bypasses_cooldown(self):
+        """The final flush (tracker.stop()) must bypass the run-creation
+        cooldown and mark the emission as final, so it is not dropped when a
+        run was recently failing to be created."""
+        with patch(
+            "codecarbon.output_methods.http.ApiClient._create_run"
+        ) as mock_create_run:
+            api_output = CodeCarbonAPIOutput(
+                endpoint_url=self.url,
+                experiment_id="exp-1",
+                api_key=self.api_key,
+                conf=None,
+            )
+            api_output.api.run_id = None
+
+            api_output.out(None, self.emissions_data, final=True)
+
+        mock_create_run.assert_called_once_with("exp-1", bypass_cooldown=True)
+        self.mock_add_emission.assert_called_once()
+        self.assertTrue(self.mock_add_emission.call_args.kwargs.get("final"))
+
     @patch("codecarbon.output_methods.http.logger.error")
     def test_codecarbon_out_api_call_failure(self, mock_logger):
         self.mock_add_emission.side_effect = Exception("Test exception")