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()