wehub-resource-sync
148 lines
4.8 KiB
Python
148 lines
4.8 KiB
Python
import os
|
|
import uuid
|
|
from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
|
|
|
|
import ray
|
|
from ray.air.execution import FixedResourceManager
|
|
from ray.air.execution._internal import RayActorManager
|
|
from ray.air.execution._internal.tracked_actor import TrackedActor
|
|
from ray.air.execution.resources import ResourceManager, ResourceRequest
|
|
from ray.train.tests.util import mock_storage_context
|
|
from ray.tune.execution.tune_controller import TuneController
|
|
from ray.tune.experiment import Trial
|
|
from ray.tune.utils.resource_updater import _ResourceUpdater
|
|
|
|
|
|
class NoopClassCache:
|
|
def get(self, trainable_name: str):
|
|
return trainable_name
|
|
|
|
|
|
class BudgetResourceManager(FixedResourceManager):
|
|
def __init__(self, total_resources: Dict[str, float]):
|
|
self._allow_strict_pack = True
|
|
self._total_resources = total_resources
|
|
self._requested_resources = []
|
|
self._used_resources = []
|
|
|
|
|
|
class NoopActorManager(RayActorManager):
|
|
def __init__(self, resource_manager: ResourceManager):
|
|
super().__init__(resource_manager=resource_manager)
|
|
|
|
self.added_actors = []
|
|
self.removed_actors = []
|
|
self.scheduled_futures = []
|
|
|
|
def add_actor(
|
|
self,
|
|
cls: Union[Type, ray.actor.ActorClass],
|
|
kwargs: Dict[str, Any],
|
|
resource_request: ResourceRequest,
|
|
*,
|
|
on_start: Optional[Callable[[TrackedActor], None]] = None,
|
|
on_stop: Optional[Callable[[TrackedActor], None]] = None,
|
|
on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
|
|
) -> TrackedActor:
|
|
fake_actor_ref = uuid.uuid4().int
|
|
tracked_actor = TrackedActor(
|
|
fake_actor_ref, on_start=on_start, on_stop=on_stop, on_error=on_error
|
|
)
|
|
self._live_actors_to_ray_actors_resources[tracked_actor] = (fake_actor_ref,)
|
|
self.added_actors.append((tracked_actor, cls, kwargs))
|
|
return tracked_actor
|
|
|
|
def remove_actor(
|
|
self,
|
|
tracked_actor: TrackedActor,
|
|
kill: bool = False,
|
|
stop_future: Optional[ray.ObjectRef] = None,
|
|
) -> None:
|
|
self.removed_actors.append(tracked_actor)
|
|
|
|
def schedule_actor_task(
|
|
self,
|
|
tracked_actor: TrackedActor,
|
|
method_name: str,
|
|
args: Optional[Tuple] = None,
|
|
kwargs: Optional[Dict] = None,
|
|
on_result: Optional[Callable[[TrackedActor, Any], None]] = None,
|
|
on_error: Optional[Callable[[TrackedActor, Exception], None]] = None,
|
|
_return_future: bool = False,
|
|
) -> Optional[int]:
|
|
fake_ref = uuid.uuid4().int
|
|
self.scheduled_futures.append(
|
|
(fake_ref, tracked_actor, method_name, args, kwargs, on_result, on_error)
|
|
)
|
|
return fake_ref
|
|
|
|
@property
|
|
def num_actor_tasks(self):
|
|
return len(self.scheduled_futures)
|
|
|
|
def get_live_actors_resources(self):
|
|
return {}
|
|
|
|
def next(self, timeout: Optional[Union[int, float]] = None) -> None:
|
|
pass
|
|
|
|
def set_num_pending(self, num_pending: int):
|
|
self._pending_actors_to_attrs = {i: None for i in range(num_pending)}
|
|
|
|
|
|
class _FakeResourceUpdater(_ResourceUpdater):
|
|
def __init__(self, resource_manager: BudgetResourceManager):
|
|
self._resource_manager = resource_manager
|
|
|
|
def get_num_cpus(self):
|
|
return self._resource_manager._total_resources.get("CPU", 0)
|
|
|
|
def get_num_gpus(self) -> int:
|
|
return self._resource_manager._total_resources.get("GPU", 0)
|
|
|
|
def update_avail_resources(self, *args, **kwargs):
|
|
pass
|
|
|
|
|
|
class TestingTrial(Trial):
|
|
def __init__(self, *args, **kwargs):
|
|
kwargs.setdefault("storage", mock_storage_context())
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def get_trainable_cls(self):
|
|
return self.trainable_name
|
|
|
|
def create_placement_group_factory(self):
|
|
self.placement_group_factory = self._default_placement_group_factory
|
|
|
|
def set_ray_actor(self, ray_actor):
|
|
pass
|
|
|
|
|
|
def create_execution_test_objects(
|
|
max_pending_trials: int = 8,
|
|
resources: Optional[Dict[str, float]] = None,
|
|
reuse_actors: bool = True,
|
|
tune_controller_cls: Type[TuneController] = TuneController,
|
|
**kwargs,
|
|
):
|
|
os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = str(max_pending_trials)
|
|
|
|
resources = resources or {"CPU": 4}
|
|
|
|
storage = kwargs.pop("storage", mock_storage_context())
|
|
|
|
tune_controller = tune_controller_cls(
|
|
reuse_actors=reuse_actors,
|
|
storage=storage,
|
|
**kwargs,
|
|
)
|
|
resource_manager = BudgetResourceManager(total_resources=resources)
|
|
resource_updater = _FakeResourceUpdater(resource_manager)
|
|
actor_manger = NoopActorManager(resource_manager)
|
|
tune_controller._actor_manager = actor_manger
|
|
tune_controller._class_cache = NoopClassCache()
|
|
tune_controller._resource_updater = resource_updater
|
|
|
|
return tune_controller, actor_manger, resource_manager
|