Mentors4EDU's picture
Remove push_to_hub.py; generalise cloud scripts (no account-specific defaults)
4b08de9 verified
Raw History Blame Contribute Delete
3.18 kB
param(
[string]$Region = "nyc2",
[string]$Size = "gpu-mi325x8-2048gb",
[string]$Image = "gpu-amd-base",
[string]$Name = "riemann-node",
[Parameter(Mandatory = $true)][string]$SshKeyName,
[Parameter(Mandatory = $true)][string]$StagingRepo,
[switch]$StopAfterStage2
)
$ErrorActionPreference = "Stop"
$proj = Split-Path -Parent $PSScriptRoot
$sessionFiles = "$env:RIEMANN_ROOT\cloud_state"
New-Item -ItemType Directory -Force -Path $sessionFiles | Out-Null
$existing = doctl compute droplet list --format ID,Name,PublicIPv4,Status --no-header | Where-Object { $_ -match "\b$Name\b" }
if ($existing) {
Write-Host "Droplet already exists: $existing"
$ip = ($existing -split '\s+')[2]
} else {
$keyId = (doctl compute ssh-key list --format ID,Name --no-header | Where-Object { $_ -match $SshKeyName } | ForEach-Object { ($_ -split '\s+')[0] })
if (-not $keyId) { throw "SSH key '$SshKeyName' not found on the account" }
Write-Host "Creating $Size in $Region with image $Image ..."
doctl compute droplet create $Name --region $Region --size $Size --image $Image --ssh-keys $keyId --wait --format ID,Name,PublicIPv4,Status | Out-Host
$ip = (doctl compute droplet list --format Name,PublicIPv4 --no-header | Where-Object { $_ -match "\b$Name\b" } | ForEach-Object { ($_ -split '\s+')[1] })
}
Write-Host "Droplet IP: $ip"
Set-Content -Path "$sessionFiles\node_ip.txt" -Value $ip
Write-Host "Waiting for SSH ..."
$deadline = (Get-Date).AddMinutes(10)
do {
$ok = $false
try { ssh -o StrictHostKeyChecking=accept-new -o ConnectTimeout=8 -o BatchMode=yes root@$ip "echo ready" 2>$null | Select-String ready | ForEach-Object { $ok = $true } } catch {}
if (-not $ok) { Start-Sleep 10 }
} while (-not $ok -and (Get-Date) -lt $deadline)
if (-not $ok) { throw "SSH not reachable after 10 minutes" }
Write-Host "Packaging project ..."
$bundle = "$env:TEMP\riemann_bundle.tar.gz"
Push-Location $proj
tar -czf $bundle riemann scripts configs cloud requirements-cuda.txt LICENSE
Pop-Location
Write-Host "Shipping code and HF token ..."
ssh root@$ip "mkdir -p /workspace/riemann ~/.cache/huggingface"
scp -q $bundle root@${ip}:/workspace/riemann_bundle.tar.gz
scp -q "$env:USERPROFILE\.cache\huggingface\token" root@${ip}:~/.cache/huggingface/token
ssh root@$ip "cd /workspace/riemann && tar -xzf /workspace/riemann_bundle.tar.gz && sed -i 's/\r$//' cloud/*.sh cloud/*.py && chmod +x cloud/bootstrap.sh && chmod 600 ~/.cache/huggingface/token && ls"
$stopFlag = if ($StopAfterStage2) { "STOP_AFTER_STAGE2=1" } else { "STOP_AFTER_STAGE2=0" }
Write-Host "Starting bootstrap under nohup ($stopFlag) ..."
ssh root@$ip "cd /workspace/riemann && mkdir -p /workspace/RiemannModel/logs && STAGING_REPO=$StagingRepo $stopFlag nohup bash cloud/bootstrap.sh > /workspace/RiemannModel/logs/bootstrap.log 2>&1 &"
Start-Sleep 5
ssh root@$ip "tail -5 /workspace/RiemannModel/logs/bootstrap.log"
Write-Host ""
Write-Host "Launched. Follow with: ssh root@$ip tail -f /workspace/RiemannModel/logs/bootstrap.log"
Write-Host "Destroy with: doctl compute droplet delete $Name --force"