vine: do not use temp for target constant keys in dask (#3918)

cooperative-computing-lab · Aug 16, 2024 · 0ec5c00 · 0ec5c00
1 parent f5b2d62
commit 0ec5c00
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/taskvine/src/bindings/python3/ndcctools/taskvine/dask_dag.py b/taskvine/src/bindings/python3/ndcctools/taskvine/dask_dag.py
@@ -173,9 +173,14 @@ def _flatten_graph(self):
             self.flatten_rec(key, self._working_graph[key], toplevel=True)
 
     def _add_second_targets(self, key):
-        if not DaskVineDag.listp(self._working_graph[key]):
+        v = self._working_graph[key]
+        if self.graph_keyp(v):
+            lst = [v]
+        elif DaskVineDag.listp(v):
+            lst = v
+        else:
             return
-        for c in self._working_graph[key]:
+        for c in lst:
             if self.graph_keyp(c):
                 self._targets.add(c)
                 self._add_second_targets(c)