diff --git a/codecarbon/external/geography.py b/codecarbon/external/geography.py index 075824959..000b58be8 100644 --- a/codecarbon/external/geography.py +++ b/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( diff --git a/tests/test_emissions_tracker.py b/tests/test_emissions_tracker.py index 8ab12e5d8..6c53fd54a 100644 --- a/tests/test_emissions_tracker.py +++ b/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) diff --git a/tests/test_geography.py b/tests/test_geography.py index 8f95f7f43..7f6056b08 100644 --- a/tests/test_geography.py +++ b/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(