Added a backup for the geo API (#453)

* Added backup API for geo-location * Committed code that belongs to another PR * fix: test test_emissions_tracker.py Signed-off-by: inimaz <[email protected]> --------- Signed-off-by: inimaz <[email protected]> Co-authored-by: inimaz <[email protected]> Co-authored-by: inimaz <[email protected]>
mlco2 · Feb 25, 2024 · fad87d7 · fad87d7
1 parent b2ec185
commit fad87d7
Show file tree

Hide file tree

Showing 5 changed files with 101 additions and 12 deletions.
diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py
@@ -475,11 +475,11 @@ def start_task(self, task_name=None) -> None:
         )
         self._active_task = task_name
 
-    def stop_task(self, task_name: str = None) -> EmissionsData:
+    def stop_task(self, task_name: str = None) -> float:
         """
         Stop tracking a dedicated execution task. Delta energy is computed by task, to isolate its contribution to total
         emissions.
-        :return: EmissionsData
+        :return: None
         """
         task_name = task_name if task_name else self._active_task
         self._measure_power_and_energy()

diff --git a/codecarbon/external/geography.py b/codecarbon/external/geography.py
@@ -3,6 +3,7 @@
 """
 
 import re
+import urllib.parse
 from dataclasses import dataclass
 from typing import Callable, Dict, Optional
 
@@ -85,11 +86,49 @@ def __repr__(self) -> str:
     def from_geo_js(cls, url: str) -> "GeoMetadata":
         try:
             response: Dict = requests.get(url, timeout=0.5).json()
+
+            return cls(
+                country_iso_code=response["country_code3"].upper(),
+                country_name=response["country"],
+                region=response.get("region", "").lower(),
+                latitude=float(response.get("latitude")),
+                longitude=float(response.get("longitude")),
+                country_2letter_iso_code=response.get("country_code"),
+            )
         except Exception as e:
-            # If there is a timeout, we default to Canada
+            # If there is a timeout, we will try using a backup API
+            logger.warning(
+                f"Unable to access geographical location through primary API. Will resort to using the backup API - Exception : {e} - url={url}"
+            )
+
+        geo_url_backup = "https://ip-api.com/json/"
+
+        try:
+            geo_response: Dict = requests.get(geo_url_backup, timeout=0.5).json()
+            country_name = geo_response["country"]
+
+            # The previous request does not return the three-letter country code
+            country_code_3_url = f"https://api.first.org/data/v1/countries?q={urllib.parse.quote_plus(country_name)}&scope=iso"
+            country_code_response: Dict = requests.get(
+                country_code_3_url, timeout=0.5
+            ).json()
+
+            return cls(
+                country_iso_code=next(
+                    iter(country_code_response["data"].keys())
+                ).upper(),
+                country_name=country_name,
+                region=geo_response.get("regionName", "").lower(),
+                latitude=float(geo_response.get("lat")),
+                longitude=float(geo_response.get("lon")),
+                country_2letter_iso_code=geo_response.get("countryCode"),
+            )
+        except Exception as e:
+            # If both API calls fail, default to Canada
             logger.warning(
                 f"Unable to access geographical location. Using 'Canada' as the default value - Exception : {e} - url={url}"
             )
+
             return cls(
                 country_iso_code="CAN",
                 country_name="Canada",
@@ -98,11 +137,3 @@ def from_geo_js(cls, url: str) -> "GeoMetadata":
                 longitude=-71.2,
                 country_2letter_iso_code="CA",
             )
-        return cls(
-            country_iso_code=response["country_code3"].upper(),
-            country_name=response["country"],
-            region=response.get("region", "").lower(),
-            latitude=float(response.get("latitude")),
-            longitude=float(response.get("longitude")),
-            country_2letter_iso_code=response.get("country_code"),
-        )
diff --git a/tests/test_emissions_tracker.py b/tests/test_emissions_tracker.py
@@ -126,7 +126,7 @@ def raise_timeout_exception(*args, **kwargs):
         tracker.start()
         heavy_computation(run_time_secs=2)
         emissions = tracker.stop()
-        self.assertEqual(1, mocked_requests_get.call_count)
+        self.assertEqual(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
@@ -8,8 +8,10 @@
     CLOUD_METADATA_AWS,
     CLOUD_METADATA_AZURE,
     CLOUD_METADATA_GCP,
+    COUNTRY_METADATA_USA,
     GEO_METADATA_CANADA,
     GEO_METADATA_USA,
+    GEO_METADATA_USA_BACKUP,
 )
 
 
@@ -69,6 +71,28 @@ def test_geo_metadata_USA(self):
         self.assertEqual("United States", geo.country_name)
         self.assertEqual("illinois", geo.region)
 
+    @responses.activate
+    def test_geo_metadata_USA_backup(self):
+        responses.add(
+            responses.GET, self.geo_js_url, json={"error": "not found"}, status=404
+        )
+        responses.add(
+            responses.GET,
+            "https://ip-api.com/json/",
+            json=GEO_METADATA_USA_BACKUP,
+            status=200,
+        )
+        responses.add(
+            responses.GET,
+            "https://api.first.org/data/v1/countries?q=United%20States&scope=iso",
+            json=COUNTRY_METADATA_USA,
+            status=200,
+        )
+        geo = GeoMetadata.from_geo_js(self.geo_js_url)
+        self.assertEqual("USA", geo.country_iso_code)
+        self.assertEqual("United States", geo.country_name)
+        self.assertEqual("illinois", geo.region)
+
     @responses.activate
     def test_geo_metadata_CANADA(self):
         responses.add(

diff --git a/tests/testdata.py b/tests/testdata.py
@@ -18,6 +18,40 @@
     "latitude": "0",
 }
 
+GEO_METADATA_USA_BACKUP = {
+    "organization_name": "foobar",
+    "regionName": "Illinois",
+    "accuracy": 1,
+    "asn": 0,
+    "organization": "foobar",
+    "timezone": "America/Chicago",
+    "lon": "88",
+    "area_code": "0",
+    "ip": "foobar",
+    "city": "Chicago",
+    "country": "United States",
+    "countryCode": "US",
+    "lat": "0",
+}
+
+COUNTRY_METADATA_USA = {
+    "status": "OK",
+    "status-code": 200,
+    "version": "1.0",
+    "access": "public",
+    "data": {
+        "USA": {
+            "id": "USA",
+            "country": "United States of America (the)",
+            "region": "North America",
+        },
+        "UMI": {
+            "id": "UMI",
+            "country": "United States Minor Outlying Islands (the)",
+            "region": "Oceania",
+        },
+    },
+}
 
 GEO_METADATA_CANADA = {
     "organization_name": "foobar",