diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 127d368b3..6737e83ab 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -283,6 +283,9 @@ def _initialize_runtime_state(self) -> None: self._ram_used_history: List[float] = [] # The monitor thread appends while the measurement thread clears/averages self._utilization_history_lock = threading.Lock() + # A scheduled measurement still running when stop() gave up joining it + # must not overlap the final one + self._measurement_lock = threading.RLock() self._total_cpu_energy: Energy = Energy.from_energy(kWh=0) self._total_gpu_energy: Energy = Energy.from_energy(kWh=0) self._total_ram_energy: Energy = Energy.from_energy(kWh=0) @@ -1357,8 +1360,9 @@ def _do_measurements(self) -> None: def _measure_power_and_energy_if_stale(self, min_interval_s: float = 0.05) -> None: """Measure only if the last sample is older than ``min_interval_s``.""" - if time.perf_counter() - self._last_measured_time >= min_interval_s: - self._measure_power_and_energy() + with self._measurement_lock: + if time.perf_counter() - self._last_measured_time >= min_interval_s: + self._measure_power_and_energy() def _measure_power_and_energy(self) -> None: """ @@ -1366,6 +1370,10 @@ def _measure_power_and_energy(self) -> None: every `self._measure_power_secs` seconds. :return: None """ + with self._measurement_lock: + self._measure_power_and_energy_locked() + + def _measure_power_and_energy_locked(self) -> None: try: last_duration = time.perf_counter() - self._last_measured_time except AttributeError as e: diff --git a/codecarbon/external/scheduler.py b/codecarbon/external/scheduler.py index b60aaad4d..4256e67fa 100644 --- a/codecarbon/external/scheduler.py +++ b/codecarbon/external/scheduler.py @@ -1,52 +1,111 @@ -from threading import Lock, Timer +import math +import time +from threading import Event, Lock, Thread, current_thread + +from codecarbon.external.logger import logger + +# Longest stop() waits for the in-flight call, a guess at "a measurement should +# never take longer than this", so a wedged one cannot hang the caller for long. +STOP_JOIN_TIMEOUT_SECONDS = 5.0 + + +class MonotonicClock: + """The real clock. Tests swap in a fake one to drive the loop.""" + + monotonic = staticmethod(time.monotonic) + + @staticmethod + def wait(event, timeout): + return event.wait(timeout) class PeriodicScheduler: """ - A periodic task running in threading.Timers - From https://stackoverflow.com/a/18906292/14541668 + Run ``function`` every ``interval`` seconds on a single daemon thread. + + The deadline is absolute, so the cadence does not drift with the time the + function itself takes. A tick that overruns its slot is skipped rather than + queued, so the function is never re-entered. """ def __init__(self, interval, function, *args, **kwargs): """ - Init the scheduler. You have to call start() after initialization. - ::interval:: interval in seconds to run the function. - ::function:: function to run. + ::interval:: in seconds, the delay between two calls to function. + ::function:: the function to call. ::args:: args to pass to the function. ::kwargs:: kwargs to pass to the function. """ - self._lock = Lock() - self._timer = None - self.function = function self.interval = interval + self.function = function self.args = args self.kwargs = kwargs - self._stopped = True + # Replace before start() to control time, e.g. in tests. + self.clock = MonotonicClock() + self._lock = Lock() + self._thread = None + self._stop_event = None + + @property + def _stopped(self): + return self._stop_event is None or self._stop_event.is_set() - def start(self, from_run=False): + def start(self): """ - Start the scheduler. - ::from_run:: For internal purposes to allow re-scheduling - Please do not use from_run=True until you know what you do ! + Start the scheduler. Calling it on a running scheduler is a no-op. + If a previous run is still finishing its callback (a timed-out + stop()), the new loop waits for it before its first tick, so the + function is never run concurrently and the restart is not lost. """ - self._lock.acquire() - if from_run or self._stopped: - self._stopped = False - self._timer = Timer(self.interval, self._run) - self._timer.daemon = True - self._timer.start() - self._lock.release() - - def _run(self): - self.start(from_run=True) - self.function(*self.args, **self.kwargs) + with self._lock: + if self._thread is not None and not self._stopped: + return + self._stop_event = Event() + self._thread = Thread( + target=self._loop, + args=(self._stop_event, self._thread), + daemon=True, + name=f"codecarbon-{getattr(self.function, '__name__', 'scheduler')}", + ) + self._thread.start() + + def _loop(self, stop_event, previous=None): + if previous is not None: + previous.join() + clock = self.clock + next_call = clock.monotonic() + self.interval + while not clock.wait(stop_event, max(0.0, next_call - clock.monotonic())): + try: + self.function(*self.args, **self.kwargs) + except Exception: # noqa: BLE001 - must not kill the only thread + logger.error("Scheduled measurement failed", exc_info=True) + next_call += self.interval + overrun = clock.monotonic() - next_call + if overrun > 0 and self.interval > 0: + # Skip the slots the call overran: no burst of catch-up ticks, + # and no back-to-back call either. + next_call += math.ceil(overrun / self.interval) * self.interval def stop(self): """ - Stop the scheduler. + Stop the scheduler and wait for the in-flight call to return, at most + min(interval, STOP_JOIN_TIMEOUT_SECONDS) seconds. """ - if not self._stopped: - self._lock.acquire() - self._stopped = True - self._timer.cancel() - self._lock.release() + with self._lock: + if self._stop_event is not None: + self._stop_event.set() + thread = self._thread + if thread is None or thread is current_thread(): + # Called from the callback itself: it is still running, so keep + # the reference for the next start() to wait on. + return + thread.join(min(self.interval, STOP_JOIN_TIMEOUT_SECONDS)) + if thread.is_alive(): + # Keep the reference so the next start() waits for it. + logger.warning( + "Scheduled measurement still running after stop(); it will" + " finish in the background" + ) + else: + with self._lock: + if self._thread is thread: + self._thread = None diff --git a/tests/test_emissions_tracker_constant.py b/tests/test_emissions_tracker_constant.py index ba3f4407e..6cfe85b17 100644 --- a/tests/test_emissions_tracker_constant.py +++ b/tests/test_emissions_tracker_constant.py @@ -1,5 +1,6 @@ import os import tempfile +import threading import time import unittest from unittest import mock @@ -12,6 +13,8 @@ OfflineEmissionsTracker, track_emissions, ) +from codecarbon.external import scheduler as scheduler_module +from tests.testutils import FakeClock def heavy_computation(run_time_secs: int = 3): @@ -202,3 +205,55 @@ def test_carbon_tracker_offline_region_error(self): finally: tracker.stop() + + +class TestFinalMeasurementWaitsForInFlightTick(unittest.TestCase): + @mock.patch.object(scheduler_module, "STOP_JOIN_TIMEOUT_SECONDS", 0.01) + def test_stop_does_not_overlap_a_slow_scheduled_measurement(self): + tracker = OfflineEmissionsTracker( + country_iso_code="USA", + measure_power_secs=1, + output_methods=[], + allow_multiple_runs=True, + ) + clock = tracker._scheduler.clock = FakeClock() + tracker.start() + main = threading.current_thread() + do_measurements = tracker._do_measurements + in_flight = {"now": 0, "max": 0} + entered = threading.Event() + # Set once stop() reaches its final measurement, whether it waits for + # the lock (fixed) or measures right away (the overlap). + final_measurement_reached = threading.Event() + + def slow_do_measurements(): + in_flight["now"] += 1 + in_flight["max"] = max(in_flight["max"], in_flight["now"]) + if threading.current_thread() is main: + final_measurement_reached.set() + else: + entered.set() + final_measurement_reached.wait(5) + do_measurements() + in_flight["now"] -= 1 + + class LockSpy: + def __init__(self, lock): + self.lock = lock + + def __enter__(self): + if threading.current_thread() is main: + final_measurement_reached.set() + return self.lock.__enter__() + + def __exit__(self, *exc): + return self.lock.__exit__(*exc) + + tracker._do_measurements = slow_do_measurements + tracker._measurement_lock = LockSpy(tracker._measurement_lock) + clock.advance(1, settle=False) + self.assertTrue(entered.wait(5), "the scheduled measurement never ran") + # Make the last sample stale so stop() takes its final measurement. + tracker._last_measured_time -= 1 + tracker.stop() # the join times out while the tick is still running + self.assertEqual(in_flight["max"], 1, "final measurement overlapped a tick") diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 000000000..b190f1be8 --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,246 @@ +import threading +import time +import unittest +from unittest import mock + +from codecarbon.external import scheduler as scheduler_module +from codecarbon.external.scheduler import PeriodicScheduler +from tests.testutils import FakeClock + +# Fake seconds: time only moves when a test calls clock.advance(). +INTERVAL = 1.0 +# Real seconds a negative check waits for something that must not happen. It +# only bounds the wait: a correct scheduler passes however slow the machine. +NEVER = 0.1 + + +class TestPeriodicScheduler(unittest.TestCase): + def make_scheduler(self, function, interval=INTERVAL): + scheduler = PeriodicScheduler(interval, function) + self.clock = scheduler.clock = FakeClock() + self.addCleanup(scheduler.stop) + return scheduler + + def test_ticks_run_on_a_single_thread(self): + """Regression guard: one long-lived thread, not one thread per tick.""" + names = [] + scheduler = self.make_scheduler( + lambda: names.append(threading.current_thread().name) + ) + scheduler.start() + for _ in range(5): + self.clock.advance(INTERVAL) + scheduler.stop() + self.assertEqual(len(names), 5) + self.assertEqual(len(set(names)), 1, f"expected one worker, got {names}") + + def test_slow_function_is_never_re_entered(self): + """A function slower than the interval must not overlap with itself.""" + state = {"in_flight": 0, "overlaps": 0, "calls": 0} + + def slow(): + state["calls"] += 1 + state["in_flight"] += 1 + if state["in_flight"] > 1: + state["overlaps"] += 1 + self.clock.sleep(INTERVAL * 2.4) + state["in_flight"] -= 1 + + scheduler = self.make_scheduler(slow) + scheduler.start() + for _ in range(20): + self.clock.advance(INTERVAL / 4) + scheduler.stop() + self.assertGreater(state["calls"], 1) + self.assertEqual(state["overlaps"], 0) + + def test_overrunning_call_waits_for_the_next_slot(self): + """A call longer than the interval skips the missed slots, not the wait.""" + starts = [] + + def slow(): + starts.append(self.clock.monotonic()) + self.clock.sleep(INTERVAL * 2.5) + + scheduler = self.make_scheduler(slow) + scheduler.start() + # The first slot is 4 steps away, then each 2.5 s call leaves half an + # interval, 2 steps, before the next free slot. + for _ in range(10): + self.clock.advance(INTERVAL / 4) + scheduler.stop() + self.assertEqual(starts, [1.0, 4.0, 7.0, 10.0]) + + def test_cadence_does_not_drift(self): + """Deadlines are absolute: the time a call takes is not added to them.""" + starts = [] + + def takes_a_while(): + starts.append(self.clock.monotonic()) + self.clock.sleep(INTERVAL * 0.5) + + scheduler = self.make_scheduler(takes_a_while) + scheduler.start() + for _ in range(10): + self.clock.advance(INTERVAL / 4) + scheduler.stop() + self.assertEqual(starts, [1.0, 2.0, 3.0, 4.0]) + + def test_stop_is_final(self): + calls = [] + scheduler = self.make_scheduler(lambda: calls.append(1)) + scheduler.start() + thread = scheduler._thread + self.assertTrue(self.clock.wait_parked()) + scheduler.stop() + self.assertTrue(scheduler._stopped) + self.assertFalse(thread.is_alive(), "stop() returned before the loop ended") + self.assertIsNone(scheduler._thread) + self.clock.sleep(INTERVAL * 10) + self.assertEqual(calls, []) + + def test_stop_before_start_and_double_stop_do_not_raise(self): + scheduler = self.make_scheduler(lambda: None) + scheduler.stop() + scheduler.start() + scheduler.stop() + scheduler.stop() + self.assertTrue(scheduler._stopped) + + def test_start_is_idempotent_and_restartable(self): + calls = [] + scheduler = self.make_scheduler(lambda: calls.append(1)) + scheduler.start() + first = scheduler._thread + scheduler.start() + self.assertIs(scheduler._thread, first, "start() armed a second thread") + self.clock.advance(INTERVAL) + scheduler.stop() + self.assertFalse(first.is_alive()) + self.assertEqual(len(calls), 1) + + scheduler.start() + second = scheduler._thread + self.clock.advance(INTERVAL) + scheduler.stop() + self.assertEqual(len(calls), 2, "restart did not resume ticking") + self.assertFalse(second.is_alive()) + + @mock.patch.object(scheduler_module, "STOP_JOIN_TIMEOUT_SECONDS", 0.01) + def test_start_after_a_timed_out_stop_resumes_without_overlap(self): + """A restart during a wedged callback waits for it, then ticks again.""" + release = threading.Event() + blocked_once = threading.Event() + threads = [] + + def wedged(): + threads.append(threading.current_thread()) + if not blocked_once.is_set(): + blocked_once.set() + release.wait(5) + + scheduler = self.make_scheduler(wedged) + scheduler.start() + first = scheduler._thread + self.clock.advance(INTERVAL, settle=False) + self.assertTrue(blocked_once.wait(2), "callback never ran") + + with self.assertLogs("codecarbon", level="WARNING"): + scheduler.stop() # join times out, callback still blocked + self.assertTrue(first.is_alive()) + self.assertIs(scheduler._thread, first, "stop() dropped the wedged thread") + + scheduler.start() # must not be refused + second = scheduler._thread + self.assertIsNot(second, first) + self.assertFalse(scheduler._stopped) + self.assertFalse( + self.clock.wait_parked(NEVER), "the new loop runs during the old call" + ) + + release.set() + first.join(2) + self.assertFalse(first.is_alive(), "the wedged thread never exited") + self.clock.advance(INTERVAL) + scheduler.stop() + self.assertEqual(threads, [first, second]) + + @mock.patch.object(scheduler_module, "STOP_JOIN_TIMEOUT_SECONDS", 0.01) + def test_stop_from_the_callback_then_start_does_not_overlap(self): + """stop() called by the callback itself must not let start() overlap it.""" + release = threading.Event() + stopped_itself = threading.Event() + threads = [] + + def stops_itself(): + threads.append(threading.current_thread()) + if not stopped_itself.is_set(): + scheduler.stop() + stopped_itself.set() + release.wait(5) + + scheduler = self.make_scheduler(stops_itself) + scheduler.start() + first = scheduler._thread + self.clock.advance(INTERVAL, settle=False) + self.assertTrue(stopped_itself.wait(2), "callback never ran") + + scheduler.start() + second = scheduler._thread + self.assertFalse( + self.clock.wait_parked(NEVER), "the new loop runs during the old call" + ) + + release.set() + first.join(2) + self.clock.advance(INTERVAL) + scheduler.stop() + self.assertEqual(threads, [first, second]) + + def test_concurrent_starts_arm_a_single_loop(self): + """Two threads calling start() at once must not create two loops.""" + + class SlowEvent(threading.Event): + def __init__(self): + time.sleep(NEVER) # widen the check-then-act window + super().__init__() + + def concurrent_start_target(): + pass + + scheduler = self.make_scheduler(concurrent_start_target) + barrier = threading.Barrier(2) + + def start(): + barrier.wait() + scheduler.start() + + with mock.patch.object(scheduler_module, "Event", SlowEvent): + starters = [threading.Thread(target=start) for _ in range(2)] + for starter in starters: + starter.start() + for starter in starters: + starter.join() + scheduler.stop() + alive = [ + t.name + for t in threading.enumerate() + if t.name == "codecarbon-concurrent_start_target" + ] + self.assertEqual(alive, [], "a loop survived stop(): two loops were armed") + + def test_exception_does_not_kill_the_loop(self): + calls = [] + + def flaky(): + calls.append(1) + if len(calls) == 1: + raise ValueError("boom") + + scheduler = self.make_scheduler(flaky) + scheduler.start() + with self.assertLogs("codecarbon", level="ERROR"): + self.clock.advance(INTERVAL) + self.clock.advance(INTERVAL) + scheduler.stop() + self.assertEqual(len(calls), 2) diff --git a/tests/testutils.py b/tests/testutils.py index e3d1dc2d1..732ce20e1 100644 --- a/tests/testutils.py +++ b/tests/testutils.py @@ -1,4 +1,5 @@ import builtins +import threading import unittest from pathlib import Path @@ -33,3 +34,55 @@ def conditional_open_func(path, *args, **kwargs): return conditional_open_func return mocked_open + + +class FakeClock: + """ + Fake time for ``PeriodicScheduler.clock``: the loop only moves when the + test calls advance(), so tests assert on exact ticks instead of sleeping. + """ + + def __init__(self): + self.now = 0.0 + self._cond = threading.Condition() + self._parked = set() + + def monotonic(self): + with self._cond: + return self.now + + def sleep(self, seconds): + """Let a scheduled function take ``seconds`` of fake time.""" + with self._cond: + self.now += seconds + + def wait(self, event, timeout): + me = threading.current_thread() + with self._cond: + deadline = self.now + timeout + while not event.is_set() and self.now < deadline: + self._parked.add(me) + self._cond.notify_all() + # stop() sets the event without notifying us, so poll it. + self._cond.wait(0.001) + self._parked.discard(me) + self._cond.notify_all() + return event.is_set() + + def wait_parked(self, timeout=2.0): + """True once a loop is waiting for its next deadline.""" + with self._cond: + return self._cond.wait_for(lambda: self._parked, timeout) + + def advance(self, seconds, settle=True): + """ + Move time forward once a loop is waiting. With ``settle``, return only + when the loop has run what was due and is waiting again. + """ + assert self.wait_parked(), "no scheduler loop is waiting" + with self._cond: + self.now += seconds + self._parked.clear() + self._cond.notify_all() + if settle: + assert self.wait_parked(), "the scheduler loop did not settle"