Fix/detached forecasting sensors

documentation/changelog.rst

@@ -110,6 +110,7 @@ Bugfixes
 * Distinguish data sources with duplicate names in chart legends without showing source IDs [see `PR #2185 <https://www.github.com/FlexMeasures/flexmeasures/pull/2185>`_]
 * Fix broken API endpoint links in Read the Docs output by post-processing both known API index layouts and rewriting ``(id)`` anchors [see `PR #1753 <https://www.github.com/FlexMeasures/flexmeasures/pull/1753>`_]
 * Wrap Darts training ``ValueError`` about no common feature times as a clearer ``NotEnoughDataException`` when there is insufficient history for the requested forecast horizon [see `PR #2192 <https://www.github.com/FlexMeasures/flexmeasures/pull/2192>`_]
+* Fix queued train-predict forecasting jobs losing their resolved forecast window or failing on detached database objects in workers [see `PR #2035 <https://www.github.com/FlexMeasures/flexmeasures/pull/2035>`_]
 * Make Sentry also ignore 404 events, if we only log them without exc_info [`PR #2212 <https://www.github.com/FlexMeasures/flexmeasures/pull/2212>`_]
 
 

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

flexmeasures/data/models/forecasting/pipelines/train_predict.py

@@ -8,21 +8,156 @@
 from datetime import datetime, timedelta
 
 from rq.job import Job
-from sqlalchemy import inspect as sa_inspect
 
 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,
 )
 from flexmeasures.utils.flexmeasures_inflection import p
 
 
+def _sensor_id(sensor: Sensor | int | None) -> int | None:
+    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:
+    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:
+    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 _make_job_config_payload(config: dict[str, Any]) -> dict[str, Any]:
+    """Build the worker config payload without ORM objects."""
+    return {
+        "model": config["model"],
+        "future_regressor_ids": [
+            _sensor_id(sensor) for sensor in config["future_regressors"]
+        ],
+        "past_regressor_ids": [
+            _sensor_id(sensor) for sensor in config["past_regressors"]
+        ],
+        "missing_threshold": config["missing_threshold"],
+        "ensure_positive": config["ensure_positive"],
+        "train_start": config.get("train_start"),
+        "train_period": config["train_period"],
+        "max_training_period": config["max_training_period"],
+        "retrain_frequency": config["retrain_frequency"],
+        "train_period_in_hours": config["train_period_in_hours"],
+    }
+
+
+def _load_job_config_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    """Load the worker config through the active worker session."""
+    return {
+        "model": payload["model"],
+        "future_regressors": [
+            _get_attached_sensor(sensor_id)
+            for sensor_id in payload["future_regressor_ids"]
+        ],
+        "past_regressors": [
+            _get_attached_sensor(sensor_id)
+            for sensor_id in payload["past_regressor_ids"]
+        ],
+        "missing_threshold": payload["missing_threshold"],
+        "ensure_positive": payload["ensure_positive"],
+        "train_start": payload["train_start"],
+        "train_period": payload["train_period"],
+        "max_training_period": payload["max_training_period"],
+        "retrain_frequency": payload["retrain_frequency"],
+        "train_period_in_hours": payload["train_period_in_hours"],
+    }
+
+
+def _make_job_parameters_payload(parameters: dict[str, Any]) -> dict[str, Any]:
+    """Build the worker parameter payload without ORM objects."""
+    sensor_id = _sensor_id(parameters["sensor"])
+    sensor_to_save_id = _sensor_id(parameters.get("sensor_to_save"))
+    if sensor_id is None:
+        raise ValueError("Cannot enqueue a forecasting job without a target sensor.")
+    return {
+        "sensor_id": sensor_id,
+        "sensor_to_save_id": sensor_to_save_id or sensor_id,
+        "model_save_dir": parameters["model_save_dir"],
+        "output_path": parameters["output_path"],
+        "end_date": parameters["end_date"],
+        "predict_start": parameters["predict_start"],
+        "predict_period_in_hours": parameters["predict_period_in_hours"],
+        "max_forecast_horizon": parameters["max_forecast_horizon"],
+        "forecast_frequency": parameters["forecast_frequency"],
+        "probabilistic": parameters["probabilistic"],
+        "save_belief_time": parameters["save_belief_time"],
+        "beliefs_before": parameters.get("beliefs_before"),
+        "m_viewpoints": parameters["m_viewpoints"],
+    }
+
+
+def _load_job_parameters_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    """Load the worker parameters through the active worker session."""
+    return {
+        "sensor": _get_attached_sensor(payload["sensor_id"]),
+        "sensor_to_save": _get_attached_sensor(payload["sensor_to_save_id"]),
+        "model_save_dir": payload["model_save_dir"],
+        "output_path": payload["output_path"],
+        "end_date": payload["end_date"],
+        "predict_start": payload["predict_start"],
+        "predict_period_in_hours": payload["predict_period_in_hours"],
+        "max_forecast_horizon": payload["max_forecast_horizon"],
+        "forecast_frequency": payload["forecast_frequency"],
+        "probabilistic": payload["probabilistic"],
+        "save_belief_time": payload["save_belief_time"],
+        "beliefs_before": payload.get("beliefs_before"),
+        "m_viewpoints": payload["m_viewpoints"],
+    }
+
+
+def run_train_predict_cycle_job(
+    config: dict,
+    parameters: dict,
+    data_source_id: int,
+    delete_model: bool,
+    **cycle_params,
+):
+    """Reconstruct pipeline state inside the worker to avoid pickling ORM objects."""
+    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 +181,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 +202,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"],
@@ -285,7 +382,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 +395,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,
             }
+            # NOTE: `job_kwargs` is intentionally not forwarded into queued cycle jobs;
+            # only plain, serialized pipeline state + per-cycle timings are enqueued.
             for cycle_params in cycles_job_params:
 
                 job = Job.create(
-                    self.run_cycle,
-                    # Some cycle job params override job kwargs
-                    kwargs={**job_kwargs, **cycle_params},
+                    run_train_predict_cycle_job,
+                    kwargs={
+                        "config": job_config,
+                        "parameters": job_parameters,
+                        "data_source_id": data_source_id,
+                        "delete_model": self.delete_model,
+                        **cycle_params,
+                    },
                     connection=connection,
                     ttl=int(
                         current_app.config.get(
@@ -340,18 +448,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(