WIP

tcjennings · tcjennings · commit da4a966b8ad6 · 2025-07-08T09:58:37.000-05:00
diff --git a/src/lsst/cmservice/common/daemon_v2.py b/src/lsst/cmservice/common/daemon_v2.py
@@ -1,4 +1,6 @@
 import pickle
+from asyncio import Task as AsyncTask
+from asyncio import TaskGroup
 from collections.abc import Mapping
 
 from sqlmodel import col, select
@@ -41,7 +43,7 @@ async def consider_campaigns(session: AsyncSession) -> None:
         # table should be to ensure the same *node* is not added twice if
         # the same campaign is "considered" by multiple daemons.
         for node in graph.processable_graph_nodes(campaign_graph):
-            logger.info("Daemon considering node", id=node.id)
+            logger.info("Daemon considering node", id=str(node.id))
             node_task = Task(
                 namespace=campaign_id,
                 node=node.id,
@@ -71,61 +73,59 @@ async def consider_nodes(session: AsyncSession) -> None:
 
     # Using a TaskGroup context manager means all "tasks" added to the group
     # are awaited when the CM exits, giving us concurrency for all the nodes
-    # being considered in the current iteration, but complicating the wrap-up
-    # operations that involve repickling the node's machine for a future
-    # iteration.
-    # . async with TaskGroup() as tg:
-    for task in tasks:
-        node = await session.get_one(Node, task.node)
-        # We expunge the node from *this* session because it will be added to
-        # whatever session the node_machine acquires during its transition
-        session.expunge(node)
-
-        node_machine: NodeMachine
-        node_machine_pickle: Machine | None
-        if node.machine is None:
-            # create a new machine for the node
-            node_machine = node_machine_factory(node.kind)(o=node)
-            node_machine_pickle = None
-        else:
-            # unpickle the node's machine and rehydrate the Node Stateful Model
-            node_machine_pickle = await session.get_one(Machine, node.machine)
-            node_machine = (pickle.loads(node_machine_pickle.state)).model
-            node_machine.node = node
-            # discard the pickled machine from this session and context
-            session.expunge(node_machine_pickle)
-            del node_machine_pickle
-
-        # the task's status field is the target status for the node, so the
-        # daemon intends to evolve the node machine to that state.
-        try:
-            assert node.status is task.previous_status
-        except AssertionError:
-            logger.error("Node status out of sync with Machine", id=node.id)
-            continue
-
-        # check possible triggers for state
-        # TODO how to pick the "best" trigger from multiple available?
-        # - Add a caller-backed conditional to the triggers, to identify which
-        # .  triggers the daemon is "allowed" to use
-        # - Determine the "desired" trigger from the (source, dest) task state
-        # .  tuple
-        if (trigger := trigger_for_transition(task, node_machine.machine.events)) is None:
-            logger.warning(
-                "No trigger available for desired state transition",
-                source=task.previous_status,
-                dest=task.status,
-            )
-            continue
-
-        # Add the node transition trigger method to the task group
-        _ = await node_machine.trigger(trigger)
-
-        # wrap up - the task is removed from the session
-        # TODO if the node has been transitioned to a terminal state, we
-        # should not need to keep its machine around.
-        await session.delete(task)
-        await session.commit()
+    # being considered in the current iteration.
+    async with TaskGroup() as tg:
+        for task in tasks:
+            node = await session.get_one(Node, task.node)
+
+            # the task's status field is the target status for the node, so the
+            # daemon intends to evolve the node machine to that state.
+            try:
+                assert node.status is task.previous_status
+            except AssertionError:
+                logger.error("Node status out of sync with Machine", id=str(node.id))
+                continue
+
+            # Expunge the node from *this* session because it will be added to
+            # whatever session the node_machine acquires during its transition
+            session.expunge(node)
+
+            node_machine: NodeMachine
+            node_machine_pickle: Machine | None
+            if node.machine is None:
+                # create a new machine for the node
+                node_machine = node_machine_factory(node.kind)(o=node)
+                node_machine_pickle = None
+            else:
+                # unpickle the node's machine and rehydrate the Stateful Model
+                node_machine_pickle = await session.get_one(Machine, node.machine)
+                node_machine = (pickle.loads(node_machine_pickle.state)).model
+                node_machine.node = node
+                # discard the pickled machine from this session and context
+                session.expunge(node_machine_pickle)
+                del node_machine_pickle
+
+            # check possible triggers for state
+            # TODO how to pick the "best" trigger from multiple available?
+            # - Add a caller-backed conditional to the triggers, to identify
+            # .  triggers the daemon is "allowed" to use
+            # - Determine the "desired" trigger from the task (source, dest)
+            if (trigger := trigger_for_transition(task, node_machine.machine.events)) is None:
+                logger.warning(
+                    "No trigger available for desired state transition",
+                    source=task.previous_status,
+                    dest=task.status,
+                )
+                continue
+
+            # Add the node transition trigger method to the task group
+            # TODO give this a name and a callback
+            task_ = tg.create_task(node_machine.trigger(trigger), name=str(node.id))
+            task_.add_done_callback(task_runner_callback)
+
+            # wrap up - the task is removed from the db.
+            await session.delete(task)
+            await session.commit()
 
 
 async def daemon_iteration(session: AsyncSession) -> None:
@@ -138,6 +138,7 @@ async def daemon_iteration(session: AsyncSession) -> None:
         await consider_campaigns(session)
     if config.daemon.process_nodes:
         await consider_nodes(session)
+    await session.close()
 
 
 def trigger_for_transition(task: Task, events: Mapping[str, Event]) -> str | None:
@@ -156,3 +157,8 @@ def trigger_for_transition(task: Task, events: Mapping[str, Event]) -> str | Non
                 ):
                     return trigger
     return None
