Run ML training off the event loop and let a shutdown abandon it - #5112
springfall2008 wants to merge 1 commit into
Conversation
| # so calling it directly would pin this component's event loop for its whole duration: | ||
| # api_stop would go unnoticed, the run_timeout watchdog could never fire, and a shutdown | ||
| # landing mid-training would sit waiting on this component's thread (#5075). | ||
| val_mae = await asyncio.to_thread( |
There was a problem hiding this comment.
Bug (high): watchdog-cancel now orphans the training thread and can start a second concurrent training on the same predictor. Pre-PR the synchronous train_curriculum pinned the loop, so the 2h run_timeout watchdog (load_ml_component.py:98) could never fire mid-training; post-PR the loop is free, so component_base.py:266-277 wait_for(shield(task), run_timeout) + task.cancel() can fire while await asyncio.to_thread(...) is pending. CancelledError lands in the coroutine but the worker thread keeps training; api_stop is False so stop_callback never trips. 60s later start() re-runs run(), should_train is still true, and a second train_curriculum starts on the same LoadPredictor — both threads mutate self.weights/self.biases/Adam/normalisation state in place (_adam_update), and both may save() to model_filepath. Consider making the stop hook also trip on a per-run cancel flag set by the timeout path.
| stop_callback=self.training_stop_requested, | ||
| ) | ||
|
|
||
| if val_mae is None and self.training_stop_requested(): |
There was a problem hiding this comment.
Bug (high, same watchdog path): a cancelled run() never stamps last_train_time/initial_training_done, so every retry relaunches the full initial curriculum — a self-sustaining runaway. The statements after the await asyncio.to_thread(...) (lines 1028-1029) don't execute when the watchdog cancels the coroutine, even though the orphaned thread finishes its full training. On re-entry is_initial is still True and retrain_age is still >= RETRAIN_INTERVAL, so each 60s-tick run() launches another full 100-epoch/9-pass initial curriculum, each again exceeding run_timeout and leaving another orphaned thread racing on self.predictor.
| ) | ||
| ) | ||
|
|
||
| if aborted: |
There was a problem hiding this comment.
Bug (medium): the abort path restores only weights — the in-memory model is left a hybrid, contradicting the comment. train() has already mutated self.feature_mean/self.feature_std (_fit_windowed_normalisation, EMA blend on fine-tune) and target statistics and _reset_adam_optimizer() before the first _should_stop() check; the post-loop restore only puts back best_weights. The comment claims the model is 'left consistent' but only the weights are. Harmless on a real shutdown, but live on the watchdog-cancel path (cancelled run(), component keeps running): every subsequent predict() normalises inputs with statistics the restored weights were never trained against, systematically skewing the published load forecast until the next successful training.
| # not finish: report failure so the caller neither publishes nor saves a partial model, | ||
| # and skip the AR rollout diagnostic, which is itself minutes of the work we just abandoned. | ||
| return None | ||
|
|
There was a problem hiding this comment.
Bug (medium-low): an abandoned initial curriculum leaves earlier passes' stamps, so the partial model reads as active. train() stamps self.training_timestamp/self.validation_mae/self.epochs_trained at the end of EVERY completed pass, including intermediate curriculum passes. When a later pass aborts, those stamps survive (only the component-side flags are skipped). _update_model_status() -> is_valid() judges purely on training_timestamp age and validation_mae, neither reflecting the abandonment, so a model trained on only the first curriculum window (e.g. 7 of 28 days) reports 'active', and the fresh timestamp suppresses the ml_max_model_age_hours staleness retrain.
| # Weights are restored above so the in-memory model is left consistent, but the run did | ||
| # not finish: report failure so the caller neither publishes nor saves a partial model, | ||
| # and skip the AR rollout diagnostic, which is itself minutes of the work we just abandoned. | ||
| return None |
There was a problem hiding this comment.
Bug (medium-low): no stop check around _ar_rollout_diagnostic, and the final pass still publishes/saves after a stop lands. The abort path skips the rollout only when the run already aborted at an epoch check; a stop landing after the final pass's last epoch check is waited out through the full rollout (the comment here calls it 'itself minutes of the work'), and because val_mae is then not None, the component's abandon gate (load_ml_component.py:1024) is bypassed: last_train_time is stamped, model_valid set and predictor.save() does disk I/O mid-shutdown. Also: train_curriculum calls train() once per intermediate pass, so an initial curriculum pays up to 9 rollouts whose results are overwritten each time — only the final pass's self.rollout_mae/pattern_mae survive; threading run_diagnostic=False through for intermediate passes would remove most of that cost.
| # Training abandoned for shutdown - a re-fetch and a prediction cycle here | ||
| # would just add their own delay to the stop we are already answering | ||
| self.log("ML Component: Stopping, skipping the post-training prediction cycle") | ||
| return True |
There was a problem hiding this comment.
Bug (low, watchdog path): run()'s finally clears load_ml_calculating while the orphaned training thread is still running NumPy work, so is_calculating() lies. When the run_timeout watchdog cancels run() at the to_thread await, this finally still executes while the worker thread keeps training (api_stop is False, so its stop hook never trips). is_calculating() then returns False during genuine heavy NumPy work, and the flag's own comment in this block documents the reason it matters (fork() is unsafe with live threads) — the same hazard the flag was kept for.
| # Training abandoned for shutdown - a re-fetch and a prediction cycle here | ||
| # would just add their own delay to the stop we are already answering | ||
| self.log("ML Component: Stopping, skipping the post-training prediction cycle") | ||
| return True |
There was a problem hiding this comment.
Bug (low): the new early return also fires when training COMPLETED successfully, skipping _update_model_status() and run()'s trailing update_success_timestamp(). A stop landing during the final epoch's work (after the last per-epoch check) lets train_curriculum return a valid val_mae: the model is saved and last_train_time stamped here in _do_training, but run() returns before the status update — the fresh model's status is never published and last_success_timestamp can sit 2h stale at the moment is_alive() is next consulted during shutdown/restart reporting.
| # Training abandoned for shutdown - a re-fetch and a prediction cycle here | ||
| # would just add their own delay to the stop we are already answering | ||
| self.log("ML Component: Stopping, skipping the post-training prediction cycle") | ||
| return True |
There was a problem hiding this comment.
Bug (low): the abandon-vs-failure classification by re-reading api_stop misreports a genuine failure. train()/train_curriculum() return None for both 'all passes failed' and 'aborted'; this check sends the all-passes-failed case down the 'Training abandoned because the component is stopping' branch whenever api_stop happens to be set, the 'Warn: ML Component: Training failed' branch and its error accounting never fire, and (because run() returns early) _update_model_status() never runs — so a recurring data problem is invisible in that cycle's log and model_status.
| # Training abandoned for shutdown - a re-fetch and a prediction cycle here | ||
| # would just add their own delay to the stop we are already answering | ||
| self.log("ML Component: Stopping, skipping the post-training prediction cycle") | ||
| return True |
There was a problem hiding this comment.
Gap (low): shutdown responsiveness is not re-established for the pre-training wait loop. The while self.base.prediction_started: await asyncio.sleep(0.5) loop just above (line 771, inside this same run()) never checks api_stop, and it sits upstream of this early return — a stop() landing while Predbat's prediction cycle is running still spins here until that cycle completes, the same delay class #5075 set out to remove.
| try: | ||
| return bool(stop_callback()) | ||
| except Exception as e: | ||
| self.log("Warn: ML Predictor: stop callback raised {}, continuing training".format(e)) |
There was a problem hiding this comment.
Design note (low): the exception polarity makes a broken stop hook silently void the whole fix. A raising stop_callback is logged and treated as 'do not stop' — the same swallow policy as progress_callback, but here the failure cost is a full-length, unaborted minutes-long run. Unlike progress_callback (whose failure only costs a liveness pulse), any future refactor that makes training_stop_requested raise (e.g. api_stop moved behind a property that raises once the component registry is torn down) turns shutdown back into a full-curriculum hang with only a per-epoch warn to show for it, while progress_callback keeps stamping last_success_timestamp so the health monitor reports the component alive. Worth at least counting/escalating repeated hook failures rather than swallowing forever.
| # Training abandoned for shutdown - a re-fetch and a prediction cycle here | ||
| # would just add their own delay to the stop we are already answering | ||
| self.log("ML Component: Stopping, skipping the post-training prediction cycle") | ||
| return True |
There was a problem hiding this comment.
Cosmetic (low): the early return re-asserts api_started=True during teardown. A True return flows into component_base.start()'s 'if run_result: if not self.api_started: self.api_started = True', so the in-flight run() sets api_started back to True after Components.stop() already cleared it — is_alive() reports load_ml healthy during the shutdown window until start()'s loop exits and re-clears it. Benign but contradicts stop()'s bookkeeping.
| ) | ||
|
|
||
| if pass_mae is None and self._should_stop(stop_callback): | ||
| self.log("ML Predictor: Curriculum training aborted during pass {}/{} - stop requested".format(pass_idx + 1, total_passes)) |
There was a problem hiding this comment.
Cleanup: this after-pass check is the redundant one of the five _should_stop sites. train() already returns None on abort and api_stop is sticky, so the top-of-next-pass check (line 1630) — or the before-final-pass check (line 1678) when this is the last intermediate pass — returns None with an identical outcome; this site only changes which of three near-identical messages fires. The minimal set is 1574, 1630, 1678 and the after-final check (1701), which must stay: deleting it lets an aborted final pass fall through to 'if final_mae is not None' and return the previous pass's val_mae as a completed curriculum, which the component would save and stamp last_train_time from.
|
|
||
| def training_stop_requested(self): | ||
| """Return True when the component is stopping, so an in-flight training run abandons itself.""" | ||
| return self.api_stop |
There was a problem hiding this comment.
Cleanup/altitude (low): this hook serves exactly one cancellation source, and the same file now spells the predicate two ways. run() (line 780) still tests if self.api_stop: directly while the hook goes through this wrapper, and ComponentBase.start() also exits on fatal_error (component_base.py:248) — a fatal_error mid-training leaves the thread running with no abandon. The repo already has several hand-rolled cancel mechanisms (api_stop polling in every long-lived component, db_manager's threading.Event pair, chat.py's run_coroutine_threadsafe bridge); this is a third. A single per-run cancellation handle that stop(), the watchdog cancel and fatal_error all feed would have made the watchdog orphan problem (see comment on line 1005) impossible by construction.
| """Release the trainer, which only happens if the loop is still being serviced.""" | ||
| loop_made_progress.set() | ||
|
|
||
| task = asyncio.create_task(releaser()) |
There was a problem hiding this comment.
Test robustness: this handshake can pass spuriously. The single-shot threading.Event only proves the loop ran at SOME point before the trainer returned, not that the loop stayed free during the trainer call: any earlier await in _do_training that yields once (today the async with self.data_lock, which happens not to yield) lets the releaser fire, after which a synchronous loop-pinning train_curriculum call still passes the assert. A regression that added await asyncio.sleep(0) before a direct synchronous trainer call would reintroduce #5075 with this test green. A loop-tick assertion collected after the trainer returns (or a second event set after the trainer's own return) would close that hole. Also worth knowing: if the regression under test occurs, the executor thread blocks here for the full 10s and asyncio.run's shutdown_default_executor() then joins it, adding a silent 10s penalty to the shared-process suite.
| ("training_progress_callback", _test_training_progress_callback, "Training reports liveness through progress_callback"), | ||
| ("component_alive_during_training", _test_component_marks_itself_alive_during_training, "Component keeps its success timestamp fresh while training"), | ||
| ("training_stop_callback", _test_training_stops_on_stop_callback, "Training abandons a run when the stop hook trips"), | ||
| ("training_off_event_loop", _test_component_training_does_not_block_event_loop, "Component training leaves its event loop free to run"), |
There was a problem hiding this comment.
Test-coverage gap: the mid-curriculum abort paths are untested. This test's curriculum case stubs train() and only exercises the pre-loop check (train_curriculum's top-of-function check returns None before any pass), so the 'aborted before pass N', 'aborted during pass N' and 'aborted before final pass' paths — and the component-side distinction between an aborted run and an all-passes-failed run — are never executed. A hook that returns True only after the first pass, with a real (stubbed-returning) train, would cover the between-passes path and pin that an intermediate val_mae is not leaked as a completed curriculum.
1 participant
This is an automated draft PR generated from issue #5075 — a maintainer should review it before merging.
Fixes #5075
Summary
LoadMLComponent._do_training()calledpredictor.train_curriculum()synchronously from inside anasync def, so a curriculum run (many passes of many epochs, minutes long) pinned the component's event loop for its whole duration. Two changes:asyncio.to_thread(...). With the loop free,api_stopcan be noticed and therun_timeoutwatchdog incomponent_base.py:266can actually fire — as the triage comment noted,asyncio.wait_forcould never evaluate its timeout while the loop was blocked, so the existing safety net was inert for exactly this case.train()andtrain_curriculum()take a newstop_callback(defaultedNone, so every existing caller is unchanged), checked at the top of each epoch and between curriculum passes._do_trainingwires it to the component'sapi_stopvia a newtraining_stop_requested(), and reads the resultingNoneas abandoned rather than failed — no partial model is published or saved, andlast_train_timeis not stamped, so the retrain interval does not swallow the next real run.Why part 2 as well, not just the
to_threadthe issue asks forThe issue expected part 1 alone to fix the symptom. Reading the shutdown path, it does not. Components are not tasks on a shared loop:
components.py:848starts each one throughbase.create_task(), which inhass.py:224spawns a realthreading.Threadrunningasyncio.run(component.start())— so each component has its own loop in its own non-daemon thread.stop_all()(hass.py:238) callsterminate()(which logs "Predbat terminated") and then joins every component thread with a 5-minute timeout. Whether the curriculum blocks that thread's loop or blocks a worker it is awaiting, the thread stays alive either way, so the join still waits — and after the join times out, the interpreter's own non-daemon-thread shutdown waits again. That matches the report exactly: "Predbat terminated" at 21:09:11, training still logging at 21:13:58, no restart until a manual one at 21:14:51.So the cooperative stop check is the part that actually lets the thread exit;
to_threadis what makes the loop (and the watchdog) responsive while it runs.run()also returns early after an abandoned run rather than doing the post-training re-fetch and prediction cycle, which would otherwise add their own delay to the stop being answered.A hook that raises is treated as "do not stop" and logged, mirroring the existing rule for
progress_callback— the guard atload_predictor.py:1385is about a failing hook never killing a run, which a deliberate stop request does not contradict.Testing
coverage/run_pre_commit— all hooks pass, bundled suite green (4 slow tests skipped).Three new sub-tests in
tests/test_load_ml.py, all registered in that module'ssub_testslist:training_stop_callback—train()stops at the epoch the hook first trips (2 epochs complete, returnsNone); a raising hook does not abort a run;train_curriculum()starts no pass at all once the hook has tripped; omitting the hook is unchanged.training_off_event_loop— drives the real_do_trainingwith a trainer that blocks until a task scheduled on the loop releases it. Deterministic rather than timed: a loop that cannot run at all fails on the fake trainer's timeout, not on a stopwatch.training_abandoned_on_stop— the component must pass a stop hook that tracksapi_stop, and must not publish, stamp or count an abandoned run.Red/green, source stashed with
git stash push -uand the test file left in place:load_mlfails, 34/37, with all three new tests red for the right reasons —LoadPredictor.train() got an unexpected keyword argument 'stop_callback'; "the event loop made no progress while training ran"; "_do_trainingmust pass a stop_callback to train_curriculum".load_mlpasses, 37/37.Also re-ran the adjacent modules that exercise the changed trainer, all green:
load_ml_rollout,ml_memory,ml_training_perf.Blast radius via GitNexus
impact():train_curriculumLOW (4 callers, all tests),_do_trainingLOW,trainMEDIUM but all 17 affected symbols are in tests — and every signature change is a defaulted keyword argument, so no existing call site changes.detect_changes()reports low risk, 3 files, no affected execution flows.Notes
_do_traininghad two near-identical 16-linetrain_curriculum(...)call sites differing only incurriculum_step_daysandmax_intermediate_passes. Rather than duplicate the newto_thread+stop_callbackwiring into both, those two values are now selected in the existing snapshot block and there is one call. Behaviour is identical (the initial path still uses the hardcoded5/8). Say the word if you would rather keep the explicit branch.LoadMLComponent; it is now_make_training_component(), shared by the existing liveness test and the two new component tests, rather than copied a third and fourth time.run()waits onwhile self.base.prediction_started: await asyncio.sleep(0.5)before ML work without consultingapi_stop. It yields, so it does not block the loop, but it is another place a stop is not noticed promptly. Left alone as out of scope for this ticket.val_mae=inf, sincebest_val_lossseeds to infinity when there is no fine-tune baseline. Harmless (the restore is a no-op there) and only reachable on the shutdown path, so left as is rather than duplicating the restore block to dodge one log line.asyncio.to_threaduses the loop's default executor. Only one training call is ever in flight per component, so at most one worker thread is created; the triage comment raisedrun_in_executorwith an explicit single-thread executor as an alternative if you would prefer that given the past NumPy thread-pool interaction.