fix: improve current-timestep impl

apps/hiopy/_data_handler.py

@@ -35,7 +35,7 @@ class DataHandler:
         self.var[self.t_flushed : self.t, ..., self.cell_slice] = self.buf[fro:to, ...]
         self.t_flushed = self.t
         if self.loco_server is not None:
-            self.loco_server.set_timestep(self.t - 1)
+            self.loco_server.set_timestep(self.var.name, self.t - 1)
 
     def __call__(self, event):
         self.buf[self.t % self.time_chunk_size, ...] = np.transpose(event.data)

apps/hiopy/loco.py

@@ -6,11 +6,15 @@ from aiohttp import web
 
 class LocoServer:
     def __init__(self, store, host="0.0.0.0", port=8080):
-        self.timestep = None
+        self.timestep = dict()
         self.port = port
 
         def _serve_current_timestep(request):
-            return web.Response(body=str(self.timestep))
+            if len(self.timestep) > 0:
+                t_idx = min(self.timestep.values())
+            else:
+                return web.Response(status=204)
+            return web.json_response({"timestep": int(t_idx)})
 
         def _serve(request):
             try:
@@ -29,14 +33,14 @@ class LocoServer:
             loop.run_forever()
 
         app = web.Application()
-        app.add_routes([web.get("/current_timestep", _serve_current_timestep)])
+        app.add_routes([web.get("/current-timestep", _serve_current_timestep)])
         app.add_routes([web.get("/{key:.*}", _serve)])
         runner = web.AppRunner(app)
         t = Thread(target=_run_server, args=(runner, host, port), daemon=True)
         t.start()
 
-    def set_timestep(self, timestep):
-        self.timestep = timestep
+    def set_timestep(self, varname, timestep):
+        self.timestep[varname] = timestep
 
     def write_nginx_config(self, distributed_data_vars, rank):
         import platform
@@ -68,7 +72,7 @@ class LocoServer:
                 # use last rank for timestep assumption: this rank has
                 # the healpix order 0 grid where the information
                 # arrives last
-                f.write(f'  "/current_timestep" loco_backend{r};\n')
+                f.write(f'  "/current-timestep" loco_backend{r};\n')
                 f.write("}\n\n")
 
                 for i in range(len(distributed_data_vars)):