wehub-resource-sync
108 lines
4.0 KiB
Python
108 lines
4.0 KiB
Python
# ABOUTME: Launches nvidia-smi monitoring on GPU worker nodes as they join the cluster.
|
|
# ABOUTME: Polls for new GPU nodes and starts nvidia-smi dmon on each via long-lived Ray actors.
|
|
|
|
import os
|
|
import threading
|
|
import time
|
|
|
|
import ray
|
|
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
|
|
|
|
|
# Module-level list of NvidiaSmiActor handles. Holding them at module scope
|
|
# keeps the actors alive for the lifetime of the driver process (Ray drops
|
|
# actors whose last handle is GC'd). Local-in-function holding broke when
|
|
# the launcher thread exited and the local frame was reclaimed.
|
|
_actor_handles = []
|
|
|
|
|
|
@ray.remote(num_cpus=0, num_gpus=0)
|
|
class _NvidiaSmiActor:
|
|
"""Long-lived actor that owns an nvidia-smi dmon subprocess on its node.
|
|
|
|
Holding nvidia-smi inside an actor (rather than launching it from a Ray
|
|
task) keeps the parent process alive for the lifetime of the actor handle.
|
|
Tasks return immediately and the worker is reaped after ~3 min idle,
|
|
taking nvidia-smi with it; the actor never goes idle, so the subprocess
|
|
survives the whole run.
|
|
"""
|
|
|
|
def __init__(self, outdir):
|
|
import subprocess
|
|
|
|
os.makedirs(outdir, exist_ok=True)
|
|
node_ip = os.environ.get("ANYSCALE_NODE_IP", "unknown").replace(".", "_")
|
|
outfile = f"{outdir}/gpu_usage_{node_ip}.txt"
|
|
self._proc = subprocess.Popen(
|
|
["stdbuf", "-oL", "nvidia-smi", "dmon", "-s", "u", "-o", "T"],
|
|
stdout=open(outfile, "w"),
|
|
stderr=subprocess.STDOUT,
|
|
)
|
|
|
|
def get_pid(self):
|
|
return self._proc.pid
|
|
|
|
|
|
def _gpu_monitor_loop(outdir, num_gpu_nodes):
|
|
"""Poll for GPU nodes and launch nvidia-smi on each as it joins.
|
|
|
|
Monitors continuously until num_gpu_nodes are found or no new nodes
|
|
appear for 60 seconds, whichever comes first. This avoids blocking
|
|
forever when the cluster has fewer GPU nodes than expected.
|
|
"""
|
|
monitored_node_ids = set()
|
|
stale_polls = 0
|
|
max_stale_polls = 30 # 30 * 2s = 60s with no new nodes
|
|
|
|
while len(monitored_node_ids) < num_gpu_nodes:
|
|
gpu_nodes = [
|
|
n for n in ray.nodes() if n["Alive"] and n["Resources"].get("GPU", 0) > 0
|
|
]
|
|
new_nodes = [n for n in gpu_nodes if n["NodeID"] not in monitored_node_ids]
|
|
if not new_nodes:
|
|
stale_polls += 1
|
|
if monitored_node_ids and stale_polls >= max_stale_polls:
|
|
print(
|
|
f"nvidia-smi: no new GPU nodes for {max_stale_polls * 2}s, "
|
|
f"proceeding with {len(monitored_node_ids)}/{num_gpu_nodes}"
|
|
)
|
|
break
|
|
else:
|
|
stale_polls = 0
|
|
for node in new_nodes:
|
|
try:
|
|
handle = _NvidiaSmiActor.options(
|
|
scheduling_strategy=NodeAffinitySchedulingStrategy(
|
|
node_id=node["NodeID"], soft=False
|
|
)
|
|
).remote(outdir)
|
|
child_pid = ray.get(handle.get_pid.remote())
|
|
_actor_handles.append(handle)
|
|
monitored_node_ids.add(node["NodeID"])
|
|
print(
|
|
f"Started nvidia-smi on GPU node {node['NodeManagerAddress']} "
|
|
f"(pid={child_pid}, {len(monitored_node_ids)}/{num_gpu_nodes})"
|
|
)
|
|
except Exception as e:
|
|
monitored_node_ids.add(node["NodeID"])
|
|
print(
|
|
f"Failed to start nvidia-smi on {node['NodeManagerAddress']}: {e}"
|
|
)
|
|
time.sleep(2)
|
|
if len(monitored_node_ids) >= num_gpu_nodes:
|
|
print(f"nvidia-smi monitoring active on all {num_gpu_nodes} GPU nodes")
|
|
|
|
|
|
def start(outdir, num_gpu_nodes):
|
|
"""Start nvidia-smi monitoring on GPU nodes in a background thread.
|
|
|
|
Args:
|
|
outdir: Shared storage directory for output files.
|
|
num_gpu_nodes: Expected number of GPU nodes to monitor.
|
|
"""
|
|
thread = threading.Thread(
|
|
target=_gpu_monitor_loop, args=(outdir, num_gpu_nodes), daemon=True
|
|
)
|
|
thread.start()
|
|
return thread
|