fix: propagate the API run id to the tracker

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

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,