mirror of
https://github.com/vale981/ray
synced 2025-03-13 22:56:38 -04:00
549527687f
This PR migrates scalability tests to the new infra. I had to copy the benchmarks folder to the release folder to make it work. I will remove some unnecessary files (e.g., benchmark.yaml or wait_for_cluster file) Alternatively we can support a different path than /release from the tool, but I think this way is cleaner. I am open to suggestion though cc @krfricke
24 lines
510 B
Python
24 lines
510 B
Python
import click
|
|
import ray
|
|
import time
|
|
|
|
|
|
def num_alive_nodes():
|
|
n = 0
|
|
for node in ray.nodes():
|
|
if node["Alive"]:
|
|
n += 1
|
|
return n
|
|
|
|
|
|
@click.command()
|
|
@click.option("--num-nodes", required=True, type=int, help="The target number of nodes")
|
|
def wait_cluster(num_nodes: int):
|
|
ray.init(address="auto")
|
|
while num_alive_nodes() != num_nodes:
|
|
print(f"Waiting for nodes: {num_alive_nodes()}/{num_nodes}")
|
|
time.sleep(5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
wait_cluster()
|