Merge pull request #72 from tiqi-group/7-hardware-worker-prevent-processing-stale-json-after-parameter-updates

7 hardware worker prevent processing stale json after parameter updates

Author: g-braeunlich

src/icon/server/data_access/db_context/sqlite/alembic/versions/b18002636885_jobrun_add_parameter_update_timestamp_.py

@@ -0,0 +1,35 @@
+"""JobRun: add parameter_update_timestamp column
+
+Revision ID: b18002636885
+Revises: fc9af856df20
+Create Date: 2025-11-20 10:28:24.835587
+
+"""
+
+from collections.abc import Sequence
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'b18002636885'
+down_revision: str | None = 'fc9af856df20'
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('job_runs', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('parameter_update_timestamp', sa.TIMESTAMP(timezone=True), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('job_runs', schema=None) as batch_op:
+        batch_op.drop_column('parameter_update_timestamp')
+
+    # ### end Alembic commands ###

src/icon/server/data_access/models/sqlite/job_run.py

@@ -65,6 +65,11 @@ class JobRun(Base):
     log: sqlalchemy.orm.Mapped[str | None] = sqlalchemy.orm.mapped_column(default=None)
     """Optional log message for this run (e.g., cancellation reason)."""
 
+    parameter_update_timestamp: sqlalchemy.orm.Mapped[datetime.datetime | None] = (
+        sqlalchemy.orm.mapped_column(default=None)
+    )
+    """Timestamp of the last parameter update."""
+
     def __repr__(self) -> str:
         return (
             f"<JobRun id={self.id} "

src/icon/server/data_access/repositories/job_run_repository.py

@@ -1,6 +1,6 @@
 import logging
 from collections.abc import Sequence
-from datetime import datetime
+from datetime import UTC, datetime
 
 import sqlalchemy.orm
 from sqlalchemy import select, update
@@ -185,3 +185,43 @@ def get_scheduled_time_by_job_id(*, job_id: int) -> datetime:
             scheduled_time = session.execute(stmt).scalar_one()
             logger.debug("Got scheduled time for job_id %s", job_id)
         return scheduled_time
+
+    @staticmethod
+    def set_parameter_update_timestamp(*, run_id: int, timestamp: datetime) -> None:
+        """Set the paramter update timestamp.
+
+        Args:
+            job_id: ID of the job.
+            timestamp: New parameter update timestamp.
+        """
+
+        with sqlalchemy.orm.Session(engine) as session:
+            stmt = (
+                update(JobRun)
+                .where(JobRun.id == run_id)
+                .values(parameter_update_timestamp=timestamp.astimezone(UTC))
+                .returning(JobRun)
+            )
+
+            run = session.execute(stmt).scalar_one()
+            session.commit()
+
+            logger.debug("Updated parameter update timestam for run %s", run)
+
+    @staticmethod
+    def get_parameter_update_timestamp(*, run_id: int) -> datetime:
+        """Get the paramter update timestamp.
+
+        Args:
+            job_id: ID of the job.
+        Returns:
+            The parameter update timestamp.
+        """
+
+        with sqlalchemy.orm.Session(engine) as session:
+            stmt = select(JobRun.parameter_update_timestamp).where(JobRun.id == run_id)
+
+            timestamp = session.execute(stmt).scalar_one()
+            logger.debug("Got parameter update timestamp for run %s", run_id)
+
+        return timestamp.replace(tzinfo=UTC)

src/icon/server/hardware_processing/task.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from datetime import datetime
-from queue import Queue
+from queue import PriorityQueue, Queue
 from typing import TYPE_CHECKING, Any
 
 import pydantic
@@ -22,15 +22,16 @@ class HardwareProcessingTask(pydantic.BaseModel):
     sequence_json: str
     src_dir: str
     created: datetime
-
     if TYPE_CHECKING:
         processed_data_points: Queue[HardwareProcessingTask]
         data_points_to_process: Queue[tuple[int, dict[str, DatabaseValueType]]]
+        outdated_tasks: PriorityQueue[HardwareProcessingTask]
     else:
         # must be Any as the queues are AutoProxy instances, which I didn't figure out
         # how to type
         processed_data_points: Any
         data_points_to_process: Any
+        outdated_tasks: Any
 
     def __lt__(self, other: HardwareProcessingTask) -> bool:
         return self.priority < other.priority or self.created < other.created

src/icon/server/hardware_processing/worker.py

@@ -160,6 +160,14 @@ def run(self) -> None:
             ):
                 continue
 
+            parameter_update_timestamp = (
+                JobRunRepository.get_parameter_update_timestamp(
+                    run_id=task.pre_processing_task.job_run.id,
+                )
+            )
+            if task.created < parameter_update_timestamp:
+                task.outdated_tasks.put(task)
+                continue
             try:
                 self._set_pydase_service_values(scanned_params=task.scanned_params)
 

src/icon/server/pre_processing/worker.py

@@ -41,7 +41,7 @@
 from icon.server.hardware_processing.task import HardwareProcessingTask
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
+    from collections.abc import Iterable, Iterator
 
     from icon.server.data_access.models.sqlite.job import Job
     from icon.server.pre_processing.task import PreProcessingTask
@@ -192,6 +192,9 @@ def __init__(
         ]
         self._processed_data_points: queue.Queue[HardwareProcessingTask]
         self._parameter_dict: dict[str, DatabaseValueType] = {}
+        self._outdated_tasks: queue.PriorityQueue[HardwareProcessingTask] = (
+            manager.PriorityQueue()
+        )
 
     def run(self) -> None:
         with tempfile.TemporaryDirectory() as tmp_dir:
@@ -289,14 +292,17 @@ def _process_task(
             readout_metadata=readout_metadata,
         )
 
-        if contains_realtime_parameter(job.scan_parameters):
+        jobs = (
             self._handle_realtime_scan(
                 pre_processing_task, src_dir=src_dir, namespace=namespace
             )
-        else:
-            self._handle_regular_scan(
+            if contains_realtime_parameter(job.scan_parameters)
+            else self._handle_regular_scan(
                 pre_processing_task, src_dir=src_dir, namespace=namespace
             )
+        )
+        for _ in jobs:
+            self._regenerate_outdated_jobs(namespace)
 
     def _update_parameter_dict(
         self,
@@ -321,6 +327,10 @@ def _update_parameter_dict(
 
         self._global_parameter_timestamp = datetime.now(timezone)
 
+        JobRunRepository.set_parameter_update_timestamp(
+            run_id=pre_processing_task.job_run.id,
+            timestamp=self._global_parameter_timestamp,
+        )
         if mode == ParamUpdateMode.ONLY_NEW_PARAMETERS:
             if new_parameters:
                 self._parameter_dict.update(new_parameters)
@@ -396,7 +406,7 @@ def _handle_regular_scan(
         pre_processing_task: PreProcessingTask,
         namespace: ExperimentIdentifier,
         src_dir: str,
-    ) -> None:
+    ) -> Iterable[None]:
         scan_parameter_value_combinations = get_scan_combinations(
             pre_processing_task.job
         )
@@ -421,6 +431,7 @@ def _handle_regular_scan(
             ):
                 break
 
+            yield
             self._submit_task_to_hw_worker(
                 task=self._create_hardware_task(
                     pre_processing_task=pre_processing_task,
@@ -454,15 +465,25 @@ def _create_hardware_task(
             sequence_json=sequence_json,
             processed_data_points=self._processed_data_points,
             data_points_to_process=self._data_points_to_process,
+            outdated_tasks=self._outdated_tasks,
             created=datetime.now(timezone),
         )
 
+    def _regenerate_outdated_jobs(self, namespace: ExperimentIdentifier) -> None:
+        for task in consume_queue(self._outdated_tasks):
+            task.sequence_json = generate_sequence_json(
+                n_shots=task.pre_processing_task.job.number_of_shots,
+                parameter_dict={**self._parameter_dict, **task.scanned_params},
+                namespace=namespace,
+            )
+            self._submit_task_to_hw_worker(task=task)
+
     def _handle_realtime_scan(
         self,
         pre_processing_task: PreProcessingTask,
         namespace: ExperimentIdentifier,
         src_dir: str,
-    ) -> None:
+    ) -> Iterable[None]:
         params = pre_processing_task.job.scan_parameters
         realtime_param = next(p for p in params if p.realtime)
         n_scan_values = len(realtime_param.scan_values)
@@ -507,6 +528,7 @@ def _handle_realtime_scan(
                     hardware_tasks[frozen_data_point] = hardware_task
                 hardware_task.created = datetime.now(timezone)
                 hardware_task.data_point_index = index
+                yield
                 self._submit_task_to_hw_worker(task=hardware_task)