fix: propagate the API run id to the tracker

codecarbon/emissions_tracker.py

@@ -266,7 +266,27 @@ def _resolve_output_methods(
         if self._save_to_logfire:
             self._output_methods.append(OutputMethod.LOGFIRE)
 
+    @property
+    def run_id(self):
+        """
+        Id of the current run. It is the API run id as soon as the API output
+        method has created one, and a locally generated uuid otherwise.
+        """
+        if self._api_output is not None and self._api_output.run_id is not None:
+            return self._api_output.run_id
+        return self._run_id
+
+    @run_id.setter
+    def run_id(self, value) -> None:
+        """
+        `run_id` used to be a plain attribute ; keep it writable for callers
+        that set their own id. An API run id, once created, still wins.
+        """
+        self._run_id = value
+
     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 +632,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 +664,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 +727,15 @@ def start(self) -> None:
             return
 
         self._ensure_hardware_ready()
+
+        if self._api_output is not None:
+            # Create the run now, so that every record of this run, whatever the
+            # output method, carries the API run id.
+            try:
+                self._api_output._ensure_api_run()
+            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,98 @@ 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
+
+            # Before the run is created, the tracker falls back on a local uuid,
+            # which stays writable for callers that set their own id.
+            self.assertIsNotNone(tracker.run_id)
+            tracker.run_id = "caller-provided"
+            self.assertEqual(tracker.run_id, "caller-provided")
+
+            tracker.start()
+            heavy_computation(1)
+            tracker.stop()
+
+        # Once the API created the run, the tracker exposes the API run id...
+        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()
+            # The API failure is reported but does not abort the tracking.
+            self.assertTrue(any("API is down" in line for line in logs.output))
+            self.assertIsNotNone(tracker._start_time)
+
+            heavy_computation(1)
+            tracker.stop()
+
+        # No API run id available: the tracker keeps its local uuid, and it is
+        # what gets persisted.
+        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,