+
+
+def task_runner_callback(task: AsyncTask):
+    """Callback function for `asyncio.TaskGroup` tasks."""
+    logger.info("Transition complete", id=task.get_name())
diff --git a/src/lsst/cmservice/machines/campaign.py b/src/lsst/cmservice/machines/campaign.py
@@ -213,8 +213,7 @@ async def update_persistent_status(self, event: EventData) -> None:
             self.activity_log_entry.detail["finished_at"] = timestamp.element_time()
 
         # Ensure database record for transitioned object is updated
-        self.session.add(self.campaign)
-        await self.session.refresh(self.campaign)
+        self.campaign = await self.session.merge(self.campaign, load=False)
         self.campaign.status = self.state
         await self.session.commit()
 
diff --git a/src/lsst/cmservice/machines/node.py b/src/lsst/cmservice/machines/node.py
@@ -124,7 +124,9 @@ async def error_handler(self, event: EventData) -> None:
                 ...
 
     async def prepare_session(self, event: EventData) -> None:
-        """Prepares the machine by acquiring a database session."""
+        """Prepares the machine by acquiring a database session and merging
+        the reference object.
+        """
         # This positive assertion concerning the Node member will prevent
         # any callback from proceeding if no such member is defined, but type
         # checkers don't know this, which is why it repeated in a TYPE_CHECKING
@@ -172,8 +174,7 @@ async def update_persistent_status(self, event: EventData) -> None:
             self.activity_log_entry.detail["finished_at"] = timestamp.element_time()
 
         # Ensure database record for transitioned object is updated
-        self.session.add(self.node)
-        await self.session.refresh(self.node)
+        self.node = await self.session.merge(self.node, load=False)
         self.node.status = self.state
         await self.session.commit()
 
diff --git a/tests/v2/conftest.py b/tests/v2/conftest.py
@@ -101,6 +101,9 @@ async def testdb(rawdb: DatabaseSessionDependency) -> AsyncGenerator[DatabaseSes
                 namespace=str(NAMESPACE_DNS),
                 name="DEFAULT",
                 owner="root",
+                status="accepted",
+                metadata={},
+                configuration={},
             )
         )
         await aconn.commit()
diff --git a/tests/v2/test_machines.py b/tests/v2/test_machines.py
@@ -29,36 +29,31 @@ async def test_campaign_machine(test_campaign: str, session: AsyncSession) -> No
     x = campaign.status
     assert x is StatusEnum.waiting
 
-    # This instance is going to be "owned" by the session managed by the State-
-    # ful Model, so we expunge it from the session managed by this test
-    session.expunge(campaign)
     campaign_machine = CampaignMachine(o=campaign)
-
     assert campaign_machine.is_waiting()
     assert await campaign_machine.may_prepare()
 
     did_prepare = await campaign_machine.prepare()
     assert did_prepare
     assert await campaign_machine.may_start()
 
-    ref_campaign = await session.get_one(Campaign, campaign_id)
-    await session.refresh(ref_campaign, attribute_names=["status"])
-    x = ref_campaign.status
+    await session.refresh(campaign, attribute_names=["status"])
+    x = campaign.status
     assert x is StatusEnum.ready
 
     did_start = await campaign_machine.start()
     assert did_start
     assert await campaign_machine.may_pause()
     assert await campaign_machine.may_finish()
 
-    await session.refresh(ref_campaign, attribute_names=["status"])
-    x = ref_campaign.status
+    await session.refresh(campaign, attribute_names=["status"])
+    x = campaign.status
     assert x is StatusEnum.running
 
     did_finish = await campaign_machine.trigger("finish")
     assert did_finish
