wehub-resource-sync
105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
import queue
|
|
import time
|
|
from unittest.mock import create_autospec
|
|
|
|
import pytest
|
|
|
|
from ray.actor import ActorHandle
|
|
from ray.train.v2._internal.constants import ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR
|
|
from ray.train.v2._internal.execution.context import (
|
|
DistributedContext,
|
|
TrainRunContext,
|
|
get_train_context,
|
|
)
|
|
from ray.train.v2._internal.execution.preemption import PreemptionInfo
|
|
from ray.train.v2._internal.execution.storage import StorageContext
|
|
from ray.train.v2._internal.execution.worker_group.worker import RayTrainWorker
|
|
from ray.train.v2._internal.util import ObjectRefWrapper
|
|
|
|
|
|
@pytest.mark.parametrize("created_nested_threads", [True, False])
|
|
def test_worker_finished_after_all_threads_finish(monkeypatch, created_nested_threads):
|
|
# Disable this to avoid TypeError from logging MagicMock
|
|
monkeypatch.setenv(ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR, False)
|
|
|
|
# Initialize RayTrainWorker state
|
|
worker = RayTrainWorker()
|
|
worker.init_train_context(
|
|
train_run_context=create_autospec(TrainRunContext, instance=True),
|
|
distributed_context=DistributedContext(
|
|
world_rank=0,
|
|
world_size=1,
|
|
local_rank=0,
|
|
local_world_size=1,
|
|
node_rank=0,
|
|
),
|
|
synchronization_actor=create_autospec(ActorHandle, instance=True),
|
|
storage_context=create_autospec(StorageContext, instance=True),
|
|
worker_callbacks=[],
|
|
controller_actor=create_autospec(ActorHandle, instance=True),
|
|
)
|
|
global_queue = queue.Queue()
|
|
|
|
def train_fn():
|
|
tc = get_train_context()
|
|
|
|
def target():
|
|
# Intentionally sleep longer than poll interval to test that we wait
|
|
# for nested threads to finish
|
|
time.sleep(0.1)
|
|
global_queue.put("nested")
|
|
|
|
if created_nested_threads:
|
|
tc.checkpoint_upload_threadpool.submit(target)
|
|
else:
|
|
global_queue.put("main")
|
|
|
|
# Run train fn and wait for it to finish
|
|
train_fn_ref = create_autospec(ObjectRefWrapper, instance=True)
|
|
train_fn_ref.get.return_value = train_fn
|
|
worker.run_train_fn(train_fn_ref)
|
|
while worker.poll_status().running:
|
|
time.sleep(0.01)
|
|
|
|
# Verify queue contents
|
|
queue_contents = []
|
|
while not global_queue.empty():
|
|
queue_contents.append(global_queue.get())
|
|
if created_nested_threads:
|
|
assert queue_contents == ["nested"]
|
|
else:
|
|
assert queue_contents == ["main"]
|
|
|
|
|
|
def test_mark_preempt_stores_info(monkeypatch):
|
|
"""mark_preempt stores the signal in the worker's PreemptionContext."""
|
|
# Disable this to avoid TypeError from logging MagicMock
|
|
monkeypatch.setenv(ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR, False)
|
|
|
|
worker = RayTrainWorker()
|
|
worker.init_train_context(
|
|
train_run_context=create_autospec(TrainRunContext, instance=True),
|
|
distributed_context=DistributedContext(
|
|
world_rank=0,
|
|
world_size=1,
|
|
local_rank=0,
|
|
local_world_size=1,
|
|
node_rank=0,
|
|
),
|
|
synchronization_actor=create_autospec(ActorHandle, instance=True),
|
|
storage_context=create_autospec(StorageContext, instance=True),
|
|
worker_callbacks=[],
|
|
controller_actor=create_autospec(ActorHandle, instance=True),
|
|
)
|
|
|
|
info = PreemptionInfo(deadline_ms=30_000, preempted_node_to_ranks={"node-a": [0]})
|
|
worker.mark_preempt(info)
|
|
|
|
assert get_train_context().preemption_context.get() is info
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|