first commit

python3-vckonline/lib/python3.8/site-packages/pop/mods/proc/run.py

@@ -0,0 +1,180 @@
+"""
+Execute functions or load subs on the workers in the named worker pool
+"""
+# import python libs
+import asyncio
+import os
+
+# Import third party libs
+import msgpack
+import pop.hub
+
+
+async def add_sub(hub: "pop.hub.Hub", worker_name, *args, **kwargs):
+    """
+    Tell all of the worker in the named pool to load the given sub,
+
+    This function takes all of the same arguments as hub.pop.sub.add
+    """
+    ret = {}
+    workers = hub.proc.Workers[worker_name]
+    for ind in workers:
+        payload = {"fun": "sub", "args": args, "kwargs": kwargs}
+        # TODO: Make these futures to the run at the same time
+        async for chunk in hub.proc.run.send(workers[ind], payload):
+            ret[ind] = chunk
+    hub.proc.WorkersTrack[worker_name]["subs"].append({"args": args, "kwargs": kwargs})
+    return ret
+
+
+async def add_proc(hub: "pop.hub.Hub", worker_name):
+    """
+    Add a single process to the worker pool, also make sure that
+    """
+    # grab and extrapolate the data we need
+    ret_ref = hub.proc.WorkersTrack[worker_name]["ret_ref"]
+    sock_dir = hub.proc.WorkersTrack[worker_name]["sock_dir"]
+    workers = hub.proc.Workers[worker_name]
+    ind = len(workers) + 1
+    for s_ind in range(len(workers) + 1):
+        if s_ind not in workers:
+            ind = s_ind
+    hub.proc.init.mk_proc(ind, workers, ret_ref, sock_dir)
+    # Make sure the process is up with a live socket
+    while True:
+        if os.path.exists(workers[ind]["path"]):
+            break
+        await asyncio.sleep(0.01)
+    # Add all of the subs that have been added to processes in this pool
+    for sub in hub.proc.WorkersTrack[worker_name]["subs"]:
+        payload = {"fun": "sub", "args": sub["args"], "kwargs": sub["kwargs"]}
+        async for chunk in hub.proc.run.send(workers[ind], payload):
+            pass
+    return ind
+
+
+async def pub(hub: "pop.hub.Hub", worker_name, func_ref, *args, **kwargs):
+    """
+    Execute the given function reference on ALL the workers in the given
+    worker pool and return the return data from each.
+
+    Pass in the arguments for the function, keep in mind that the sub needs
+    to be loaded into the workers for a function to be available via
+    hub.proc.run.add_sub
+    """
+    workers = hub.proc.Workers[worker_name]
+    ret = {}
+    for ind in workers:
+        payload = {"fun": "run", "ref": func_ref, "args": args, "kwargs": kwargs}
+        # TODO: Make these futures to the run at the same time
+        async for chunk in hub.proc.run.send(workers[ind], payload):
+            ret[ind] = chunk
+    return ret
+
+
+async def set_attr(hub: "pop.hub.Hub", worker_name, ref, value):
+    """
+    Set the given attribute to the given location on the hub of all
+    worker procs
+    """
+    workers = hub.proc.Workers[worker_name]
+    ret = {}
+    for ind in workers:
+        payload = {"fun": "setattr", "ref": ref, "value": value}
+        # TODO: Make these futures to the run at the same time
+        async for chunk in hub.proc.run.send(workers[ind], payload):
+            ret[ind] = chunk
+    return ret
+
+
+async def ind_func(hub: "pop.hub.Hub", worker_name, _ind, func_ref, *args, **kwargs):
+    """
+    Execute the function on the indexed process within the named worker pool
+    """
+    workers = hub.proc.Workers[worker_name]
+    worker = workers[_ind]
+    payload = {"fun": "run", "ref": func_ref, "args": args, "kwargs": kwargs}
+    async for ret in hub.proc.run.send(worker, payload):
+        return ret
+
+
+async def func(hub: "pop.hub.Hub", worker_name, func_ref, *args, **kwargs):
+    """
+    Execute the given function reference on one worker in the given worker
+    pool and return the return data.
+
+    Pass in the arguments for the function, keep in mind that the sub needs
+    to be loaded into the workers for a function to be available via
+    hub.proc.run.add_sub
+    """
+    ind, coro = await hub.proc.run.track_func(worker_name, func_ref, *args, **kwargs)
+    return await coro
+
+
+async def track_func(hub: "pop.hub.Hub", worker_name, func_ref, *args, **kwargs):
+    """
+    Run a function and return the index of the worker that the function was
+    executed on and a coroutine to track
+    """
+    w_iter = hub.proc.WorkersIter[worker_name]
+    ind = next(w_iter)
+    coro = hub.proc.run.ind_func(worker_name, ind, func_ref, *args, **kwargs)
+    return ind, coro
+
+
+async def gen(hub: "pop.hub.Hub", worker_name, func_ref, *args, **kwargs):
+    """
+    Execute a generator function reference within one worker within the given
+    worker pool.
+
+    Like `func` the sub needs to be made available to all workers first
+    """
+    ind, coro = await hub.proc.run.track_gen(worker_name, func_ref, *args, **kwargs)
+    async for chunk in coro:
+        yield chunk
+
+
+async def track_gen(hub: "pop.hub.Hub", worker_name, func_ref, *args, **kwargs):
+    """
+    Return an iterable coroutine and the index executed on
+    """
+    w_iter = hub.proc.WorkersIter[worker_name]
+    ind = next(w_iter)
+    coro = hub.proc.run.ind_gen(worker_name, ind, func_ref, *args, **kwargs)
+    return ind, coro
+
+
+async def ind_gen(hub: "pop.hub.Hub", worker_name, _ind, func_ref, *args, **kwargs):
+    """
+    run the given iterator on the defined index
+    """
+    workers = hub.proc.Workers[worker_name]
+    worker = workers[_ind]
+    payload = {"fun": "gen", "ref": func_ref, "args": args, "kwargs": kwargs}
+    async for chunk in hub.proc.run.send(worker, payload):
+        yield chunk
+
+
+async def send(hub: "pop.hub.Hub", worker, payload):
+    """
+    Send the given payload to the given worker, yield iterations based on the
+    returns from the remote.
+    """
+    mp = msgpack.dumps(payload, use_bin_type=True)
+    mp += hub.proc.DELIM
+    reader, writer = await asyncio.open_unix_connection(path=worker["path"])
+    writer.write(mp)
+    await writer.drain()
+    final_ret = True
+    while True:
+        ret = await reader.readuntil(hub.proc.DELIM)
+        p_ret = ret[: -len(hub.proc.DELIM)]
+        i_flag = p_ret[-1:]
+        ret = msgpack.loads(p_ret[:-1], raw=False)
+        if i_flag == hub.proc.D_FLAG:
+            # break for the end of the sequence
+            break
+        yield ret
+        final_ret = False
+    if final_ret:
+        yield ret