src/libertem/executor/pipelined.py
import itertools
import os
import sys
import logging
import functools
import contextlib
import multiprocessing as mp
from typing import (
TYPE_CHECKING, Any, Callable, NamedTuple, Optional, TypeVar, Union,
)
from collections.abc import Iterable, Generator
from typing_extensions import TypedDict, Literal
import uuid
import warnings
import time
from tblib import pickling_support
import cloudpickle
from opentelemetry import trace
from libertem.common.backend import set_use_cpu, set_use_cuda, set_file_limit
from libertem.common.executor import (
Environment, TaskProtocol, WorkerContext, WorkerQueue,
WorkerQueueEmpty, TaskCommHandler, SimpleMPWorkerQueue,
JobCancelledError, ResourceDef,
)
from libertem.common.scheduler import Worker, WorkerSet, Scheduler
from libertem.common.tracing import add_partition_to_span, attach_to_parent, maybe_setup_tracing
from .utils import assign_cudas
from .base import BaseJobExecutor, ResourceError
try:
import prctl
except ImportError:
prctl = None
if TYPE_CHECKING:
from opentelemetry.trace import SpanContext
tracer = trace.get_tracer(__name__)
logger = logging.getLogger(__name__)
T = TypeVar('T')
class WorkerQueues(NamedTuple):
request: "WorkerQueue"
response: "WorkerQueue"
message: "WorkerQueue"
class WorkerSpec(TypedDict):
name: str
device_id: int
device_kind: Union[Literal['CPU'], Literal['CUDA']]
worker_idx: int
has_cupy: bool
class PoolWorkerInfo(NamedTuple):
queues: WorkerQueues
process: mp.Process
spec: WorkerSpec
class PoolStateError(Exception):
"""
The worker pool is not in the expected state to perform the requested operation
"""
pass
def _resources_for_spec(spec: WorkerSpec) -> ResourceDef:
worker_resources: "ResourceDef" = {}
if spec['device_kind'] == 'CPU':
worker_resources["compute"] = 1
worker_resources["CPU"] = 1
worker_resources["ndarray"] = 1
if spec["device_kind"] == "CUDA":
worker_resources["compute"] = 1
worker_resources["CUDA"] = 1
if spec["has_cupy"]:
worker_resources["ndarray"] = 1
return worker_resources
class WorkerPool:
"""
Combination of worker processes and matching request queues,
and a single response queue.
Processes are started using the spawn method, so they need
to use primitives that are compatible with spawn.
Take care to properly close queues and join workers at shutdown.
Note
----
We are not using the vanilla :class:`multiprocesing.Pool` here, because
we need to coordinate the sending and receiving side for each task,
and we need to keep state on the worker, pin processes to cores etc.
"""
def __init__(self, worker_fn: Callable, spec: list[WorkerSpec]):
self._worker_q_cls = SimpleMPWorkerQueue
self._workers: Optional[list[PoolWorkerInfo]] = []
self._worker_fn = worker_fn
self._response_q: Optional[SimpleMPWorkerQueue] = self._worker_q_cls()
self._message_q = self._worker_q_cls()
self._mp_ctx = mp.get_context("spawn")
self._spec = spec
self._start_workers()
@property
def response_queue(self) -> "WorkerQueue":
if self._response_q is None:
raise PoolStateError("Response queue not available")
return self._response_q
@property
def message_queue(self) -> "WorkerQueue":
return self._message_q
@property
def size(self) -> int:
return len(self._spec)
def _start_workers(self):
with tracer.start_as_current_span("WorkerPool._start_workers") as span:
span_context = span.get_span_context()
for spec_item in self._spec:
queues = self._make_worker_queues()
p = self._mp_ctx.Process(target=self._worker_fn, kwargs={
"queues": queues,
"spec": spec_item,
"span_context": span_context,
})
p.start()
self._workers.append(
PoolWorkerInfo(
queues=queues,
process=p,
spec=spec_item,
)
)
def kill_worker(self, worker_info: PoolWorkerInfo, timeout: float = 5.0):
worker_info.queues.request.close(drain=False, force=True)
worker_info.process.terminate()
worker_info.process.join(timeout)
if worker_info.process.exitcode is None:
worker_info.process.kill()
# reap the dead process:
worker_info.process.join(30)
def kill(self, timeout: float = 5):
if self._workers is None:
raise PoolStateError("Cannot kill workers, not running")
for worker in self._workers:
self.kill_worker(worker, timeout=timeout)
exitcodes = [
worker.process.exitcode
for worker in self._workers
]
self._workers = None
if self._response_q is None:
raise PoolStateError("Response queue not available")
self._response_q.close(drain=False, force=True)
self._response_q = None
if any([e is None for e in exitcodes]):
raise PoolStateError("At least one worker failed to be killed")
@property
def workers(self) -> list[PoolWorkerInfo]:
if self._workers is None:
raise PoolStateError("No workers are running")
return self._workers
def all_alive(self) -> bool:
if self._workers is None:
raise PoolStateError("No workers are running")
return all(qp.process.is_alive() for qp in self._workers)
def assert_all_alive(self):
"""
Check if any workers are dead, if so, kill the whole pool and raise a
`RuntimeError`.
"""
if not self.all_alive():
self.kill()
raise RuntimeError(
"One or more workers failed to start"
)
def close_resp_queue(self):
self._response_q.close()
def close_mesg_queue(self):
self._message_q.close()
def get_worker_queues(self, worker_idx: int) -> WorkerQueues:
if self._workers is None:
raise PoolStateError("No workers are running")
return self._workers[worker_idx].queues
def _make_worker_queues(self):
return WorkerQueues(
request=self._worker_q_cls(),
response=self._response_q,
message=self._message_q,
)
WorkerSelector = Callable[
[list[PoolWorkerInfo], WorkerPool, TaskProtocol, int],
PoolWorkerInfo
]
def _select_by_queue_size(
eligible: list[PoolWorkerInfo],
pool: WorkerPool,
task: TaskProtocol,
task_idx: int,
) -> PoolWorkerInfo:
# FIXME: I think this introduces some latency currently,
# as the workers `get` from their queue as soon as possible,
# so the following could happen:
# - task 0 becomes available
# - all queues have size 0, so task gets assigned to worker 0
# - worker 0 picks the task from its queue, so its queue size is 0 again
# - task 1 becomes available
# - all queues have size 0 (again), so task gets assigned to worker 0
# => this should only happen in the offline case, as in the live
# processing case, we push more than one message per task into the
# queue, so there should be more than one message in the backlog
# when the `TaskCommHandler` is finished pushing messages for one task.
return min(eligible, key=lambda w: w.queues.request.size())
def _select_by_round_robin(
eligible: list[PoolWorkerInfo],
pool: WorkerPool,
task: TaskProtocol,
task_idx: int,
) -> PoolWorkerInfo:
return eligible[task_idx % len(eligible)]
def schedule_task(
task_idx: int,
task: TaskProtocol,
pool: WorkerPool,
scheduler: Scheduler,
pool_info_for_worker: dict[Worker, PoolWorkerInfo],
selector: WorkerSelector,
) -> tuple[int, WorkerQueues]:
"""
Returns the worker index and its queues of the worker the given task should
be scheduled on.
Workers are first filtered based on requested resources, then the concrete
worker is selected by the function specified in the :code:`selector`
parameter.
"""
eligible = [
pool_info_for_worker[w]
for w in scheduler.workers_for_task(task)
]
if len(eligible) == 0:
resources = task.get_resources()
raise ResourceError(f"No worker available with resources {resources}")
worker = selector(eligible, pool, task, task_idx)
idx = pool.workers.index(worker)
return idx, worker.queues
def set_thread_name(name: str):
"""
Set a thread name; mostly useful for using system tools for profiling
Parameters
----------
name : str
The thread name
"""
if prctl is None:
return
prctl.set_name(name)
def worker_run_task(
header: dict,
work_mem: dict,
queues: WorkerQueues,
worker_idx: int,
env: Environment,
):
"""
Called from the worker main loop when a RUN_TASK message is received
Parameters
----------
header
The header of the message that was received
work_mem
The worker's working memory, for accessing scattered data
queues
request and response queues
worker_idx
This worker's index
env
The Environment for preparing thread counts etc. for the UDF run
"""
with tracer.start_as_current_span("RUN_TASK") as span:
try:
span.set_attributes({
"libertem.task_size_pickled": len(header["task"]),
"libertem.os.pid": os.getpid(),
})
task: TaskProtocol = cloudpickle.loads(header["task"])
params_handle = header["params_handle"]
params = work_mem[params_handle]
partition = task.get_partition()
add_partition_to_span(partition)
result = task(params, env)
queues.response.put({
"type": "RESULT",
"result": result,
"task_id": header["task_id"],
"uuid": header["uuid"],
"worker_id": worker_idx,
})
except Exception as e:
logger.exception("failure in RUN_TASK")
pickling_support.install(e)
queues.response.put({
"type": "ERROR",
"error": e,
"exception": e,
"worker_id": worker_idx,
"uuid": header["uuid"],
})
finally:
pass
def worker_run_function(header, queues, worker_idx):
"""
Called from the worker main loop when a RUN_FUNCTION message is received
Parameters
----------
header
The header of the message that was received
queues
request and response queues
worker_idx
This worker's index
"""
with tracer.start_as_current_span("RUN_FUNCTION"):
try:
fn = cloudpickle.loads(header["fn"])
result = fn()
queues.response.put({
"type": "RUN_FUNCTION_RESULT",
"result": result,
"worker_id": worker_idx,
})
except Exception as e:
logger.exception("failure in RUN_FUNCTION")
pickling_support.install(e)
queues.response.put({
"type": "ERROR",
"error": e,
"exception": e,
"worker_id": worker_idx,
})
def _drain_after_task(queues: WorkerQueues):
"""
In case of an error, need to drain the request queue
(anything that is left over from the detector-specific
data that was sent in `TaskCommHandler.handle_task`):
"""
with tracer.start_as_current_span("drain after task") as span:
while True:
with queues.request.get() as msg:
header, payload = msg
header_type = header["type"]
span.add_event("msg", {"type": header_type})
if header_type in (
"RUN_TASK", "SCATTER", "SCATTER_UPDATE",
"SCATTER_UPDATE_PATCH", "RUN_FUNCTION",
"DELETE", "SHUTDOWN", "WARMUP",
):
raise RuntimeError(
f"unexpected message type {header_type}"
)
if header_type == "END_TASK":
break
def worker_loop(
queues: WorkerQueues,
work_mem: dict,
worker_idx: int,
env: Environment
):
"""
The worker main loop, called when the worker setup is done.
Waits for messages on the request queue until a SHUTDOWN message
is received.
Parameters
----------
work_mem
The worker's working memory, for accessing scattered data
queues
request and response queues
worker_idx
This worker's index
env
The Environment for preparing thread counts etc. for the UDF run
"""
while True:
try:
with queues.request.get() as msg:
header, payload = msg
header_type = header["type"]
if header_type == "RUN_TASK":
with attach_to_parent(header["span_context"]):
worker_run_task(header, work_mem, queues, worker_idx, env)
_drain_after_task(queues)
elif header_type == "SCATTER":
# FIXME: array data could be transferred and stored in SHM instead
key = header["key"]
if key in work_mem:
queues.response.put({
"type": "ERROR",
"error": f"key {key} already stored in worker memory",
"worker_id": worker_idx,
})
continue
work_mem[key] = header["value"]
continue
elif header_type == "SCATTER_UPDATE":
# FIXME: array data could be transferred and stored in SHM instead
key = header["key"]
if key not in work_mem:
queues.response.put({
"type": "ERROR",
"error": f"key {key} not stored in worker memory, can't update",
"worker_id": worker_idx,
})
continue
work_mem[key] = header["value"]
continue
elif header_type == "SCATTER_UPDATE_PATCH":
# FIXME: array data could be transferred and stored in SHM instead
key = header["key"]
if key not in work_mem:
queues.response.put({
"type": "ERROR",
"error": f"key {key} not stored in worker memory, can't update",
"worker_id": worker_idx,
})
continue
if not hasattr(work_mem[key], 'patch'):
queues.response.put({
"type": "ERROR",
"error": f"key {key} is not patcheable",
"worker_id": worker_idx,
})
continue
work_mem[key].patch(header["patch"])
continue
elif header_type == "RUN_FUNCTION":
with attach_to_parent(header["span_context"]):
worker_run_function(header, queues, worker_idx)
elif header_type == "DELETE":
key = header["key"]
if key in work_mem:
del work_mem[key]
continue
elif header_type == "SHUTDOWN":
with attach_to_parent(header["span_context"]):
with tracer.start_as_current_span("SHUTDOWN"):
queues.request.close()
queues.response.close(drain=False)
queues.message.close(drain=False)
break
elif header_type == "WARMUP":
with attach_to_parent(header["span_context"]):
with tracer.start_as_current_span("WARMUP"):
import libertem.udf.base # NOQA
import libertem.api # NOQA
import libertem.preload # NOQA
with env.enter():
pass
else:
queues.response.put({
"type": "ERROR",
"error": f"unknown message {header}",
"worker_id": worker_idx,
})
# probably desynchronized with the main process, so give up:
raise RuntimeError(f"unknown message, shutting down worker {worker_idx}")
except KeyboardInterrupt:
queues.response.put({
"type": "ERROR",
"error": "KeyboardInterrupt",
"worker_id": worker_idx,
})
raise
DeviceT = tuple[int, Union[Literal['CUDA'], Literal['CPU']]]
def _setup_device(spec: WorkerSpec, pin: bool):
"""
Set up this worker for its given task - either CPU or GPU comptation,
and maybe pin CPU workers to a given CPU core.
"""
if spec["device_kind"].lower() == "cpu":
if hasattr(os, 'sched_setaffinity') and pin:
os.sched_setaffinity(0, [spec["device_id"]])
set_use_cpu(spec["device_id"])
elif spec["device_kind"].lower() == "cuda":
set_use_cuda(spec["device_id"])
else:
raise RuntimeError(f"unknown device kind: {spec['device_kind']}")
def pipelined_worker(
queues: WorkerQueues,
pin: bool,
spec: WorkerSpec,
span_context: "SpanContext",
early_setup: Optional[Callable] = None,
):
"""
Main pipelined worker function.
Parameters
----------
queues
request and response queues
pin
Whether or not the CPU worker should be pinned to a specific CPU
spec
This worker's spec, containing name, worker index, device kind etc.
span_context
The tracing span we should attach to for the setup code
early_setup
Function that will be called very early in the setup code,
allowing to inject custom functionality or specific warmup code
"""
# FIXME: propagate to parent process with a pipe or similar?
sys.stderr.close()
sys.stderr = open(os.open(os.devnull, os.O_RDWR), closefd=False)
try:
if early_setup:
early_setup()
set_file_limit()
# FIXME: explicitly propagate exporter settings to the workers?
# right now taken care of by environment variables...
worker_idx = spec["worker_idx"]
maybe_setup_tracing(service_name="pipelined_worker", service_id=f"worker-{worker_idx}")
# attach to parent span context for startup:
with attach_to_parent(span_context), \
tracer.start_as_current_span("pipelined_worker.startup") as span:
_setup_device(spec, pin)
work_mem: dict[str, Any] = {}
span.set_attributes({
"libertem.spec.name": spec["name"],
"libertem.spec.device_id": spec["device_id"],
"libertem.spec.devide_kind": spec["device_kind"],
"libertem.spec.worker_idx": spec["worker_idx"],
})
set_thread_name(f"worker-{worker_idx}")
worker_context = PipelinedWorkerContext(queues.request, queues.message)
env = Environment(
threaded_executor=False,
threads_per_worker=1,
worker_context=worker_context
)
queues.response.put({
"type": "STARTUP_DONE",
"worker_id": worker_idx,
})
return worker_loop(queues, work_mem, worker_idx, env)
except Exception as e:
queues.response.put({
"type": "ERROR",
"worker_id": worker_idx,
"error": e,
"exception": e,
})
# drain, close and join queues:
queues.request.close()
queues.response.close(drain=False)
queues.message.close(drain=False)
sys.exit(1)
class PipelinedWorkerContext(WorkerContext):
"""
A context object that is made available to the Partition
for custom communication to its matching DataSet class
(currently uni-directional)
"""
def __init__(self, queue: "WorkerQueue", msg_queue: "WorkerQueue"):
self._queue = queue
self._msg_queue = msg_queue
def get_worker_queue(self) -> WorkerQueue:
return self._queue
def signal(self, ident: str, topic: str, msg_dict: dict[str, Any]):
msg_dict.update({'ident': ident})
self._msg_queue.put((topic, msg_dict))
ResultT = Generator[tuple[Any, TaskProtocol], None, None]
ResultWithID = Generator[tuple[Any, TaskProtocol, int], None, None]
def _order_results(results_in: ResultWithID) -> ResultT:
"""
Order the `results_in` generator by the result id, yielding ordered results.
Requires indexes to be without gaps.
"""
last_sent_id = -1
# for results that are received out-of-order, keep sorted:
result_stack: list[tuple[Any, TaskProtocol, int]] = []
span = trace.get_current_span()
for result, task, task_id in results_in:
if task_id == last_sent_id + 1:
span.add_event("_order_results.yield", {
'task_id': task_id,
})
yield (result, task)
last_sent_id = task_id
else:
span.add_event("_order_results.postpone", {
"expect": last_sent_id + 1,
"is": task_id,
})
result_stack = sorted(
result_stack + [(result, task, task_id)],
key=lambda x: x[-1]
)
# top of the stack looks good, yield as long as it matches:
while len(result_stack) > 0 and result_stack[0][2] == last_sent_id + 1:
res = result_stack.pop(0)
span.add_event("_order_results.yield", {
'task_id': res[2],
})
yield res[0], res[1]
last_sent_id = res[2]
for result, task, task_id in result_stack:
if task_id != last_sent_id + 1:
raise RuntimeError(
f"missing tasks? end of result but next id on result_stack is "
f"{task_id}, was expecting {last_sent_id + 1}"
)
span.add_event("_order_results.yield")
yield (result, task)
last_sent_id = task_id
def _make_spec(
cpus: Union[int, Iterable[int]],
cudas: Union[int, Iterable[int]],
has_cupy: bool = False,
) -> list[WorkerSpec]:
"""
Takes the output of :func:`libertem.utils.devices.detect`
and makes a plan for starting workers on them.
Parameters
----------
cpus: int | Iterable
Iterable of integer CPU identifiers or an integer number of workers to create.
If pinning is enabled, each worker processe is pinned to one of these identifiers,
as accepted by :func:`python:os.sched_setaffinity`. Pinning is currently only
supported on platforms that implement :func:`python:os.sched_setaffinity`.
cudas: int | Iterable
Interable of CUDA device identifiers for which workers should be started or
an integer number of GPU workers to create across the available devices.
Identifiers can be repeated to start multiple workers per GPU, which can
result in better device utilization.
has_cupy
Is cupy available?
"""
spec = []
worker_idx = 0
if isinstance(cpus, int):
cpus = tuple(range(cpus))
cudas = assign_cudas(cudas)
grouped_cudas = itertools.groupby(cudas, lambda x: x)
for device_id, group in grouped_cudas:
for i in range(len(list(group))):
spec.append(WorkerSpec(
name=f"cuda-{device_id}-{i}",
device_id=device_id,
device_kind="CUDA",
worker_idx=worker_idx,
has_cupy=has_cupy,
))
worker_idx += 1
for device_id in cpus:
spec.append(WorkerSpec(
name=f"cpu-{device_id}",
device_id=device_id,
device_kind="CPU",
worker_idx=worker_idx,
has_cupy=False,
))
worker_idx += 1
return spec
def _raise_from_msg(msg: dict, err_prefix: str):
if "exception" in msg:
raise msg["exception"]
else:
raise RuntimeError(f"{err_prefix}: {msg['error']}")
def _inspect_startup(msg, span):
if msg["type"] == "ERROR":
_raise_from_msg(msg, "error on startup")
if msg["type"] != "STARTUP_DONE":
raise RuntimeError(
f"unknown message type {msg['type']}, expected STARTUP_DONE"
)
span.add_event("worker startup done", {"worker_id": msg["worker_id"]})
class PipelinedExecutor(BaseJobExecutor):
"""
Multi-process pipelined executor. Useful for live processing using
`LiberTEM-live <https://libertem.github.io/LiberTEM-live/>`_
if your processing function is not able to keep up with the incoming data
stream in a single process, but also works for offline processing.
Parameters
----------
spec
Specification for the worker processes - can be generated
by :meth:`make_spec`.
pin_workers
Pin each CPU worker to a specific CPU, as defined by :func:`python:os.sched_setaffinity`.
Only works on OSes that implement :func:`python:os.sched_setaffinity`.
startup_timeout
Startup of the executor is cancelled if it doesn't finish within
this limit (in detail: each worker's startup time is limited by this timeout).
In seconds.
cleanup_timeout
When cleaning up using :meth:`close`, give up after
this limit. In seconds.
early_setup
Callable that will be run as early as possible on each worker process.
Useful for custom warmup code or testing.
Note
----
This executor is not thread-safe - concurrent calls into :meth:`run_tasks` or
:meth:`run_function` are not supported.
"""
def __init__(
self,
spec: Optional[list[WorkerSpec]] = None,
pin_workers: bool = True,
startup_timeout: float = 30.0,
cleanup_timeout: float = 10.0,
early_setup: Optional[Callable] = None,
) -> None:
self._pin_workers = pin_workers
if spec is None:
spec = self._default_spec()
self._spec = spec
self._closed = True
self._early_setup = early_setup
# for testing via monkeypatching, for example:
self._worker_selector: Optional[WorkerSelector] = None
# timeout for cleanup, either from exception or when joining processes
self._cleanup_timeout = cleanup_timeout
# timeout for starting a single worker
self._startup_timeout = startup_timeout
# keep this at the bottom:
self._pool = self._start_pool()
def _start_pool(self) -> WorkerPool:
with tracer.start_as_current_span("PipelinedExecutor.start_pool") as span:
pool = WorkerPool(
worker_fn=functools.partial(
pipelined_worker,
pin=self._pin_workers,
early_setup=self._early_setup,
),
spec=self._spec,
)
# if any processes are already dead here, raise an exception so we
# don't have to run into a timeout below:
pool.assert_all_alive()
warn_time = 5.0
check_interval = 0.1
warn_deadline = time.monotonic() + warn_time
startup_deadline = time.monotonic() + self._startup_timeout
warned = False
num_started = 0
while num_started < pool.size:
pool.assert_all_alive()
try:
with pool.response_queue.get(timeout=check_interval) as (msg, _):
_inspect_startup(msg, span)
num_started += 1
continue
except WorkerQueueEmpty:
if time.monotonic() > startup_deadline:
pool.assert_all_alive()
pool.kill()
# break possibly confusing exception chain using "from None":
raise RuntimeError(
f"Timeout while starting workers, might need to increase "
f"`startup_timeout` (is {self._startup_timeout}s)"
) from None
if time.monotonic() > warn_deadline and not warned:
warnings.warn('Slow worker startup, please be patient...', RuntimeWarning)
warned = True
for qp in pool.workers:
qp.queues.request.put({
"type": "WARMUP",
"span_context": span.get_span_context(),
})
# set here, so we don't try to close the pool if it doesn't exist
self._closed = False
return pool
def _restart_pool(self):
self._pool.kill()
self._pool = self._start_pool()
@classmethod
def _default_spec(cls):
from libertem.utils.devices import detect
detected = detect()
return _make_spec(**detected)
@classmethod
def make_local(cls, **kwargs):
"""
Create a :code:`PipelinedExecutor` with the default spec.
"""
spec = cls._default_spec()
return cls(spec=spec, **kwargs)
@classmethod
def make_spec(
cls,
cpus: Union[int, Iterable[int]],
cudas: Union[int, Iterable[int]],
has_cupy: bool = False,
):
return _make_spec(cpus=cpus, cudas=cudas, has_cupy=has_cupy)
make_spec.__doc__ = _make_spec.__doc__
def _validate_worker_state(self):
if not self._pool.all_alive():
raise RuntimeError("some workers are stopped, cannot continue")
def _run_tasks_inner(
self,
tasks: Iterable[TaskProtocol],
params_handle: Any,
cancel_id: Any,
task_comm_handler: "TaskCommHandler",
) -> ResultWithID:
# In theory, `in_flight` could be calculated from `id_to_task`, but in case
# of exceptions, it becomes a bit harder to keep attribution of messages to
# tasks, which is why we have a separate counter for now.
in_flight = 0
id_to_task: dict[int, TaskProtocol] = {}
tasks_uuid = str(uuid.uuid4())
all_workers = self.get_available_workers()
scheduler = Scheduler(all_workers=all_workers)
pool_info_for_worker = dict(zip(all_workers, self._pool.workers))
selector: WorkerSelector
# feature detect: can we access the queue size? support for this can be
# missing depending on platform
if self._worker_selector is None:
try:
self._pool.workers[0].queues.request.size()
selector = _select_by_queue_size
except NotImplementedError:
# fall back to simple round robin worker selection
selector = _select_by_round_robin
else:
selector = self._worker_selector
try:
self._validate_worker_state()
task_comm_handler.start()
span = trace.get_current_span()
span_context = span.get_span_context()
def yield_result_if_found(block, timeout):
nonlocal in_flight
try:
with self._pool.response_queue.get(block=block, timeout=timeout) as (result, _):
if result.get('uuid') != tasks_uuid:
# mismatch, log and ignore:
logger.warning(
"mismatched result, ignoring: %s != %s",
result.get('uuid'), tasks_uuid,
)
return
in_flight -= 1
if result["type"] == "ERROR":
_raise_from_msg(result, "failed to run tasks")
result_task_id = result["task_id"]
yield (result["result"], id_to_task[result_task_id], result_task_id)
del id_to_task[result_task_id]
if len(id_to_task) != in_flight:
raise RuntimeError(
"state mismatch; `id_to_task` mapping should match `in_flight`"
)
except WorkerQueueEmpty:
self._validate_worker_state()
with tracer.start_as_current_span('schedule_tasks') as span:
for task_idx, task in enumerate(tasks):
# important: call `schedule_task` before incrementing number of
# in-flight tasks, otherwise, if we run into an exception, we
# will wait for a response that never comes.
# select a worker for `task` according to `selector`:
_worker_idx, worker_queues = schedule_task(
task_idx,
task,
self._pool,
scheduler,
pool_info_for_worker,
selector=selector,
)
worker_queues.request.put({
"type": "RUN_TASK",
"uuid": tasks_uuid,
"task": cloudpickle.dumps(task),
"task_id": task_idx,
"params_handle": params_handle,
"span_context": span_context,
})
span.add_event('RUN_TASK', {
'task_id': task_idx,
})
in_flight += 1
id_to_task[task_idx] = task
if len(id_to_task) != in_flight:
raise RuntimeError(
"state mismatch; `id_to_task` mapping should match `in_flight`"
)
# FIXME: semantics of this - is this enough?
# does it matter if this is enough? we can change it in the future if not
# could be: the function returns once it has forwarded
# all the data necessary for the given task,
# (or, in the offline case, immediately)
try:
task_comm_handler.handle_task(task, worker_queues.request)
except JobCancelledError:
worker_queues.request.put({
"type": "END_TASK",
"task_id": task_idx,
"params_handle": params_handle,
"span_context": span_context,
})
raise
# NOTE: sentinel message; in case of errors, the worker
# needs to discard the data from the queue until it receives
# this message:
worker_queues.request.put({
"type": "END_TASK",
"task_id": task_idx,
"params_handle": params_handle,
"span_context": span_context,
})
yield from yield_result_if_found(block=False, timeout=None)
# at the end, block to get the remaining results:
while in_flight > 0:
yield from yield_result_if_found(block=True, timeout=0.1)
except JobCancelledError:
# don't drain here, as the next acquisition could start
# very soon; instead, we just ignore the mismatched
# responses in the next run
raise
except Exception as e:
# In case of an exception, we need to drain the response queue,
# so the next `run_tasks` call isn't polluted by old responses.
# -> just like in the happy case, we need to wait for responses
# for all in-flight tasks. In case the error happened between incrementing
# `in_flight` and actually sending the task to the queue, we should
# have a timeout here to not wait infinitely long.
try:
self._drain_response_queue(in_flight=in_flight)
except RuntimeError as e2:
raise e2 from e
# if from a worker, this is the first exception that got put into the queue
raise
finally:
task_comm_handler.done()
def _drain_response_queue(self, in_flight: int) -> None:
"""
Drain response queue and log any errors returned from workers
Parameters
----------
in_flight : int
The number of requests that are still in flight
"""
t0 = time.time()
while in_flight > 0:
try:
timeout = max(0.010, self._cleanup_timeout - (time.time() - t0))
with self._pool.response_queue.get(
timeout=timeout,
) as (result, _):
t0 = time.time()
in_flight -= 1
# we only raise the first exception; log the others here:
if result["type"] == "ERROR":
logger.error(f"Error response from worker: {result['error']}")
except WorkerQueueEmpty:
# kill and restart workers:
self._restart_pool()
raise RuntimeError(
f'Worker or queue presumably in a bad state, lost {in_flight} in-flight tasks.'
)
def run_tasks(
self,
tasks: Iterable[TaskProtocol],
params_handle: Any,
cancel_id: Any,
task_comm_handler: "TaskCommHandler",
) -> ResultT:
with tracer.start_as_current_span("PipelinedExecutor.run_tasks"):
with task_comm_handler.monitor(self._pool.message_queue):
yield from _order_results(self._run_tasks_inner(
tasks, params_handle, cancel_id, task_comm_handler,
))
def get_available_workers(self) -> WorkerSet:
return WorkerSet([
Worker(
name=worker_info.spec["name"],
host="localhost",
resources=_resources_for_spec(worker_info.spec),
nthreads=1,
)
for worker_info in self._pool.workers
])
def close(self):
with tracer.start_as_current_span("PipelinedExecutor.close") as span:
if self._closed:
return
for idx, worker_info in enumerate(self._pool.workers):
span.add_event("sending SHUTDOWN", {"idx": idx})
worker_info.queues.request.put({
"type": "SHUTDOWN",
"span_context": span.get_span_context(),
})
span.add_event("SHUTDOWN sent", {"idx": idx})
while True:
try:
with worker_info.queues.response.get(block=False) as msg:
logger.warning(f"got message on close: {msg[0]}")
except WorkerQueueEmpty:
break
worker_info.process.join(timeout=self._cleanup_timeout)
if worker_info.process.exitcode is None:
self._pool.kill_worker(worker_info)
worker_info.queues.request.close(force=True)
self._pool.close_resp_queue()
self._pool.close_mesg_queue()
self._closed = True
def __del__(self):
# `self` may be already partially garbage-collected; only close
# if "enough" of `self` still exists:
if hasattr(self, '_closed') and not self._closed:
self.close()
def _run_function(self, fn: Callable[..., T], worker_idx, *args, **kwargs) -> T:
self._validate_worker_state()
qs = self._pool.get_worker_queues(worker_idx)
f = functools.partial(fn, *args, **kwargs)
pickled = cloudpickle.dumps(f)
qs.request.put({
"type": "RUN_FUNCTION",
"fn": pickled,
"span_context": trace.get_current_span().get_span_context(),
})
# FIXME: timeout?
with qs.response.get() as (response, _):
if response["type"] == "ERROR":
if "exception" in response:
raise response["exception"]
else:
raise RuntimeError(f"failed to run function: {response['error']}")
if not response["type"] == "RUN_FUNCTION_RESULT":
raise RuntimeError(f"invalid response type: {response['type']}")
result: T = response["result"]
return result
def run_function(self, fn: Callable[..., T], *args, **kwargs) -> T:
# FIXME: this is not concurrency-safe currently! beware!
with tracer.start_as_current_span("PipelinedExecutor.run_function"):
return self._run_function(fn, 0, *args, **kwargs)
def run_each_worker(self, fn, *args, **kwargs):
# FIXME: not as fast as it could be, but also not really perf-sensitive?
result = {}
for idx, worker_info in enumerate(self._pool.workers):
result[worker_info.spec["name"]] = self._run_function(
fn, worker_idx=idx, *args, **kwargs
)
return result
def run_each_host(self, fn, *args, **kwargs):
return {"localhost": self.run_function(fn, *args, **kwargs)}
@contextlib.contextmanager
def scatter(self, obj):
self._validate_worker_state()
key = str(uuid.uuid4())
with tracer.start_as_current_span("scatter_put_to_workers"):
for worker_info in self._pool.workers:
worker_info.queues.request.put({
"type": "SCATTER",
"key": key,
"value": obj,
})
try:
yield key
finally:
if not self._closed:
for worker_info in self._pool.workers:
worker_info.queues.request.put({
"type": "DELETE",
"key": key,
})
def scatter_update(self, handle: str, obj):
self._validate_worker_state()
for worker_info in self._pool.workers:
worker_info.queues.request.put({
"type": "SCATTER_UPDATE",
"key": handle,
"value": obj,
})
def scatter_update_patch(self, handle: str, patch):
self._validate_worker_state()
for worker_info in self._pool.workers:
worker_info.queues.request.put({
"type": "SCATTER_UPDATE_PATCH",
"key": handle,
"patch": patch,
})
def map(self, fn, iterable):
# FIXME: replace with efficient impl if needed
with tracer.start_as_current_span("PipelinedExecutor.map"):
return [
self._run_function(fn=lambda: fn(item), worker_idx=0)
for item in iterable
]