chore: import upstream snapshot with attribution

python/ray/train/tests/test_train_head.py

@@ -0,0 +1,74 @@
+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__]))