[hiopy] Subtasks

Creating subtasks for the flush operation in hiopy.

apps/hiopy/_data_handler.py

@@ -30,30 +30,33 @@ class DataHandler:
             self.buf = np.empty([self.time_chunk_size, *var.shape[1:-1], ma - mi], dtype=var.dtype)
 
     def _create_flush_task(self):
-        old_buffer = self.buf.copy()
+        fro = self.t_flushed % self.time_chunk_size
+        to = fro + (self.t - self.t_flushed)
+        buf = self.buf[fro:to, ...].copy()
+        t_flushed = int(self.t_flushed)
+        t = int(self.t)
 
         def _flush():
+            nonlocal t, t_flushed, buf, self
             if self.prune_offset is not None:
-                fro = max(0, self.t_flushed - self.prune_offset)
-                to = max(0, self.t - self.prune_offset)
-                with Timer() as t:
-                    self.var[fro:to, ..., self.cell_slice] = self.var.fill_value
+                prune_from = max(0, t_flushed - self.prune_offset)
+                prune_to = max(0, t - self.prune_offset)
+                with Timer() as timer:
+                    self.var[prune_from:prune_to, ..., self.cell_slice] = self.var.fill_value
                 logging.info(
-                    f"pruned {self.var.name} for timesteps {fro}:{to}. Took {t.elapsed_time}s"
+                    f"pruned {self.var.name} for timesteps {fro}:{to}. Took {timer.elapsed_time}s"
                 )
 
-            fro = self.t_flushed % self.time_chunk_size
-            to = fro + (self.t - self.t_flushed)
-            with Timer() as t:
-                self.var[self.t_flushed : self.t, ..., self.cell_slice] = old_buffer[fro:to, ...]
+            with Timer() as timer:
+                self.var[t_flushed:t, ..., self.cell_slice] = buf
             logging.info(
-                f"wrote {self.var.name} for timesteps {self.t_flushed}:{self.t}."
-                f" Took {t.elapsed_time}s"
+                f"wrote {self.var.name} for timesteps {t_flushed}:{t}."
+                f" Took {timer.elapsed_time}s"
             )
-            self.t_flushed = self.t
             if self.loco_server is not None:
-                self.loco_server.set_timestep(self.var.name, self.t - 1)
+                self.loco_server.set_timestep(self.var.name, t - 1)
 
+        self.t_flushed = self.t
         return _flush
 
     def flush(self):