class FuncXWorker:
"""The FuncX worker
Parameters
----------
worker_id : str
Worker id string
address : str
Address at which the manager might be reached. This is usually 127.0.0.1
port : int
Port at which the manager can be reached
result_size_limit : int
Maximum result size allowed in Bytes
Default = 10 MB
Funcx worker will use the REP sockets to:
task = recv ()
result = execute(task)
send(result)
"""
def __init__(
self,
worker_id,
address,
port,
worker_type="RAW",
result_size_limit=DEFAULT_RESULT_SIZE_LIMIT_B,
):
self.worker_id = worker_id
self.address = address
self.port = port
self.worker_type = worker_type
self.serializer = FuncXSerializer()
self.serialize = self.serializer.serialize
self.deserialize = self.serializer.deserialize
self.result_size_limit = result_size_limit
log.info(f"Initializing worker {worker_id}")
log.info(f"Worker is of type: {worker_type}")
self.context = zmq.Context()
self.poller = zmq.Poller()
self.identity = worker_id.encode()
self.task_socket = self.context.socket(zmq.DEALER)
self.task_socket.setsockopt(zmq.IDENTITY, self.identity)
log.info(f"Trying to connect to : tcp://{self.address}:{self.port}")
self.task_socket.connect(f"tcp://{self.address}:{self.port}")
self.poller.register(self.task_socket, zmq.POLLIN)
signal.signal(signal.SIGTERM, self.handler)
def handler(self, signum, frame):
log.error("Signal handler called with signal", signum)
sys.exit(1)
def registration_message(self):
return {"worker_id": self.worker_id, "worker_type": self.worker_type}
def start(self):
log.info("Starting worker")
result = self.registration_message()
task_type = b"REGISTER"
log.debug("Sending registration")
self.task_socket.send_multipart([task_type,
pickle.dumps(result)] # Byte encoded
)
while True:
log.debug("Waiting for task")
p_task_id, p_container_id, msg = self.task_socket.recv_multipart()
task_id = pickle.loads(p_task_id)
container_id = pickle.loads(p_container_id)
log.debug(f"Received task_id:{task_id} with task:{msg}")
result = None
task_type = None
if task_id == "KILL":
task = Message.unpack(msg)
if task.task_buffer.decode("utf-8") == "KILL":
log.info("[KILL] -- Worker KILL message received! ")
task_type = b"WRKR_DIE"
else:
log.exception(
"Caught an exception of non-KILL message for KILL task"
)
continue
else:
log.debug("Executing task...")
try:
result = self.execute_task(msg)
serialized_result = self.serialize(result)
if len(serialized_result) > self.result_size_limit:
raise MaxResultSizeExceeded(len(serialized_result),
self.result_size_limit)
except Exception as e:
log.exception(f"Caught an exception {e}")
result_package = {
"task_id":
task_id,
"container_id":
container_id,
"exception":
self.serialize(
RemoteExceptionWrapper(*sys.exc_info())),
}
else:
log.debug("Execution completed without exception")
result_package = {
"task_id": task_id,
"container_id": container_id,
"result": serialized_result,
}
result = result_package
task_type = b"TASK_RET"
log.debug("Sending result")
self.task_socket.send_multipart([task_type,
pickle.dumps(result)
] # Byte encoded
)
if task_type == b"WRKR_DIE":
log.info(f"*** WORKER {self.worker_id} ABOUT TO DIE ***")
# Kill the worker after accepting death in message to manager.
sys.exit()
# We need to return here to allow for sys.exit mocking in tests
return
log.warning("Broke out of the loop... dying")
def execute_task(self, message):
"""Deserialize the buffer and execute the task.
Returns the result or throws exception.
"""
task = Message.unpack(message)
f, args, kwargs = self.serializer.unpack_and_deserialize(
task.task_buffer.decode("utf-8"))
return f(*args, **kwargs)
class FuncXWorker(object):
""" The FuncX worker
Parameters
----------
worker_id : str
Worker id string
address : str
Address at which the manager might be reached. This is usually 127.0.0.1
port : int
Port at which the manager can be reached
logdir : str
Logging directory
debug : Bool
Enables debug logging
Funcx worker will use the REP sockets to:
task = recv ()
result = execute(task)
send(result)
"""
def __init__(self,
worker_id,
address,
port,
logdir,
debug=False,
worker_type='RAW'):
self.worker_id = worker_id
self.address = address
self.port = port
self.logdir = logdir
self.debug = debug
self.worker_type = worker_type
self.serializer = FuncXSerializer()
self.serialize = self.serializer.serialize
self.deserialize = self.serializer.deserialize
global logger
logger = set_file_logger(
'{}/funcx_worker_{}.log'.format(logdir, worker_id),
name="worker_log",
level=logging.DEBUG if debug else logging.INFO)
logger.info('Initializing worker {}'.format(worker_id))
logger.info('Worker is of type: {}'.format(worker_type))
if debug:
logger.debug('Debug logging enabled')
self.context = zmq.Context()
self.poller = zmq.Poller()
self.identity = worker_id.encode()
self.task_socket = self.context.socket(zmq.DEALER)
self.task_socket.setsockopt(zmq.IDENTITY, self.identity)
logger.info('Trying to connect to : tcp://{}:{}'.format(
self.address, self.port))
self.task_socket.connect('tcp://{}:{}'.format(self.address, self.port))
self.poller.register(self.task_socket, zmq.POLLIN)
def registration_message(self):
return {'worker_id': self.worker_id, 'worker_type': self.worker_type}
def start(self):
logger.info("Starting worker")
result = self.registration_message()
task_type = b'REGISTER'
while True:
logger.debug("Sending result")
self.task_socket.send_multipart([
task_type, # Byte encoded
pickle.dumps(result)
])
if task_type == b'WRKR_DIE':
logger.info("*** WORKER {} ABOUT TO DIE ***".format(
self.worker_id))
# Kill the worker after accepting death in message to manager.
exit()
logger.debug("Waiting for task")
p_task_id, msg = self.task_socket.recv_multipart()
task_id = pickle.loads(p_task_id)
logger.debug("Received task_id:{} with task:{}".format(
task_id, msg))
if msg == b"KILL":
logger.info("[KILL] -- Worker KILL message received! ")
task_type = b'WRKR_DIE'
result = None
continue
logger.debug("Executing task...")
try:
result = self.execute_task(msg)
serialized_result = self.serialize(result)
except Exception as e:
logger.exception(f"Caught an exception {e}")
result_package = {
'task_id':
task_id,
'exception':
self.serialize(RemoteExceptionWrapper(*sys.exc_info()))
}
else:
logger.debug("Execution completed without exception")
result_package = {
'task_id': task_id,
'result': serialized_result
}
result = result_package
task_type = b'TASK_RET'
logger.warning("Broke out of the loop... dying")
def execute_task(self, message):
"""Deserialize the buffer and execute the task.
Returns the result or throws exception.
"""
user_ns = locals()
user_ns.update({'__builtins__': __builtins__})
decoded = message.decode()
f, args, kwargs = self.serializer.unpack_and_deserialize(decoded)
return f(*args, **kwargs)
class FuncXWorker(object):
""" The FuncX worker
Parameters
----------
worker_id : str
Worker id string
address : str
Address at which the manager might be reached. This is usually 127.0.0.1
port : int
Port at which the manager can be reached
logdir : str
Logging directory
debug : Bool
Enables debug logging
Funcx worker will use the REP sockets to:
task = recv ()
result = execute(task)
send(result)
"""
def __init__(self,
worker_id,
address,
port,
logdir,
debug=False,
worker_type='RAW'):
self.worker_id = worker_id
self.address = address
self.port = port
self.logdir = logdir
self.debug = debug
self.worker_type = worker_type
self.serializer = FuncXSerializer()
self.serialize = self.serializer.serialize
self.deserialize = self.serializer.deserialize
global logger
logger = set_file_logger(
'{}/funcx_worker_{}.log'.format(logdir, worker_id),
name="worker_log",
level=logging.DEBUG if debug else logging.INFO)
logger.info('Initializing worker {}'.format(worker_id))
if debug:
logger.debug('Debug logging enabled')
self.context = zmq.Context()
self.poller = zmq.Poller()
self.identity = worker_id.encode()
self.task_socket = self.context.socket(zmq.DEALER)
self.task_socket.setsockopt(zmq.IDENTITY, self.identity)
logger.info('Trying to connect to : tcp://{}:{}'.format(
self.address, self.port))
self.task_socket.connect('tcp://{}:{}'.format(self.address, self.port))
self.poller.register(self.task_socket, zmq.POLLIN)
def registration_message(self):
return {'worker_id': self.worker_id, 'worker_type': self.worker_type}
def start(self):
logger.info("Starting worker")
result = self.registration_message()
task_type = b'REGISTER'
while True:
logger.debug("Sending result")
# TODO : Swap for our serialization methods
self.task_socket.send_multipart([
task_type, # Byte encoded
pickle.dumps(result)
])
if task_type == b'WRKR_DIE':
logger.info("*** WORKER {} ABOUT TO DIE ***".format(
self.worker_id))
exit(
) # Kill the worker after accepting death in message to manager.
logger.debug("Waiting for task")
p_task_id, msg = self.task_socket.recv_multipart()
task_id = pickle.loads(p_task_id)
logger.debug("Received task_id:{} with task:{}".format(
task_id, msg))
if task_id == "KILL":
logger.info("[KILL] -- Worker KILL message received! ")
task_type = b'WRKR_DIE'
result = None
logger.debug("Executing task...")
try:
result = self.execute_task(msg)
logger.debug("Executed result: {}".format(result))
serialized_result = serialize_object(result)
except Exception:
logger.exception("Caught an exception {}")
result_package = {
'task_id':
task_id,
'exception':
serialize_object(RemoteExceptionWrapper(*sys.exc_info()))
}
else:
logger.debug("Execution completed without exception")
result_package = {
'task_id': task_id,
'result': serialized_result
}
# TODO: Change this to serialize_object to match IX?
result = result_package
task_type = b'TASK_RET'
logger.warning("Broke out of the loop... dying")
def execute_task(self, message):
"""Deserialize the buffer and execute the task.
Returns the result or throws exception.
"""
logger.debug("Inside execute_task function")
user_ns = locals()
user_ns.update({'__builtins__': __builtins__})
logger.info("Trying to pickle load the message {}".format(message))
decoded = message.decode()
f, args, kwargs = self.serializer.unpack_and_deserialize(decoded)
logger.debug("Message unpacked")
# We might need to look into callability of the function from itself
# since we change it's name in the new namespace
"""
prefix = "parsl_"
fname = prefix + "f"
argname = prefix + "args"
kwargname = prefix + "kwargs"
resultname = prefix + "result"
user_ns.update({fname: f,
argname: args,
kwargname: kwargs,
resultname: resultname})
logger.debug("Namespace updated")
code = "{0} = {1}(*{2}, **{3})".format(resultname, fname,
argname, kwargname)
try:
exec(code, user_ns, user_ns)
except Exception as e:
raise e
else:
return user_ns.get(resultname)
"""
return f(*args, **kwargs)
