radical.asyncflow.backends.execution.dask_parallel¶

def __await__(self):
    """Make DaskExecutionBackend awaitable like Dask Client."""
    return self._async_init().__await__()

def register_callback(self, callback: Callable) -> None:
    """Register a callback for task state changes.

    Args:
        callback: Function to be called when task states change. Should accept
            task and state parameters.
    """
    self._callback_func = callback

def get_task_states_map(self):
    """Retrieve a mapping of task IDs to their current states.

    Returns:
        StateMapper: Object containing the mapping of task states for this backend.
    """
    return StateMapper(backend=self)

async def cancel_task(self, uid: str) -> bool:
    """Cancel a task by its UID.

    Args:
        uid (str): The UID of the task to cancel.

    Returns:
        bool: True if the task was found and cancellation was attempted,
        False otherwise.
    """
    self._ensure_initialized()
    if uid in self.tasks:
        task = self.tasks[uid]
        future = task.get("future")
        if future:
            return await future.cancel()
    return False

async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None:
    """Submit async tasks to Dask cluster.

    Processes a list of tasks and submits them to the Dask cluster for execution.
    Filters out future objects from arguments and validates that all functions
    are async coroutine functions.

    Args:
        tasks: List of task dictionaries containing:
            - uid: Unique task identifier
            - function: Async callable to execute
            - args: Positional arguments
            - kwargs: Keyword arguments
            - executable: Optional executable path (not supported)
            - task_backend_specific_kwargs: Backend-specific parameters

    Note:
        Executable tasks are not supported and will result in FAILED state.
        Only async functions are supported - sync functions will result in
        FAILED state.
        Future objects are filtered out from arguments as they are not picklable.
    """
    self._ensure_initialized()

    for task in tasks:
        is_func_task = bool(task.get("function"))
        is_exec_task = bool(task.get("executable"))

        if is_exec_task:
            error_msg = "DaskExecutionBackend does not support executable tasks"
            task["stderr"] = ValueError(error_msg)
            self._callback_func(task, "FAILED")
            continue

        # Validate that function is async
        if is_func_task and not asyncio.iscoroutinefunction(task["function"]):
            error_msg = "DaskExecutionBackend only supports async functions"
            task["exception"] = ValueError(error_msg)
            self._callback_func(task, "FAILED")
            continue

        self.tasks[task["uid"]] = task

        # Filter out future objects as they are not picklable
        filtered_args = [
            arg for arg in task["args"] if not isinstance(arg, asyncio.Future)
        ]
        task["args"] = tuple(filtered_args)
        try:
            await self._submit_async_function(task)
        except Exception as e:
            task["exception"] = e
            self._callback_func(task, "FAILED")

async def cancel_all_tasks(self) -> int:
    """Cancel all currently running/pending tasks.

    Returns:
        Number of tasks that were successfully cancelled
    """
    self._ensure_initialized()
    cancelled_count = 0
    task_uids = list(self.tasks.keys())

    for task_uid in task_uids:
        if await self.cancel_task(task_uid):
            cancelled_count += 1

    return cancelled_count

def link_explicit_data_deps(
    self, src_task=None, dst_task=None, file_name=None, file_path=None
):
    """Handle explicit data dependencies between tasks.

    Args:
        src_task: The source task that produces the dependency.
        dst_task: The destination task that depends on the source.
        file_name: Name of the file that represents the dependency.
        file_path: Full path to the file that represents the dependency.
    """
    pass

def link_implicit_data_deps(self, src_task, dst_task):
    """Handle implicit data dependencies for a task.

    Args:
        src_task: The source task that produces data.
        dst_task: The destination task that depends on the source task's output.
    """
    pass

async def state(self) -> str:
    """Get the current state of the Dask execution backend.

    Returns:
        Current state of the backend as a string.
    """
    if not self._initialized or self._client is None:
        return "DISCONNECTED"

    try:
        # Check if client is still connected
        await self._client.scheduler_info()
        return "CONNECTED"
    except Exception:
        return "DISCONNECTED"

async def task_state_cb(self, task: dict, state: str) -> None:
    """Callback function invoked when a task's state changes.

    Args:
        task: Dictionary containing task information and metadata.
        state: The new state of the task.
    """
    pass

async def build_task(self, task: dict) -> None:
    """Build or prepare a task for execution.

    Args:
        task: Dictionary containing task definition, parameters, and metadata
            required for task construction.
    """
    pass

async def shutdown(self) -> None:
    """Shutdown the Dask client and clean up resources.

    Closes the Dask client connection, clears task storage, and handles
    any cleanup exceptions gracefully.
    """
    if self._client is not None:
        try:
            # Cancel all running tasks first
            await self.cancel_all_tasks()

            # Close the client
            await self._client.close()
            logger.info("Dask client shutdown complete")
        except Exception as e:
            logger.exception(f"Error during shutdown: {str(e)}")
        finally:
            self._client = None
            self.tasks.clear()
            self._initialized = False
            logger.info("Dask execution backend shutdown complete")

async def __aenter__(self):
    """Async context manager entry."""
    if not self._initialized:
        await self._async_init()
    return self

async def __aexit__(self, exc_type, exc_val, exc_tb):
    """Async context manager exit."""
    await self.shutdown()

@classmethod
async def create(cls, resources: Optional[dict] = None):
    """Alternative factory method for creating initialized backend.

    Args:
        resources: Configuration parameters for Dask client initialization.

    Returns:
        Fully initialized DaskExecutionBackend instance.
    """
    backend = cls(resources)
    return await backend