wehub-resource-sync
107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
from abc import ABC, abstractmethod
|
|
from typing import ContextManager, Iterator, List, Optional
|
|
|
|
from .execution_options import ExecutionOptions
|
|
from .physical_operator import PhysicalOperator
|
|
from .ref_bundle import RefBundle
|
|
from ray.data._internal.stats import DatasetStats
|
|
|
|
|
|
class OutputIterator(Iterator[RefBundle], ABC):
|
|
"""Iterator used to access the output of an Executor execution.
|
|
|
|
This is a blocking iterator. Datasets guarantees that all its iterators are
|
|
thread-safe (i.e., multiple threads can block on them at the same time).
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
|
|
"""Can be used to pull outputs by a specified output index.
|
|
|
|
This is used to support the streaming_split() API, where the output of a
|
|
streaming execution is to be consumed by multiple processes.
|
|
|
|
Args:
|
|
output_split_idx: The output split index to get results for. This arg is
|
|
only allowed for iterators created by `Dataset.streaming_split()`.
|
|
|
|
Returns:
|
|
The next ``RefBundle`` of outputs for the given split index.
|
|
|
|
Raises:
|
|
StopIteration: If there are no more outputs to return.
|
|
"""
|
|
...
|
|
|
|
def __next__(self) -> RefBundle:
|
|
return self.get_next()
|
|
|
|
|
|
class Executor(ContextManager, ABC):
|
|
"""Abstract class for executors, which implement physical operator execution.
|
|
|
|
Subclasses:
|
|
StreamingExecutor
|
|
"""
|
|
|
|
def __init__(self, options: ExecutionOptions):
|
|
"""Create the executor."""
|
|
options.validate()
|
|
self._options = options
|
|
|
|
@abstractmethod
|
|
def execute(
|
|
self,
|
|
dag: PhysicalOperator,
|
|
initial_stats: Optional[DatasetStats] = None,
|
|
callbacks: Optional[List] = None,
|
|
) -> OutputIterator:
|
|
"""Start execution.
|
|
|
|
Args:
|
|
dag: The operator graph to execute.
|
|
initial_stats: The DatasetStats to prepend to the stats returned by the
|
|
executor. These stats represent actions done to compute inputs.
|
|
callbacks: A list of ExecutionCallbacks to run during execution.
|
|
This method keeps and uses the exact list you pass in, so do not
|
|
pass an empty list like ``[]`` directly. Create the list first,
|
|
then pass it.
|
|
|
|
Returns:
|
|
An ``OutputIterator`` over the execution's output ref bundles.
|
|
"""
|
|
...
|
|
|
|
def shutdown(self, force: bool, exception: Optional[Exception] = None):
|
|
"""Shutdown an executor, which may still be running.
|
|
|
|
This should interrupt execution and clean up any used resources.
|
|
|
|
Args:
|
|
force: Controls whether shutdown should forcefully terminate all execution
|
|
activity (making sure that upon returning from this method all
|
|
activities are stopped). When force=False, some activities could be
|
|
terminated asynchronously (ie this method won't provide guarantee
|
|
that they stop executing before returning from this method)
|
|
exception: The exception that causes the executor to shut down, or None if
|
|
the executor finishes successfully.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def get_stats(self) -> DatasetStats:
|
|
"""Return stats for the execution so far.
|
|
|
|
This is generally called after `execute` has completed, but may be called
|
|
while iterating over `execute` results for streaming execution.
|
|
"""
|
|
...
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback, /):
|
|
# NOTE: ``ContextManager`` semantic must guarantee that executor
|
|
# fully shutdown upon returning from this method
|
|
self.shutdown(force=True, exception=exc_value)
|