fix: raise geolocation API timeout and retry the primary API once

codecarbon/external/geography.py

@@ -12,6 +12,9 @@
 from codecarbon.core.cloud import get_env_cloud_details
 from codecarbon.external.logger import logger
 
+GEO_API_TIMEOUT: float = 5
+GEO_API_RETRIES: int = 1
+
 
 @dataclass
 class CloudMetadata:
@@ -88,10 +91,29 @@ def __repr__(self) -> str:
             self.region,
         )
 
+    @staticmethod
+    def _get_geo_json(url: str, retries: int = 0) -> Dict:
+        """
+        Query a geolocation API, retrying only when the network itself fails,
+        so a slow or busy connection does not send us straight to the fallback.
+        """
+        for attempt in range(retries + 1):
+            try:
+                return requests.get(url, timeout=GEO_API_TIMEOUT).json()
+            except (
+                requests.exceptions.Timeout,
+                requests.exceptions.ConnectionError,
+            ) as e:
+                if attempt == retries:
+                    raise
+                logger.debug(
+                    f"Could not reach {url}, retrying ({attempt + 1}/{retries}) - Exception : {e}"
+                )
+
     @classmethod
     def from_geo_js(cls, url: str) -> "GeoMetadata":
         try:
-            response: Dict = requests.get(url, timeout=0.5).json()
+            response: Dict = cls._get_geo_json(url, retries=GEO_API_RETRIES)
 
             region = response.get("region", "").lower()
             if not region:
@@ -114,7 +136,7 @@ def from_geo_js(cls, url: str) -> "GeoMetadata":
         geo_url_backup = "https://ipinfo.io/json"
 
         try:
-            geo_response: Dict = requests.get(geo_url_backup, timeout=0.5).json()
+            geo_response: Dict = cls._get_geo_json(geo_url_backup)
 
             # extract latitude and longitude from loc (e.g., "loc": "37.4056,-122.0775")
             loc = geo_response.get("loc", "").split(",")
@@ -140,7 +162,7 @@ def from_geo_js(cls, url: str) -> "GeoMetadata":
         except Exception as e:
             # If both API calls fail, default to Canada
             logger.warning(
-                f"Unable to access geographical location through fallback API. Using 'Canada' as the default value - Exception : {e} - url={geo_url_backup}"
+                f"Unable to access geographical location through fallback API. Defaulting to Canada, so emissions will be computed with the Canadian carbon intensity - Exception : {e} - url={geo_url_backup}"
             )
 
             return cls(

tests/test_emissions_tracker.py

@@ -17,7 +17,7 @@
     OfflineEmissionsTracker,
     track_emissions,
 )
-from codecarbon.external.geography import CloudMetadata
+from codecarbon.external.geography import GEO_API_RETRIES, CloudMetadata
 from codecarbon.output import BoAmpsOutput, CodeCarbonAPIOutput, OutputMethod
 from tests.fake_modules import pynvml as fake_pynvml
 from tests.testdata import (
@@ -289,7 +289,8 @@ def raise_timeout_exception(*args, **kwargs):
         tracker.start()
         heavy_computation(run_time_secs=2)
         emissions = tracker.stop()
-        self.assertEqual(2, mocked_requests_get.call_count)
+        # The primary API is tried once more before the backup one is called.
+        self.assertEqual(GEO_API_RETRIES + 2, mocked_requests_get.call_count)
         self.assertIsInstance(emissions, float)
         self.assertAlmostEqual(1.1037980397280433e-05, emissions, places=2)
 

tests/test_geography.py

@@ -1,9 +1,14 @@
 import unittest
 from unittest import mock
 
+import requests
 import responses
 
-from codecarbon.external.geography import CloudMetadata, GeoMetadata
+from codecarbon.external.geography import (
+    GEO_API_TIMEOUT,
+    CloudMetadata,
+    GeoMetadata,
+)
 from tests.testdata import (
     CLOUD_METADATA_AWS,
     CLOUD_METADATA_AZURE,
@@ -118,6 +123,44 @@ def test_geo_metadata_empty_region_fallback(self):
         self.assertEqual("United States", geo.country_name)
         self.assertEqual("illinois", geo.region)
 
+    @responses.activate
+    def test_geo_metadata_retries_primary_api_on_timeout(self):
+        responses.add(
+            responses.GET,
+            self.geo_js_url,
+            body=requests.exceptions.Timeout("Read timed out"),
+        )
+        responses.add(responses.GET, self.geo_js_url, json=GEO_METADATA_USA, status=200)
+        responses.add(
+            responses.GET,
+            "https://ipinfo.io/json",
+            json=GEO_METADATA_USA_BACKUP,
+            status=200,
+        )
+
+        geo = GeoMetadata.from_geo_js(self.geo_js_url)
+
+        self.assertEqual("USA", geo.country_iso_code)
+        self.assertEqual("illinois", geo.region)
+        # The primary API answered on the second try, so the backup is never called.
+        self.assertEqual(
+            [self.geo_js_url, self.geo_js_url],
+            [call.request.url for call in responses.calls],
+        )
+
+    def test_geo_metadata_uses_configured_timeout(self):
+        mocked_response = mock.Mock()
+        mocked_response.json.return_value = GEO_METADATA_USA
+
+        with mock.patch(
+            "codecarbon.external.geography.requests.get", return_value=mocked_response
+        ) as mocked_get:
+            geo = GeoMetadata.from_geo_js(self.geo_js_url)
+
+        self.assertEqual("USA", geo.country_iso_code)
+        self.assertGreater(GEO_API_TIMEOUT, 0.5)
+        mocked_get.assert_called_once_with(self.geo_js_url, timeout=GEO_API_TIMEOUT)
+
     @responses.activate
     def test_geo_metadata_CANADA(self):
         responses.add(