Use priority to send queue all jobs on workers

Author: domfournier

simpeg/dask/objective_function.py

@@ -132,7 +132,9 @@ def __call__(self, m, f=None):
 
         values = []
         count = 0
-        for futures in self._futures:
+        for ind, futures in enumerate(self._futures):
+
+            priority = len(self._futures) - ind  # reverse order for priority
             for objfct, worker in zip(futures, self._workers, strict=True):
 
                 if self.multipliers[count] == 0.0:
@@ -145,6 +147,7 @@ def __call__(self, m, f=None):
                         self.multipliers[count],
                         m_future,
                         workers=worker,
+                        priority=priority,
                     )
                 )
                 count += 1
@@ -193,11 +196,12 @@ def deriv(self, m, f=None):
         client = self.client
         m_future = self._m_as_future
 
-        derivs = 0.0
         count = 0
+        future_deriv = []
+        for ind, futures in enumerate(self._futures):
+
+            priority = len(self._futures) - ind  # reverse order for priority
 
-        for futures in self._futures:
-            future_deriv = []
             for objfct, worker in zip(futures, self._workers):
                 if self.multipliers[count] == 0.0:  # don't evaluate the fct
                     continue
@@ -209,13 +213,14 @@ def deriv(self, m, f=None):
                         self.multipliers[count],
                         m_future,
                         workers=worker,
+                        priority=priority,
                     )
                 )
 
                 count += 1
-            future_deriv = client.gather(future_deriv)
+        future_deriv = client.gather(future_deriv)
 
-            derivs += np.sum(future_deriv, axis=0)
+        derivs = np.sum(future_deriv, axis=0)
 
         return derivs
 
@@ -234,12 +239,11 @@ def deriv2(self, m, v=None, f=None):
         m_future = self._m_as_future
         [v_future] = client.scatter([v], broadcast=True)
 
-        derivs = 0.0
         count = 0
+        future_derivs = []
+        for ind, futures in enumerate(self._futures):
 
-        for futures in self._futures:
-
-            future_derivs = []
+            priority = len(self._futures) - ind  # reverse order for priority
             for objfct, worker in zip(futures, self._workers):
                 if self.multipliers[count] == 0.0:  # don't evaluate the fct
                     continue
@@ -253,12 +257,13 @@ def deriv2(self, m, v=None, f=None):
                         v_future,
                         # field,
                         workers=worker,
+                        priority=priority,
                     )
                 )
                 count += 1
 
-            future_derivs = self.client.gather(future_derivs)
-            derivs += np.sum(future_derivs, axis=0)
+        future_derivs = self.client.gather(future_derivs)
+        derivs = np.sum(future_derivs, axis=0)
 
         return derivs
 
@@ -270,20 +275,22 @@ def get_dpred(self, m, f=None):
 
         client = self.client
         m_future = self._m_as_future
-        dpred = []
+        # dpred = []
+        future_preds = []
+        for ind, futures in enumerate(self._futures):
 
-        for futures in self._futures:
-            future_preds = []
+            priority = len(self._futures) - ind  # reverse order for priority
             for objfct, worker in zip(futures, self._workers):
                 future_preds.append(
                     client.submit(
                         _calc_dpred,
                         objfct,
                         m_future,
                         workers=worker,
+                        priority=priority,
                     )
                 )
-            dpred += client.gather(future_preds)
+        dpred = client.gather(future_preds)
 
         return dpred
 
@@ -294,25 +301,24 @@ def getJtJdiag(self, m, f=None):
         self.model = m
         m_future = self._m_as_future
         if getattr(self, "_jtjdiag", None) is None:
-
-            jtj_diag = 0.0
             client = self.client
+            work = []
+            for ind, futures in enumerate(self._futures):
 
-            for futures in self._futures:
-                work = []
-
+                priority = len(self._futures) - ind  # reverse order for priority
                 for objfct, worker in zip(futures, self._workers):
                     work.append(
                         client.submit(
                             _get_jtj_diag,
                             objfct,
                             m_future,
                             workers=worker,
+                            priority=priority,
                         )
                     )
 
-                work = client.gather(work)
-                jtj_diag += np.sum(work, axis=0)
+            work = client.gather(work)
+            jtj_diag = np.sum(work, axis=0)
 
             self._jtjdiag = jtj_diag
 
@@ -332,15 +338,18 @@ def fields(self, m):
         # The above should pass the model to all the internal simulations.
         f = []
 
-        for futures in self._futures:
-            f.append([])
+        for ind, futures in enumerate(self._futures):
+
+            priority = len(self._futures) - ind  # reverse order for priority
+
             for objfct, worker in zip(futures, self._workers):
-                f[-1].append(
+                f.append(
                     client.submit(
                         _calc_fields,
                         objfct,
                         m_future,
                         workers=worker,
+                        priority=priority,
                     )
                 )
         self._stashed_fields = f
@@ -367,14 +376,17 @@ def model(self, value):
         [self._m_as_future] = client.scatter([value], broadcast=True)
 
         stores = []
-        for futures in self._futures:
+        for ind, futures in enumerate(self._futures):
+
+            priority = len(self._futures) - ind  # reverse order for priority
             for objfct, worker in zip(futures, self._workers):
                 stores.append(
                     client.submit(
                         _store_model,
                         objfct,
                         self._m_as_future,
                         workers=worker,
+                        priority=priority,
                     )
                 )
         self.client.gather(stores)  # blocking call to ensure all models were stored
@@ -418,19 +430,22 @@ def residuals(self, m, f=None):
         client = self.client
         m_future = self._m_as_future
         residuals = []
+        future_residuals = []
+        for ind, futures in enumerate(self._futures):
+
+            priority = len(self._futures) - ind  # reverse order for priority
 
-        for futures in self._futures:
-            future_residuals = []
             for objfct, worker in zip(futures, self._workers):
                 future_residuals.append(
                     client.submit(
                         _calc_residual,
                         objfct,
                         m_future,
                         workers=worker,
+                        priority=priority,
                     )
                 )
-            residuals += client.gather(future_residuals)
+        residuals = client.gather(future_residuals)
 
         return residuals