Use lazy pirate pattern in tomato-job (#159)

* lazy pirate in jobs

* ruff

* fix test & nicer debug

* higher frequency ketchup checks

* one more timeout

* retries in npoints handling

* Fix metadata passing.

* ruff & longer timeout

* rewrite assert

* io None

* tidy up debug

* lpp in manager

Author: PeterKraus

src/tomato/daemon/driver.py

@@ -26,6 +26,7 @@
 from tomato.driverinterface_2_1 import ModelInterface as MI_2_1
 from tomato.drivers import driver_to_interface
 from tomato.models import Reply, Daemon
+from tomato.daemon import lpp
 
 logger = logging.getLogger(__name__)
 ModelInterface = TypeVar("ModelInterface", MI_1_0, MI_2_0, MI_2_1)
@@ -371,25 +372,30 @@ def manager(port: int, timeout: int = 1000):
     logger.info("launched successfully")
     req = context.socket(zmq.REQ)
     req.connect(f"tcp://127.0.0.1:{port}")
-    poller = zmq.Poller()
-    poller.register(req, zmq.POLLIN)
-    to = timeout
+    lppargs = dict(
+        endpoint=f"tcp://127.0.0.1:{port}",
+        context=context,
+        sender=sender,
+        timeout=timeout,
+    )
 
     spawned_drivers = dict()
     driver_retries = defaultdict(int)
     component_retries = defaultdict(int)
 
     while getattr(thread, "do_run"):
-        req.send_pyobj(dict(cmd="status", sender=sender))
-        events = dict(poller.poll(to))
-        if req not in events:
-            logger.warning("could not contact tomato-daemon in %d ms", to)
-            to = to * 2
-            continue
-        elif to > timeout:
-            to = timeout
+        msg = dict(cmd="status", sender=sender)
+        ret, req = lpp.comm(req, msg, **lppargs)
+        if req.closed:
+            thread.do_run = False
+            break
+        if ret.success:
+            daemon: Daemon = ret.data
+        else:
+            logger.critical(ret.msg)
+            thread.do_run = False
+            break
 
-        daemon = req.recv_pyobj().data
         for driver in daemon.drvs.keys():
             args = [
                 daemon.port,

src/tomato/daemon/io.py

@@ -50,13 +50,13 @@ def merge_netcdfs(job: Job, snapshot=False):
     using the Component `role` as the group label.
     """
     logger = logging.getLogger(f"{__name__}.merge_netcdf")
-    logger.debug("opening datasets")
+    logger.debug("opening datasets in '%s'", job.jobpath)
     datasets = []
-    logger.debug(f"{job=}")
-    logger.debug(f"{job.jobpath=}")
     for fn in Path(job.jobpath).glob("*.pkl"):
-        with pickle.load(fn.open("rb")) as ds:
-            datasets.append(ds)
+        with fn.open("rb") as pkl:
+            ds = pickle.load(pkl)
+            if ds is not None:
+                datasets.append(ds)
     logger.debug("creating a DataTree from %d groups", len(datasets))
     dt = xr.DataTree.from_dict({ds.attrs["role"]: ds for ds in datasets})
     logger.debug(f"{dt=}")
@@ -66,10 +66,9 @@ def merge_netcdfs(job: Job, snapshot=False):
     }
     dt.attrs = root_attrs
     outpath = job.snappath if snapshot else job.respath
-    logger.debug("saving DataTree into '%s'", outpath)
+    logger.debug("saving DataTree into a NetCDF file at '%s'", outpath)
     dt.to_netcdf(outpath)
     dt.close()
-    logger.debug(f"{dt=}")
 
 
 def data_to_pickle(ds: xr.Dataset, path: Path, role: str):
@@ -81,9 +80,11 @@ def data_to_pickle(ds: xr.Dataset, path: Path, role: str):
     ds.attrs["role"] = role
     logger.debug("checking for existing pickle at '%s'", path)
     if path.exists():
-        with pickle.load(path.open("rb")) as oldds:
-            logger.debug("concatenating Dataset with existing data")
-            ds = xr.concat([oldds, ds], dim="uts")
+        with path.open("rb") as old:
+            oldds = pickle.load(old)
+            if oldds is not None:
+                logger.debug("concatenating Dataset with existing data")
+                ds = xr.concat([oldds, ds], dim="uts")
     logger.debug("dumping Dataset into pickle at '%s'", path)
     with path.open("wb") as out:
         pickle.dump(ds, out, protocol=5)

src/tomato/daemon/job.py

@@ -25,10 +25,11 @@
 import zmq
 import psutil
 import sys
+import xarray as xr
 
 from tomato.daemon.io import merge_netcdfs, data_to_pickle
-from tomato.daemon import jobdb
-from tomato.models import Pipeline, Daemon, Component, Device, Driver, Job, Reply
+from tomato.daemon import jobdb, lpp
+from tomato.models import Pipeline, Daemon, Component, Device, Driver, Job
 from dgbowl_schemas.tomato import to_payload
 from dgbowl_schemas.tomato.payload import Task
 
@@ -56,10 +57,14 @@ def method_validate(
                     address=cmps[cmp].address,
                     channel=cmps[cmp].channel,
                 )
-                req.send_pyobj(dict(cmd="task_validate", params=params))
-                ret = req.recv_pyobj()
-                req.close()
+                ret, req = lpp.comm(
+                    req,
+                    dict(cmd="task_validate", params=params),
+                    f"tcp://127.0.0.1:{drv.port}",
+                    context,
+                )
                 if ret.success:
+                    req.close()
                     break
         else:
             return False
@@ -295,19 +300,15 @@ def manager(port: int, timeout: int = 500):
     thread = current_thread()
     logger.info("launched successfully")
     req: zmq.Socket = context.socket(zmq.REQ)
-    req.RCVTIMEO = 1000
     req.connect(f"tcp://127.0.0.1:{port}")
-    poller = zmq.Poller()
-    poller.register(req, zmq.POLLIN)
+    lppargs = dict(endpoint=f"tcp://127.0.0.1:{port}", context=context)
     while getattr(thread, "do_run"):
         logger.debug("tick")
-        try:
-            req.send_pyobj(dict(cmd="status", sender=f"{__name__}.manager"))
-            ret: Reply = req.recv_pyobj()
-        except zmq.ZMQError:
-            logger.critical("could not contact tomato-daemon in 1 s", exc_info=True)
+        msg = dict(cmd="status", sender=f"{__name__}.manager")
+        ret, req = lpp.comm(req, msg, **lppargs)
+        if req.closed:
             break
-        if ret.success is False:
+        elif ret.success is False:
             logger.critical("tomato-daemon is not running: %s", ret.msg)
             break
         daemon: Daemon = ret.data
@@ -466,81 +467,85 @@ def job_thread(
     thread = current_thread()
     sender = f"{__name__}.job_thread({thread.ident})"
     logger = logging.getLogger(sender)
-    logger.debug(f"in job thread of {component.role!r}")
-
     context = zmq.Context()
     req = context.socket(zmq.REQ)
-    req.RCVTIMEO = 1000
     req.connect(f"tcp://127.0.0.1:{driver.port}")
-    logger.info(f"job thread of {component.role!r} connected to tomato-daemon")
+    lppargs = dict(
+        endpoint=f"tcp://127.0.0.1:{driver.port}", context=context, sender=sender
+    )
+
+    logger.info(
+        "%s: job thread of %s attached to tomato-daemon", component.role, component.name
+    )
 
     kwargs = dict(address=component.address, channel=component.channel)
 
     datapath = Path(jobpath) / f"{component.role}.pkl"
-    logger.debug("distributing tasks:")
+    logger.debug("%s: processing tasks on component %s", component.role, component.name)
     for ti, task in enumerate(tasks):
+        taskid = f"{component.role}:{ti}"
+        if task.task_name is not None:
+            taskid += f":{task.task_name!r}"
         thread.current_task = task
-        logger.info("processing task %s:%d", component.role, ti)
+        logger.info("%s: processing task", taskid)
         while True:
             time.sleep(1e-1)
             if task.start_with_task_name is None:
                 pass
             elif task.start_with_task_name in thread.started_task_names:
                 pass
             else:
-                logger.debug("waiting for task_name '%s'", task.start_with_task_name)
+                logger.debug(
+                    "%s: waiting for task_name '%s'", taskid, task.start_with_task_name
+                )
                 continue
-            logger.debug("polling component '%s' for task readiness", component.role)
-            try:
-                req.send_pyobj(dict(cmd="task_status", params={**kwargs}))
-                ret = req.recv_pyobj()
-            except zmq.ZMQError as e:
-                logger.critical(e, exc_info=True)
-                thread.crashed = True
-                sys.exit(e)
+            logger.debug(
+                "%s: polling component %s for task readiness", taskid, component.name
+            )
+            ret, req = lpp.comm(
+                req, dict(cmd="task_status", params={**kwargs}), **lppargs
+            )
             if ret.success and ret.data["can_submit"]:
                 break
-            logger.warning("cannot submit onto component '%s', waiting", component.role)
+            elif req.closed:
+                thread.crashed = True
+                sys.exit()
+            logger.warning(
+                "%s: cannot submit onto component %s, waiting", taskid, component.name
+            )
 
-        logger.info("sending task %s:%d to component", component.role, ti)
-        try:
-            req.send_pyobj(dict(cmd="task_start", params={"task": task, **kwargs}))
-            ret = req.recv_pyobj()
-        except zmq.ZMQError as e:
-            logger.critical(e, exc_info=True)
+        logger.info("%s: sending task to component %s", taskid, component.name)
+        msg = dict(cmd="task_start", params={"task": task, **kwargs})
+        ret, req = lpp.comm(req, msg, **lppargs)
+        if req.closed:
             thread.crashed = True
-            sys.exit(e)
+            sys.exit()
 
         t0 = time.perf_counter()
         while True:
             tN = time.perf_counter()
             if tN - t0 > device.pollrate:
-                logger.debug("polling task %s:%d for data", component.role, ti)
-                try:
-                    req.send_pyobj(dict(cmd="task_data", params={**kwargs}))
-                    ret = req.recv_pyobj()
-                except zmq.ZMQError as e:
-                    logger.critical(e, exc_info=True)
+                logger.debug("%s: polling task for data", taskid)
+                msg = dict(cmd="task_data", params={**kwargs})
+                ret, req = lpp.comm(req, msg, **lppargs, timeout=5000)
+                if req.closed:
                     thread.crashed = True
-                    sys.exit(e)
-                if ret.success:
-                    logger.debug("pickling received data")
-                    ds = ret.data
+                    sys.exit()
+                elif ret.success and ret.data is not None:
+                    logger.debug("%s: pickling received data", taskid)
+                    ds: xr.Dataset = ret.data
                     ds.attrs["tomato_Component"] = component.model_dump_json()
                     data_to_pickle(ds, datapath, role=component.role)
                 t0 += device.pollrate
 
-            logger.debug("polling task %s:%d for completion", component.role, ti)
-            try:
-                req.send_pyobj(dict(cmd="task_status", params={**kwargs}))
-                ret = req.recv_pyobj()
-            except zmq.ZMQError as e:
-                logger.critical(e, exc_info=True)
+            logger.debug("%s: polling task for completion", taskid)
+            msg = dict(cmd="task_status", params={**kwargs})
+            ret, req = lpp.comm(req, msg, **lppargs)
+            if req.closed:
                 thread.crashed = True
-                sys.exit(e)
-
-            if ret.success and not ret.data["running"]:
-                logger.info("task %s:%d no longer running, break", component.role, ti)
+                sys.exit()
+            elif ret.success and not ret.data["running"]:
+                logger.info("%s: task no longer running, break", taskid)
                 break
             elif ret.success is False:
                 logger.critical(f"{ret=}")
@@ -550,50 +555,49 @@ def job_thread(
                 task.stop_with_task_name is not None
                 and task.stop_with_task_name in thread.started_task_names
             ):
-                logger.info("task %s:%d stop trigger met", component.role, ti)
-                try:
-                    req.RCVTIMEO = 10000
-                    req.send_pyobj(dict(cmd="task_stop", params={**kwargs}))
-                    ret = req.recv_pyobj()
-                    req.RCVTIMEO = 1000
-                except zmq.ZMQError as e:
-                    logger.critical(e, exc_info=True)
+                logger.info("%s: task stop trigger met", taskid)
+                msg = dict(cmd="task_stop", params={**kwargs})
+                ret, req = lpp.comm(req, msg, **lppargs, timeout=5000)
+                if req.closed:
                     thread.crashed = True
-                    sys.exit(e)
-                if ret.success and ret.data is not None:
-                    data_to_pickle(ret.data, datapath, role=component.role)
+                    sys.exit()
+                elif ret.success and ret.data is not None:
+                    logger.debug("%s: pickling received data", taskid)
+                    ds: xr.Dataset = ret.data
+                    ds.attrs["tomato_Component"] = component.model_dump_json()
+                    data_to_pickle(ds, datapath, role=component.role)
                 break
 
             time.sleep(max(1e-1, (device.pollrate - (tN - t0)) / 2))
-        logger.info("task %s:%d fetching final data", component.role, ti)
-        try:
-            req.send_pyobj(dict(cmd="task_data", params={**kwargs}))
-            ret = req.recv_pyobj()
-        except zmq.ZMQError as e:
-            logger.critical(e, exc_info=True)
+        logger.info("%s: task fetching final data", taskid)
+        msg = dict(cmd="task_data", params={**kwargs})
+        ret, req = lpp.comm(req, msg, **lppargs, timeout=5000)
+        if req.closed:
             thread.crashed = True
-            sys.exit(e)
-        if ret.success:
-            data_to_pickle(ret.data, datapath, role=component.role)
+            sys.exit()
+        elif ret.success and ret.data is not None:
+            logger.debug("%s: pickling received data", taskid)
+            ds: xr.Dataset = ret.data
+            ds.attrs["tomato_Component"] = component.model_dump_json()
+            data_to_pickle(ds, datapath, role=component.role)
         thread.completed_tasks.append(task)
         thread.current_task = None
 
-    logger.info("all tasks done on component '%s', resetting", component.role)
-    try:
-        if driver.version == "1.0":
-            req.send_pyobj(dict(cmd="dev_reset", params={**kwargs}))
-        else:
-            req.send_pyobj(dict(cmd="cmp_reset", params={**kwargs}))
-        req.RCVTIMEO = 10000
-        ret = req.recv_pyobj()
-    except zmq.ZMQError as e:
-        logger.critical(e, exc_info=True)
+    logger.info(
+        "%s: all tasks done on component %s, resetting", component.role, component.name
+    )
+    if driver.version == "1.0":
+        msg = dict(cmd="dev_reset", params={**kwargs})
+    else:
+        msg = dict(cmd="cmp_reset", params={**kwargs})
+    ret, req = lpp.comm(req, msg, **lppargs, timeout=5000)
+    if req.closed:
         thread.crashed = True
-        sys.exit(e)
-    if not ret.success:
-        logger.warning("could not reset component '%s': %s", component.role, ret.msg)
+        sys.exit()
+    elif not ret.success:
+        logger.warning("%s: could not reset component: %s", component.role, ret.msg)
     else:
-        logger.info("reset of component '%s' complete", component.role)
+        logger.info("%s: reset of component %s done", component.role, component.name)
     req.close()
 
 
@@ -613,15 +617,14 @@ def job_main_loop(
 
     req = context.socket(zmq.REQ)
     req.connect(f"tcp://127.0.0.1:{port}")
-    req.RCVTIMEO = 1000
+    lppargs = dict(endpoint=f"tcp://127.0.0.1:{port}", context=context)
 
     while True:
-        req.send_pyobj(dict(cmd="status", sender=sender))
-        try:
-            daemon: Daemon = req.recv_pyobj().data
-        except zmq.ZMQError as e:
-            logger.critical(e, exc_info=True)
-            sys.exit(e)
+        ret, req = lpp.comm(req, dict(cmd="status", sender=sender), **lppargs)
+        if ret.success:
+            daemon: Daemon = ret.data
+        else:
+            sys.exit()
         if all([drv.port is not None for drv in daemon.drvs.values()]):
             break
         else: