diff --git a/codecarbon/core/cloud.py b/codecarbon/core/cloud.py index 405bc3528..466726909 100644 --- a/codecarbon/core/cloud.py +++ b/codecarbon/core/cloud.py @@ -1,3 +1,4 @@ +from concurrent.futures import ThreadPoolExecutor from typing import Any, Dict, Optional import requests @@ -51,21 +52,37 @@ def get_env_cloud_details(timeout: int = 1) -> Optional[Any]: 'region': 'us-east-1', 'version': '2017-09-30'}} """ - for provider in CLOUD_METADATA_MAPPING.keys(): - try: - params = CLOUD_METADATA_MAPPING[provider] - response = requests.get( - params["url"], headers=params["headers"], timeout=timeout - ) - response.raise_for_status() - response_data = response.json() + providers = list(CLOUD_METADATA_MAPPING.keys()) - postprocess_function = params.get("postprocess_function") - if postprocess_function is not None: - response_data = postprocess_function(response_data) - - return {"provider": provider, "metadata": response_data} - except requests.exceptions.RequestException: - logger.debug("Not running on %s", provider) + # All providers answer on the same link-local address, so probing them + # sequentially costs one timeout each on a machine that is not on a cloud. + # `or 1`: ThreadPoolExecutor rejects max_workers=0. + with ThreadPoolExecutor(max_workers=len(providers) or 1) as executor: + futures = [executor.submit(_probe_provider, p, timeout) for p in providers] + # Resolve in mapping order, not completion order, so detection stays + # deterministic if more than one provider answers. + for provider, future in zip(providers, futures): + response_data = future.result() + if response_data is not None: + return {"provider": provider, "metadata": response_data} return None + + +def _probe_provider(provider: str, timeout: int) -> Optional[Dict[str, Any]]: + params = CLOUD_METADATA_MAPPING[provider] + try: + response = requests.get( + params["url"], headers=params["headers"], timeout=timeout + ) + response.raise_for_status() + response_data = response.json() + except requests.exceptions.RequestException: + logger.debug("Not running on %s", provider) + return None + + postprocess_function = params.get("postprocess_function") + if postprocess_function is not None: + response_data = postprocess_function(response_data) + + return response_data diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 96ed00c91..56cd2630b 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -424,6 +424,7 @@ def __init__( allow_multiple_runs: Optional[bool] = _sentinel, rapl_include_dram: Optional[bool] = _sentinel, rapl_prefer_psys: Optional[bool] = _sentinel, + cloud_detection: Optional[bool] = _sentinel, ): """ :param project_name: Project name for current experiment run, default name @@ -519,6 +520,10 @@ def __init__( (CPU + chipset + PCIe). When False, uses package domains which are more reliable. Note: psys can report higher values than CPU TDP and may be unreliable on older systems. + :param cloud_detection: Query the cloud provider metadata service at startup to + detect AWS/Azure/GCP, defaults to True. Set to False on + machines that are not on a cloud, or in air-gapped and + egress-filtered environments, to skip the network probe. """ # logger.info("base tracker init") @@ -588,6 +593,7 @@ def __init__( self._set_from_conf(force_mode_cpu_load, "force_mode_cpu_load", False, bool) self._set_from_conf(rapl_include_dram, "rapl_include_dram", False, bool) self._set_from_conf(rapl_prefer_psys, "rapl_prefer_psys", False, bool) + self._set_from_conf(cloud_detection, "cloud_detection", True, bool) self._set_from_conf( experiment_id, "experiment_id", "5b0fa12a-3dd7-45bb-9766-cc326314d9f1" ) @@ -1434,7 +1440,11 @@ def _get_cloud_metadata(self) -> CloudMetadata: from codecarbon.external.geography import CloudMetadata if self._cloud is None: - self._cloud = CloudMetadata.from_utils() + if self._cloud_detection: + self._cloud = CloudMetadata.from_utils() + else: + logger.debug("Cloud detection is disabled, skipping metadata probe.") + self._cloud = CloudMetadata(provider=None, region=None) return self._cloud diff --git a/docs/how-to/configuration.md b/docs/how-to/configuration.md index 9f6766aa1..4bec581ac 100644 --- a/docs/how-to/configuration.md +++ b/docs/how-to/configuration.md @@ -178,6 +178,38 @@ Despite its name, this option applies to every counter-based CPU interface: It has no effect when CodeCarbon falls back to TDP/CPU-load estimation, since that mode models the CPU package only. +## Cloud Detection + +At startup, `EmissionsTracker` queries the cloud instance metadata service on the +link-local address `169.254.169.254` to find out whether it runs on AWS, Azure or +GCP. All providers are probed in parallel, so the check costs at most one +network timeout. + +On a machine that is not on a cloud the probe always fails, and in air-gapped or +egress-filtered environments it is unwanted network noise. Set `cloud_detection` +to `false` to skip it entirely: + +``` ini +[codecarbon] +cloud_detection = false +``` + +Or in code: + +``` python +EmissionsTracker(cloud_detection=False) +``` + +Or as an environment variable: + +``` shell +export CODECARBON_CLOUD_DETECTION=false +``` + +With detection disabled, CodeCarbon reports no cloud provider or region and uses +the usual geolocation path to pick a carbon intensity. `OfflineEmissionsTracker` +never probes the metadata service, so the option has no effect there. + ## Access internet through proxy server If you need a proxy to access internet, which is needed to call a Web diff --git a/tests/test_cloud.py b/tests/test_cloud.py index 2353e25e2..dbe1b84fb 100644 --- a/tests/test_cloud.py +++ b/tests/test_cloud.py @@ -19,9 +19,12 @@ # OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR # OTHER DEALINGS IN THE SOFTWARE. +from unittest import mock + import responses from codecarbon.core.cloud import CLOUD_METADATA_MAPPING, get_env_cloud_details +from codecarbon.emissions_tracker import EmissionsTracker def setup_cloud_details_responses(tested_provider, provider_metadata): @@ -79,3 +82,44 @@ def test_get_env_cloud_details_mapping_nothing(): setup_cloud_details_responses("localhost", metadata) assert get_env_cloud_details() is None + + +@responses.activate +def test_get_env_cloud_details_probes_all_providers_concurrently(): + """Every provider is probed, and detection stays deterministic if several answer.""" + for params in CLOUD_METADATA_MAPPING.values(): + responses.add(responses.GET, params["url"], json={"region": "here"}, status=200) + + first_provider = next(iter(CLOUD_METADATA_MAPPING)) + assert get_env_cloud_details() == { + "provider": first_provider, + "metadata": {"region": "here"}, + } + assert len(responses.calls) == len(CLOUD_METADATA_MAPPING) + + +def test_cloud_detection_disabled_makes_no_request(): + with mock.patch( + "codecarbon.external.geography.get_env_cloud_details" + ) as mocked_get_env_cloud_details: + tracker = EmissionsTracker(cloud_detection=False, allow_multiple_runs=True) + cloud = tracker._get_cloud_metadata() + + mocked_get_env_cloud_details.assert_not_called() + assert cloud.is_on_private_infra + + +def test_cloud_detection_enabled_by_default(): + with mock.patch( + "codecarbon.external.geography.get_env_cloud_details", return_value=None + ) as mocked_get_env_cloud_details: + tracker = EmissionsTracker(allow_multiple_runs=True) + tracker._get_cloud_metadata() + + mocked_get_env_cloud_details.assert_called_once() + + +def test_get_env_cloud_details_with_no_provider(): + """An empty mapping must not blow up on ThreadPoolExecutor(max_workers=0).""" + with mock.patch("codecarbon.core.cloud.CLOUD_METADATA_MAPPING", {}): + assert get_env_cloud_details() is None