mars-project / mars-project/mars

[PROPOSAL] The Execution API

Open
#2,893 0 comments 2 reactions 1 assignee Claimed by @fyrestone View on GitHub
proposal
Dominant language
Python
Stars
2.7k
Forks
325
PR merge metrics
No merged PRs in 30d

Description

# The Execution API

## Background

Mars provides some interfaces for third party backends, but it is difficult to implement them and can not take the advantages of third party engine. For example, the Ray backend has to implement the Mars actor pool for Ray, the channel for Ray, the storage backend for Ray. However, the Mars on Ray can not use the great features of Ray: data-locality scheduling, data prefetching, distributed object management, ...

This proposal introduces an Execution API to split the graph construction and execution, it gives great freedom for the third party execution backends. The execution backend is able to choose the best way to execute each subtask graph and does not need to care about how the graph is constructed or optimized.

## Execution API

A classic execution logic of Mars is mixing tiling and execution, that is the iterative tiling. So, the execution API should handle this case.
![image](https://user-images.githubusercontent.com/6308809/161222445-31a7ae99-154a-4d23-98ae-d334cf561e1d.png)

Then, we can get the Execution API.

```python
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import List, Dict, Any, Type

from ..core import ChunkGraph
from ..typing import BandType
from ..services.task import Task
from ..services.subtask import SubtaskGraph

@dataclass
class ExecutionChunkResult:
key: str # The chunk key for fetching the result.
meta: Dict # The chunk meta for iterative tiling.
context: Any # The context info, e.g. ray.ObjectRef.

class TaskExecutor(ABC):
name = None

@classmethod
@abstractmethod
async def create(
cls, config: Dict, *, session_id: str, address: str, task: Task, **kwargs
) -> "TaskExecutor":
"""Create TaskExecutor instance from config.

An example of config:
{
"backend": "mars",
"mars": {...},
"ray": {...},
}
"""

async def __aenter__(self):
"""Called when begin to execute the task."""

@abstractmethod
async def execute_subtask_graph(
self,
stage_id: str,
subtask_graph: SubtaskGraph,
chunk_graph: ChunkGraph,
context: Any = None,
) -> List[ExecutionChunkResult]:
"""Execute a subtask graph and returns result."""

async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Called when finish the task."""

@abstractmethod
async def get_available_band_slots(self) -> Dict[BandType, int]:
"""Get available band slots."""

@abstractmethod
async def get_progress(self) -> float:
"""Get the execution progress."""

@abstractmethod
async def cancel(self):
"""Cancel execution."""
```

The execution backend can be a third party Python package, it may include some Mars services. For example, an execution backend on Ray looks like:

```shell
ray # An execution backend for ray.
├── __init__.py # [Optional] May defines a _init_extension for mars_extensions.
├── core.py # [Required] Implements above APIs.
└── services # [Optional] May contains some services.
└── __init__.py
```
We can define a `_init_extension()` in the `__init__.py` for the backend registration and initialization through the `mars_extensions` entrypoint.

Then, Mars task service only relies on the Execution API.
![image](https://user-images.githubusercontent.com/6308809/161222757-452db3b1-bc45-4b7f-94ba-d6cc8fbc4a46.png)

## Dev Plans

- Introduce the Execution API, the Mars execution backend follows the API. https://github.com/mars-project/mars/pull/2894
- Extract the subtask execution ~~logic to `mars.execution.core`~~ utils. https://github.com/mars-project/mars/pull/2921
- Add the `output_keys()` API for the shuffle operands for constructing the subtask DAG on other engines before execution.
- Make the fetch logic general for the Execution API. https://github.com/mars-project/mars/pull/2921
- Implement Mars on Ray DAG execution backend.
- Support execution backend registration and initialization through the `mars_extensions` entrypoint.
- Support task progress for ray executor backend. #3008
- Support task cancel for ray executor backend. #3044
- Support task context gc for ray executor backend. #3061

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.