ray/release/golden_notebook_tests/workloads/util.py

from pathlib import Path
import importlib.util
import ray
import time


def import_and_execute_test_script(relative_path_to_test_script: str):
    """Imports and executes a module from a path relative to Ray repo root."""
    # get the ray folder
    ray_path = next(
        x for x in Path(__file__).resolve().parents if str(x).endswith("/ray")
    )
    notebook_path = ray_path.joinpath(relative_path_to_test_script)
    assert notebook_path.exists()

    spec = importlib.util.spec_from_file_location("notebook_test", notebook_path)
    notebook_test_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(notebook_test_module)


def wait_for_cluster_client(
    num_nodes: int, max_time_s: int, feedback_interval_s: int = 10
):
    assert ray.is_initialized()
    curr_nodes = 0
    start = time.time()
    next_feedback = start
    max_time = start + max_time_s
    while not curr_nodes >= num_nodes:
        now = time.time()

        if now >= max_time:
            raise RuntimeError(
                f"Maximum wait time reached, but only "
                f"{curr_nodes}/{num_nodes} nodes came up. Aborting."
            )

        if now >= next_feedback:
            passed = now - start
            print(
                f"Waiting for more nodes to come up: "
                f"{curr_nodes}/{num_nodes} "
                f"({passed:.0f} seconds passed)"
            )
            next_feedback = now + feedback_interval_s

        time.sleep(5)
        curr_nodes = len(ray.nodes())

    passed = time.time() - start
    print(
        f"Cluster is up: {curr_nodes}/{num_nodes} nodes online after "
        f"{passed:.0f} seconds"
    )