fix: resume the scheduler in stop_task only if start_task paused it

codecarbon/emissions_tracker.py

@@ -296,6 +296,7 @@ def _initialize_runtime_state(self) -> None:
         self._tasks: Dict[str, Task] = {}
         self._active_task: Optional[str] = None
         self._active_task_emissions_at_start: Optional[EmissionsData] = None
+        self._scheduler_paused_by_task = False
         self._hardware = []
         self._hardware_initialized = False
 
@@ -754,6 +755,13 @@ def start_task(self, task_name=None) -> None:
 
         # Stop scheduler as we do not want it to interfere with the task measurement
         if self._scheduler:
+            # Only resume it in stop_task if it was actually running, i.e. the tracker
+            # was started with start(). Pure start_task/stop_task usage must not leave
+            # a periodic measurement running behind. The flag is sticky: a second
+            # start_task call sees an already stopped scheduler and must not clear it.
+            self._scheduler_paused_by_task = (
+                self._scheduler_paused_by_task or not self._scheduler._stopped
+            )
             self._scheduler.stop()
 
         # Task background thread for measuring power
@@ -794,6 +802,16 @@ def start_task(self, task_name=None) -> None:
         )
         self._active_task = task_name
 
+    def _resume_scheduler_if_paused_by_task(self) -> None:
+        """
+        Restart the periodic scheduler if, and only if, start_task paused a running
+        one. No-op for pure start_task/stop_task usage, and when called from stop()
+        where the scheduler has already been released.
+        """
+        if self._scheduler is not None and self._scheduler_paused_by_task:
+            self._scheduler_paused_by_task = False
+            self._scheduler.start()
+
     def stop_task(self, task_name: str = None) -> EmissionsData:
         """
         Stop tracking a dedicated execution task. Delta energy is computed by task, to isolate its contribution to total
@@ -806,6 +824,9 @@ def stop_task(self, task_name: str = None) -> EmissionsData:
         task_name = task_name if task_name else self._active_task
         if self._tasks.get(task_name) is None:
             logger.warning("stop_task : No active task to stop.")
+            # Still resume, so an unknown task name does not leave the periodic
+            # scheduler paused for the rest of the run.
+            self._resume_scheduler_if_paused_by_task()
             return None
         self._measure_power_and_energy()
         emissions_data = (
@@ -850,6 +871,8 @@ def stop_task(self, task_name: str = None) -> EmissionsData:
         self._active_task = None
         self._active_task_emissions_at_start = None  # Clear task-specific start data
 
+        self._resume_scheduler_if_paused_by_task()
+
         return task_emission_data
 
     @suppress(Exception)

tests/test_emissions_tracker.py

@@ -765,6 +765,51 @@ def test_start_task_returns_when_engine_initialization_fails(
             any("Tracker not initialized" in message for message in logs.output)
         )
 
+    def test_stop_task_resumes_scheduler_only_if_start_task_paused_it(
+        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,
+    ):
+        # Pure start_task/stop_task usage: no periodic scheduler must be left running.
+        tracker = EmissionsTracker(save_to_file=False)
+        tracker.start_task("task-only")
+        tracker.stop_task()
+        self.assertTrue(tracker._scheduler._stopped)
+        tracker.stop()
+
+        # start() then start_task/stop_task: the paused scheduler must be resumed.
+        tracker = EmissionsTracker(save_to_file=False)
+        tracker.start()
+        tracker.start_task("task-in-run")
+        self.assertTrue(tracker._scheduler._stopped)
+        tracker.stop_task()
+        self.assertFalse(tracker._scheduler._stopped)
+        tracker.stop()
+
+        # A second start_task on an already measuring tracker must not clear the
+        # "paused by task" flag, otherwise the scheduler is never resumed.
+        tracker = EmissionsTracker(save_to_file=False)
+        tracker.start()
+        tracker.start_task("first-task")
+        tracker.start_task("ignored-second-task")
+        self.assertTrue(tracker._scheduler._stopped)
+        tracker.stop_task()
+        self.assertFalse(tracker._scheduler._stopped)
+        tracker.stop()
+
+        # An unknown task name must not leave the scheduler paused either.
+        tracker = EmissionsTracker(save_to_file=False)
+        tracker.start()
+        tracker.start_task("known-task")
+        self.assertIsNone(tracker.stop_task("unknown-task"))
+        self.assertFalse(tracker._scheduler._stopped)
+        tracker.stop()
+
     @mock.patch("codecarbon.external.ram.RAM.measure_power_and_energy")
     @mock.patch("codecarbon.external.hardware.CPU.measure_power_and_energy")
     @mock.patch(