chore: import upstream snapshot with attribution

python/ray/data/_internal/execution/__init__.py

@@ -0,0 +1,57 @@
+import enum
+import logging
+import os
+from typing import TYPE_CHECKING, Optional
+
+from .ranker import DefaultRanker, Ranker
+
+if TYPE_CHECKING:
+    from ... import DataContext
+    from .resource_manager import OpResourceAllocator, ResourceManager
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_USE_OP_RESOURCE_ALLOCATOR_VERSION = os.environ.get(
+    "RAY_DATA_USE_OP_RESOURCE_ALLOCATOR_VERSION", "V1"
+)
+
+
+class OpResourceAllocatorVersion(str, enum.Enum):
+    V1 = "V1"  # ReservationOpResourceAllocator
+
+
+def create_resource_allocator(
+    resource_manager: "ResourceManager",
+    data_context: "DataContext",
+) -> Optional["OpResourceAllocator"]:
+    """Creates ``OpResourceAllocator`` instances.``"""
+
+    if not data_context.op_resource_reservation_enabled:
+        # This is a historical kill-switch to disable resource allocator, that
+        # will be soon deprecated and removed.
+        return None
+
+    logger.debug(
+        f"Using op resource allocator version: "
+        f"{DEFAULT_USE_OP_RESOURCE_ALLOCATOR_VERSION!r}"
+    )
+
+    if DEFAULT_USE_OP_RESOURCE_ALLOCATOR_VERSION == OpResourceAllocatorVersion.V1:
+        from .resource_manager import ReservationOpResourceAllocator
+
+        return ReservationOpResourceAllocator(
+            resource_manager,
+            reservation_ratio=data_context.op_resource_reservation_ratio,
+        )
+
+    else:
+        raise ValueError(
+            "Resource allocator version of "
+            f"'{DEFAULT_USE_OP_RESOURCE_ALLOCATOR_VERSION}' is not supported"
+        )
+
+
+def create_ranker() -> Ranker:
+    """Create a ranker instance based on environment and configuration."""
+    return DefaultRanker()