wehub-resource-sync
259 lines
7.8 KiB
Python
259 lines
7.8 KiB
Python
import logging
|
|
import threading
|
|
import typing
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any, List, Optional
|
|
|
|
import ray
|
|
from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
|
|
from ray.data._internal.progress.utils import truncate_operator_name
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from ray.data._internal.execution.resource_manager import ResourceManager
|
|
from ray.data._internal.execution.streaming_executor_state import OpState, Topology
|
|
from ray.types import ObjectRef
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Used a signal to cancel execution.
|
|
_canceled_threads = set()
|
|
_canceled_threads_lock = threading.Lock()
|
|
|
|
|
|
def _extract_num_rows(result: Any) -> int:
|
|
"""Extract the number of rows from a result object.
|
|
|
|
Args:
|
|
result: The result object from which to extract the number of rows.
|
|
|
|
Returns:
|
|
The number of rows, defaulting to 1 if it cannot be determined.
|
|
"""
|
|
if hasattr(result, "num_rows"):
|
|
return result.num_rows
|
|
elif hasattr(result, "__len__"):
|
|
# For output is DataFrame,i.e. sort_sample
|
|
return len(result)
|
|
else:
|
|
return 1
|
|
|
|
|
|
class BaseProgressBar(ABC):
|
|
"""Base class to define a progress bar."""
|
|
|
|
def block_until_complete(self, remaining: List["ObjectRef"]) -> None:
|
|
t = threading.current_thread()
|
|
while remaining:
|
|
done, remaining = ray.wait(
|
|
remaining, num_returns=len(remaining), fetch_local=False, timeout=0.1
|
|
)
|
|
total_rows_processed = 0
|
|
for _, result in zip(done, ray.get(done)):
|
|
num_rows = _extract_num_rows(result)
|
|
total_rows_processed += num_rows
|
|
self.update(total_rows_processed)
|
|
|
|
with _canceled_threads_lock:
|
|
if t in _canceled_threads:
|
|
break
|
|
|
|
def fetch_until_complete(self, refs: List["ObjectRef"]) -> List[Any]:
|
|
ref_to_result = {}
|
|
remaining = refs
|
|
t = threading.current_thread()
|
|
# Triggering fetch_local redundantly for the same object is slower.
|
|
# We only need to trigger the fetch_local once for each object,
|
|
# raylet will persist these fetch requests even after ray.wait returns.
|
|
# See https://github.com/ray-project/ray/issues/30375.
|
|
fetch_local = True
|
|
while remaining:
|
|
done, remaining = ray.wait(
|
|
remaining,
|
|
num_returns=len(remaining),
|
|
fetch_local=fetch_local,
|
|
timeout=0.1,
|
|
)
|
|
if fetch_local:
|
|
fetch_local = False
|
|
total_rows_processed = 0
|
|
for ref, result in zip(done, ray.get(done)):
|
|
ref_to_result[ref] = result
|
|
num_rows = _extract_num_rows(result)
|
|
total_rows_processed += num_rows
|
|
self.update(total_rows_processed)
|
|
|
|
with _canceled_threads_lock:
|
|
if t in _canceled_threads:
|
|
break
|
|
|
|
return [ref_to_result[ref] for ref in refs]
|
|
|
|
@abstractmethod
|
|
def set_description(self, name: str) -> None:
|
|
...
|
|
|
|
@abstractmethod
|
|
def get_description(self) -> str:
|
|
...
|
|
|
|
@abstractmethod
|
|
def update(self, increment: int = 0, total: Optional[int] = None) -> None:
|
|
...
|
|
|
|
def refresh(self):
|
|
pass
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
|
|
class BaseExecutionProgressManager(ABC):
|
|
"""Base Data Execution Progress Display Manager"""
|
|
|
|
# If the name/description of the progress bar exceeds this length,
|
|
# it will be truncated.
|
|
MAX_NAME_LENGTH = 100
|
|
|
|
# Total progress refresh rate (update interval in scheduling step)
|
|
# refer to `streaming_executor.py::StreamingExecutor::_scheduling_loop_step`
|
|
TOTAL_PROGRESS_REFRESH_EVERY_N_STEPS = 50
|
|
|
|
@abstractmethod
|
|
def __init__(
|
|
self,
|
|
dataset_id: str,
|
|
topology: "Topology",
|
|
show_op_progress: bool,
|
|
verbose_progress: bool,
|
|
):
|
|
"""Initialize the progress manager, create all necessary progress bars
|
|
and sub-progress bars for the given topology. Sub-progress bars are
|
|
created for operators that implement the SubProgressBarMixin.
|
|
|
|
Args:
|
|
dataset_id: id of Dataset
|
|
topology: operation topology built via `build_streaming_topology`
|
|
show_op_progress: whether to show individual operator progress
|
|
(only for non-AllToAll by default).
|
|
verbose_progress: whether to show individual operator progress for
|
|
non-AllToAll operators as well.
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def start(self) -> None:
|
|
"""Start the progress manager."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def refresh(self) -> None:
|
|
"""Refresh displayed progress."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def close_with_finishing_description(self, desc: str, success: bool) -> None:
|
|
"""Close the progress manager with a finishing message.
|
|
|
|
Args:
|
|
desc: description to display
|
|
success: whether the dataset execution was successful
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def update_total_progress(self, new_rows: int, total_rows: Optional[int]) -> None:
|
|
"""Update the total progress rows.
|
|
|
|
Args:
|
|
new_rows: new rows processed by the streaming_executor
|
|
total_rows: total rows to be processed (if known)
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def update_total_resource_status(self, resource_status: str) -> None:
|
|
"""Update the total resource usage statistics.
|
|
|
|
Args:
|
|
resource_status: resource status information string.
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def update_operator_progress(
|
|
self, opstate: "OpState", resource_manager: "ResourceManager"
|
|
) -> None:
|
|
"""Update individual operator progress.
|
|
|
|
Args:
|
|
opstate: opstate of the operator.
|
|
resource_manager: the ResourceManager.
|
|
"""
|
|
...
|
|
|
|
|
|
class NoopSubProgressBar(BaseProgressBar):
|
|
"""Sub-Progress Bar for Noop (Disabled) Progress Manager"""
|
|
|
|
def __init__(self, name: str, max_name_length: int):
|
|
self._max_name_length = max_name_length
|
|
self._desc = truncate_operator_name(name, self._max_name_length)
|
|
|
|
def set_description(self, name: str) -> None:
|
|
self._desc = truncate_operator_name(name, self._max_name_length)
|
|
|
|
def get_description(self) -> str:
|
|
return self._desc
|
|
|
|
def update(self, increment: int = 0, total: Optional[int] = None) -> None:
|
|
pass
|
|
|
|
def refresh(self):
|
|
pass
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
|
|
class NoopExecutionProgressManager(BaseExecutionProgressManager):
|
|
"""Noop Data Execution Progress Display Manager (Progress Display Disabled)"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataset_id: str,
|
|
topology: "Topology",
|
|
show_op_progress: bool,
|
|
verbose_progress: bool,
|
|
):
|
|
for state in topology.values():
|
|
op = state.op
|
|
if not isinstance(op, SubProgressBarMixin):
|
|
continue
|
|
sub_pg_names = op.get_sub_progress_bar_names()
|
|
if sub_pg_names is not None:
|
|
for name in sub_pg_names:
|
|
pg = NoopSubProgressBar(
|
|
name=name, max_name_length=self.MAX_NAME_LENGTH
|
|
)
|
|
op.set_sub_progress_bar(name, pg)
|
|
|
|
def start(self) -> None:
|
|
pass
|
|
|
|
def refresh(self) -> None:
|
|
pass
|
|
|
|
def close_with_finishing_description(self, desc: str, success: bool) -> None:
|
|
pass
|
|
|
|
def update_total_progress(self, new_rows: int, total_rows: Optional[int]) -> None:
|
|
pass
|
|
|
|
def update_total_resource_status(self, resource_status: str) -> None:
|
|
pass
|
|
|
|
def update_operator_progress(
|
|
self, opstate: "OpState", resource_manager: "ResourceManager"
|
|
) -> None:
|
|
pass
|