diff --git a/documentation/changelog.rst b/documentation/changelog.rst index 0f41c0a1d4..1634f41a9b 100644 --- a/documentation/changelog.rst +++ b/documentation/changelog.rst @@ -34,6 +34,7 @@ Bugfixes ----------- * Let storage scheduling treat missing constant SoC bounds as unconstrained lower or upper bounds [see `PR #2221 `_] * Allow root assets belonging to different accounts to share the same name, while keeping asset names unique among root assets within the same account and among children of the same parent [see `PR #2226 `_] +* Fix queued train-predict forecasting jobs losing their resolved forecast window or failing on detached database objects in workers [see `PR #2035 `_] v0.33.1 | July 1, 2026 diff --git a/flexmeasures/data/models/forecasting/pipelines/base.py b/flexmeasures/data/models/forecasting/pipelines/base.py index 4de0287600..e512371a7b 100644 --- a/flexmeasures/data/models/forecasting/pipelines/base.py +++ b/flexmeasures/data/models/forecasting/pipelines/base.py @@ -8,9 +8,9 @@ import pandas as pd from darts import TimeSeries from darts.dataprocessing.transformers import MissingValuesFiller -from flexmeasures.data.models.time_series import Sensor from timely_beliefs import utils as tb_utils +from flexmeasures.data.models.time_series import Sensor from flexmeasures.data.models.forecasting.exceptions import NotEnoughDataException @@ -90,11 +90,11 @@ def __init__( self.target = f"{target_sensor.name} (ID: {target_sensor.id})_target" self.future_regressors = [ f"{sensor.name} (ID: {sensor.id})_FR-{idx}" - for idx, sensor in enumerate(future_regressors) + for idx, sensor in enumerate(self.future) ] self.past_regressors = [ f"{sensor.name} (ID: {sensor.id})_PR-{idx}" - for idx, sensor in enumerate(past_regressors) + for idx, sensor in enumerate(self.past) ] self.predict_start = predict_start if predict_start else None self.predict_end = predict_end if predict_end else None diff --git a/flexmeasures/data/models/forecasting/pipelines/train_predict.py b/flexmeasures/data/models/forecasting/pipelines/train_predict.py index 88883e132c..26c91542b3 100644 --- a/flexmeasures/data/models/forecasting/pipelines/train_predict.py +++ b/flexmeasures/data/models/forecasting/pipelines/train_predict.py @@ -13,9 +13,11 @@ from flask import current_app from flexmeasures.data import db +from flexmeasures.data.models.data_sources import DataSource from flexmeasures.data.models.forecasting import Forecaster from flexmeasures.data.models.forecasting.pipelines.predict import PredictPipeline from flexmeasures.data.models.forecasting.pipelines.train import TrainPipeline +from flexmeasures.data.models.time_series import Sensor from flexmeasures.data.schemas.forecasting.pipeline import ( ForecasterParametersSchema, TrainPredictPipelineConfigSchema, @@ -23,6 +25,134 @@ from flexmeasures.utils.flexmeasures_inflection import p +def _sensor_id(sensor: Sensor | int | None) -> int | None: + """Return the sensor ID from a Sensor object or already-serialized ID.""" + if sensor is None: + return None + return sensor.id if isinstance(sensor, Sensor) else sensor + + +def _get_attached_sensor(sensor_id: int | None) -> Sensor | None: + """Load a sensor in the current session from a queued job payload ID.""" + if sensor_id is None: + return None + attached_sensor = db.session.get(Sensor, sensor_id) + if attached_sensor is None: + raise ValueError(f"Could not load sensor with id {sensor_id}.") + return attached_sensor + + +def _get_attached_data_source(data_source_id: int | None) -> DataSource | None: + """Load a data source in the current session from a queued job payload ID.""" + if data_source_id is None: + return None + attached_source = db.session.get(DataSource, data_source_id) + if attached_source is None: + raise ValueError(f"Could not load data source with id {data_source_id}.") + return attached_source + + +def _assert_no_orm_objects(value: Any, path: str = "payload") -> None: + """Reject ORM objects before they can be pickled into an RQ job.""" + inspection = sa_inspect(value, raiseerr=False) + if inspection is not None and hasattr(inspection, "object"): + raise ValueError( + f"Queued forecasting job {path} contains a " + f"{value.__class__.__name__} ORM object. Pass its ID instead." + ) + + if isinstance(value, dict): + for key, nested_value in value.items(): + _assert_no_orm_objects(nested_value, f"{path}.{key}") + elif isinstance(value, (list, tuple, set)): + for index, nested_value in enumerate(value): + _assert_no_orm_objects(nested_value, f"{path}[{index}]") + + +def _make_job_config_payload(config: dict[str, Any]) -> dict[str, Any]: + """Build the queued worker config payload. + + ORM-backed fields are replaced by IDs, while plain config fields are preserved. + """ + # Preserve plain config fields, but replace ORM-backed regressors by IDs. + payload = dict(config) + future_regressors = payload.pop("future_regressors", []) + past_regressors = payload.pop("past_regressors", []) + payload["future_regressor_ids"] = [ + _sensor_id(sensor) for sensor in future_regressors + ] + payload["past_regressor_ids"] = [_sensor_id(sensor) for sensor in past_regressors] + _assert_no_orm_objects(payload) + return payload + + +def _load_job_config_payload(payload: dict[str, Any]) -> dict[str, Any]: + """Restore worker config and reload regressors in the worker session.""" + config = dict(payload) + config["future_regressors"] = [ + _get_attached_sensor(sensor_id) + for sensor_id in config.pop("future_regressor_ids", []) + ] + config["past_regressors"] = [ + _get_attached_sensor(sensor_id) + for sensor_id in config.pop("past_regressor_ids", []) + ] + return config + + +def _make_job_parameters_payload(parameters: dict[str, Any]) -> dict[str, Any]: + """Build the queued worker parameter payload. + + ORM-backed fields are replaced by IDs, while plain parameter fields are preserved. + """ + # Preserve plain parameters, but replace ORM-backed sensors by IDs. + payload = dict(parameters) + sensor_id = _sensor_id(payload.pop("sensor")) + sensor_to_save_id = _sensor_id(payload.pop("sensor_to_save", None)) + if sensor_id is None: + raise ValueError("Cannot enqueue a forecasting job without a target sensor.") + payload["sensor_id"] = sensor_id + payload["sensor_to_save_id"] = sensor_to_save_id or sensor_id + _assert_no_orm_objects(payload) + return payload + + +def _load_job_parameters_payload(payload: dict[str, Any]) -> dict[str, Any]: + """Restore worker parameters and reload sensors in the worker session.""" + parameters = dict(payload) + parameters["sensor"] = _get_attached_sensor(parameters.pop("sensor_id")) + parameters["sensor_to_save"] = _get_attached_sensor( + parameters.pop("sensor_to_save_id") + ) + return parameters + + +def run_train_predict_cycle_job( + config: dict, + parameters: dict, + data_source_id: int, + delete_model: bool, + **cycle_params, +): + """Run one train-predict cycle after reconstructing worker-local ORM state.""" + pipeline = TrainPredictPipeline(delete_model=delete_model) + pipeline._config = _load_job_config_payload(config) + for key, value in pipeline._config.items(): + setattr(pipeline, key, value) + pipeline._parameters = _load_job_parameters_payload(parameters) + pipeline._data_source = _get_attached_data_source(data_source_id) + return pipeline.run_cycle(**cycle_params) + + +def run_train_predict_wrap_up_job(cycle_job_ids: list[str], queue: str = "forecasting"): + """Log the status of all cycle jobs after completion.""" + connection = current_app.queues[queue].connection + + for index, job_id in enumerate(cycle_job_ids): + status = Job.fetch(job_id, connection=connection).get_status() + logging.info(f"{queue} job-{index}: {job_id} status: {status}") + + class TrainPredictPipeline(Forecaster): __version__ = "1" @@ -46,28 +176,9 @@ def __init__( self.delete_model = delete_model self.return_values = [] # To store forecasts and jobs - @staticmethod - def _reattach_if_needed(obj): - """Re-merge a SQLAlchemy object into the current session if it is detached or expired. - - After ``db.session.commit()``, all objects in the session are expired. - When RQ pickles ``self.run_cycle`` for a worker, expired or detached - objects may raise ``DetachedInstanceError`` on attribute access. This - helper merges such objects back into the active session so they are - usable when the worker executes the job. - """ - insp = sa_inspect(obj) - if insp.detached or insp.expired: - return db.session.merge(obj) - return obj - def run_wrap_up(self, cycle_job_ids: list[str], queue: str = "forecasting"): """Log the status of all cycle jobs after completion.""" - connection = current_app.queues[queue].connection - - for index, job_id in enumerate(cycle_job_ids): - status = Job.fetch(job_id, connection=connection).get_status() - logging.info(f"{queue} job-{index}: {job_id} status: {status}") + run_train_predict_wrap_up_job(cycle_job_ids, queue) def run_cycle( self, @@ -86,25 +197,6 @@ def run_cycle( f"Starting Train-Predict cycle from {train_start} to {predict_end}" ) - # Re-attach sensor objects if they are detached after RQ pickles/unpickles self - # (this can happen when a commit expires objects before RQ serializes the job). - self._parameters["sensor"] = self._reattach_if_needed( - self._parameters["sensor"] - ) - sensor_to_save = self._parameters.get("sensor_to_save") - if sensor_to_save is not None: - self._parameters["sensor_to_save"] = self._reattach_if_needed( - sensor_to_save - ) - # Also re-attach regressor sensors stored in _config - self._config["future_regressors"] = [ - self._reattach_if_needed(s) - for s in self._config.get("future_regressors", []) - ] - self._config["past_regressors"] = [ - self._reattach_if_needed(s) for s in self._config.get("past_regressors", []) - ] - # Train model train_pipeline = TrainPipeline( future_regressors=self._config["future_regressors"], @@ -187,8 +279,8 @@ def run_cycle( return total_runtime def _compute_forecast(self, as_job: bool = False, **kwargs) -> list[dict[str, Any]]: - # Run the train-and-predict pipeline - return self.run(as_job=as_job, **kwargs) + # DataGenerator.compute already loaded kwargs into self._parameters. + return self.run(as_job=as_job) def _derive_training_period(self) -> tuple[datetime, datetime]: """Derive the effective training period for model fitting. @@ -234,7 +326,6 @@ def run( self, as_job: bool = False, queue: str = "forecasting", - **job_kwargs, ): logging.info( f"Starting Train-Predict Pipeline to predict for {self._parameters['predict_period_in_hours']} hours." @@ -285,7 +376,6 @@ def run( cycle_runtime = self.run_cycle(**train_predict_params) cumulative_cycles_runtime += cycle_runtime else: - train_predict_params["target_sensor_id"] = self._parameters["sensor"].id cycles_job_params.append(train_predict_params) train_end += cycle_frequency @@ -299,27 +389,39 @@ def run( if as_job: cycle_job_ids = [] - # Ensure the data source is attached to the current session before - # committing. get_or_create_source() only flushes (does not commit), so - # without this merge the data source would not be found by the worker. - db.session.merge(self.data_source) + job_config = _make_job_config_payload(self._config) + job_parameters = _make_job_parameters_payload(self._parameters) + sensor_id = job_parameters["sensor_id"] + sensor_to_save_id = job_parameters["sensor_to_save_id"] + + # Ensure the data source ID is available in the database when the job runs. + self._data_source = db.session.merge(self.data_source) + db.session.flush() + data_source_id = self._data_source.id db.session.commit() # job metadata for tracking # Serialize start and end to ISO format strings # Workaround for https://github.com/Parallels/rq-dashboard/issues/510 job_metadata = { - "data_source_info": {"id": self.data_source.id}, + "data_source_info": {"id": data_source_id}, "start": self._parameters["predict_start"].isoformat(), "end": self._parameters["end_date"].isoformat(), - "sensor_id": self._parameters["sensor_to_save"].id, + "sensor_id": sensor_to_save_id, } for cycle_params in cycles_job_params: + job_kwargs = { + "config": job_config, + "parameters": job_parameters, + "data_source_id": data_source_id, + "delete_model": self.delete_model, + **cycle_params, + } + _assert_no_orm_objects(job_kwargs) job = Job.create( - self.run_cycle, - # Some cycle job params override job kwargs - kwargs={**job_kwargs, **cycle_params}, + run_train_predict_cycle_job, + kwargs=job_kwargs, connection=connection, ttl=int( current_app.config.get( @@ -340,18 +442,18 @@ def run( current_app.queues[queue].enqueue_job(job) current_app.job_cache.add( - self._parameters["sensor"].id, + sensor_id, job_id=job.id, queue=queue, asset_or_sensor_type="sensor", ) wrap_up_job = Job.create( - self.run_wrap_up, + run_train_predict_wrap_up_job, kwargs={ - "cycle_job_ids": cycle_job_ids, # cycles jobs IDs to wait for + "cycle_job_ids": cycle_job_ids, "queue": queue, - }, + }, # cycles jobs IDs to wait for connection=connection, depends_on=cycle_job_ids, # wrap-up job depends on all cycle jobs ttl=int( diff --git a/flexmeasures/data/tests/test_forecasting_pipeline.py b/flexmeasures/data/tests/test_forecasting_pipeline.py index ac63643606..91f3020b25 100644 --- a/flexmeasures/data/tests/test_forecasting_pipeline.py +++ b/flexmeasures/data/tests/test_forecasting_pipeline.py @@ -7,6 +7,7 @@ from datetime import datetime, timedelta from marshmallow import ValidationError +from sqlalchemy import inspect as sa_inspect from flexmeasures.data.models.forecasting.custom_models.lgbm_model import CustomLGBM from flexmeasures.data.models.data_sources import DataSource @@ -18,12 +19,117 @@ GenericAssetType, ) from flexmeasures.data.models.forecasting.pipelines import TrainPredictPipeline +from flexmeasures.data.models.forecasting.pipelines.train_predict import ( + _load_job_config_payload, + _load_job_parameters_payload, + _make_job_config_payload, + _make_job_parameters_payload, + run_train_predict_cycle_job, +) from flexmeasures.data.models.time_series import Sensor, TimedBelief from flexmeasures.data.queries.utils import simplify_index from flexmeasures.utils.job_utils import work_on_rq from flexmeasures.data.services.forecasting import handle_forecasting_exception +def _contains_orm_instance(value) -> bool: + inspection = sa_inspect(value, raiseerr=False) + if inspection is not None and hasattr(inspection, "object"): + return True + + if isinstance(value, dict): + return any(_contains_orm_instance(v) for v in value.values()) + if isinstance(value, (list, tuple, set)): + return any(_contains_orm_instance(v) for v in value) + return False + + +def test_train_predict_job_config_payload_preserves_plain_fields( + setup_fresh_test_forecast_data, +): + future_regressor = setup_fresh_test_forecast_data["irradiance-sensor"] + past_regressor = setup_fresh_test_forecast_data["solar-sensor-1"] + + config = { + "model": "CustomLGBM", + "future_regressors": [future_regressor], + "past_regressors": [past_regressor], + "missing_threshold": 0.25, + "plain_future_option": { + "lower": "0 kW", + "upper": "20 kW", + "snap": {"0 kW": ["0 kW", "4 kW"]}, + }, + } + + payload = _make_job_config_payload(config) + + assert "future_regressors" not in payload + assert "past_regressors" not in payload + assert payload["future_regressor_ids"] == [future_regressor.id] + assert payload["past_regressor_ids"] == [past_regressor.id] + assert payload["plain_future_option"] == config["plain_future_option"] + + restored_config = _load_job_config_payload(payload) + + assert restored_config["future_regressors"] == [future_regressor] + assert restored_config["past_regressors"] == [past_regressor] + assert restored_config["plain_future_option"] == config["plain_future_option"] + + +def test_train_predict_job_parameters_payload_preserves_plain_fields( + setup_fresh_test_forecast_data, +): + sensor = setup_fresh_test_forecast_data["solar-sensor"] + sensor_to_save = setup_fresh_test_forecast_data["solar-sensor-1"] + parameters = { + "sensor": sensor, + "sensor_to_save": sensor_to_save, + "model_save_dir": "flexmeasures/data/models/forecasting/artifacts/models", + "plain_future_parameter": {"labels": ["expected", "preserved"]}, + } + + payload = _make_job_parameters_payload(parameters) + + assert "sensor" not in payload + assert "sensor_to_save" not in payload + assert payload["sensor_id"] == sensor.id + assert payload["sensor_to_save_id"] == sensor_to_save.id + assert payload["plain_future_parameter"] == parameters["plain_future_parameter"] + + restored_parameters = _load_job_parameters_payload(payload) + + assert restored_parameters["sensor"] == sensor + assert restored_parameters["sensor_to_save"] == sensor_to_save + assert ( + restored_parameters["plain_future_parameter"] + == parameters["plain_future_parameter"] + ) + + +def test_train_predict_job_payload_rejects_unexpected_orm_objects( + setup_fresh_test_forecast_data, +): + sensor = setup_fresh_test_forecast_data["solar-sensor"] + + with pytest.raises(ValueError, match="payload.unexpected_sensor.*Sensor"): + _make_job_config_payload( + { + "future_regressors": [], + "past_regressors": [], + "unexpected_sensor": sensor, + } + ) + + with pytest.raises(ValueError, match="payload.unexpected_sensor.*Sensor"): + _make_job_parameters_payload( + { + "sensor": sensor, + "unexpected_sensor": sensor, + } + ) + + def test_custom_lgbm_falls_back_when_daily_lag_is_under_sampled(): """Short histories should drop daily lags only where they are under-sampled.""" under_sampled_model = CustomLGBM( @@ -335,6 +441,50 @@ def test_train_predict_pipeline( # noqa: C901 assert hasattr(pipeline, attr) if as_job: + queued_job = app.queues["forecasting"].fetch_job(pipeline_returns["job_id"]) + assert queued_job is not None + if not queued_job.dependency_ids: + queued_cycle_job_ids = [queued_job.id] + else: + queued_cycle_job_ids = queued_job.kwargs.get("cycle_job_ids", []) + + for job_id in queued_cycle_job_ids: + queued_cycle_job = app.queues["forecasting"].fetch_job(job_id) + assert queued_cycle_job is not None + assert queued_cycle_job.func == run_train_predict_cycle_job + assert not _contains_orm_instance(queued_cycle_job.kwargs) + assert isinstance( + queued_cycle_job.kwargs["parameters"]["sensor_id"], int + ) + assert isinstance( + queued_cycle_job.kwargs["parameters"]["sensor_to_save_id"], int + ) + assert all( + isinstance(sensor_id, int) + for sensor_id in queued_cycle_job.kwargs["config"][ + "future_regressor_ids" + ] + ) + assert all( + isinstance(sensor_id, int) + for sensor_id in queued_cycle_job.kwargs["config"][ + "past_regressor_ids" + ] + ) + for timing_field in ( + "predict_start", + "end_date", + "predict_period_in_hours", + "max_forecast_horizon", + "forecast_frequency", + "save_belief_time", + "m_viewpoints", + ): + assert ( + queued_cycle_job.kwargs["parameters"][timing_field] + == pipeline._parameters[timing_field] + ) + work_on_rq( app.queues["forecasting"], exc_handler=handle_forecasting_exception )