chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,796 @@
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import typing
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from ray.data._internal.actor_autoscaler import (
|
||||
create_actor_autoscaler,
|
||||
)
|
||||
from ray.data._internal.cluster_autoscaler import create_cluster_autoscaler
|
||||
from ray.data._internal.execution import create_ranker
|
||||
from ray.data._internal.execution.backpressure_policy import (
|
||||
BackpressurePolicy,
|
||||
get_backpressure_policies,
|
||||
)
|
||||
from ray.data._internal.execution.block_ref_counter import BlockRefCounter
|
||||
from ray.data._internal.execution.dataset_state import DatasetState
|
||||
from ray.data._internal.execution.execution_callback import ExecutionCallback
|
||||
from ray.data._internal.execution.interfaces import (
|
||||
Executor,
|
||||
OutputIterator,
|
||||
PhysicalOperator,
|
||||
RefBundle,
|
||||
)
|
||||
from ray.data._internal.execution.metadata_fetcher import make_metadata_fetcher
|
||||
from ray.data._internal.execution.operators.base_physical_operator import (
|
||||
InternalQueueOperatorMixin,
|
||||
)
|
||||
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
|
||||
from ray.data._internal.execution.resource_manager import (
|
||||
ResourceManager,
|
||||
)
|
||||
from ray.data._internal.execution.streaming_executor_state import (
|
||||
OpState,
|
||||
OutputBackpressureGuard,
|
||||
Topology,
|
||||
build_streaming_topology,
|
||||
format_op_state_summary,
|
||||
process_completed_tasks,
|
||||
select_operator_to_run,
|
||||
update_operator_states,
|
||||
)
|
||||
from ray.data._internal.logging import (
|
||||
get_log_directory,
|
||||
register_dataset_logger,
|
||||
unregister_dataset_logger,
|
||||
)
|
||||
from ray.data._internal.metadata_exporter import (
|
||||
Topology as TopologyMetadata,
|
||||
sanitize_for_struct,
|
||||
)
|
||||
from ray.data._internal.operator_schema_exporter import (
|
||||
OperatorSchema,
|
||||
get_operator_schema_exporter,
|
||||
)
|
||||
from ray.data._internal.progress import get_progress_manager
|
||||
from ray.data._internal.stats import DatasetStats, Timer, _StatsManager
|
||||
from ray.data.context import OK_PREFIX, WARN_PREFIX, DataContext
|
||||
from ray.util.debug import log_once
|
||||
from ray.util.metrics import Gauge
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from ray.data._internal.issue_detection.issue_detector_manager import (
|
||||
IssueDetectorManager,
|
||||
)
|
||||
from ray.data._internal.progress.base_progress import BaseExecutionProgressManager
|
||||
from ray.data.block import Schema
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Interval for logging execution progress updates and operator metrics.
|
||||
DEBUG_LOG_INTERVAL_SECONDS = 5
|
||||
|
||||
# Maximum string/sequence length for DataContext logging. Set high to avoid truncation
|
||||
# while still protecting against pathological cases.
|
||||
DATA_CONTEXT_LOG_TRUNCATE_LENGTH = 10000
|
||||
|
||||
# Visible for testing.
|
||||
_num_shutdown = 0
|
||||
|
||||
|
||||
# Extra environment variables to log that don't start with RAY_DATA.
|
||||
_EXTRA_ENV_VARS_TO_LOG = (
|
||||
# We historically recommended users configure this value. If a Ray Data job uses
|
||||
# more object store memory than expected, it's worth checking how this environment
|
||||
# variable has been configured.
|
||||
"RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION",
|
||||
)
|
||||
|
||||
|
||||
def _log_ray_data_env_vars() -> None:
|
||||
env_vars = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k.startswith("RAY_DATA") or k in _EXTRA_ENV_VARS_TO_LOG
|
||||
}
|
||||
if env_vars:
|
||||
formatted = ", ".join(f"{k}={v}" for k, v in sorted(env_vars.items()))
|
||||
logger.debug(f"RAY_DATA environment variables: {formatted}")
|
||||
else:
|
||||
logger.debug("No RAY_DATA environment variables set.")
|
||||
|
||||
|
||||
class StreamingExecutor(Executor, threading.Thread):
|
||||
"""A streaming Dataset executor.
|
||||
|
||||
This implementation executes Dataset DAGs in a fully streamed way. It runs
|
||||
by setting up the operator topology, and then routing blocks through operators in
|
||||
a way that maximizes throughput under resource constraints.
|
||||
"""
|
||||
|
||||
UPDATE_METRICS_INTERVAL_S: float = 5.0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_context: DataContext,
|
||||
dataset_id: str = "unknown_dataset",
|
||||
):
|
||||
self._data_context = data_context
|
||||
self._ranker = create_ranker()
|
||||
self._start_time: Optional[float] = None
|
||||
self._initial_stats: Optional[DatasetStats] = None
|
||||
self._final_stats: Optional[DatasetStats] = None
|
||||
self._progress_manager: Optional["BaseExecutionProgressManager"] = None
|
||||
self._callbacks: List["ExecutionCallback"] = []
|
||||
|
||||
# The executor can be shutdown while still running.
|
||||
self._shutdown_lock = threading.RLock()
|
||||
self._execution_started = False
|
||||
self._shutdown = False
|
||||
|
||||
# Internal execution state shared across thread boundaries. We run the control
|
||||
# loop on a separate thread so that it doesn't become stalled between
|
||||
# generator `yield`s.
|
||||
self._topology: Optional[Topology] = None
|
||||
self._output_node: Optional[Tuple[PhysicalOperator, OpState]] = None
|
||||
self._backpressure_policies: List[BackpressurePolicy] = []
|
||||
self._op_schema: Dict[PhysicalOperator, Schema] = {}
|
||||
|
||||
self._dataset_id = dataset_id
|
||||
# Set by IssueDetectionExecutionCallback when issue detection is registered;
|
||||
# otherwise remains None. Access via the issue_detector_manager property.
|
||||
self._issue_detector_manager: Optional["IssueDetectorManager"] = None
|
||||
# Stores if an operator is completed,
|
||||
# used for marking when an op has just completed.
|
||||
self._has_op_completed: Optional[Dict[PhysicalOperator, bool]] = None
|
||||
self._max_errored_blocks = self._data_context.max_errored_blocks
|
||||
self._num_errored_blocks = 0
|
||||
|
||||
self._last_debug_log_time = 0
|
||||
|
||||
self._data_context.set_dataset_logger_id(
|
||||
register_dataset_logger(self._dataset_id)
|
||||
)
|
||||
|
||||
# This stores the last time we updated the metrics.
|
||||
# This allows us to update metrics on some interval,
|
||||
# by comparing it with the current timestamp.
|
||||
self._metrics_last_updated: float = 0.0
|
||||
|
||||
self._sched_loop_duration_s = Gauge(
|
||||
"data_sched_loop_duration_s",
|
||||
description="Duration of the scheduling loop in seconds",
|
||||
tag_keys=("dataset",),
|
||||
)
|
||||
|
||||
# Resolves pulled (block_ref, meta_ref) pairs into emitted RefBundles.
|
||||
# The threaded fetcher (default) fetches metadata on a background thread
|
||||
# so the scheduling loop never blocks on ``ray.get(meta_refs)``; the
|
||||
# inline fetcher reproduces the synchronous, master-identical path.
|
||||
self._metadata_fetcher = make_metadata_fetcher()
|
||||
|
||||
Executor.__init__(self, self._data_context.execution_options)
|
||||
thread_name = f"StreamingExecutor-{self._dataset_id}"
|
||||
threading.Thread.__init__(self, daemon=True, name=thread_name)
|
||||
|
||||
@property
|
||||
def issue_detector_manager(self) -> Optional["IssueDetectorManager"]:
|
||||
"""The issue detector manager, or None if issue detection isn't registered."""
|
||||
return self._issue_detector_manager
|
||||
|
||||
def execute(
|
||||
self,
|
||||
dag: PhysicalOperator,
|
||||
initial_stats: Optional[DatasetStats] = None,
|
||||
callbacks: Optional[List] = None,
|
||||
) -> OutputIterator:
|
||||
"""Executes the DAG using a streaming execution strategy.
|
||||
|
||||
We take an event-loop approach to scheduling. We block on the next scheduling
|
||||
event using `ray.wait`, updating operator state and dispatching new tasks.
|
||||
"""
|
||||
|
||||
if callbacks is not None:
|
||||
self._callbacks = callbacks
|
||||
else:
|
||||
self._callbacks = []
|
||||
|
||||
self._initial_stats = initial_stats
|
||||
self._start_time = time.perf_counter()
|
||||
|
||||
if logger.isEnabledFor(logging.DEBUG):
|
||||
_log_ray_data_env_vars()
|
||||
|
||||
if not isinstance(dag, InputDataBuffer):
|
||||
if self._data_context.print_on_execution_start:
|
||||
message = f"Starting execution of Dataset {self._dataset_id}."
|
||||
log_path = get_log_directory()
|
||||
if log_path is not None:
|
||||
message += f" Full logs are in {log_path}"
|
||||
logger.info(message)
|
||||
logger.info(
|
||||
f"Execution plan of Dataset {self._dataset_id}: {dag.dag_str}"
|
||||
)
|
||||
|
||||
# Log the full DataContext for traceability
|
||||
if logger.isEnabledFor(logging.DEBUG) and log_once(
|
||||
f"ray_data_log_context_{self._dataset_id}"
|
||||
):
|
||||
logger.debug(
|
||||
f"Data Context for dataset {self._dataset_id}:\n%s",
|
||||
sanitize_for_struct(
|
||||
self._data_context,
|
||||
truncate_length=DATA_CONTEXT_LOG_TRUNCATE_LENGTH,
|
||||
),
|
||||
)
|
||||
|
||||
# Setup the streaming DAG topology and start the runner thread.
|
||||
self._block_ref_counter = BlockRefCounter()
|
||||
self._topology = build_streaming_topology(
|
||||
dag, self._options, self._block_ref_counter
|
||||
)
|
||||
|
||||
self._resource_manager = ResourceManager(
|
||||
self._topology,
|
||||
self._options,
|
||||
lambda: self._cluster_autoscaler.get_total_resources(),
|
||||
self._data_context,
|
||||
self._block_ref_counter,
|
||||
)
|
||||
|
||||
# Constructed once per executor (not per scheduling iteration) so the
|
||||
# guard's idle-detection state accumulates across scheduling iterations.
|
||||
self._output_backpressure_guard = OutputBackpressureGuard(
|
||||
self._topology, self._resource_manager
|
||||
)
|
||||
|
||||
# Setup progress manager
|
||||
self._progress_manager = get_progress_manager(
|
||||
self._data_context,
|
||||
self._dataset_id,
|
||||
self._topology,
|
||||
self._options.verbose_progress,
|
||||
)
|
||||
self._progress_manager.start()
|
||||
|
||||
self._backpressure_policies = get_backpressure_policies(
|
||||
self._data_context, self._topology, self._resource_manager
|
||||
)
|
||||
self._cluster_autoscaler = create_cluster_autoscaler(
|
||||
self._topology,
|
||||
self._resource_manager,
|
||||
self._data_context,
|
||||
execution_id=self._dataset_id,
|
||||
)
|
||||
self._actor_autoscaler = create_actor_autoscaler(
|
||||
self._topology,
|
||||
self._resource_manager,
|
||||
config=self._data_context.autoscaling_config,
|
||||
)
|
||||
|
||||
self._has_op_completed = dict.fromkeys(self._topology, False)
|
||||
|
||||
self._output_node = dag, self._topology[dag]
|
||||
|
||||
op_to_id = {
|
||||
op: self._get_operator_id(op, i) for i, op in enumerate(self._topology)
|
||||
}
|
||||
_StatsManager.register_dataset_to_stats_actor(
|
||||
self._dataset_id,
|
||||
self._get_operator_tags(),
|
||||
TopologyMetadata.create_topology_metadata(dag, op_to_id),
|
||||
self._data_context,
|
||||
)
|
||||
for callback in self._callbacks:
|
||||
callback.before_execution_starts(self)
|
||||
|
||||
self.start()
|
||||
self._execution_started = True
|
||||
|
||||
return _ClosingIterator(self)
|
||||
|
||||
def __del__(self):
|
||||
# NOTE: Upon garbage-collection we're allowing running tasks
|
||||
# to be terminated asynchronously (ie avoid unnecessary
|
||||
# synchronization on their completion)
|
||||
self.shutdown(force=False)
|
||||
|
||||
def shutdown(self, force: bool, exception: Optional[Exception] = None):
|
||||
global _num_shutdown
|
||||
|
||||
with self._shutdown_lock:
|
||||
if not self._execution_started or self._shutdown:
|
||||
return
|
||||
|
||||
start = time.perf_counter()
|
||||
|
||||
status_detail = (
|
||||
f"failed with {exception}" if exception else "completed successfully"
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"Shutting down executor for dataset {self._dataset_id} "
|
||||
f"({status_detail})"
|
||||
)
|
||||
|
||||
_num_shutdown += 1
|
||||
self._shutdown = True
|
||||
# Give the scheduling loop some time to finish processing.
|
||||
self.join(timeout=2.0)
|
||||
# Stop the metadata fetcher (after the loop thread that feeds it has
|
||||
# been joined). No-op for the inline fetcher.
|
||||
self._metadata_fetcher.stop()
|
||||
self._update_stats_metrics(
|
||||
state=DatasetState.FINISHED.name
|
||||
if exception is None
|
||||
else DatasetState.FAILED.name,
|
||||
force_update=True,
|
||||
)
|
||||
# Freeze the stats and save it.
|
||||
self._final_stats = self._generate_stats()
|
||||
stats_summary_string = self._final_stats.to_summary().to_string(
|
||||
include_parent=False
|
||||
)
|
||||
# Reset the scheduling loop duration gauge + resource manager budgets/usages.
|
||||
self._resource_manager.update_usages()
|
||||
self.update_metrics(0)
|
||||
if self._data_context.enable_auto_log_stats:
|
||||
logger.info(stats_summary_string)
|
||||
# Close the progress manager with a finishing message.
|
||||
if exception is None:
|
||||
desc = (
|
||||
f"{OK_PREFIX} Dataset {self._dataset_id} execution finished in "
|
||||
f"{self._final_stats.time_total_s:.2f} seconds"
|
||||
)
|
||||
else:
|
||||
desc = f"{WARN_PREFIX} Dataset {self._dataset_id} execution failed"
|
||||
|
||||
self._progress_manager.close_with_finishing_description(
|
||||
desc, exception is None
|
||||
)
|
||||
logger.info(desc)
|
||||
|
||||
timer = Timer()
|
||||
|
||||
for op in self._topology.keys():
|
||||
op.shutdown(timer, force=force)
|
||||
|
||||
self._clear_topology_queues_post_shutdown(force, exception)
|
||||
# Queues have been drained; any remaining Ray Core callbacks that fire
|
||||
# after this point should be no-ops.
|
||||
self._block_ref_counter.clear()
|
||||
|
||||
min_ = round(timer.min(), 3)
|
||||
max_ = round(timer.max(), 3)
|
||||
total = round(timer.get(), 3)
|
||||
logger.debug(
|
||||
f"Shut down operator hierarchy for dataset {self._dataset_id}"
|
||||
f" (min/max/total={min_}/{max_}/{total}s)"
|
||||
)
|
||||
|
||||
if exception is None:
|
||||
for callback in self._callbacks:
|
||||
callback.after_execution_succeeds(self)
|
||||
else:
|
||||
for callback in self._callbacks:
|
||||
callback.after_execution_fails(self, exception)
|
||||
|
||||
self._cluster_autoscaler.on_executor_shutdown()
|
||||
|
||||
dur = time.perf_counter() - start
|
||||
|
||||
logger.debug(
|
||||
f"Shut down executor for dataset {self._dataset_id} "
|
||||
f"(took {round(dur, 3)}s)"
|
||||
)
|
||||
|
||||
# Unregister should be called after all operators are shut down to
|
||||
# capture as many logs as possible.
|
||||
self._data_context.set_dataset_logger_id(
|
||||
unregister_dataset_logger(self._dataset_id)
|
||||
)
|
||||
|
||||
def _clear_topology_queues_post_shutdown(
|
||||
self, force: bool, exception: Optional[Exception] = None
|
||||
) -> None:
|
||||
"""Drain topology queues after operator shutdown (releases block refs)."""
|
||||
for op, state in self._topology.items():
|
||||
if isinstance(op, InternalQueueOperatorMixin):
|
||||
op.clear_internal_input_queue()
|
||||
op.clear_internal_output_queue()
|
||||
# Input queues alias upstream output queues; clears the DAG except the sink.
|
||||
for inqueue in state.input_queues:
|
||||
inqueue.clear()
|
||||
|
||||
output_op, _ = self._output_node
|
||||
# Clear sink output unless cooperative multi-split success (splits may still read).
|
||||
is_live_multi_split_sink = (
|
||||
output_op.num_output_splits() > 1 and not force and exception is None
|
||||
)
|
||||
if not is_live_multi_split_sink:
|
||||
self._topology[output_op].output_queue.clear()
|
||||
|
||||
def run(self):
|
||||
"""Run the control loop in a helper thread.
|
||||
|
||||
Results are returned via the output node's outqueue.
|
||||
"""
|
||||
exc: Optional[Exception] = None
|
||||
self._metadata_fetcher.start()
|
||||
try:
|
||||
# Run scheduling loop until complete.
|
||||
while True:
|
||||
# Use `perf_counter` rather than `process_time` to ensure we include
|
||||
# time spent on IO, like RPCs to Ray Core.
|
||||
t_start = time.perf_counter()
|
||||
continue_sched = self._scheduling_loop_step(self._topology)
|
||||
|
||||
sched_loop_duration = time.perf_counter() - t_start
|
||||
|
||||
self.update_metrics(sched_loop_duration)
|
||||
if self._initial_stats:
|
||||
self._initial_stats.streaming_exec_schedule_s.add(
|
||||
sched_loop_duration
|
||||
)
|
||||
|
||||
for callback in self._callbacks:
|
||||
callback.on_execution_step(self)
|
||||
if not continue_sched or self._shutdown:
|
||||
break
|
||||
except Exception as e:
|
||||
# Propagate it to the result iterator.
|
||||
exc = e
|
||||
finally:
|
||||
# Mark state of outputting operator as finished
|
||||
_, state = self._output_node
|
||||
state.mark_finished(exc)
|
||||
|
||||
def update_metrics(self, sched_loop_duration: int):
|
||||
self._sched_loop_duration_s.set(
|
||||
sched_loop_duration, tags={"dataset": self._dataset_id}
|
||||
)
|
||||
|
||||
def get_stats(self):
|
||||
"""Return the stats object for the streaming execution.
|
||||
|
||||
The stats object will be updated as streaming execution progresses.
|
||||
"""
|
||||
if self._final_stats:
|
||||
return self._final_stats
|
||||
else:
|
||||
return self._generate_stats()
|
||||
|
||||
def set_external_consumer_bytes(self, num_bytes: int) -> None:
|
||||
"""Set the bytes buffered by external consumers."""
|
||||
if self._resource_manager is not None:
|
||||
self._resource_manager.set_external_consumer_bytes(num_bytes)
|
||||
|
||||
def _generate_stats(self) -> DatasetStats:
|
||||
"""Create a new stats object reflecting execution status so far."""
|
||||
stats = self._initial_stats or DatasetStats(metadata={}, parent=None)
|
||||
for op in self._topology:
|
||||
if isinstance(op, InputDataBuffer):
|
||||
continue
|
||||
builder = stats.child_builder(op.name, override_start_time=self._start_time)
|
||||
stats = builder.build_multioperator(op.get_stats())
|
||||
stats.extra_metrics = op.metrics.as_dict(skip_internal_metrics=True)
|
||||
# Always assign a ``Timer`` so downstream consumers can call
|
||||
# ``.get()`` / ``.avg()`` / ``.max()`` / ``.percentile()``
|
||||
# unconditionally. When ``_initial_stats`` is absent we hand
|
||||
# back an empty Timer; zero-sample semantics yield 0 across all
|
||||
# four.
|
||||
stats.streaming_exec_schedule_s = (
|
||||
self._initial_stats.streaming_exec_schedule_s
|
||||
if self._initial_stats
|
||||
else Timer()
|
||||
)
|
||||
return stats
|
||||
|
||||
def _scheduling_loop_step(self, topology: Topology) -> bool:
|
||||
"""Run one step of the scheduling loop.
|
||||
|
||||
This runs a few general phases:
|
||||
1. Waiting for the next task completion using `ray.wait()`.
|
||||
2. Pulling completed refs into operator outqueues.
|
||||
3. Selecting and dispatching new inputs to operators.
|
||||
|
||||
Args:
|
||||
topology: The :class:`Topology` of operators being executed.
|
||||
|
||||
Returns:
|
||||
True if we should continue running the scheduling loop.
|
||||
"""
|
||||
self._resource_manager.update_usages()
|
||||
# Note: calling process_completed_tasks() is expensive since it incurs
|
||||
# ray.wait() overhead, so make sure to allow multiple dispatch per call for
|
||||
# greater parallelism.
|
||||
num_errored_blocks = process_completed_tasks(
|
||||
topology,
|
||||
self._backpressure_policies,
|
||||
self._max_errored_blocks,
|
||||
output_backpressure_guard=self._output_backpressure_guard,
|
||||
metadata_fetcher=self._metadata_fetcher,
|
||||
)
|
||||
if self._max_errored_blocks > 0:
|
||||
self._max_errored_blocks -= num_errored_blocks
|
||||
self._num_errored_blocks += num_errored_blocks
|
||||
|
||||
self._resource_manager.update_usages()
|
||||
# Dispatch as many operators as we can for completed tasks.
|
||||
self._report_current_usage()
|
||||
|
||||
i = 0
|
||||
while True:
|
||||
op = select_operator_to_run(
|
||||
topology,
|
||||
self._resource_manager,
|
||||
self._backpressure_policies,
|
||||
# If consumer is idling (there's nothing for it to consume)
|
||||
# enforce liveness, ie that at least a single task gets scheduled
|
||||
ensure_liveness=self._consumer_idling(),
|
||||
ranker=self._ranker,
|
||||
)
|
||||
|
||||
if op is None:
|
||||
break
|
||||
|
||||
topology[op].dispatch_next_task()
|
||||
|
||||
self._resource_manager.update_usages()
|
||||
|
||||
i += 1
|
||||
if i % self._progress_manager.TOTAL_PROGRESS_REFRESH_EVERY_N_STEPS == 0:
|
||||
self._refresh_progress_manager(topology)
|
||||
|
||||
# Trigger autoscaling
|
||||
self._cluster_autoscaler.try_trigger_scaling()
|
||||
self._actor_autoscaler.try_trigger_scaling()
|
||||
|
||||
update_operator_states(topology)
|
||||
self._refresh_progress_manager(topology)
|
||||
|
||||
self._update_stats_metrics(state=DatasetState.RUNNING.name)
|
||||
if time.time() - self._last_debug_log_time >= DEBUG_LOG_INTERVAL_SECONDS:
|
||||
_log_op_metrics(topology)
|
||||
_debug_dump_topology(topology, self._resource_manager)
|
||||
self._last_debug_log_time = time.time()
|
||||
|
||||
for op, state in topology.items():
|
||||
# Export operator schema if it's updated
|
||||
if state._schema is not None and self._op_schema.get(op) != state._schema:
|
||||
self._op_schema[op] = state._schema
|
||||
self._export_operator_schema(op)
|
||||
|
||||
# Log metrics of newly completed operators.
|
||||
if not op.has_completed():
|
||||
op.refresh_state()
|
||||
elif not self._has_op_completed[op]:
|
||||
log_str = (
|
||||
f"Operator {op} completed. "
|
||||
f"Operator Metrics:\n{op._metrics.as_dict(skip_internal_metrics=True)}"
|
||||
)
|
||||
logger.debug(log_str)
|
||||
self._has_op_completed[op] = True
|
||||
self._validate_operator_queues_empty(op, state)
|
||||
|
||||
# Keep going until all operators run to completion.
|
||||
return not all(op.has_completed() for op in topology)
|
||||
|
||||
def _refresh_progress_manager(self, topology: Topology):
|
||||
# Update the progress manager to reflect scheduling decisions.
|
||||
if self._progress_manager:
|
||||
for op_state in topology.values():
|
||||
if not isinstance(op_state.op, InputDataBuffer):
|
||||
self._progress_manager.update_operator_progress(
|
||||
op_state, self._resource_manager
|
||||
)
|
||||
self._progress_manager.refresh()
|
||||
|
||||
def _consumer_idling(self) -> bool:
|
||||
"""Returns whether the user thread is blocked on topology execution."""
|
||||
_, state = self._output_node
|
||||
return len(state.output_queue) == 0
|
||||
|
||||
def _export_operator_schema(self, op: PhysicalOperator) -> None:
|
||||
schema = self._op_schema.get(op)
|
||||
operator_schema_exporter = get_operator_schema_exporter()
|
||||
if (
|
||||
operator_schema_exporter is not None
|
||||
and hasattr(schema, "names")
|
||||
and hasattr(schema, "types")
|
||||
):
|
||||
names = [str(n) for n in schema.names]
|
||||
types = [str(t) for t in schema.types]
|
||||
operator_schema = OperatorSchema(
|
||||
operator_uuid=op.id,
|
||||
schema_fields=dict(zip(names, types)),
|
||||
)
|
||||
operator_schema_exporter.export_operator_schema(operator_schema)
|
||||
|
||||
def _validate_operator_queues_empty(
|
||||
self, op: PhysicalOperator, state: OpState
|
||||
) -> None:
|
||||
"""Validate that all queues are empty when an operator completes.
|
||||
|
||||
Args:
|
||||
op: The completed operator to validate.
|
||||
state: The operator's execution state.
|
||||
"""
|
||||
error_msg = "Expected {} Queue for {} to be empty, but found {} bundles"
|
||||
|
||||
if isinstance(op, InternalQueueOperatorMixin):
|
||||
# 1) Check Internal Input Queue is empty
|
||||
assert op.internal_input_queue_num_blocks() == 0, error_msg.format(
|
||||
"Internal Input", op.name, op.internal_input_queue_num_blocks()
|
||||
)
|
||||
|
||||
# 2) Check Internal Output Queue is empty
|
||||
assert op.internal_output_queue_num_blocks() == 0, error_msg.format(
|
||||
"Internal Output",
|
||||
op.name,
|
||||
op.internal_output_queue_num_blocks(),
|
||||
)
|
||||
|
||||
# 3) Check that External Input Queue is empty
|
||||
for input_q in state.input_queues:
|
||||
assert len(input_q) == 0, error_msg.format(
|
||||
"External Input", op.name, len(input_q)
|
||||
)
|
||||
|
||||
def _report_current_usage(self) -> None:
|
||||
# running_usage is the amount of resources that have been requested but
|
||||
# not necessarily available
|
||||
# TODO(sofian) https://github.com/ray-project/ray/issues/47520
|
||||
# We need to split the reported resources into running, pending-scheduling,
|
||||
# pending-node-assignment.
|
||||
running_usage = self._resource_manager.get_global_running_usage()
|
||||
pending_usage = self._resource_manager.get_global_pending_usage()
|
||||
limits = self._resource_manager.get_global_limits()
|
||||
resources_status = (
|
||||
f"Active & requested resources: "
|
||||
f"{running_usage.cpu:.4g}/{limits.cpu:.4g} CPU, "
|
||||
)
|
||||
if running_usage.memory > 0:
|
||||
resources_status += (
|
||||
f"{running_usage.memory_str()}/{limits.memory_str()} memory, "
|
||||
)
|
||||
if running_usage.gpu > 0:
|
||||
resources_status += f"{running_usage.gpu:.4g}/{limits.gpu:.4g} GPU, "
|
||||
resources_status += (
|
||||
f"{running_usage.object_store_memory_str()}/"
|
||||
f"{limits.object_store_memory_str()} object store"
|
||||
)
|
||||
|
||||
# Only include pending section when there are pending resources.
|
||||
pending_parts = []
|
||||
if pending_usage.cpu:
|
||||
pending_parts.append(f"{pending_usage.cpu:.4g} CPU")
|
||||
if pending_usage.memory:
|
||||
pending_parts.append(f"{pending_usage.memory_str()} memory")
|
||||
if pending_usage.gpu:
|
||||
pending_parts.append(f"{pending_usage.gpu:.4g} GPU")
|
||||
if pending_parts:
|
||||
resources_status += f" (pending: {', '.join(pending_parts)})"
|
||||
|
||||
self._progress_manager.update_total_resource_status(resources_status)
|
||||
|
||||
def _get_operator_id(self, op: PhysicalOperator, topology_index: int) -> str:
|
||||
return f"{op.name}_{topology_index}"
|
||||
|
||||
def _get_operator_tags(self):
|
||||
"""Returns a list of operator tags."""
|
||||
return [
|
||||
f"{self._get_operator_id(op, i)}" for i, op in enumerate(self._topology)
|
||||
]
|
||||
|
||||
def _get_state_dict(self, state):
|
||||
last_op, last_state = list(self._topology.items())[-1]
|
||||
return {
|
||||
"state": state,
|
||||
"progress": last_state.num_completed_tasks,
|
||||
"total": last_op.num_outputs_total(),
|
||||
"total_rows": last_op.num_output_rows_total(),
|
||||
"end_time": time.time()
|
||||
if state in (DatasetState.FINISHED.name, DatasetState.FAILED.name)
|
||||
else None,
|
||||
"operators": {
|
||||
f"{self._get_operator_id(op, i)}": {
|
||||
"name": op.name,
|
||||
"progress": op_state.num_completed_tasks,
|
||||
"total": op.num_outputs_total(),
|
||||
"total_rows": op.num_output_rows_total(),
|
||||
"queued_blocks": op_state.total_enqueued_input_blocks(),
|
||||
"state": DatasetState.FINISHED.name
|
||||
if op.has_execution_finished()
|
||||
else state,
|
||||
}
|
||||
for i, (op, op_state) in enumerate(self._topology.items())
|
||||
},
|
||||
}
|
||||
|
||||
def _update_stats_metrics(self, state: str, force_update: bool = False):
|
||||
now = time.time()
|
||||
if (
|
||||
force_update
|
||||
or (now - self._metrics_last_updated) > self.UPDATE_METRICS_INTERVAL_S
|
||||
):
|
||||
_StatsManager.update_execution_metrics(
|
||||
self._dataset_id,
|
||||
[op.metrics for op in self._topology],
|
||||
self._get_operator_tags(),
|
||||
self._get_state_dict(state=state),
|
||||
)
|
||||
self._metrics_last_updated = now
|
||||
|
||||
|
||||
def _debug_dump_topology(topology: Topology, resource_manager: ResourceManager) -> None:
|
||||
"""Log current execution state for the topology for debugging.
|
||||
|
||||
Args:
|
||||
topology: The topology to debug.
|
||||
resource_manager: The resource manager for this topology.
|
||||
"""
|
||||
logger.debug("Execution Progress:")
|
||||
for i, (op, state) in enumerate(topology.items()):
|
||||
summary_str = format_op_state_summary(state, resource_manager, verbose=True)
|
||||
logger.debug(
|
||||
f"{i}: {op.name} - {summary_str}, "
|
||||
f"Blocks Outputted: {state.num_completed_tasks}/{op.num_outputs_total()}"
|
||||
)
|
||||
|
||||
|
||||
def _log_op_metrics(topology: Topology) -> None:
|
||||
"""Logs the metrics of each operator.
|
||||
|
||||
Args:
|
||||
topology: The topology to debug.
|
||||
"""
|
||||
log_str = "Operator Metrics:\n"
|
||||
for op in topology:
|
||||
metrics_dict = op.metrics.as_dict(skip_internal_metrics=True)
|
||||
log_str += f"{op.name}: {metrics_dict}\n"
|
||||
logger.debug(log_str)
|
||||
|
||||
|
||||
class _ClosingIterator(OutputIterator):
|
||||
"""Iterator automatically shutting down executor upon exhausting the
|
||||
iterable sequence.
|
||||
|
||||
NOTE: If this iterator isn't fully exhausted, executor still have to
|
||||
be closed manually by the caller!
|
||||
"""
|
||||
|
||||
def __init__(self, executor: StreamingExecutor):
|
||||
self._executor = executor
|
||||
|
||||
def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
|
||||
try:
|
||||
op, state = self._executor._output_node
|
||||
bundle = state.get_output_blocking(output_split_idx)
|
||||
|
||||
# Update progress-bars
|
||||
if self._executor._progress_manager:
|
||||
self._executor._progress_manager.update_total_progress(
|
||||
bundle.num_rows() or 0, op.num_output_rows_total()
|
||||
)
|
||||
|
||||
return bundle
|
||||
|
||||
# Have to be BaseException to catch ``KeyboardInterrupt``
|
||||
#
|
||||
# NOTE: This also handles ``StopIteration``
|
||||
except BaseException as e:
|
||||
# Asynchronously shutdown the executor (ie avoid unnecessary
|
||||
# synchronization on tasks termination)
|
||||
self._executor.shutdown(
|
||||
force=False, exception=e if not isinstance(e, StopIteration) else None
|
||||
)
|
||||
raise
|
||||
|
||||
def __del__(self):
|
||||
# NOTE: Upon garbage-collection we're allowing running tasks
|
||||
# to be terminated asynchronously (ie avoid unnecessary
|
||||
# synchronization on their completion)
|
||||
self._executor.shutdown(force=False)
|
||||
Reference in New Issue
Block a user