Don't monitor lost drops until serviceDrops are completed.

ICRAR · May 22, 2024 · 69b4da5 · 69b4da5
1 parent 2f0f31d
commit 69b4da5
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 22 deletions.
diff --git a/daliuge-engine/dlg/droputils.py b/daliuge-engine/dlg/droputils.py
@@ -228,7 +228,8 @@ def getLeafNodes(drops):
     return [
         drop
         for drop, _ in breadFirstTraverse(drops)
-        if not getDownstreamObjects(drop) and drop.CategoryType != "dropclass"
+        if not getDownstreamObjects(drop) and (drop.CategoryType != "dropclass" and
+                                               drop.CategoryType != "Service")
     ]
 
 

diff --git a/daliuge-engine/dlg/lifecycle/dlm.py b/daliuge-engine/dlg/lifecycle/dlm.py
@@ -124,6 +124,7 @@
 import threading
 import time
 
+from typing import Dict
 from dlg.common import CategoryType
 from . import registry
 from .hsm import manager
@@ -229,8 +230,8 @@ def __init__(
         # instead of _drops.itervalues() to get a full, thread-safe copy of the
         # dictionary values. Maybe there's a better approach for thread-safety
         # here
-        self._drops: dict[str, AbstractDROP] = {}
-
+        self._drops: Dict[str, AbstractDROP] = {}
+        self._serviceDrops: Dict[str, AbstractDROP] = {}
         self._check_period = check_period
         self._cleanup_period = cleanup_period
         self._drop_checker = None
@@ -334,6 +335,12 @@ def _disappeared(self, drop):
 
     def deleteLostDrops(self):
 
+        for sd in self._serviceDrops.values():
+            if sd.status == DROPStates.INITIALIZED:
+                logger.info("No need to delete lost drops whilst Service Drops are "
+                            "starting up...")
+                return
+
         toRemove = []
         for drop in self._drops.values():
 
@@ -465,15 +472,9 @@ def addDrop(self, drop: AbstractDROP):
         drop.phase = DROPPhases.GAS
         drop.subscribe(self._listener)
 
-        # if drop.CategoryType == CategoryType.SERVICE:
-        #     connection = drop.getIO().exists()
-        #     continue
-        #TODO LOOK HERE FOR SETTING UP SERVICES BASED ON THE DROP
-        # if drop.persist:
-        #     self._updatePersistentStore(drop)
-
-        if drop.CategoryType == CategoryType.SERVICE: # and self._hsm:
-            print("We are in the DLM!")
+        if drop.CategoryType == CategoryType.SERVICE:  # and self._hsm:
+            print("Tracking ServiceDROPS in the DLM...")
+            self._serviceDrops[drop.uid] = drop
             # self._hsm.addStore(drop.store)
 
         self._reg.addDrop(drop)
@@ -509,16 +510,18 @@ def handleCompletedDrop(self, uid):
 
         if not self._enable_drop_replication:
             return
-
-        drop = self._drops[uid]
-        if drop.persist and self.isReplicable(drop):
-            logger.debug(
-                "Replicating %r because it's marked to be persisted", drop
-            )
-            try:
-                self.replicateDrop(drop)
-            except:
-                logger.exception("Problem while replicating %r", drop)
+        try:
+            drop = self._drops[uid]
+            if drop.persist and self.isReplicable(drop):
+                logger.debug(
+                    "Replicating %r because it's marked to be persisted", drop
+                )
+                try:
+                    self.replicateDrop(drop)
+                except:
+                    logger.exception("Problem while replicating %r", drop)
+        except KeyError:
+            logger.warning("Drop %s was removed from DLM early!", uid)
 
     def isReplicable(self, drop):
         return not isinstance(drop, ContainerDROP)