181 lines
6.1 KiB
Python
181 lines
6.1 KiB
Python
"""
|
|
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
|