163 lines
4.1 KiB
Python
163 lines
4.1 KiB
Python
import cloudpickle
|
|
import grpc
|
|
import proto.task_pb2
|
|
import proto.task_pb2_grpc
|
|
import uuid
|
|
|
|
|
|
class ObjectID:
|
|
def __init__(self, id):
|
|
self.id = id
|
|
|
|
def __repr__(self):
|
|
return "ObjectID(%s)" % self.id.decode()
|
|
|
|
|
|
worker_registry = {}
|
|
|
|
|
|
def set_global_worker(worker):
|
|
global _global_worker
|
|
_global_worker = worker
|
|
|
|
|
|
def register_worker(worker):
|
|
id = uuid.uuid4()
|
|
worker_registry[id] = worker
|
|
return id
|
|
|
|
|
|
def get_worker_registry(id):
|
|
out = worker_registry.get(id)
|
|
if out is None:
|
|
return _global_worker
|
|
return out
|
|
|
|
|
|
class Worker:
|
|
def __init__(self, conn_str="", stub=None):
|
|
if stub is None:
|
|
self.channel = grpc.insecure_channel(conn_str)
|
|
self.server = proto.task_pb2_grpc.TaskServerStub(self.channel)
|
|
else:
|
|
self.server = stub
|
|
self.uuid = register_worker(self)
|
|
|
|
def get(self, ids):
|
|
to_get = []
|
|
single = False
|
|
if isinstance(ids, list):
|
|
to_get = [x.id for x in ids]
|
|
elif isinstance(ids, ObjectID):
|
|
to_get = [ids.id]
|
|
single = True
|
|
else:
|
|
raise Exception("Can't get something that's not a list of IDs or just an ID")
|
|
|
|
out = [self._get(x) for x in to_get]
|
|
if single:
|
|
out = out[0]
|
|
return out
|
|
|
|
def _get(self, id: bytes):
|
|
req = proto.task_pb2.GetRequest(id=id)
|
|
data = self.server.GetObject(req)
|
|
return cloudpickle.loads(data.data)
|
|
|
|
def put(self, vals):
|
|
to_put = []
|
|
single = False
|
|
if isinstance(vals, list):
|
|
to_put = vals
|
|
else:
|
|
single = True
|
|
to_put.append(vals)
|
|
|
|
out = [self._put(x) for x in to_put]
|
|
if single:
|
|
out = out[0]
|
|
return out
|
|
|
|
def _put(self, val):
|
|
data = cloudpickle.dumps(val)
|
|
#print("val: %s\ndata: %s"%(val, data))
|
|
req = proto.task_pb2.PutRequest(data=data)
|
|
resp = self.server.PutObject(req)
|
|
return ObjectID(resp.id)
|
|
|
|
def remote(self, func):
|
|
return RemoteFunc(self, func)
|
|
|
|
def schedule(self, task):
|
|
return self.server.Schedule(task)
|
|
|
|
def close(self):
|
|
self.channel.close()
|
|
|
|
|
|
class RemoteFunc:
|
|
def __init__(self, worker, f):
|
|
self._func = f
|
|
self._name = f.__name__
|
|
self.id = None
|
|
self._worker_id = worker.uuid
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
raise Exception("Matching the old API")
|
|
|
|
def remote(self, *args):
|
|
if self.id is None:
|
|
self._push_func()
|
|
t = proto.task_pb2.Task()
|
|
t.name = self._name
|
|
t.payload_id = self.id.id
|
|
for a in args:
|
|
arg = proto.task_pb2.Arg()
|
|
if isinstance(a, ObjectID):
|
|
arg.local = proto.task_pb2.Arg.Locality.REFERENCE
|
|
arg.reference_id = a.id
|
|
else:
|
|
arg.local = proto.task_pb2.Arg.Locality.INTERNED
|
|
arg.data = cloudpickle.dumps(a)
|
|
t.args.append(arg)
|
|
worker = get_worker_registry(self._worker_id)
|
|
ticket = worker.schedule(t)
|
|
return ObjectID(ticket.return_id)
|
|
|
|
def _push_func(self):
|
|
worker = get_worker_registry(self._worker_id)
|
|
self.id = worker.put(self._func)
|
|
|
|
def __repr__(self):
|
|
return "RemoteFunc(%s, %s)" % (self._name, self.id)
|
|
|
|
|
|
_global_worker = None
|
|
|
|
|
|
def connect(*args, **kwargs):
|
|
global _global_worker
|
|
if _global_worker is not None:
|
|
raise Exception("Can't connect a second global worker")
|
|
_global_worker = Worker(*args, **kwargs)
|
|
|
|
|
|
def get(*args, **kwargs):
|
|
global _global_worker
|
|
if _global_worker is None:
|
|
raise Exception("Need a connection before calling")
|
|
return _global_worker.get(*args, **kwargs)
|
|
|
|
|
|
def put(*args, **kwargs):
|
|
global _global_worker
|
|
if _global_worker is None:
|
|
raise Exception("Need a connection before calling")
|
|
return _global_worker.put(*args, **kwargs)
|
|
|
|
|
|
def remote(*args, **kwargs):
|
|
global _global_worker
|
|
if _global_worker is None:
|
|
raise Exception("Need a connection before calling")
|
|
return _global_worker.remote(*args, **kwargs)
|