def fake_remote_fct(
docker_auth: DockerBasicAuth,
service_key: str,
service_version: str,
input_data: TaskInputData,
output_data_keys: TaskOutputDataSchema,
log_file_url: AnyUrl,
command: List[str],
) -> TaskOutputData:
# get the task data
worker = get_worker()
task = worker.tasks.get(worker.get_current_task())
assert task is not None
print(f"--> task {task=} started")
cancel_event = Event(TaskCancelEventName.format(task.key))
# tell the client we are started
start_event = Event(_DASK_EVENT_NAME)
start_event.set()
# sleep a bit in case someone is aborting us
print("--> waiting for task to be aborted...")
cancel_event.wait(timeout=10)
if cancel_event.is_set():
# NOTE: asyncio.CancelledError is not propagated back to the client...
print("--> raising cancellation error now")
raise TaskCancelledError
return TaskOutputData.parse_obj({"some_output_key": 123})
def _maybe_run(event_name: str, fn: Callable, *args: Any,
**kwargs: Any) -> Any:
"""Check if the task should run against a `distributed.Event` before
starting the task. This offers stronger guarantees than distributed's
current cancellation mechanism, which only cancels pending tasks."""
import dask
from distributed import Event, get_client
try:
# Explicitly pass in the timeout from dask's config_dict. Some versions of
# distributed hardcode this rather than using the value from the
# config_dict. Can be removed once we bump our min requirements for
# distributed to >= 2.31.0.
timeout = dask.config.get("distributed.comm.timeouts.connect")
event = Event(event_name, client=get_client(timeout=timeout))
should_run = event.is_set()
except Exception:
# Failure to create an event is usually due to connection errors. These
# are either due to flaky behavior in distributed's comms under high
# loads, or due to the scheduler shutting down. Either way, the safest
# course here is to assume we *should* run the task still. If we guess
# wrong, we're either doing a bit of unnecessary work, or the cluster
# is shutting down and the task will be cancelled anyway.
should_run = True
if should_run:
return fn(*args, **kwargs)
def wait_for_it_ok(x):
event = Event("x")
# Event is set in another task
assert event.wait(timeout=0.5)
assert event.is_set()
def wait_for_it_failing(x):
event = Event("x")
# Event is not set in another task so far
assert not event.wait(timeout=0.05)
assert not event.is_set()
