diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 96ed00c91..d1e7af6df 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -267,6 +267,8 @@ def _resolve_output_methods( self._output_methods.append(OutputMethod.LOGFIRE) def _initialize_runtime_state(self) -> None: + self._api_output = None + self.run_id = uuid.uuid4() self._start_time: Optional[float] = None self._last_measured_time: float = time.perf_counter() self._total_energy: Energy = Energy.from_energy(kWh=0) @@ -612,7 +614,6 @@ def _init_output_methods(self, *, api_key: str = None): methods = set(self._output_methods) if self._output_methods else set() if not methods and not self._emissions_endpoint: - self.run_id = uuid.uuid4() return from codecarbon.output_methods.boamps import BoAmpsOutput @@ -645,10 +646,8 @@ def _init_output_methods(self, *, api_key: str = None): api_key=api_key, conf=self._conf, ) - self.run_id = cc_api__out.run_id + self._api_output = cc_api__out self._output_handlers.append(cc_api__out) - else: - self.run_id = uuid.uuid4() if OutputMethod.PROMETHEUS in methods: self._output_handlers.append( @@ -710,6 +709,15 @@ def start(self) -> None: return self._ensure_hardware_ready() + + if self._api_output is not None: + # Create the run now, so every record carries the API run id. + try: + self._api_output._ensure_api_run() + self.run_id = self._api_output.run_id or self.run_id + except Exception as e: + logger.error(e, exc_info=True) + self._last_measured_time = self._start_time = time.perf_counter() # Clear utilization history for fresh measurements diff --git a/tests/test_emissions_tracker.py b/tests/test_emissions_tracker.py index 8ab12e5d8..badeac1f5 100644 --- a/tests/test_emissions_tracker.py +++ b/tests/test_emissions_tracker.py @@ -356,6 +356,92 @@ def test_output_methods_boamps_adds_boamps_output_handler( ) ) + def test_run_id_with_api_output_is_never_none( + self, + mock_cli_setup, + mock_log_values, + mocked_get_gpu_details, + mocked_env_cloud_details, + mocked_get_gpu_utilization_list, + mocked_is_gpu_details_available, + mocked_is_nvidia_system, + ): + with ( + mock.patch( + "codecarbon.output_methods.http.ApiClient._create_run" + ) as mock_create_run, + mock.patch("codecarbon.output_methods.http.ApiClient.add_emission"), + ): + tracker = EmissionsTracker( + output_dir=self.temp_path, + output_handlers=[], + output_methods=[OutputMethod.CSV, OutputMethod.API], + experiment_id="test-experiment-id", + api_key="test-api-key", + ) + api_output = next( + handler + for handler in tracker._output_handlers + if isinstance(handler, CodeCarbonAPIOutput) + ) + + def create_run(experiment_id): + api_output.api.run_id = "run-created" + return "run-created" + + mock_create_run.side_effect = create_run + + self.assertIsNotNone(tracker.run_id) + tracker.run_id = "caller-provided" + self.assertEqual(tracker.run_id, "caller-provided") + + tracker.start() + heavy_computation(1) + tracker.stop() + + self.assertEqual(tracker.run_id, "run-created") + # ...and it is what got persisted, instead of the string "None". + emissions_df = pd.read_csv(self.emissions_file_path) + self.assertEqual(emissions_df["run_id"].iloc[0], "run-created") + + def test_run_id_falls_back_to_uuid_when_api_run_creation_fails( + self, + mock_cli_setup, + mock_log_values, + mocked_get_gpu_details, + mocked_env_cloud_details, + mocked_get_gpu_utilization_list, + mocked_is_gpu_details_available, + mocked_is_nvidia_system, + ): + with ( + mock.patch( + "codecarbon.output_methods.http.ApiClient._create_run", + side_effect=Exception("API is down"), + ), + mock.patch("codecarbon.output_methods.http.ApiClient.add_emission"), + ): + tracker = EmissionsTracker( + output_dir=self.temp_path, + output_handlers=[], + output_methods=[OutputMethod.CSV, OutputMethod.API], + experiment_id="test-experiment-id", + api_key="test-api-key", + ) + local_run_id = tracker.run_id + + with self.assertLogs("codecarbon", level="ERROR") as logs: + tracker.start() + self.assertTrue(any("API is down" in line for line in logs.output)) + self.assertIsNotNone(tracker._start_time) + + heavy_computation(1) + tracker.stop() + + self.assertEqual(tracker.run_id, local_run_id) + emissions_df = pd.read_csv(self.emissions_file_path) + self.assertEqual(emissions_df["run_id"].iloc[0], str(local_run_id)) + def test_default_output_methods_is_csv( self, mock_cli_setup,