345 lines
14 KiB
C#
345 lines
14 KiB
C#
|
using System;
|
||
|
|
using System.IO;
|
||
|
|
using System.Reflection;
|
||
|
|
using System.Text;
|
||
|
|
using System.Threading.Tasks;
|
||
|
|
using UnifiedFarmLauncher.Models;
|
||
|
|
|
||
|
|
namespace UnifiedFarmLauncher.Services
|
||
|
|
{
|
||
|
|
public class WorkerControllerService
|
||
|
|
{
|
||
|
|
private readonly SshService _sshService;
|
||
|
|
private readonly ConfigService _configService;
|
||
|
|
private byte[]? _controllerScriptBytes;
|
||
|
|
private byte[]? _attachHelperScriptBytes;
|
||
|
|
|
||
|
|
public WorkerControllerService(SshService sshService, ConfigService configService)
|
||
|
|
{
|
||
|
|
_sshService = sshService;
|
||
|
|
_configService = configService;
|
||
|
|
}
|
||
|
|
|
||
|
|
private byte[] GetControllerScriptBytes()
|
||
|
|
{
|
||
|
|
if (_controllerScriptBytes != null)
|
||
|
|
return _controllerScriptBytes;
|
||
|
|
|
||
|
|
var assembly = Assembly.GetExecutingAssembly();
|
||
|
|
var resourceName = "UnifiedFarmLauncher.Scripts.remote_worker_controller.ps1";
|
||
|
|
|
||
|
|
using var stream = assembly.GetManifestResourceStream(resourceName);
|
||
|
|
if (stream == null)
|
||
|
|
throw new InvalidOperationException($"Resource {resourceName} not found");
|
||
|
|
|
||
|
|
using var reader = new BinaryReader(stream);
|
||
|
|
_controllerScriptBytes = reader.ReadBytes((int)stream.Length);
|
||
|
|
return _controllerScriptBytes;
|
||
|
|
}
|
||
|
|
|
||
|
|
private byte[] GetAttachHelperScriptBytes()
|
||
|
|
{
|
||
|
|
if (_attachHelperScriptBytes != null)
|
||
|
|
return _attachHelperScriptBytes;
|
||
|
|
|
||
|
|
var assembly = Assembly.GetExecutingAssembly();
|
||
|
|
var resourceName = "UnifiedFarmLauncher.Scripts.remote_worker_attach.ps1";
|
||
|
|
|
||
|
|
using var stream = assembly.GetManifestResourceStream(resourceName);
|
||
|
|
if (stream == null)
|
||
|
|
throw new InvalidOperationException($"Resource {resourceName} not found");
|
||
|
|
|
||
|
|
using var reader = new BinaryReader(stream);
|
||
|
|
_attachHelperScriptBytes = reader.ReadBytes((int)stream.Length);
|
||
|
|
return _attachHelperScriptBytes;
|
||
|
|
}
|
||
|
|
|
||
|
|
public async Task DeployControllerAsync(WorkerConfig worker)
|
||
|
|
{
|
||
|
|
var controllerBase64 = Convert.ToBase64String(GetControllerScriptBytes());
|
||
|
|
var script = $@"
|
||
|
|
$ProgressPreference = 'SilentlyContinue'
|
||
|
|
$dataRoot = Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'
|
||
|
|
New-Item -ItemType Directory -Path $dataRoot -Force | Out-Null
|
||
|
|
$controllerPath = Join-Path $dataRoot 'controller.ps1'
|
||
|
|
[IO.File]::WriteAllBytes($controllerPath, [Convert]::FromBase64String('{controllerBase64}'))
|
||
|
|
";
|
||
|
|
|
||
|
|
await _sshService.ExecuteRemoteCommandAsync(worker, $"powershell -NoLogo -NoProfile -NonInteractive -OutputFormat Text -ExecutionPolicy Bypass -EncodedCommand {Convert.ToBase64String(Encoding.Unicode.GetBytes(script))}");
|
||
|
|
}
|
||
|
|
|
||
|
|
public async Task DeployAttachHelperAsync(WorkerConfig worker)
|
||
|
|
{
|
||
|
|
var helperBase64 = Convert.ToBase64String(GetAttachHelperScriptBytes());
|
||
|
|
var script = $@"
|
||
|
|
$ProgressPreference = 'SilentlyContinue'
|
||
|
|
$dataRoot = Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'
|
||
|
|
New-Item -ItemType Directory -Path $dataRoot -Force | Out-Null
|
||
|
|
$attachPath = Join-Path $dataRoot 'attach-helper.ps1'
|
||
|
|
[IO.File]::WriteAllBytes($attachPath, [Convert]::FromBase64String('{helperBase64}'))
|
||
|
|
";
|
||
|
|
|
||
|
|
await _sshService.ExecuteRemoteCommandAsync(worker, $"powershell -NoLogo -NoProfile -NonInteractive -OutputFormat Text -ExecutionPolicy Bypass -EncodedCommand {Convert.ToBase64String(Encoding.Unicode.GetBytes(script))}");
|
||
|
|
}
|
||
|
|
|
||
|
|
public string GenerateSheepItPayload(WorkerConfig worker)
|
||
|
|
{
|
||
|
|
if (worker.WorkerTypes.SheepIt == null)
|
||
|
|
throw new InvalidOperationException("Worker does not have SheepIt configuration");
|
||
|
|
|
||
|
|
var config = _configService.Load();
|
||
|
|
var sheepIt = worker.WorkerTypes.SheepIt;
|
||
|
|
var safeKey = sheepIt.RenderKey.Replace("'", "''");
|
||
|
|
var safeUser = sheepIt.Username.Replace("'", "''");
|
||
|
|
var urls = config.GlobalSettings.SheepItJarUrls;
|
||
|
|
var urlLiteral = "@(" + string.Join(", ", Array.ConvertAll(urls.ToArray(), url => $"'{url}'")) + ")";
|
||
|
|
|
||
|
|
return $@"
|
||
|
|
$ProgressPreference = 'SilentlyContinue'
|
||
|
|
$ErrorActionPreference = 'Stop'
|
||
|
|
|
||
|
|
# Write startup message to stderr so controller can capture it
|
||
|
|
[Console]::Error.WriteLine('[SHEEPIT] Payload script starting...')
|
||
|
|
|
||
|
|
try {{
|
||
|
|
$appData = [Environment]::GetFolderPath('ApplicationData')
|
||
|
|
$sheepDir = Join-Path $appData 'sheepit'
|
||
|
|
if (-not (Test-Path $sheepDir)) {{
|
||
|
|
New-Item -Path $sheepDir -ItemType Directory -Force | Out-Null
|
||
|
|
}}
|
||
|
|
|
||
|
|
$jarPath = Join-Path $sheepDir 'sheepit-client.jar'
|
||
|
|
$urls = {urlLiteral}
|
||
|
|
$headers = @{{ 'User-Agent' = 'Mozilla/5.0' }}
|
||
|
|
|
||
|
|
if (Test-Path $jarPath) {{
|
||
|
|
Write-Host ""SheepIt client already present at $jarPath. Skipping download."" -ForegroundColor Green
|
||
|
|
}}
|
||
|
|
else {{
|
||
|
|
$downloaded = $false
|
||
|
|
|
||
|
|
foreach ($url in $urls) {{
|
||
|
|
Write-Host ""Downloading SheepIt client from $url..."" -ForegroundColor Cyan
|
||
|
|
try {{
|
||
|
|
Invoke-WebRequest -Uri $url -OutFile $jarPath -UseBasicParsing -Headers $headers
|
||
|
|
$downloaded = $true
|
||
|
|
Write-Host ""Download complete."" -ForegroundColor Green
|
||
|
|
break
|
||
|
|
}}
|
||
|
|
catch {{
|
||
|
|
Write-Host (""Download failed from {{0}}: {{1}}"" -f $url, $_.Exception.Message) -ForegroundColor Yellow
|
||
|
|
}}
|
||
|
|
}}
|
||
|
|
|
||
|
|
if (-not $downloaded) {{
|
||
|
|
throw 'Unable to download SheepIt client from any known URL.'
|
||
|
|
}}
|
||
|
|
}}
|
||
|
|
|
||
|
|
[Console]::Error.WriteLine('[SHEEPIT] Starting Java with SheepIt client...')
|
||
|
|
Set-Location $sheepDir
|
||
|
|
|
||
|
|
$javaArgs = @('-XX:+IgnoreUnrecognizedVMOptions', '-jar', $jarPath,
|
||
|
|
'-ui', 'text', '--log-stdout', '--verbose',
|
||
|
|
'-gpu', '{sheepIt.Gpu}', '-login', '{safeUser}', '-password', '{safeKey}')
|
||
|
|
|
||
|
|
try {{
|
||
|
|
& java @javaArgs
|
||
|
|
}}
|
||
|
|
catch {{
|
||
|
|
Write-Host ('Java execution error: {{0}}' -f $_.Exception.Message) -ForegroundColor Red
|
||
|
|
Write-Host ""If the error persists, try reinstalling Java (Temurin 21 recommended)."" -ForegroundColor Yellow
|
||
|
|
[Console]::Error.WriteLine(""Java execution error: $($_.Exception.Message)"")
|
||
|
|
throw
|
||
|
|
}}
|
||
|
|
}}
|
||
|
|
catch {{
|
||
|
|
$errorMsg = ('Error: {{0}}' -f $_.Exception.Message)
|
||
|
|
$stackMsg = ('Stack trace: {{0}}' -f $_.ScriptStackTrace)
|
||
|
|
Write-Host $errorMsg -ForegroundColor Red
|
||
|
|
Write-Host $stackMsg -ForegroundColor DarkRed
|
||
|
|
[Console]::Error.WriteLine($errorMsg)
|
||
|
|
[Console]::Error.WriteLine($stackMsg)
|
||
|
|
exit 1
|
||
|
|
}}
|
||
|
|
";
|
||
|
|
}
|
||
|
|
|
||
|
|
public string GenerateFlamencoPayload(WorkerConfig worker)
|
||
|
|
{
|
||
|
|
if (worker.WorkerTypes.Flamenco == null)
|
||
|
|
throw new InvalidOperationException("Worker does not have Flamenco configuration");
|
||
|
|
|
||
|
|
var flamenco = worker.WorkerTypes.Flamenco;
|
||
|
|
var drives = string.Join(", ", Array.ConvertAll(flamenco.NetworkDrives.ToArray(), d => $"'{d}'"));
|
||
|
|
var paths = string.Join(", ", Array.ConvertAll(flamenco.NetworkPaths.ToArray(), p => $"'{p.Replace("\\", "\\\\")}'"));
|
||
|
|
|
||
|
|
return $@"
|
||
|
|
$ProgressPreference = 'SilentlyContinue'
|
||
|
|
$ErrorActionPreference = 'Stop'
|
||
|
|
|
||
|
|
# Write startup message to stderr so controller can capture it
|
||
|
|
[Console]::Error.WriteLine('[FLAMENCO] Payload script starting...')
|
||
|
|
|
||
|
|
try {{
|
||
|
|
Write-Host ""Setting up network connections..."" -ForegroundColor Cyan
|
||
|
|
|
||
|
|
$drives = @({drives})
|
||
|
|
$networkPaths = @({paths})
|
||
|
|
|
||
|
|
# Disconnect all existing connections
|
||
|
|
Write-Host ""Disconnecting existing network connections..."" -ForegroundColor Yellow
|
||
|
|
foreach ($path in $networkPaths) {{ net use $path /delete /y 2>$null }}
|
||
|
|
foreach ($drive in $drives) {{ net use $drive /delete /y 2>$null }}
|
||
|
|
Write-Host ""All network connections cleared."" -ForegroundColor Green
|
||
|
|
|
||
|
|
# Connect to network shares (simplified - credentials should be stored securely)
|
||
|
|
Write-Host ""Establishing network connections..."" -ForegroundColor Cyan
|
||
|
|
# TODO: Add credential handling for network shares
|
||
|
|
|
||
|
|
# Start worker
|
||
|
|
Write-Host ""Starting Flamenco worker..."" -ForegroundColor Cyan
|
||
|
|
Set-Location '{flamenco.WorkerPath}'
|
||
|
|
if (Test-Path 'flamenco-worker.exe') {{
|
||
|
|
Write-Host ""Running flamenco-worker.exe..."" -ForegroundColor Green
|
||
|
|
$workerProcess = Start-Process -FilePath '.\flamenco-worker.exe' -NoNewWindow -PassThru -Wait
|
||
|
|
$exitCode = $workerProcess.ExitCode
|
||
|
|
Write-Host ""Flamenco worker process has terminated with exit code: $exitCode"" -ForegroundColor Yellow
|
||
|
|
exit $exitCode
|
||
|
|
}} else {{
|
||
|
|
Write-Host ""Error: flamenco-worker.exe not found in {flamenco.WorkerPath}"" -ForegroundColor Red
|
||
|
|
[Console]::Error.WriteLine(""Error: flamenco-worker.exe not found in {flamenco.WorkerPath}"")
|
||
|
|
exit 1
|
||
|
|
}}
|
||
|
|
}}
|
||
|
|
catch {{
|
||
|
|
$errorMsg = ('Error: {{0}}' -f $_.Exception.Message)
|
||
|
|
$stackMsg = ('Stack trace: {{0}}' -f $_.ScriptStackTrace)
|
||
|
|
Write-Host $errorMsg -ForegroundColor Red
|
||
|
|
Write-Host $stackMsg -ForegroundColor DarkRed
|
||
|
|
[Console]::Error.WriteLine($errorMsg)
|
||
|
|
[Console]::Error.WriteLine($stackMsg)
|
||
|
|
exit 1
|
||
|
|
}}
|
||
|
|
";
|
||
|
|
}
|
||
|
|
|
||
|
|
public async Task StartWorkerAsync(WorkerConfig worker, string workerType)
|
||
|
|
{
|
||
|
|
await DeployControllerAsync(worker);
|
||
|
|
|
||
|
|
string payloadScript;
|
||
|
|
if (workerType == "sheepit")
|
||
|
|
{
|
||
|
|
payloadScript = GenerateSheepItPayload(worker);
|
||
|
|
}
|
||
|
|
else if (workerType == "flamenco")
|
||
|
|
{
|
||
|
|
payloadScript = GenerateFlamencoPayload(worker);
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
throw new ArgumentException($"Unknown worker type: {workerType}", nameof(workerType));
|
||
|
|
}
|
||
|
|
|
||
|
|
var payloadBase64 = Convert.ToBase64String(Encoding.Unicode.GetBytes(payloadScript));
|
||
|
|
var remoteBasePath = await _sshService.GetWorkerBasePathAsync(worker);
|
||
|
|
|
||
|
|
var ensureScript = $@"
|
||
|
|
$ProgressPreference = 'SilentlyContinue'
|
||
|
|
$params = ConvertFrom-Json ([Text.Encoding]::Unicode.GetString([Convert]::FromBase64String('{Convert.ToBase64String(Encoding.Unicode.GetBytes($@"{{""WorkerName"":""{worker.Name}"",""WorkerType"":""{workerType}"",""PayloadBase64"":""{payloadBase64}""}}"@))}')))
|
||
|
|
$workerName = $params.WorkerName
|
||
|
|
$workerType = $params.WorkerType
|
||
|
|
$payloadBase64 = $params.PayloadBase64
|
||
|
|
$dataRoot = Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'
|
||
|
|
$instanceRoot = Join-Path (Join-Path $dataRoot $workerType) $workerName
|
||
|
|
$logsRoot = Join-Path $instanceRoot 'logs'
|
||
|
|
$stateRoot = Join-Path $instanceRoot 'state'
|
||
|
|
New-Item -ItemType Directory -Path $logsRoot -Force | Out-Null
|
||
|
|
New-Item -ItemType Directory -Path $stateRoot -Force | Out-Null
|
||
|
|
$logPath = Join-Path $logsRoot 'worker.log'
|
||
|
|
$commandPath = Join-Path $stateRoot 'commands.txt'
|
||
|
|
$payloadPath = Join-Path $stateRoot 'payload.ps1'
|
||
|
|
$payloadBase64Path = Join-Path $stateRoot 'payload.b64'
|
||
|
|
if (-not (Test-Path $logPath)) {{ New-Item -Path $logPath -ItemType File -Force | Out-Null }}
|
||
|
|
if (-not (Test-Path $commandPath)) {{ New-Item -Path $commandPath -ItemType File -Force | Out-Null }}
|
||
|
|
[IO.File]::WriteAllText($payloadBase64Path, $payloadBase64, [System.Text.Encoding]::UTF8)
|
||
|
|
$metaPath = Join-Path $instanceRoot 'state\worker-info.json'
|
||
|
|
$controllerPath = Join-Path $dataRoot 'controller.ps1'
|
||
|
|
|
||
|
|
if (-not (Test-Path $controllerPath)) {{
|
||
|
|
throw ""Controller missing at $controllerPath""
|
||
|
|
}}
|
||
|
|
|
||
|
|
$shouldStart = $true
|
||
|
|
if (Test-Path $metaPath) {{
|
||
|
|
try {{
|
||
|
|
$meta = Get-Content $metaPath -Raw | ConvertFrom-Json
|
||
|
|
if ($meta.Status -eq 'running' -and $meta.WorkerPid) {{
|
||
|
|
if (Get-Process -Id $meta.WorkerPid -ErrorAction SilentlyContinue) {{
|
||
|
|
Write-Host ""Worker $workerName already running (PID $($meta.WorkerPid)).""
|
||
|
|
$shouldStart = $false
|
||
|
|
}}
|
||
|
|
}}
|
||
|
|
}} catch {{
|
||
|
|
Write-Host ""Failed to read metadata. Controller will restart worker."" -ForegroundColor Yellow
|
||
|
|
}}
|
||
|
|
}}
|
||
|
|
|
||
|
|
if ($shouldStart) {{
|
||
|
|
$initialMeta = [pscustomobject]@{{
|
||
|
|
WorkerName = $workerName
|
||
|
|
WorkerType = $workerType
|
||
|
|
Status = 'launching'
|
||
|
|
ControllerPid = $null
|
||
|
|
WorkerPid = $null
|
||
|
|
Restarts = 0
|
||
|
|
LastExitCode = $null
|
||
|
|
LogPath = $logPath
|
||
|
|
CommandPath = $commandPath
|
||
|
|
PayloadPath = $payloadPath
|
||
|
|
UpdatedAtUtc = (Get-Date).ToUniversalTime()
|
||
|
|
}} | ConvertTo-Json -Depth 5
|
||
|
|
$initialMeta | Set-Content -Path $metaPath -Encoding UTF8
|
||
|
|
|
||
|
|
$pwsh = Get-Command pwsh -ErrorAction SilentlyContinue
|
||
|
|
if ($pwsh) {{
|
||
|
|
$psExe = $pwsh.Source
|
||
|
|
}}
|
||
|
|
else {{
|
||
|
|
$psExe = (Get-Command powershell -ErrorAction Stop).Source
|
||
|
|
}}
|
||
|
|
|
||
|
|
$controllerArgs = @(
|
||
|
|
'-NoLogo','-NoProfile','-ExecutionPolicy','Bypass',
|
||
|
|
'-File',""$controllerPath"",
|
||
|
|
'-WorkerName',""$workerName"",
|
||
|
|
'-WorkerType',""$workerType"",
|
||
|
|
'-PayloadBase64Path',""$payloadBase64Path""
|
||
|
|
)
|
||
|
|
|
||
|
|
Start-Process -FilePath $psExe -ArgumentList $controllerArgs -WindowStyle Hidden | Out-Null
|
||
|
|
Write-Host ""Worker $workerName started under controller."" -ForegroundColor Green
|
||
|
|
}}
|
||
|
|
";
|
||
|
|
|
||
|
|
await _sshService.ExecuteRemoteCommandAsync(worker, $"powershell -NoLogo -NoProfile -NonInteractive -OutputFormat Text -ExecutionPolicy Bypass -EncodedCommand {Convert.ToBase64String(Encoding.Unicode.GetBytes(ensureScript))}");
|
||
|
|
}
|
||
|
|
|
||
|
|
public async Task StopWorkerAsync(WorkerConfig worker, string workerType)
|
||
|
|
{
|
||
|
|
var script = $@"
|
||
|
|
$ProgressPreference = 'SilentlyContinue'
|
||
|
|
$dataRoot = Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'
|
||
|
|
$instanceRoot = Join-Path (Join-Path $dataRoot '{workerType}') '{worker.Name}'
|
||
|
|
$commandPath = Join-Path $instanceRoot 'state\commands.txt'
|
||
|
|
[IO.File]::WriteAllText($commandPath, 'quit', [System.Text.Encoding]::UTF8)
|
||
|
|
Write-Host ""Quit command sent to worker {worker.Name}.""
|
||
|
|
";
|
||
|
|
|
||
|
|
await _sshService.ExecuteRemoteCommandAsync(worker, $"powershell -NoLogo -NoProfile -NonInteractive -OutputFormat Text -ExecutionPolicy Bypass -EncodedCommand {Convert.ToBase64String(Encoding.Unicode.GetBytes(script))}");
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|