wehub-resource-sync
89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
import contextlib
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data._internal.execution.operators.map_transformer import MapTransformer
|
|
from ray.data._internal.progress.base_progress import BaseProgressBar
|
|
|
|
|
|
_thread_local = threading.local()
|
|
|
|
|
|
@dataclass
|
|
class TaskContext:
|
|
"""This describes the information of a task running block transform."""
|
|
|
|
# The index of task. Each task has a unique task index within the same
|
|
# operator.
|
|
task_idx: int
|
|
|
|
# Name of the operator that this task belongs to.
|
|
op_name: str
|
|
|
|
# The dictionary of sub progress bar to update. The key is name of sub progress
|
|
# bar. Note this is only used on driver side.
|
|
# TODO(chengsu): clean it up from TaskContext with new optimizer framework.
|
|
sub_progress_bar_dict: Optional[Dict[str, "BaseProgressBar"]] = None
|
|
|
|
# NOTE(hchen): `upstream_map_transformer` and `upstream_map_ray_remote_args`
|
|
# are only used for `RandomShuffle`. DO NOT use them for other operators.
|
|
# Ideally, they should be handled by the optimizer, and should be transparent
|
|
# to the specific operators.
|
|
# But for `RandomShuffle`, the AllToAllOperator doesn't do the shuffle itself.
|
|
# It uses `ExchangeTaskScheduler` to launch new tasks to do the shuffle.
|
|
# That's why we need to pass them to `ExchangeTaskScheduler`.
|
|
# TODO(hchen): Use a physical operator to do the shuffle directly.
|
|
|
|
# The underlying function called in a MapOperator; this is used when fusing
|
|
# an AllToAllOperator with an upstream MapOperator.
|
|
upstream_map_transformer: Optional["MapTransformer"] = None
|
|
|
|
# The Ray remote arguments of the fused upstream MapOperator.
|
|
# This should be set if upstream_map_transformer is set.
|
|
upstream_map_ray_remote_args: Optional[Dict[str, Any]] = None
|
|
|
|
# Override of the target max-block-size for the task
|
|
target_max_block_size_override: Optional[int] = None
|
|
|
|
# Additional keyword arguments passed to the task.
|
|
kwargs: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
@classmethod
|
|
def get_current(cls) -> Optional["TaskContext"]:
|
|
"""Get the TaskContext for the current thread.
|
|
Returns None if no TaskContext has been set.
|
|
"""
|
|
|
|
return getattr(_thread_local, "task_context", None)
|
|
|
|
@classmethod
|
|
def set_current(cls, context: "TaskContext") -> None:
|
|
"""Set the TaskContext for the current thread.
|
|
|
|
Args:
|
|
context: The TaskContext instance to set for this thread
|
|
"""
|
|
|
|
_thread_local.task_context = context
|
|
|
|
@classmethod
|
|
def reset_current(cls):
|
|
"""Clear the current thread's TaskContext."""
|
|
|
|
if hasattr(_thread_local, "task_context"):
|
|
delattr(_thread_local, "task_context")
|
|
|
|
@classmethod
|
|
@contextlib.contextmanager
|
|
def current(cls, context: "TaskContext") -> Iterator["TaskContext"]:
|
|
"""Sets this TaskContext as current for the scope
|
|
of the context block and resets it on exit.
|
|
"""
|
|
cls.set_current(context)
|
|
try:
|
|
yield context
|
|
finally:
|
|
cls.reset_current()
|