wehub-resource-sync
75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
import sys
|
|
import time
|
|
|
|
import pytest
|
|
import requests
|
|
|
|
import ray
|
|
from ray.train import RunConfig, ScalingConfig
|
|
from ray.train.torch import TorchTrainer
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_8_cpus(monkeypatch):
|
|
monkeypatch.setenv("RAY_TRAIN_ENABLE_STATE_TRACKING", "1")
|
|
address_info = ray.init(num_cpus=8)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
def test_get_train_runs(ray_start_8_cpus):
|
|
def train_func():
|
|
print("Training Starts")
|
|
time.sleep(0.5)
|
|
|
|
datasets = {"train": ray.data.range(100), "val": ray.data.range(100)}
|
|
|
|
trainer = TorchTrainer(
|
|
train_func,
|
|
run_config=RunConfig(name="my_train_run", storage_path="/tmp/cluster_storage"),
|
|
scaling_config=ScalingConfig(num_workers=4, use_gpu=False),
|
|
datasets=datasets,
|
|
)
|
|
trainer.fit()
|
|
|
|
# Call the train run api
|
|
url = ray._private.worker.get_dashboard_url()
|
|
resp = requests.get("http://" + url + "/api/train/v2/runs")
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert len(body["train_runs"]) == 1
|
|
assert body["train_runs"][0]["name"] == "my_train_run"
|
|
assert len(body["train_runs"][0]["workers"]) == 4
|
|
|
|
|
|
def test_add_actor_status(ray_start_8_cpus):
|
|
from ray.train._internal.state.schema import ActorStatusEnum
|
|
|
|
def check_actor_status(expected_actor_status):
|
|
url = ray._private.worker.get_dashboard_url()
|
|
resp = requests.get("http://" + url + "/api/train/v2/runs")
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
|
|
for worker_info in body["train_runs"][0]["workers"]:
|
|
assert worker_info["status"] == expected_actor_status
|
|
|
|
def train_func():
|
|
print("Training Starts")
|
|
time.sleep(0.5)
|
|
check_actor_status(expected_actor_status=ActorStatusEnum.ALIVE)
|
|
|
|
trainer = TorchTrainer(
|
|
train_func,
|
|
run_config=RunConfig(name="my_train_run", storage_path="/tmp/cluster_storage"),
|
|
scaling_config=ScalingConfig(num_workers=4, use_gpu=False),
|
|
)
|
|
trainer.fit()
|
|
|
|
check_actor_status(expected_actor_status=ActorStatusEnum.DEAD)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|