wehub-resource-sync
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
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()
|