diff --git a/codecarbon/external/scheduler.py b/codecarbon/external/scheduler.py index b60aaad4d..107e6c643 100644 --- a/codecarbon/external/scheduler.py +++ b/codecarbon/external/scheduler.py @@ -1,10 +1,16 @@ -from threading import Lock, Timer +import time +from threading import Event, Thread, current_thread + +from codecarbon.external.logger import logger 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): @@ -15,38 +21,65 @@ def __init__(self, interval, function, *args, **kwargs): ::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 + self._stop_event = Event() + self._thread = None + + @property + def _stopped(self): + return self._thread is None or not self._thread.is_alive() - 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. """ - 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) - - def stop(self): + if not self._stopped: + return + self._stop_event.clear() + self._thread = Thread( + target=self._loop, + daemon=True, + name=f"codecarbon-{getattr(self.function, '__name__', 'scheduler')}", + ) + self._thread.start() + + def _loop(self): + next_call = time.monotonic() + self.interval + while not self._stop_event.wait(max(0.0, next_call - time.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) + # Absolute deadline so the cadence does not drift, but if we + # overran (or the process was suspended) skip ahead instead of + # firing a burst of catch-up ticks. + next_call += self.interval + now = time.monotonic() + if next_call <= now: + next_call = now + self.interval + + def stop(self, timeout=None): """ - Stop the scheduler. + Stop the scheduler and wait for the in-flight call to return. + ::timeout:: seconds to wait for the running function, bounded by + default so a wedged measurement cannot hang the caller for long. """ - if not self._stopped: - self._lock.acquire() - self._stopped = True - self._timer.cancel() - self._lock.release() + self._stop_event.set() + thread = self._thread + if thread is not None and thread is not current_thread(): + # ponytail: 5s cap is a guess at "a measurement should never take + # longer than this"; make it configurable if a slow hardware + # backend ever needs more. + thread.join(min(self.interval, 5.0) if timeout is None else timeout) + if thread.is_alive(): + logger.warning( + "Scheduled measurement did not return in time; the " + "scheduler thread is still running and will exit on its " + "own. A new start() is ignored until then." + ) + # `_thread` is deliberately kept: `_stopped` tests `is_alive()`, so a + # still-running thread keeps `start()` a no-op instead of letting it + # clear `_stop_event` and release the old loop alongside a new one. diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 000000000..1d3cc81f1 --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,124 @@ +import threading +import time +import unittest + +from codecarbon.external.scheduler import PeriodicScheduler + +INTERVAL = 0.05 + + +class TestPeriodicScheduler(unittest.TestCase): + def test_ticks_run_on_a_single_thread(self): + """Regression guard: one long-lived thread, not one thread per tick.""" + names = set() + scheduler = PeriodicScheduler( + INTERVAL, lambda: names.add(threading.current_thread().name) + ) + scheduler.start() + time.sleep(INTERVAL * 10) + scheduler.stop() + # A one-element set also proves at least one tick happened. + self.assertEqual(len(names), 1, f"expected one worker thread, 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 + time.sleep(INTERVAL * 2.4) + state["in_flight"] -= 1 + + scheduler = PeriodicScheduler(INTERVAL, slow) + scheduler.start() + time.sleep(INTERVAL * 20) + scheduler.stop() + self.assertGreater(state["calls"], 1) + self.assertEqual(state["overlaps"], 0) + + def test_stop_is_prompt_and_final(self): + calls = [] + scheduler = PeriodicScheduler(10.0, lambda: calls.append(1)) + scheduler.start() + before = time.monotonic() + scheduler.stop() + self.assertLess(time.monotonic() - before, 1.0) + self.assertTrue(scheduler._stopped) + time.sleep(0.1) + self.assertEqual(calls, []) + + def test_stop_before_start_and_double_stop_do_not_raise(self): + scheduler = PeriodicScheduler(INTERVAL, lambda: None) + scheduler.stop() + scheduler.start() + scheduler.stop() + scheduler.stop() + self.assertTrue(scheduler._stopped) + + def test_start_is_idempotent_and_restartable(self): + calls = [] + scheduler = PeriodicScheduler(INTERVAL, lambda: calls.append(1)) + scheduler.start() + first = scheduler._thread + scheduler.start() + self.assertIs(scheduler._thread, first, "start() armed a second thread") + scheduler.stop() + self.assertFalse(first.is_alive()) + stopped_at = len(calls) + + scheduler.start() + second = scheduler._thread + time.sleep(INTERVAL * 4) + scheduler.stop() + self.assertGreater(len(calls), stopped_at, "restart did not resume ticking") + self.assertFalse(second.is_alive()) + + def test_start_after_a_timed_out_stop_does_not_run_two_loops(self): + """A wedged callback must not be released back into its loop. + + `stop()` gives up after a bounded join. If it forgot the thread anyway, + the next `start()` would clear `_stop_event` and the still-blocked + thread would resume ticking alongside the freshly started one. + """ + release = threading.Event() + idents = set() + blocked_once = threading.Event() + + def wedged(): + idents.add(threading.get_ident()) + if not blocked_once.is_set(): + blocked_once.set() + release.wait(5) + + scheduler = PeriodicScheduler(INTERVAL, wedged) + scheduler.start() + first = scheduler._thread + self.assertTrue(blocked_once.wait(2), "callback never ran") + + scheduler.stop() # join times out, callback still blocked + self.assertTrue(first.is_alive()) + + scheduler.start() # must not arm a second loop + release.set() + time.sleep(INTERVAL * 6) + scheduler.stop() + + self.assertEqual(len(idents), 1, f"two scheduler loops ran: {idents}") + self.assertFalse(first.is_alive()) + + def test_exception_does_not_kill_the_loop(self): + calls = [] + + def flaky(): + calls.append(1) + if len(calls) == 1: + raise ValueError("boom") + + scheduler = PeriodicScheduler(INTERVAL, flaky) + scheduler.start() + time.sleep(INTERVAL * 6) + scheduler.stop() + self.assertGreater(len(calls), 1)