fix: allow restarting a tracker after stop()

codecarbon/emissions_tracker.py

@@ -268,6 +268,12 @@ def _resolve_output_methods(
 
     def _initialize_runtime_state(self) -> None:
         self._start_time: Optional[float] = None
+        # Timestamp of the last stop(); None while the tracker is running.
+        # This is the single start/stop state flag: `_start_time is None` means
+        # never started, `_stopped_at is not None` means stopped.
+        self._stopped_at: Optional[float] = None
+        self.final_emissions: Optional[float] = None
+        self.final_emissions_data: Optional[EmissionsData] = None
         self._last_measured_time: float = time.perf_counter()
         self._total_energy: Energy = Energy.from_energy(kWh=0)
         self._total_emissions: float = 0.0
@@ -705,12 +711,33 @@ def start(self) -> None:
                 "Another instance of codecarbon is already running. Exiting."
             )
             return
-        if self._start_time is not None:
+        if self._start_time is not None and self._stopped_at is None:
             logger.warning("Already started tracking")
             return
 
         self._ensure_hardware_ready()
-        self._last_measured_time = self._start_time = time.perf_counter()
+        now = time.perf_counter()
+        if self._stopped_at is None:
+            self._start_time = now
+        else:
+            # Restarting after a stop(). Undo what stop() tore down, in reverse
+            # order: take the lock back, then rebuild the schedulers it dropped.
+            if not self._allow_multiple_runs:
+                try:
+                    self._lock.acquire()
+                except FileExistsError:
+                    logger.error(
+                        f"Error: Another instance of codecarbon is probably running as we find `{self._lock.lockfile_path}`. Turn off the other instance to be able to run this one or use `allow_multiple_runs` or delete the file. Exiting."
+                    )
+                    self._another_instance_already_running = True
+                    return
+            self._initialize_scheduler_state()
+            # The energy accumulators are kept across the restart, so shift the
+            # start time by the stopped interval to keep `duration` and the
+            # accumulators on the same clock (active time only).
+            self._start_time += now - self._stopped_at
+            self._stopped_at = None
+        self._last_measured_time = now
 
         # Clear utilization history for fresh measurements
         self._cpu_utilization_history.clear()
@@ -899,21 +926,24 @@ def stop(self) -> Optional[float]:
                 "Another instance of codecarbon is already running. Exiting."
             )
             return
-        if not self._allow_multiple_runs:
-            # Release the lock
-            self._lock.release()
         if self._start_time is None:
             logger.error("You first need to start the tracker.")
             return None
+        if self._stopped_at is not None:
+            logger.warning("Tracker already stopped !")
+            return self.final_emissions
+        self._stopped_at = time.perf_counter()
+
+        if not self._allow_multiple_runs:
+            # Release the lock
+            self._lock.release()
 
         if self._scheduler:
             self._scheduler.stop()
             self._scheduler = None
         if self._scheduler_monitor_power:
             self._scheduler_monitor_power.stop()
             self._scheduler_monitor_power = None
-        else:
-            logger.warning("Tracker already stopped !")
         for task_name in self._tasks:
             if self._tasks[task_name].is_active:
                 self.stop_task(task_name=task_name)
@@ -934,6 +964,9 @@ def stop(self) -> Optional[float]:
         self.final_emissions_data = emissions_data
         self.final_emissions = emissions_data.emissions
 
+        # Every stop() ends a run, so the handlers are torn down here even
+        # though start() may resume the tracker later: they are re-usable, a
+        # later out() re-creates whatever exit() cleaned up.
         for handler in self._output_handlers:
             handler.exit()
 

codecarbon/lock.py

@@ -61,6 +61,7 @@ def release(self):
             try:
                 # Remove the lock file only if it was created by this instance
                 if self._has_created_lock:
+                    self._has_created_lock = False
                     os.remove(LOCKFILE)
             except OSError as e:
                 logger.debug(f"Error: {e}")

tests/output_methods/metrics/test_prometheus.py

@@ -57,6 +57,21 @@ def test_exit_method(self, mock_delete):
         output.exit()
         mock_delete.assert_called_once_with("url", job="custom_job")
 