-    await session.refresh(ref_campaign, attribute_names=["status"])
-    x = ref_campaign.status
+    await session.refresh(campaign, attribute_names=["status"])
+    x = campaign.status
     assert x is StatusEnum.accepted
 
     await session.close()
@@ -74,7 +69,6 @@ async def test_bad_transition(test_campaign: str, session: AsyncSession) -> None
     campaign = await session.get_one(Campaign, campaign_id)
     x = campaign.status
     assert x is StatusEnum.waiting
-    session.expunge(campaign)
     campaign_machine = CampaignMachine(o=campaign)
 
     with patch(
@@ -97,26 +91,25 @@ async def test_bad_transition(test_campaign: str, session: AsyncSession) -> None
 
     # Both the state machine and db object should reflect the failed state
     assert campaign_machine.is_failed()
-    ref_campaign = await session.get_one(Campaign, campaign_id)
-    await session.refresh(ref_campaign, attribute_names=["status"])
-    x = ref_campaign.status
+    await session.refresh(campaign, attribute_names=["status"])
+    x = campaign.status
     assert x is StatusEnum.failed
 
     # retry by rolling back to ready
     assert await campaign_machine.may_retry()
     await campaign_machine.trigger("retry")
 
     assert campaign_machine.is_ready()
-    await session.refresh(ref_campaign, attribute_names=["status"])
-    x = ref_campaign.status
+    await session.refresh(campaign, attribute_names=["status"])
+    x = campaign.status
     assert x is StatusEnum.ready
 
     # start and finish without further error
     assert await campaign_machine.start()
     assert await campaign_machine.finish()
 
-    await session.refresh(ref_campaign, attribute_names=["status"])
-    x = ref_campaign.status
+    await session.refresh(campaign, attribute_names=["status"])
+    x = campaign.status
     assert x is StatusEnum.accepted
     await session.close()
 
@@ -130,8 +123,6 @@ async def test_machine_pickle(test_campaign: str, session: AsyncSession) -> None
     x = campaign.status
     assert x is StatusEnum.waiting
     campaign_machine = CampaignMachine(o=campaign)
-    session.expunge(campaign)
-    ref_campaign = await session.get_one(Campaign, campaign_id)
 
     # evolve the machine to "prepared"
     await campaign_machine.prepare()
@@ -145,15 +136,13 @@ async def test_machine_pickle(test_campaign: str, session: AsyncSession) -> None
     # not necessarily deterministic (i.e., they do not have to be namespaced.)
     session.add(machine_db)
     await session.commit()
-    ref_campaign = await session.get_one(Campaign, campaign_id)
-    ref_campaign.machine = machine_db.id
+    campaign.machine = machine_db.id
     await session.commit()
-    await session.refresh(ref_campaign, attribute_names=["status", "machine"])
+    await session.refresh(campaign, attribute_names=["status", "machine"])
 
     await session.close()
     del campaign_machine
     del campaign
-    del ref_campaign
     del machine_db
     del machine_pickle
 
@@ -169,18 +158,16 @@ async def test_machine_pickle(test_campaign: str, session: AsyncSession) -> None
     assert machine_unpickle.state is not None
     new_campaign_machine: CampaignMachine = (pickle.loads(machine_unpickle.state)).model
     new_campaign_machine.campaign = new_campaign
-    session.expunge(new_campaign)
-    new_ref_campaign = await session.get_one(Campaign, campaign_id)
 
     # test the rehydrated machine and continue to evolve it
-    assert new_campaign_machine.state == new_ref_campaign.status
+    assert new_campaign_machine.state == new_campaign.status
 
     await new_campaign_machine.start()
-    await session.refresh(new_ref_campaign, attribute_names=["status"])
-    assert new_campaign_machine.state == new_ref_campaign.status
+    await session.refresh(new_campaign, attribute_names=["status"])
+    assert new_campaign_machine.state == new_campaign.status
 
     await new_campaign_machine.finish()
-    await session.refresh(new_ref_campaign, attribute_names=["status"])
-    assert new_campaign_machine.state == new_ref_campaign.status
+    await session.refresh(new_campaign, attribute_names=["status"])
+    assert new_campaign_machine.state == new_campaign.status
 
     await session.close()

Original file line number	Diff line number	Diff line change
`@@ -101,6 +101,9 @@ async def testdb(rawdb: DatabaseSessionDependency) -> AsyncGenerator[DatabaseSes`
`101`	`101`	`namespace=str(NAMESPACE_DNS),`
`102`	`102`	`name="DEFAULT",`
`103`	`103`	`owner="root",`
	`104`	`+ status="accepted",`
	`105`	`+ metadata={},`
	`106`	`+ configuration={},`
`104`	`107`	`)`
`105`	`108`	`)`
`106`	`109`	`await aconn.commit()`