wehub-resource-sync
404 lines
12 KiB
Python
404 lines
12 KiB
Python
import os
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
from unittest.mock import MagicMock
|
|
|
|
import ray
|
|
from ray.train import BackendConfig, Checkpoint
|
|
from ray.train._internal.data_config import DataConfig
|
|
from ray.train.backend import Backend
|
|
from ray.train.context import TrainContext
|
|
from ray.train.v2._internal.execution.context import (
|
|
DistributedContext,
|
|
TrainRunContext,
|
|
)
|
|
from ray.train.v2._internal.execution.failure_handling import (
|
|
FailureDecision,
|
|
FailurePolicy,
|
|
)
|
|
from ray.train.v2._internal.execution.scaling_policy import (
|
|
NoopDecision,
|
|
ScalingDecision,
|
|
ScalingPolicy,
|
|
)
|
|
from ray.train.v2._internal.execution.storage import StorageContext
|
|
from ray.train.v2._internal.execution.training_report import _TrainingReport
|
|
from ray.train.v2._internal.execution.worker_group import (
|
|
WorkerGroup,
|
|
WorkerGroupContext,
|
|
WorkerGroupPollStatus,
|
|
WorkerGroupState,
|
|
WorkerStatus,
|
|
)
|
|
from ray.train.v2._internal.execution.worker_group.execution_group import ReplicaGroup
|
|
from ray.train.v2._internal.state.schema import (
|
|
ActorStatus,
|
|
BackendConfig as BackendConfigSchema,
|
|
CheckpointConfig as CheckpointConfigSchema,
|
|
DataConfig as DataConfigSchema,
|
|
DataExecutionOptions,
|
|
FailureConfig as FailureConfigSchema,
|
|
RunAttemptStatus,
|
|
RunConfig as RunConfigSchema,
|
|
RunSettings,
|
|
RunStatus,
|
|
ScalingConfig as ScalingConfigSchema,
|
|
TrainResources,
|
|
TrainRun,
|
|
TrainRunAttempt,
|
|
TrainWorker,
|
|
)
|
|
from ray.train.v2._internal.state.util import execution_options_to_model
|
|
from ray.train.v2._internal.util import ObjectRefWrapper, time_monotonic
|
|
from ray.train.v2.api.exceptions import TrainingFailedError
|
|
from ray.train.v2.api.validation_config import ValidationTaskConfig
|
|
|
|
|
|
class MockReplicaGroupBackend(Backend):
|
|
has_replica_groups = True
|
|
|
|
|
|
class MockReplicaGroupBackendConfig(BackendConfig):
|
|
@property
|
|
def backend_cls(self):
|
|
return MockReplicaGroupBackend
|
|
|
|
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self):
|
|
self.count = 0
|
|
|
|
def increment(self):
|
|
self.count += 1
|
|
return self.count
|
|
|
|
def get_count(self):
|
|
return self.count
|
|
|
|
|
|
class DummyWorkerGroup(WorkerGroup):
|
|
|
|
_start_failure = None
|
|
_poll_failure = None
|
|
|
|
# TODO: Clean this up and use Mocks instead.
|
|
def __init__(
|
|
self,
|
|
train_run_context: TrainRunContext,
|
|
worker_group_context: WorkerGroupContext,
|
|
callbacks=None,
|
|
):
|
|
self._num_workers = worker_group_context.num_workers
|
|
self._worker_group_state = None
|
|
self._worker_statuses = {}
|
|
self._replaced_replica_groups: List[int] = []
|
|
self._replica_groups = None
|
|
self._latest_poll_status: Optional[WorkerGroupPollStatus] = None
|
|
|
|
def poll_status(self, *args, **kwargs) -> WorkerGroupPollStatus:
|
|
if self._poll_failure:
|
|
raise self._poll_failure
|
|
return WorkerGroupPollStatus(
|
|
worker_statuses=self._worker_statuses,
|
|
)
|
|
|
|
def _start(self):
|
|
num_workers = self._num_workers
|
|
if self._start_failure:
|
|
raise self._start_failure
|
|
|
|
workers = [MagicMock() for i in range(num_workers)]
|
|
self._worker_group_state = WorkerGroupState(
|
|
start_time=time_monotonic(),
|
|
workers=workers,
|
|
placement_group_handle=MagicMock(),
|
|
sync_actor=None,
|
|
)
|
|
|
|
self._worker_statuses = {
|
|
i: WorkerStatus(running=True, error=None) for i in range(num_workers)
|
|
}
|
|
|
|
self._replica_groups = [
|
|
ReplicaGroup([workers[i]], resources_per_worker={})
|
|
for i in range(num_workers)
|
|
]
|
|
|
|
def shutdown(self):
|
|
self._worker_group_state = None
|
|
|
|
def abort(self):
|
|
pass
|
|
|
|
def replace_replica_group(self, replica_group_index: int):
|
|
self._replaced_replica_groups.append(replica_group_index)
|
|
|
|
# === Test methods ===
|
|
def clear_worker(self):
|
|
for worker_status in self._worker_statuses.values():
|
|
worker_status.error = None
|
|
worker_status.running = True
|
|
|
|
def error_worker(self, worker_index):
|
|
status = self._worker_statuses[worker_index]
|
|
status.error = RuntimeError(f"Worker {worker_index} failed")
|
|
|
|
def finish_worker(self, worker_index):
|
|
status = self._worker_statuses[worker_index]
|
|
status.running = False
|
|
|
|
@classmethod
|
|
def set_start_failure(cls, start_failure):
|
|
cls._start_failure = start_failure
|
|
|
|
@classmethod
|
|
def set_poll_failure(cls, poll_failure):
|
|
cls._poll_failure = poll_failure
|
|
|
|
|
|
class MockScalingPolicy(ScalingPolicy):
|
|
def __init__(self, scaling_config):
|
|
self._recovery_decision_queue = []
|
|
self._monitor_decision_queue = []
|
|
|
|
super().__init__(scaling_config)
|
|
|
|
def _get_num_workers_for_resource_request(self) -> int:
|
|
return self.scaling_config.num_workers
|
|
|
|
def make_decision_for_non_running_worker_group(self) -> ScalingDecision:
|
|
if self._recovery_decision_queue:
|
|
return self._recovery_decision_queue.pop(0)
|
|
return NoopDecision()
|
|
|
|
def make_decision_for_running_worker_group(
|
|
self,
|
|
worker_group_state: WorkerGroupState,
|
|
worker_group_status: WorkerGroupPollStatus,
|
|
) -> ScalingDecision:
|
|
if self._monitor_decision_queue:
|
|
return self._monitor_decision_queue.pop(0)
|
|
return NoopDecision()
|
|
|
|
# === Test methods ===
|
|
def queue_recovery_decision(self, decision):
|
|
self._recovery_decision_queue.append(decision)
|
|
|
|
def queue_monitor_decision(self, decision):
|
|
self._monitor_decision_queue.append(decision)
|
|
|
|
|
|
class MockFailurePolicy(FailurePolicy):
|
|
def __init__(self, failure_config):
|
|
self._decision_queue = []
|
|
|
|
super().__init__(failure_config)
|
|
|
|
def make_decision(
|
|
self, training_failed_error: TrainingFailedError
|
|
) -> FailureDecision:
|
|
if self._decision_queue:
|
|
return self._decision_queue.pop(0)
|
|
return FailureDecision.NOOP
|
|
|
|
# === Test methods ===
|
|
def queue_decision(self, decision):
|
|
self._decision_queue.append(decision)
|
|
|
|
|
|
class DummyObjectRefWrapper(ObjectRefWrapper):
|
|
"""Mock object that returns the object passed in without going through ray.put."""
|
|
|
|
def __init__(self, obj):
|
|
self._obj = obj
|
|
|
|
def get(self):
|
|
return self._obj
|
|
|
|
|
|
_RUN_ID = "mock_run_id"
|
|
|
|
|
|
def create_mock_train_run(
|
|
status: RunStatus = RunStatus.RUNNING,
|
|
controller_actor_id: Optional[str] = None,
|
|
end_time_ns: Optional[int] = None,
|
|
id: Optional[str] = None,
|
|
status_detail: Optional[str] = None,
|
|
train_loop_config: Optional[Dict] = None,
|
|
):
|
|
return TrainRun(
|
|
schema_version=0,
|
|
id=id or _RUN_ID,
|
|
name="test_run",
|
|
job_id=uuid.uuid4().hex,
|
|
controller_actor_id=controller_actor_id or uuid.uuid4().hex,
|
|
status=status,
|
|
status_detail=status_detail,
|
|
start_time_ns=time.time_ns(),
|
|
end_time_ns=end_time_ns,
|
|
controller_log_file_path="/tmp/ray/session_xxx/logs/train/ray-train-app-controller.log",
|
|
framework_versions={"ray": ray.__version__},
|
|
run_settings=RunSettings(
|
|
train_loop_config=train_loop_config,
|
|
backend_config=BackendConfigSchema(framework=None, config={}),
|
|
scaling_config=ScalingConfigSchema(
|
|
num_workers=1,
|
|
use_gpu=False,
|
|
resources_per_worker=None,
|
|
placement_strategy="PACK",
|
|
accelerator_type=None,
|
|
use_tpu=False,
|
|
topology=None,
|
|
bundle_label_selector=None,
|
|
),
|
|
datasets=["dataset_1"],
|
|
data_config=DataConfigSchema(
|
|
datasets_to_split="all",
|
|
data_execution_options=DataExecutionOptions(
|
|
default=execution_options_to_model(
|
|
DataConfig.default_ingest_options()
|
|
),
|
|
),
|
|
enable_shard_locality=True,
|
|
),
|
|
run_config=RunConfigSchema(
|
|
name="test_run",
|
|
failure_config=FailureConfigSchema(
|
|
max_failures=0, controller_failure_limit=-1
|
|
),
|
|
worker_runtime_env={"type": "conda"},
|
|
checkpoint_config=CheckpointConfigSchema(
|
|
num_to_keep=None,
|
|
checkpoint_score_attribute=None,
|
|
checkpoint_score_order="max",
|
|
),
|
|
storage_path="s3://bucket/path",
|
|
storage_filesystem=None,
|
|
),
|
|
),
|
|
)
|
|
|
|
|
|
def create_mock_train_run_attempt(
|
|
attempt_id: str = "mock_attempt_id",
|
|
status: RunAttemptStatus = RunAttemptStatus.RUNNING,
|
|
end_time_ns: Optional[int] = None,
|
|
run_id: Optional[str] = None,
|
|
worker_status: Optional[ActorStatus] = ActorStatus.ALIVE,
|
|
status_detail: Optional[str] = None,
|
|
):
|
|
worker = TrainWorker(
|
|
world_rank=0,
|
|
local_rank=0,
|
|
node_rank=0,
|
|
actor_id=uuid.uuid4().hex,
|
|
node_id=uuid.uuid4().hex,
|
|
node_ip="127.0.0.1",
|
|
pid=1234,
|
|
gpu_ids=[0],
|
|
status=worker_status,
|
|
resources=TrainResources(resources={"CPU": 1}),
|
|
log_file_path="/tmp/ray/session_xxx/logs/train/ray-train-app-worker.log",
|
|
)
|
|
|
|
return TrainRunAttempt(
|
|
schema_version=0,
|
|
attempt_id=attempt_id,
|
|
run_id=run_id or _RUN_ID,
|
|
status=status,
|
|
status_detail=status_detail,
|
|
start_time_ns=time.time_ns(),
|
|
resources=[TrainResources(resources={"CPU": 1})],
|
|
workers=[worker],
|
|
end_time_ns=end_time_ns,
|
|
)
|
|
|
|
|
|
def create_dummy_run_context(**kwargs: dict) -> TrainRunContext:
|
|
"""Create a standardized TrainRunContext for testing.
|
|
|
|
Args:
|
|
**kwargs: Optional overrides for the default configuration.
|
|
|
|
Returns:
|
|
TrainRunContext: A standardized TrainRunContext instance for testing.
|
|
"""
|
|
from ray.train import BackendConfig, DataConfig
|
|
from ray.train.v2._internal.execution.context import TrainRunContext
|
|
from ray.train.v2.api.config import RunConfig, ScalingConfig
|
|
|
|
config = dict(
|
|
run_config=RunConfig(name="test"),
|
|
train_loop_config={},
|
|
scaling_config=ScalingConfig(num_workers=1),
|
|
backend_config=BackendConfig(),
|
|
dataset_config=DataConfig(),
|
|
)
|
|
config.update(kwargs)
|
|
return TrainRunContext(**config)
|
|
|
|
|
|
class DummyTrainContext(TrainContext):
|
|
"""A dummy TrainContext subclass for testing."""
|
|
|
|
def __init__(self):
|
|
self.train_run_context = create_dummy_run_context()
|
|
self.distributed_context = DistributedContext(
|
|
world_rank=0,
|
|
world_size=1,
|
|
local_rank=0,
|
|
local_world_size=1,
|
|
node_rank=0,
|
|
)
|
|
# Mock everything else since we don't need the actual functionality
|
|
self.execution_context = MagicMock()
|
|
self.storage_context = MagicMock()
|
|
self.dataset_shards = {}
|
|
|
|
def get_run_config(self):
|
|
return self.train_run_context.run_config
|
|
|
|
|
|
def create_dummy_train_context() -> TrainContext:
|
|
"""Create a standardized TrainContext for testing.
|
|
|
|
Returns:
|
|
TrainContext: A standardized TrainContext instance for testing.
|
|
"""
|
|
return DummyTrainContext()
|
|
|
|
|
|
def create_dummy_training_reports(
|
|
num_results: int,
|
|
storage_context: StorageContext,
|
|
include_metrics: bool = True,
|
|
include_validation: bool = False,
|
|
starting_checkpoint_index: int = 0,
|
|
) -> List[_TrainingReport]:
|
|
training_results = []
|
|
for i in range(num_results):
|
|
metrics = {"score": i} if include_metrics else {}
|
|
validation = (
|
|
ValidationTaskConfig(fn_kwargs={"arg": i}) if include_validation else False
|
|
)
|
|
checkpoint_path = os.path.join(
|
|
storage_context.experiment_fs_path,
|
|
f"checkpoint_{starting_checkpoint_index + i}",
|
|
)
|
|
os.makedirs(checkpoint_path, exist_ok=True)
|
|
training_results.append(
|
|
_TrainingReport(
|
|
checkpoint=Checkpoint(
|
|
path=Path(checkpoint_path).as_posix(),
|
|
filesystem=storage_context.storage_filesystem,
|
|
),
|
|
metrics=metrics,
|
|
validation=validation,
|
|
)
|
|
)
|
|
return training_results
|