Fabian
82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Install missing ROCm math libraries for PyTorch build on BC-250."""
|
|
import paramiko
|
|
import time
|
|
|
|
ssh = paramiko.SSHClient()
|
|
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
ssh.connect('192.168.178.150', username='fabian', key_filename=r'C:\Users\fabia\.ssh\id_ed25519')
|
|
|
|
def run(cmd, timeout=300, desc=""):
|
|
if desc:
|
|
print(f"\n{'='*60}")
|
|
print(f" {desc}")
|
|
print(f"{'='*60}")
|
|
print(f"$ {cmd}")
|
|
_, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
|
|
out = stdout.read().decode()
|
|
err = stderr.read().decode()
|
|
rc = stdout.channel.recv_exit_status()
|
|
if out.strip():
|
|
lines = out.strip().split('\n')
|
|
if len(lines) > 60:
|
|
print(f" ... ({len(lines)} lines, showing last 60)")
|
|
print('\n'.join(lines[-60:]))
|
|
else:
|
|
print(out.strip())
|
|
if err.strip():
|
|
lines = err.strip().split('\n')
|
|
show = lines[-20:] if len(lines) > 20 else lines
|
|
print(f"STDERR: {chr(10).join(show)}")
|
|
print(f" Exit code: {rc}")
|
|
return rc, out, err
|
|
|
|
# Find all available ROCm packages
|
|
run("pacman -Ss rocm | grep -E '^cachyos|^extra|^core' | head -40",
|
|
desc="Available ROCm packages")
|
|
|
|
# Install all ROCm math/compute libraries needed by PyTorch
|
|
run("sudo pacman -S --needed --noconfirm "
|
|
"hiprand rocrand "
|
|
"hipblas rocblas "
|
|
"hipfft rocfft "
|
|
"hipsparse rocsparse "
|
|
"hipsolver rocsolver "
|
|
"miopen-hip "
|
|
"rocprim hipcub "
|
|
"rocthrust "
|
|
"rccl "
|
|
"hipblaslt "
|
|
"roctracer "
|
|
"2>&1 | tail -40",
|
|
desc="Install ROCm math libraries",
|
|
timeout=600)
|
|
|
|
# Verify hiprand is now available
|
|
run("find /opt/rocm -name 'hiprandConfig.cmake' -o -name 'hiprand-config.cmake' 2>/dev/null | head -5",
|
|
desc="Verify hiprand cmake config")
|
|
|
|
# Check all libraries
|
|
run("ls /opt/rocm/lib/libhiprand.so /opt/rocm/lib/librocblas.so /opt/rocm/lib/libhipblas.so /opt/rocm/lib/librocfft.so /opt/rocm/lib/libMIOpen.so 2>&1",
|
|
desc="Verify key libraries exist")
|
|
|
|
# Clean the failed build and restart
|
|
run("bash -c 'source ~/comfyui-env/bin/activate && cd ~/pytorch && python3 setup.py clean 2>&1 | tail -5'",
|
|
desc="Clean failed build")
|
|
|
|
# Restart build
|
|
run("rm -f /home/fabian/pytorch_build.log", desc="Clean old log")
|
|
run("bash -c 'nohup bash /home/fabian/build_pytorch.sh </dev/null >/dev/null 2>&1 & echo PID=$!'",
|
|
desc="Restart PyTorch build")
|
|
|
|
time.sleep(15)
|
|
run("pgrep -fa 'setup.py\\|cmake\\|ninja' | head -10",
|
|
desc="Verify build restarted")
|
|
|
|
time.sleep(45)
|
|
run("tail -40 /home/fabian/pytorch_build.log 2>/dev/null",
|
|
desc="Build progress")
|
|
|
|
ssh.close()
|
|
print("\nDone — ROCm libs installed and build restarted.")
|