Fabian
78 lines
2.7 KiB
Python
78 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Try pre-built PyTorch ROCm from CachyOS repos, test 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) > 80:
|
|
print(f" ... ({len(lines)} lines, showing last 80)")
|
|
print('\n'.join(lines[-80:]))
|
|
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
|
|
|
|
# Check what architectures the pre-built packages support
|
|
run("pacman -Si python-pytorch-rocm 2>&1 | head -20",
|
|
desc="Pre-built PyTorch ROCm package info")
|
|
|
|
run("pacman -Si python-pytorch-opt-rocm 2>&1 | head -20",
|
|
desc="Pre-built PyTorch Opt ROCm package info")
|
|
|
|
# Install the pre-built package (system-wide, venv will pick it up via --system-site-packages)
|
|
run("sudo pacman -S --needed --noconfirm python-pytorch-rocm 2>&1 | tail -30",
|
|
desc="Install pre-built PyTorch ROCm",
|
|
timeout=600)
|
|
|
|
# Test if it works in venv
|
|
run("""bash -c 'source ~/comfyui-env/bin/activate && \
|
|
HSA_OVERRIDE_GFX_VERSION=10.1.0 \
|
|
HIP_VISIBLE_DEVICES=0 \
|
|
HSA_ENABLE_SDMA=0 \
|
|
python3 -c "
|
|
import torch
|
|
print(f\\"PyTorch version: {torch.__version__}\\")
|
|
print(f\\"HIP version: {torch.version.hip}\\")
|
|
print(f\\"CUDA available (HIP): {torch.cuda.is_available()}\\")
|
|
if torch.cuda.is_available():
|
|
print(f\\"Device count: {torch.cuda.device_count()}\\")
|
|
print(f\\"Device name: {torch.cuda.get_device_name(0)}\\")
|
|
print(f\\"Device arch: {torch.cuda.get_device_capability(0)}\\")
|
|
# Try a simple tensor operation on GPU
|
|
t = torch.randn(4, 4, device=\\"cuda\\")
|
|
print(f\\"Tensor device: {t.device}\\")
|
|
print(f\\"Tensor sum: {t.sum().item():.4f}\\")
|
|
# Try matmul
|
|
a = torch.randn(64, 64, device=\\"cuda\\")
|
|
b = torch.randn(64, 64, device=\\"cuda\\")
|
|
c = torch.matmul(a, b)
|
|
print(f\\"Matmul result shape: {c.shape}\\")
|
|
print(\\"GPU COMPUTE: WORKING\\")
|
|
else:
|
|
print(\\"CUDA/HIP NOT AVAILABLE\\")
|
|
" 2>&1'""",
|
|
desc="Test PyTorch on GPU",
|
|
timeout=120)
|
|
|
|
ssh.close()
|
|
print("\nDone.")
|