chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,258 @@
|
||||
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
|
||||
Reference in New Issue
Block a user