diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 96ed00c91..f7520bc71 100644 --- a/codecarbon/emissions_tracker.py +++ b/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 @@ -899,12 +905,17 @@ 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() @@ -912,8 +923,6 @@ def stop(self) -> Optional[float]: 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) diff --git a/codecarbon/lock.py b/codecarbon/lock.py index 38d112324..47ed53bb8 100644 --- a/codecarbon/lock.py +++ b/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}") diff --git a/tests/test_emissions_tracker.py b/tests/test_emissions_tracker.py index 8ab12e5d8..52f5d52e2 100644 --- a/tests/test_emissions_tracker.py +++ b/tests/test_emissions_tracker.py @@ -633,6 +633,62 @@ 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_offline_tracker_invalid_headers( self, mock_cli_setup, diff --git a/tests/test_lock.py b/tests/test_lock.py index aafb46a1b..3a6ae2c70 100644 --- a/tests/test_lock.py +++ b/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(