chore: import upstream snapshot with attribution

python/ray/data/_internal/execution/interfaces/task_context.py

@@ -0,0 +1,88 @@
+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()