+    @patch("codecarbon.output_methods.metrics.prometheus.delete_from_gateway")
+    @patch("codecarbon.output_methods.metrics.prometheus.push_to_gateway")
+    def test_exit_does_not_prevent_a_later_out(self, mock_push, mock_delete):
+        # A tracker can be restarted after stop(), which means exit() may be
+        # followed by another out(). The push re-creates the job that exit()
+        # deleted, so the handler stays usable.
+        output = prometheus.PrometheusOutput("url", job_name="custom_job")
+        output.out(total=EMISSIONS_DATA, delta=EMISSIONS_DATA)
+        output.exit()
+        output.out(total=EMISSIONS_DATA, delta=EMISSIONS_DATA)
+
+        self.assertEqual(mock_push.call_count, 2)
+        self.assertEqual(mock_push.call_args.kwargs["job"], "custom_job")
+        mock_delete.assert_called_once_with("url", job="custom_job")
+
     @patch(
         "codecarbon.output_methods.metrics.prometheus.push_to_gateway",
         side_effect=Exception("Test error"),

tests/test_emissions_tracker.py

@@ -633,6 +633,189 @@ def test_offline_tracker_country_name(
         self.assertEqual("United States", emissions_df["country_name"].values[0])
         self.assertEqual("USA", emissions_df["country_iso_code"].values[0])
 
+    def test_offline_tracker_stop_is_idempotent(
+        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,
+    ):
+        tracker = OfflineEmissionsTracker(
+            country_iso_code="USA",
+            output_dir=self.temp_path,
+            experiment_id="test",
+        )
+        tracker.start()
+        heavy_computation(run_time_secs=1)
+        first_emissions = tracker.stop()
+        second_emissions = tracker.stop()
+
+        self.assertEqual(first_emissions, second_emissions)
+        self.verify_output_file(self.emissions_file_path, 2)
+
+    def test_stop_releases_the_lock_only_once(
+        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.emissions_tracker.Lock") as mock_lock_class:
+            tracker = OfflineEmissionsTracker(
+                country_iso_code="USA",
+                output_dir=self.temp_path,
+                experiment_id="test",
+                allow_multiple_runs=False,
+            )
+            lock = mock_lock_class.return_value
+            lock.acquire.assert_called_once()
+
+            tracker.start()
+            heavy_computation(run_time_secs=1)
+            first_emissions = tracker.stop()
+            lock.release.assert_called_once()
+
+            # A second stop() is a no-op: it must not touch the lock again, which
+            # by then may belong to another tracker.
+            second_emissions = tracker.stop()
+            lock.release.assert_called_once()
+
+        self.assertEqual(first_emissions, second_emissions)
+        self.verify_output_file(self.emissions_file_path, 2)
+
+    def test_restart_reacquires_the_lock(
+        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.emissions_tracker.Lock") as mock_lock_class:
+            tracker = OfflineEmissionsTracker(
+                country_iso_code="USA",
+                output_dir=self.temp_path,
+                experiment_id="test",
+                allow_multiple_runs=False,
+            )
+            lock = mock_lock_class.return_value
+
+            tracker.start()
+            heavy_computation(run_time_secs=1)
+            tracker.stop()
+            self.assertEqual(lock.acquire.call_count, 1)
+            self.assertEqual(lock.release.call_count, 1)
+
+            # Restarting re-enters the protected section, so the lock has to be
+            # taken back, and released again by the matching stop().
+            tracker.start()
+            heavy_computation(run_time_secs=1)
+            tracker.stop()
+            self.assertEqual(lock.acquire.call_count, 2)
+            self.assertEqual(lock.release.call_count, 2)
+
+            # And a redundant stop() still touches nothing.
+            tracker.stop()
+            self.assertEqual(lock.release.call_count, 2)
+
+    def test_restart_aborts_when_the_lock_is_taken_by_another_instance(
+        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.emissions_tracker.Lock") as mock_lock_class:
+            tracker = OfflineEmissionsTracker(
+                country_iso_code="USA",
+                output_dir=self.temp_path,
+                experiment_id="test",
+                allow_multiple_runs=False,
+            )
+            tracker.start()
+            heavy_computation(run_time_secs=1)
+            tracker.stop()
+
+            mock_lock_class.return_value.acquire.side_effect = FileExistsError
+            tracker.start()
+
+            # The restart was refused, so the tracker stays stopped, nothing was
+            # rebuilt, and the clash was reported rather than swallowed.
+            self.assertTrue(tracker._another_instance_already_running)
+            self.assertIsNotNone(tracker._stopped_at)
+            self.assertIsNone(tracker._scheduler)
+
+    def test_tracker_can_be_restarted_after_stop(
+        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,
+    ):
+        # GIVEN
+        tracker = OfflineEmissionsTracker(
+            country_iso_code="USA", measure_power_secs=1, save_to_file=False
+        )
+
+        # WHEN
+        tracker.start()
+        heavy_computation(run_time_secs=2)
+        tracker.stop()
+        first_duration = tracker.final_emissions_data.duration
+
+        time.sleep(2)  # stopped: this gap must not be counted
+
+        tracker.start()
+        self.assertIsNotNone(tracker._scheduler)
+        heavy_computation(run_time_secs=2)
+        tracker.stop()
+        second_duration = tracker.final_emissions_data.duration
+
+        # THEN the second row covers the two active phases only, not the pause
+        self.assertAlmostEqual(first_duration, 2, delta=1)
+        self.assertAlmostEqual(second_duration, 4, delta=1)
+
+    def test_second_start_while_running_is_a_no_op(
+        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,
+    ):
+        tracker = OfflineEmissionsTracker(
+            country_iso_code="USA", measure_power_secs=1, save_to_file=False
+        )
+        tracker.start()
+        start_time = tracker._start_time
+        scheduler = tracker._scheduler
+
+        with self.assertLogs("codecarbon", level="WARNING") as logs:
+            tracker.start()
+
+        # A restart is only allowed after stop(): while running, start() must not
+        # move the clock nor replace the running schedulers.
+        self.assertTrue(any("Already started tracking" in line for line in logs.output))
+        self.assertEqual(start_time, tracker._start_time)
+        self.assertIs(scheduler, tracker._scheduler)
+        tracker.stop()
+
     def test_offline_tracker_invalid_headers(
         self,
         mock_cli_setup,

tests/test_lock.py

@@ -30,6 +30,17 @@ def test_release_removes_lock_file(self, mock_file, mock_remove):
         self.lock.release()
         mock_remove.assert_called_once_with(LOCKFILE)
 
+    @patch("codecarbon.lock.os.remove")
+    @patch("codecarbon.lock.open", new_callable=mock_open)
+    def test_release_is_idempotent(self, mock_file, mock_remove):
+        # A second release() must not delete the lock file again: by then it may
+        # have been re-created by another instance of codecarbon.
+        self.lock.acquire()
+        self.lock.release()
+        self.lock.release()
+        mock_remove.assert_called_once_with(LOCKFILE)
+        self.assertFalse(self.lock._has_created_lock)
+
     @patch("codecarbon.lock.os.remove")
     @patch("codecarbon.lock.open", new_callable=mock_open)
     def test_release_does_not_release_when_not_created_by_this_instance(