begin persistence/attachment overhaul
This commit is contained in:
Nathan
File diff suppressed because one or more lines are too long
29
persistent_workers.md
Normal file
29
persistent_workers.md
Normal file
@@ -0,0 +1,29 @@
|
||||
## Persistent Remote Worker Architecture
|
||||
|
||||
This repository now includes a small controller/attach framework that keeps both Flamenco and SheepIt workers running on each remote Windows host even after the SSH session closes.
|
||||
|
||||
### Components
|
||||
|
||||
| File | Purpose |
|
||||
| --- | --- |
|
||||
| `remote_worker_controller.ps1` | Runs on each remote host. Launches the real worker process (Flamenco or SheepIt), redirects its stdout/stderr to a log file, listens for commands, and restarts the worker if it crashes. |
|
||||
| `remote_worker_attach.ps1` | Also runs remotely. Streams the controller’s log file back to the local SSH session and forwards commands (`pause`, `resume`, `quit`, etc.) to the worker’s stdin. It can also be invoked in a `-CommandOnly` mode to send a single command without attaching. |
|
||||
| `unified_flamenco_launcher.ps1` / `unified_sheepit_launcher.ps1` | Local launchers that copy the helper scripts to each host, start controllers in the background, and open attach sessions as requested. |
|
||||
|
||||
All metadata, logs and helper scripts live under `C:\ProgramData\UnifiedWorkers\<worker-type>\<worker-name>\`.
|
||||
|
||||
### Typical Flow
|
||||
|
||||
1. **Start/attach**: The launcher ensures the controller + payload are present on the target host, starts the controller through `Start-Process`, then launches `remote_worker_attach.ps1` via SSH so you can see live output and type commands. Closing the SSH window only ends the attach session—the worker keeps running.
|
||||
2. **Reattach**: Run the launcher again and choose the same worker. If a controller reports that a worker is already running, the launcher simply opens another attach session.
|
||||
3. **Pause/Resume/Quit all (SheepIt)**: The SheepIt launcher now exposes menu options that iterate over every enabled worker and invoke the attach helper in `-CommandOnly` mode to send the requested command.
|
||||
|
||||
### Manual Verification Checklist
|
||||
|
||||
1. **Controller deployment**: From the launcher, start a worker once and confirm `%ProgramData%\UnifiedWorkers` appears on the remote host with `controller.ps1`, `attach-helper.ps1`, `logs\worker.log`, and `state\worker-info.json`.
|
||||
2. **Persistence**: Attach to a worker, close the SSH window, then reattach. The log stream should resume and the worker PID reported in metadata should remain unchanged.
|
||||
3. **Command channel**: While attached, type `pause`, `resume`, and `quit`. The controller should log each command and the worker should react accordingly. Repeat the same commands using the SheepIt “Pause/Resume/Quit All” menu entries to validate the automation path.
|
||||
4. **Failure recovery**: Manually terminate the worker process on a remote host. The controller should log the exit, update metadata, and restart the worker up to the configured retry limit.
|
||||
|
||||
These steps ensure both Flamenco and SheepIt workflows operate as intended with the new persistent worker model. Adjust the retry/backoff values inside `remote_worker_controller.ps1` if you need more aggressive or conservative restart behavior.
|
||||
|
||||
107
remote_worker_attach.ps1
Normal file
107
remote_worker_attach.ps1
Normal file
@@ -0,0 +1,107 @@
|
||||
param(
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$WorkerName,
|
||||
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$WorkerType,
|
||||
|
||||
[string]$DataRoot = "$env:ProgramData\UnifiedWorkers",
|
||||
|
||||
[switch]$CommandOnly,
|
||||
|
||||
[string]$Command
|
||||
)
|
||||
|
||||
$ErrorActionPreference = 'Stop'
|
||||
|
||||
function Get-WorkerPaths {
|
||||
param([string]$Root, [string]$Type, [string]$Name)
|
||||
|
||||
$instanceRoot = Join-Path -Path (Join-Path -Path $Root -ChildPath $Type) -ChildPath $Name
|
||||
return [pscustomobject]@{
|
||||
Metadata = Join-Path -Path $instanceRoot -ChildPath 'state\worker-info.json'
|
||||
Command = Join-Path -Path $instanceRoot -ChildPath 'state\commands.txt'
|
||||
Log = Join-Path -Path $instanceRoot -ChildPath 'logs\worker.log'
|
||||
}
|
||||
}
|
||||
|
||||
$paths = Get-WorkerPaths -Root $DataRoot -Type $WorkerType -Name $WorkerName
|
||||
|
||||
if (-not (Test-Path $paths.Metadata)) {
|
||||
Write-Host "No worker metadata found for $WorkerName ($WorkerType)." -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
|
||||
try {
|
||||
$metadata = Get-Content -Path $paths.Metadata -Raw | ConvertFrom-Json
|
||||
}
|
||||
catch {
|
||||
Write-Host "Unable to read worker metadata: $($_.Exception.Message)" -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
|
||||
if (Test-Path $paths.Log) {
|
||||
# ensure log file exists but do not truncate
|
||||
$null = (Get-Item $paths.Log)
|
||||
} else {
|
||||
New-Item -Path $paths.Log -ItemType File -Force | Out-Null
|
||||
}
|
||||
|
||||
if (-not (Test-Path $paths.Command)) {
|
||||
New-Item -Path $paths.Command -ItemType File -Force | Out-Null
|
||||
}
|
||||
|
||||
function Send-WorkerCommand {
|
||||
param([string]$Value)
|
||||
|
||||
$clean = $Value.Trim()
|
||||
if (-not $clean) {
|
||||
return
|
||||
}
|
||||
|
||||
Add-Content -Path $paths.Command -Value $clean -Encoding UTF8
|
||||
Write-Host "Sent command '$clean' to $WorkerName." -ForegroundColor DarkGray
|
||||
}
|
||||
|
||||
if ($CommandOnly) {
|
||||
if (-not $Command) {
|
||||
Write-Host "CommandOnly flag set but no command provided." -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
|
||||
Send-WorkerCommand -Value $Command
|
||||
exit 0
|
||||
}
|
||||
|
||||
Write-Host "Attaching to $WorkerName ($WorkerType) logs." -ForegroundColor Cyan
|
||||
Write-Host "Type commands and press Enter. Type 'detach' to exit session." -ForegroundColor Yellow
|
||||
|
||||
$logJob = Start-Job -ScriptBlock {
|
||||
param($LogPath)
|
||||
Get-Content -Path $LogPath -Tail 50 -Wait
|
||||
} -ArgumentList $paths.Log
|
||||
|
||||
try {
|
||||
while ($true) {
|
||||
$input = Read-Host "> "
|
||||
if ($null -eq $input) {
|
||||
continue
|
||||
}
|
||||
|
||||
if ($input.Trim().ToLower() -eq 'detach') {
|
||||
break
|
||||
}
|
||||
|
||||
Send-WorkerCommand -Value $input
|
||||
}
|
||||
}
|
||||
finally {
|
||||
if ($logJob) {
|
||||
Stop-Job -Job $logJob -Force -ErrorAction SilentlyContinue
|
||||
Receive-Job -Job $logJob -ErrorAction SilentlyContinue | Out-Null
|
||||
Remove-Job -Job $logJob -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
|
||||
Write-Host "Detached from worker $WorkerName." -ForegroundColor Cyan
|
||||
}
|
||||
|
||||
195
remote_worker_controller.ps1
Normal file
195
remote_worker_controller.ps1
Normal file
@@ -0,0 +1,195 @@
|
||||
param(
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$WorkerName,
|
||||
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$WorkerType,
|
||||
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$PayloadBase64,
|
||||
|
||||
[string]$DataRoot = "$env:ProgramData\UnifiedWorkers",
|
||||
|
||||
[int]$MaxRestarts = 5,
|
||||
|
||||
[int]$RestartDelaySeconds = 10
|
||||
)
|
||||
|
||||
$ErrorActionPreference = 'Stop'
|
||||
|
||||
# region Path setup
|
||||
$workerRoot = Join-Path -Path $DataRoot -ChildPath $WorkerType
|
||||
$instanceRoot = Join-Path -Path $workerRoot -ChildPath $WorkerName
|
||||
New-Item -ItemType Directory -Path $instanceRoot -Force | Out-Null
|
||||
|
||||
$logsRoot = Join-Path -Path $instanceRoot -ChildPath 'logs'
|
||||
$stateRoot = Join-Path -Path $instanceRoot -ChildPath 'state'
|
||||
New-Item -ItemType Directory -Path $logsRoot -Force | Out-Null
|
||||
New-Item -ItemType Directory -Path $stateRoot -Force | Out-Null
|
||||
|
||||
$logPath = Join-Path -Path $logsRoot -ChildPath 'worker.log'
|
||||
$metaPath = Join-Path -Path $stateRoot -ChildPath 'worker-info.json'
|
||||
$commandPath = Join-Path -Path $stateRoot -ChildPath 'commands.txt'
|
||||
$payloadPath = Join-Path -Path $stateRoot -ChildPath "payload-$([Guid]::NewGuid().ToString()).ps1"
|
||||
# endregion
|
||||
|
||||
# region Helpers
|
||||
function Write-ControllerLog {
|
||||
param([string]$Message)
|
||||
$timestamp = (Get-Date).ToString('u')
|
||||
Add-Content -Path $logPath -Value "[CTRL $timestamp] $Message" -Encoding UTF8
|
||||
}
|
||||
|
||||
function Write-Metadata {
|
||||
param(
|
||||
[string]$Status,
|
||||
[nullable[int]]$WorkerPid = $null,
|
||||
[nullable[int]]$ControllerPid = $PID,
|
||||
[int]$Restarts = 0,
|
||||
[nullable[int]]$LastExitCode = $null
|
||||
)
|
||||
|
||||
$payload = [pscustomobject]@{
|
||||
WorkerName = $WorkerName
|
||||
WorkerType = $WorkerType
|
||||
Status = $Status
|
||||
ControllerPid = $ControllerPid
|
||||
WorkerPid = $WorkerPid
|
||||
Restarts = $Restarts
|
||||
LastExitCode = $LastExitCode
|
||||
LogPath = $logPath
|
||||
CommandPath = $commandPath
|
||||
PayloadPath = $payloadPath
|
||||
UpdatedAtUtc = (Get-Date).ToUniversalTime()
|
||||
}
|
||||
|
||||
$payload | ConvertTo-Json -Depth 5 | Set-Content -Path $metaPath -Encoding UTF8
|
||||
}
|
||||
|
||||
function Get-PendingCommands {
|
||||
if (-not (Test-Path $commandPath)) {
|
||||
return @()
|
||||
}
|
||||
|
||||
try {
|
||||
$lines = Get-Content -Path $commandPath -ErrorAction Stop
|
||||
Remove-Item -Path $commandPath -Force -ErrorAction SilentlyContinue
|
||||
return $lines | Where-Object { -not [string]::IsNullOrWhiteSpace($_) }
|
||||
}
|
||||
catch {
|
||||
return @()
|
||||
}
|
||||
}
|
||||
# endregion
|
||||
|
||||
try {
|
||||
# Write payload script to disk
|
||||
$payloadBytes = [Convert]::FromBase64String($PayloadBase64)
|
||||
[IO.File]::WriteAllBytes($payloadPath, $payloadBytes)
|
||||
Write-ControllerLog "Payload written to $payloadPath"
|
||||
}
|
||||
catch {
|
||||
Write-Error "Unable to write payload: $($_.Exception.Message)"
|
||||
exit 1
|
||||
}
|
||||
|
||||
$restartCount = 0
|
||||
$controllerPid = $PID
|
||||
|
||||
while ($restartCount -le $MaxRestarts) {
|
||||
try {
|
||||
# Initialize worker process
|
||||
$psi = [System.Diagnostics.ProcessStartInfo]::new()
|
||||
$psi.FileName = (Get-Command pwsh -ErrorAction SilentlyContinue)?.Source
|
||||
if (-not $psi.FileName) {
|
||||
$psi.FileName = (Get-Command powershell -ErrorAction Stop).Source
|
||||
}
|
||||
|
||||
$psi.Arguments = "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$payloadPath`""
|
||||
$psi.UseShellExecute = $false
|
||||
$psi.RedirectStandardInput = $true
|
||||
$psi.RedirectStandardOutput = $true
|
||||
$psi.RedirectStandardError = $true
|
||||
$psi.CreateNoWindow = $true
|
||||
|
||||
$workerProcess = New-Object System.Diagnostics.Process
|
||||
$workerProcess.StartInfo = $psi
|
||||
|
||||
if (-not $workerProcess.Start()) {
|
||||
throw "Failed to start worker process."
|
||||
}
|
||||
|
||||
Write-ControllerLog "Worker process started with PID $($workerProcess.Id)"
|
||||
Write-Metadata -Status 'running' -WorkerPid $workerProcess.Id -ControllerPid $controllerPid -Restarts $restartCount
|
||||
|
||||
$logWriter = [System.IO.StreamWriter]::new($logPath, $true, [System.Text.Encoding]::UTF8)
|
||||
$logWriter.AutoFlush = $true
|
||||
|
||||
$outputHandler = [System.Diagnostics.DataReceivedEventHandler]{
|
||||
param($s, $e)
|
||||
if ($e.Data) {
|
||||
$ts = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[OUT $ts] $($e.Data)")
|
||||
}
|
||||
}
|
||||
$errorHandler = [System.Diagnostics.DataReceivedEventHandler]{
|
||||
param($s, $e)
|
||||
if ($e.Data) {
|
||||
$ts = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[ERR $ts] $($e.Data)")
|
||||
}
|
||||
}
|
||||
|
||||
$workerProcess.add_OutputDataReceived($outputHandler)
|
||||
$workerProcess.add_ErrorDataReceived($errorHandler)
|
||||
$workerProcess.BeginOutputReadLine()
|
||||
$workerProcess.BeginErrorReadLine()
|
||||
|
||||
while (-not $workerProcess.HasExited) {
|
||||
$commands = Get-PendingCommands
|
||||
foreach ($command in $commands) {
|
||||
$trimmed = $command.Trim()
|
||||
if (-not $trimmed) { continue }
|
||||
|
||||
Write-ControllerLog "Received command '$trimmed'"
|
||||
try {
|
||||
$workerProcess.StandardInput.WriteLine($trimmed)
|
||||
$workerProcess.StandardInput.Flush()
|
||||
}
|
||||
catch {
|
||||
Write-ControllerLog "Failed to forward command '$trimmed': $($_.Exception.Message)"
|
||||
}
|
||||
|
||||
if ($trimmed -ieq 'quit') {
|
||||
Write-ControllerLog "Quit command issued. Waiting for worker to exit."
|
||||
}
|
||||
}
|
||||
|
||||
Start-Sleep -Milliseconds 500
|
||||
}
|
||||
|
||||
$exitCode = $workerProcess.ExitCode
|
||||
Write-ControllerLog "Worker exited with code $exitCode"
|
||||
Write-Metadata -Status 'stopped' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount -LastExitCode $exitCode
|
||||
|
||||
if ($exitCode -eq 0) {
|
||||
break
|
||||
}
|
||||
}
|
||||
catch {
|
||||
Write-ControllerLog "Controller error: $($_.Exception.Message)"
|
||||
}
|
||||
|
||||
$restartCount++
|
||||
if ($restartCount -gt $MaxRestarts) {
|
||||
Write-ControllerLog "Maximum restart attempts reached. Controller stopping."
|
||||
break
|
||||
}
|
||||
|
||||
Write-ControllerLog "Restarting worker in $RestartDelaySeconds seconds (attempt $restartCount of $MaxRestarts)."
|
||||
Start-Sleep -Seconds $RestartDelaySeconds
|
||||
}
|
||||
|
||||
Write-Metadata -Status 'inactive' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount
|
||||
Write-ControllerLog "Controller exiting."
|
||||
|
||||
@@ -50,6 +50,205 @@ $workers = @(
|
||||
}
|
||||
)
|
||||
|
||||
$global:UnifiedWorkerDataRoot = 'C:\ProgramData\UnifiedWorkers'
|
||||
$script:ControllerScriptBase64 = $null
|
||||
$script:AttachHelperScriptBase64 = $null
|
||||
|
||||
function Get-FileBase64 {
|
||||
param([string]$Path)
|
||||
[Convert]::ToBase64String([IO.File]::ReadAllBytes($Path))
|
||||
}
|
||||
|
||||
function Get-ControllerScriptBase64 {
|
||||
if (-not $script:ControllerScriptBase64) {
|
||||
$controllerPath = Join-Path $PSScriptRoot 'remote_worker_controller.ps1'
|
||||
$script:ControllerScriptBase64 = Get-FileBase64 -Path $controllerPath
|
||||
}
|
||||
return $script:ControllerScriptBase64
|
||||
}
|
||||
|
||||
function Get-AttachHelperScriptBase64 {
|
||||
if (-not $script:AttachHelperScriptBase64) {
|
||||
$helperPath = Join-Path $PSScriptRoot 'remote_worker_attach.ps1'
|
||||
$script:AttachHelperScriptBase64 = Get-FileBase64 -Path $helperPath
|
||||
}
|
||||
return $script:AttachHelperScriptBase64
|
||||
}
|
||||
|
||||
function ConvertTo-Base64Unicode {
|
||||
param([string]$Content)
|
||||
[Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($Content))
|
||||
}
|
||||
|
||||
function Get-WorkerSshArgs {
|
||||
param([string]$RawArgs)
|
||||
if ([string]::IsNullOrWhiteSpace($RawArgs)) {
|
||||
return @()
|
||||
}
|
||||
return $RawArgs -split '\s+' | Where-Object { $_.Trim().Length -gt 0 }
|
||||
}
|
||||
|
||||
function Invoke-RemotePowerShell {
|
||||
param(
|
||||
[Parameter(Mandatory = $true)][object]$Worker,
|
||||
[Parameter(Mandatory = $true)][string]$Script
|
||||
)
|
||||
|
||||
$bytes = [Text.Encoding]::Unicode.GetBytes($Script)
|
||||
$encoded = [Convert]::ToBase64String($bytes)
|
||||
$sshArgs = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30')
|
||||
$sshArgs += (Get-WorkerSshArgs -RawArgs $Worker.SSHArgs)
|
||||
$sshArgs += "powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -EncodedCommand $encoded"
|
||||
|
||||
& ssh @sshArgs
|
||||
return $LASTEXITCODE
|
||||
}
|
||||
|
||||
function Ensure-ControllerDeployed {
|
||||
param([object]$Worker)
|
||||
$controllerBase64 = Get-ControllerScriptBase64
|
||||
$script = @'
|
||||
`$dataRoot = '{0}'
|
||||
New-Item -ItemType Directory -Path `$dataRoot -Force | Out-Null
|
||||
`$controllerPath = Join-Path `$dataRoot 'controller.ps1'
|
||||
[IO.File]::WriteAllBytes(`$controllerPath, [Convert]::FromBase64String('{1}'))
|
||||
Write-Host "Controller ready at `$controllerPath"
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $controllerBase64
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
function Ensure-AttachHelperDeployed {
|
||||
param([object]$Worker)
|
||||
$helperBase64 = Get-AttachHelperScriptBase64
|
||||
$script = @'
|
||||
`$dataRoot = '{0}'
|
||||
New-Item -ItemType Directory -Path `$dataRoot -Force | Out-Null
|
||||
`$attachPath = Join-Path `$dataRoot 'attach-helper.ps1'
|
||||
[IO.File]::WriteAllBytes(`$attachPath, [Convert]::FromBase64String('{1}'))
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $helperBase64
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
function Get-RemoteWorkerPaths {
|
||||
param(
|
||||
[string]$WorkerType,
|
||||
[string]$WorkerName
|
||||
)
|
||||
|
||||
$base = Join-Path $global:UnifiedWorkerDataRoot $WorkerType
|
||||
$instance = Join-Path $base $WorkerName
|
||||
return [pscustomobject]@{
|
||||
InstanceRoot = $instance
|
||||
MetadataPath = Join-Path $instance 'state\worker-info.json'
|
||||
CommandPath = Join-Path $instance 'state\commands.txt'
|
||||
LogPath = Join-Path $instance 'logs\worker.log'
|
||||
}
|
||||
}
|
||||
|
||||
function Get-EnsureWorkerScript {
|
||||
param(
|
||||
[string]$WorkerName,
|
||||
[string]$WorkerType,
|
||||
[string]$PayloadBase64
|
||||
)
|
||||
|
||||
$paths = Get-RemoteWorkerPaths -WorkerType $WorkerType -WorkerName $WorkerName
|
||||
$controllerPath = Join-Path $global:UnifiedWorkerDataRoot 'controller.ps1'
|
||||
|
||||
return @'
|
||||
`$controllerPath = '{0}'
|
||||
`$metaPath = '{1}'
|
||||
`$workerName = '{2}'
|
||||
`$workerType = '{3}'
|
||||
`$payloadBase64 = '{4}'
|
||||
|
||||
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 "Unable to parse metadata for `$workerName. Starting new controller." -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
|
||||
if (`$shouldStart) {
|
||||
`$psExe = (Get-Command pwsh -ErrorAction SilentlyContinue)?.Source
|
||||
if (-not `$psExe) {
|
||||
`$psExe = (Get-Command powershell -ErrorAction Stop).Source
|
||||
}
|
||||
|
||||
`$controllerArgs = @(
|
||||
'-NoLogo','-NoProfile','-ExecutionPolicy','Bypass',
|
||||
'-File',"`$controllerPath",
|
||||
'-WorkerName',"`$workerName",
|
||||
'-WorkerType',"`$workerType",
|
||||
'-PayloadBase64',"`$payloadBase64"
|
||||
)
|
||||
|
||||
Start-Process -FilePath `$psExe -ArgumentList `$controllerArgs -WindowStyle Hidden | Out-Null
|
||||
Write-Host "Worker `$workerName started under controller." -ForegroundColor Green
|
||||
}
|
||||
'@ -f $controllerPath, $paths.MetadataPath, $WorkerName, $WorkerType, $PayloadBase64
|
||||
}
|
||||
|
||||
function Ensure-PersistentWorker {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$WorkerType,
|
||||
[string]$PayloadScript
|
||||
)
|
||||
|
||||
Write-Host "[$($Worker.Name)] Ensuring worker is running under controller..." -ForegroundColor Cyan
|
||||
Ensure-ControllerDeployed -Worker $Worker
|
||||
$payloadBase64 = ConvertTo-Base64Unicode -Content $PayloadScript
|
||||
$ensureScript = Get-EnsureWorkerScript -WorkerName $Worker.Name -WorkerType $WorkerType -PayloadBase64 $payloadBase64
|
||||
$result = Invoke-RemotePowerShell -Worker $Worker -Script $ensureScript
|
||||
if ($result -ne 0) {
|
||||
Write-Host "[$($Worker.Name)] Remote ensure command exited with code $result." -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
|
||||
function Invoke-WorkerAttach {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$WorkerType,
|
||||
[switch]$CommandOnly,
|
||||
[string]$Command
|
||||
)
|
||||
|
||||
Ensure-AttachHelperDeployed -Worker $Worker
|
||||
$paramsBlock = @("-WorkerName","$($Worker.Name)","-WorkerType","$WorkerType")
|
||||
if ($CommandOnly) {
|
||||
$paramsBlock += "-CommandOnly"
|
||||
}
|
||||
if ($Command) {
|
||||
$paramsBlock += "-Command"
|
||||
$paramsBlock += $Command
|
||||
}
|
||||
|
||||
$paramsLiteral = ($paramsBlock | ForEach-Object { "'$_'" }) -join ', '
|
||||
$script = @'
|
||||
`$helperPath = Join-Path '{0}' 'attach-helper.ps1'
|
||||
if (-not (Test-Path `$helperPath)) {
|
||||
throw "Attach helper missing at `$helperPath"
|
||||
}
|
||||
`$arguments = @({1})
|
||||
& `$helperPath @arguments
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $paramsLiteral
|
||||
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
# FUNCTIONS
|
||||
|
||||
# This function generates the standard PowerShell remote command
|
||||
@@ -258,57 +457,10 @@ function Start-StandardWorker {
|
||||
[Parameter(Mandatory = $true)]
|
||||
[object]$Worker
|
||||
)
|
||||
|
||||
$retryCount = 0
|
||||
$retryDelay = 15 # seconds between retries
|
||||
$workerRestarted = $false
|
||||
|
||||
while ($true) { # Changed to infinite loop
|
||||
if ($retryCount -gt 0) {
|
||||
Write-Host "`nRestarting worker process (Attempt $($retryCount + 1))..." -ForegroundColor Yellow
|
||||
Start-Sleep -Seconds $retryDelay
|
||||
$workerRestarted = $true
|
||||
}
|
||||
|
||||
Write-Host "Connecting to $($Worker.Name)..." -ForegroundColor Cyan
|
||||
if ($workerRestarted) {
|
||||
Write-Host "Worker was restarted due to disconnection or crash." -ForegroundColor Yellow
|
||||
}
|
||||
|
||||
try {
|
||||
$remoteCommand = Get-RemoteStandardWorkerCommand
|
||||
|
||||
# Encode the command to handle special characters
|
||||
$bytes = [System.Text.Encoding]::Unicode.GetBytes($remoteCommand)
|
||||
$encodedCommand = [Convert]::ToBase64String($bytes)
|
||||
|
||||
# Execute the encoded command on the remote machine
|
||||
Write-Host "Connecting to $($Worker.Name) and executing worker script..." -ForegroundColor Yellow
|
||||
|
||||
# Add SSH keepalive settings to reduce chance of random disconnections
|
||||
$sshCommand = "ssh -o ServerAliveInterval=60 -o ServerAliveCountMax=30 $($Worker.SSHArgs) ""powershell -EncodedCommand $encodedCommand"""
|
||||
|
||||
# Execute the SSH command and capture the exit code
|
||||
Invoke-Expression $sshCommand
|
||||
$sshExitCode = $LASTEXITCODE
|
||||
|
||||
# Check if SSH command completed successfully
|
||||
if ($sshExitCode -eq 0) {
|
||||
Write-Host "`nWorker completed successfully. Restarting automatically..." -ForegroundColor Green
|
||||
Start-Sleep -Seconds 2 # Brief pause before restarting
|
||||
$retryCount = 0 # Reset counter for successful completion
|
||||
continue # Continue the loop instead of breaking
|
||||
} else {
|
||||
throw "Worker process exited with code: $sshExitCode"
|
||||
}
|
||||
}
|
||||
catch {
|
||||
$retryCount++
|
||||
Write-Host "`nAn error occurred while running worker on $($Worker.Name):" -ForegroundColor Red
|
||||
Write-Host $_.Exception.Message -ForegroundColor Red
|
||||
Write-Host "`nAttempting to restart worker in $retryDelay seconds..." -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
|
||||
$payloadScript = Get-RemoteStandardWorkerCommand
|
||||
Ensure-PersistentWorker -Worker $Worker -WorkerType 'flamenco' -PayloadScript $payloadScript
|
||||
Invoke-WorkerAttach -Worker $Worker -WorkerType 'flamenco'
|
||||
}
|
||||
|
||||
# This function launches the CMD worker
|
||||
@@ -317,57 +469,10 @@ function Start-CmdWorker {
|
||||
[Parameter(Mandatory = $true)]
|
||||
[object]$Worker
|
||||
)
|
||||
|
||||
$retryCount = 0
|
||||
$retryDelay = 5 # seconds between retries
|
||||
$workerRestarted = $false
|
||||
|
||||
while ($true) { # Changed to infinite loop
|
||||
if ($retryCount -gt 0) {
|
||||
Write-Host "`nRestarting worker process (Attempt $($retryCount + 1))..." -ForegroundColor Yellow
|
||||
Start-Sleep -Seconds $retryDelay
|
||||
$workerRestarted = $true
|
||||
}
|
||||
|
||||
Write-Host "Connecting to $($Worker.Name) (CMD mode)..." -ForegroundColor Cyan
|
||||
if ($workerRestarted) {
|
||||
Write-Host "Worker was restarted due to disconnection or crash." -ForegroundColor Yellow
|
||||
}
|
||||
|
||||
try {
|
||||
$remoteCommand = Get-RemoteCmdWorkerCommand
|
||||
|
||||
# Encode the command to handle special characters
|
||||
$bytes = [System.Text.Encoding]::Unicode.GetBytes($remoteCommand)
|
||||
$encodedCommand = [Convert]::ToBase64String($bytes)
|
||||
|
||||
# Execute the encoded command on the remote machine
|
||||
Write-Host "Connecting to $($Worker.Name) and executing CMD worker script..." -ForegroundColor Yellow
|
||||
|
||||
# Add SSH keepalive settings to reduce chance of random disconnections
|
||||
$sshCommand = "ssh -o ServerAliveInterval=60 -o ServerAliveCountMax=30 $($Worker.SSHArgs) ""powershell -EncodedCommand $encodedCommand"""
|
||||
|
||||
# Execute the SSH command and capture the exit code
|
||||
Invoke-Expression $sshCommand
|
||||
$sshExitCode = $LASTEXITCODE
|
||||
|
||||
# Check if SSH command completed successfully
|
||||
if ($sshExitCode -eq 0) {
|
||||
Write-Host "`nWorker completed successfully. Restarting automatically..." -ForegroundColor Green
|
||||
Start-Sleep -Seconds 2 # Brief pause before restarting
|
||||
$retryCount = 0 # Reset counter for successful completion
|
||||
continue # Continue the loop instead of breaking
|
||||
} else {
|
||||
throw "Worker process exited with code: $sshExitCode"
|
||||
}
|
||||
}
|
||||
catch {
|
||||
$retryCount++
|
||||
Write-Host "`nAn error occurred while running worker on $($Worker.Name):" -ForegroundColor Red
|
||||
Write-Host $_.Exception.Message -ForegroundColor Red
|
||||
Write-Host "`nAttempting to restart worker in $retryDelay seconds..." -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
|
||||
$payloadScript = Get-RemoteCmdWorkerCommand
|
||||
Ensure-PersistentWorker -Worker $Worker -WorkerType 'flamenco' -PayloadScript $payloadScript
|
||||
Invoke-WorkerAttach -Worker $Worker -WorkerType 'flamenco'
|
||||
}
|
||||
|
||||
# This function launches ALL workers in Windows Terminal tabs
|
||||
@@ -376,174 +481,22 @@ function Start-AllWorkers {
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$WorkerType
|
||||
)
|
||||
|
||||
Write-Host "Launching ALL $WorkerType workers in Windows Terminal tabs..." -ForegroundColor Cyan
|
||||
|
||||
try {
|
||||
# First, check if Windows Terminal is available
|
||||
if (-not (Get-Command wt.exe -ErrorAction SilentlyContinue)) {
|
||||
Write-Host "Windows Terminal (wt.exe) not found. Falling back to separate windows." -ForegroundColor Yellow
|
||||
$useTerminal = $false
|
||||
} else {
|
||||
$useTerminal = $true
|
||||
}
|
||||
|
||||
foreach ($worker in $workers) {
|
||||
# Create a new PowerShell script with a unique name for this worker
|
||||
$tempScriptPath = [System.IO.Path]::GetTempFileName() + ".ps1"
|
||||
|
||||
# Create different script content based on worker type
|
||||
if ($WorkerType -eq "CMD") {
|
||||
# CMD workers get retry logic at the local level
|
||||
$scriptContent = @"
|
||||
# Wrap everything in a try-catch to prevent script termination
|
||||
try {
|
||||
Write-Host 'Launching $WorkerType worker for $($worker.Name)' -ForegroundColor Cyan
|
||||
Write-Host "Ensuring ALL $WorkerType workers are running under controllers..." -ForegroundColor Cyan
|
||||
|
||||
`$retryCount = 0
|
||||
`$retryDelay = 5 # seconds between retries
|
||||
`$workerRestarted = `$false
|
||||
|
||||
while (`$true) { # Changed to infinite loop
|
||||
try {
|
||||
if (`$retryCount -gt 0) {
|
||||
Write-Host "`nRestarting worker process (Attempt `$(`$retryCount + 1))..." -ForegroundColor Yellow
|
||||
Start-Sleep -Seconds `$retryDelay
|
||||
`$workerRestarted = `$true
|
||||
}
|
||||
|
||||
Write-Host "Connecting to $($worker.Name) ($WorkerType mode)..." -ForegroundColor Cyan
|
||||
if (`$workerRestarted) {
|
||||
Write-Host "Worker was restarted due to disconnection or crash." -ForegroundColor Yellow
|
||||
}
|
||||
|
||||
# Get remote command
|
||||
`$remoteCommand = @'
|
||||
$(Get-RemoteSimplifiedCmdWorkerCommand)
|
||||
'@
|
||||
|
||||
# Encode the command
|
||||
`$bytes = [System.Text.Encoding]::Unicode.GetBytes(`$remoteCommand)
|
||||
`$encodedCommand = [Convert]::ToBase64String(`$bytes)
|
||||
|
||||
# Execute SSH command with keepalive settings and capture exit code
|
||||
ssh -o ServerAliveInterval=60 -o ServerAliveCountMax=30 $($worker.SSHArgs) "powershell -EncodedCommand `$encodedCommand"
|
||||
`$sshExitCode = `$LASTEXITCODE
|
||||
|
||||
# Check if SSH command completed successfully
|
||||
if (`$sshExitCode -eq 0) {
|
||||
Write-Host "`nWorker completed successfully. Restarting automatically..." -ForegroundColor Green
|
||||
Start-Sleep -Seconds 2 # Brief pause before restarting
|
||||
`$retryCount = 0 # Reset counter for successful completion
|
||||
continue # Continue the loop instead of breaking
|
||||
} else {
|
||||
throw "SSH command failed with exit code: `$sshExitCode"
|
||||
}
|
||||
}
|
||||
catch {
|
||||
`$retryCount++
|
||||
Write-Host "An error occurred while connecting to $($worker.Name):" -ForegroundColor Red
|
||||
Write-Host `$_.Exception.Message -ForegroundColor Red
|
||||
Write-Host "Attempting to reconnect in `$retryDelay seconds..." -ForegroundColor Yellow
|
||||
# Don't rethrow - we want to continue the retry loop
|
||||
}
|
||||
$payloadScript = if ($WorkerType -eq 'CMD') {
|
||||
Get-RemoteCmdWorkerCommand
|
||||
} else {
|
||||
Get-RemoteStandardWorkerCommand
|
||||
}
|
||||
}
|
||||
catch {
|
||||
# This outer catch block is for any unexpected errors that might terminate the script
|
||||
Write-Host "`nCRITICAL ERROR: Script encountered an unexpected error:" -ForegroundColor Red
|
||||
Write-Host `$_.Exception.Message -ForegroundColor Red
|
||||
Write-Host "`nRestarting the entire worker process in 5 seconds..." -ForegroundColor Yellow
|
||||
Start-Sleep -Seconds 5
|
||||
# Restart the script by calling itself
|
||||
& `$MyInvocation.MyCommand.Path
|
||||
}
|
||||
"@
|
||||
} else {
|
||||
# Standard workers keep the original retry logic
|
||||
$scriptContent = @"
|
||||
# Wrap everything in a try-catch to prevent script termination
|
||||
try {
|
||||
Write-Host 'Launching $WorkerType worker for $($worker.Name)' -ForegroundColor Cyan
|
||||
|
||||
`$retryCount = 0
|
||||
`$retryDelay = 5 # seconds between retries
|
||||
|
||||
while (`$true) { # Changed to infinite loop
|
||||
try {
|
||||
if (`$retryCount -gt 0) {
|
||||
Write-Host "Retry attempt `$retryCount..." -ForegroundColor Yellow
|
||||
Start-Sleep -Seconds `$retryDelay
|
||||
}
|
||||
|
||||
# Get remote command
|
||||
`$remoteCommand = @'
|
||||
$(Get-RemoteStandardWorkerCommand)
|
||||
'@
|
||||
|
||||
# Encode the command
|
||||
`$bytes = [System.Text.Encoding]::Unicode.GetBytes(`$remoteCommand)
|
||||
`$encodedCommand = [Convert]::ToBase64String(`$bytes)
|
||||
|
||||
# Execute SSH command with keepalive settings
|
||||
ssh -o ServerAliveInterval=60 -o ServerAliveCountMax=30 $($worker.SSHArgs) "powershell -EncodedCommand `$encodedCommand"
|
||||
`$sshExitCode = `$LASTEXITCODE
|
||||
|
||||
# Check SSH exit code and handle accordingly
|
||||
if (`$sshExitCode -eq 0) {
|
||||
Write-Host "`nWorker completed successfully. Restarting automatically..." -ForegroundColor Green
|
||||
Start-Sleep -Seconds 2 # Brief pause before restarting
|
||||
`$retryCount = 0 # Reset counter for successful completion
|
||||
continue # Continue the loop instead of breaking
|
||||
} else {
|
||||
throw "SSH command failed with exit code: `$sshExitCode"
|
||||
}
|
||||
}
|
||||
catch {
|
||||
`$retryCount++
|
||||
Write-Host "An error occurred while connecting to $($worker.Name):" -ForegroundColor Red
|
||||
Write-Host `$_.Exception.Message -ForegroundColor Red
|
||||
Write-Host "Attempting to reconnect in `$retryDelay seconds..." -ForegroundColor Yellow
|
||||
# Don't rethrow - we want to continue the retry loop
|
||||
}
|
||||
foreach ($worker in $workers) {
|
||||
Ensure-PersistentWorker -Worker $worker -WorkerType 'flamenco' -PayloadScript $payloadScript
|
||||
}
|
||||
}
|
||||
catch {
|
||||
# This outer catch block is for any unexpected errors that might terminate the script
|
||||
Write-Host "`nCRITICAL ERROR: Script encountered an unexpected error:" -ForegroundColor Red
|
||||
Write-Host `$_.Exception.Message -ForegroundColor Red
|
||||
Write-Host "`nRestarting the entire worker process in 5 seconds..." -ForegroundColor Yellow
|
||||
Start-Sleep -Seconds 5
|
||||
# Restart the script by calling itself
|
||||
& `$MyInvocation.MyCommand.Path
|
||||
}
|
||||
"@
|
||||
}
|
||||
|
||||
# Write the script to file
|
||||
Set-Content -Path $tempScriptPath -Value $scriptContent
|
||||
|
||||
if ($useTerminal) {
|
||||
# Launch in a new Windows Terminal tab
|
||||
$tabTitle = "$($worker.Name) - $WorkerType Worker"
|
||||
Start-Process wt.exe -ArgumentList "-w 0 new-tab --title `"$tabTitle`" powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$tempScriptPath`""
|
||||
} else {
|
||||
# Fallback to separate window if Windows Terminal is not available
|
||||
Start-Process powershell -ArgumentList "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$tempScriptPath`""
|
||||
}
|
||||
|
||||
Write-Host "Started $($worker.Name) ($WorkerType) worker in a new tab." -ForegroundColor Green
|
||||
Start-Sleep -Milliseconds 300 # Small delay between launches
|
||||
}
|
||||
|
||||
Write-Host "`nAll $WorkerType worker scripts have been launched in Windows Terminal tabs." -ForegroundColor Cyan
|
||||
}
|
||||
catch {
|
||||
Write-Host "Error launching workers: $($_.Exception.Message)" -ForegroundColor Red
|
||||
}
|
||||
|
||||
|
||||
Write-Host "All workers processed. Attach to any worker from the menu to monitor output." -ForegroundColor Green
|
||||
Write-Host "Press Enter to return to the menu..." -ForegroundColor Green
|
||||
Read-Host
|
||||
Read-Host | Out-Null
|
||||
}
|
||||
|
||||
# Main menu loop
|
||||
|
||||
@@ -32,6 +32,196 @@ $workers = @(
|
||||
@{ ID = 6; Name = "i9-13ks"; SSHArgs = "-t -p 22146 i9-13ks"; Enabled = $true }
|
||||
)
|
||||
|
||||
$global:UnifiedWorkerDataRoot = 'C:\ProgramData\UnifiedWorkers'
|
||||
$script:ControllerScriptBase64 = $null
|
||||
$script:AttachHelperScriptBase64 = $null
|
||||
|
||||
function Get-FileBase64 {
|
||||
param([string]$Path)
|
||||
[Convert]::ToBase64String([IO.File]::ReadAllBytes($Path))
|
||||
}
|
||||
|
||||
function Get-ControllerScriptBase64 {
|
||||
if (-not $script:ControllerScriptBase64) {
|
||||
$controllerPath = Join-Path $PSScriptRoot 'remote_worker_controller.ps1'
|
||||
$script:ControllerScriptBase64 = Get-FileBase64 -Path $controllerPath
|
||||
}
|
||||
return $script:ControllerScriptBase64
|
||||
}
|
||||
|
||||
function Get-AttachHelperScriptBase64 {
|
||||
if (-not $script:AttachHelperScriptBase64) {
|
||||
$helperPath = Join-Path $PSScriptRoot 'remote_worker_attach.ps1'
|
||||
$script:AttachHelperScriptBase64 = Get-FileBase64 -Path $helperPath
|
||||
}
|
||||
return $script:AttachHelperScriptBase64
|
||||
}
|
||||
|
||||
function ConvertTo-Base64Unicode {
|
||||
param([string]$Content)
|
||||
[Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($Content))
|
||||
}
|
||||
|
||||
function Get-WorkerSshArgs {
|
||||
param([string]$RawArgs)
|
||||
if ([string]::IsNullOrWhiteSpace($RawArgs)) {
|
||||
return @()
|
||||
}
|
||||
return $RawArgs -split '\s+' | Where-Object { $_.Trim().Length -gt 0 }
|
||||
}
|
||||
|
||||
function Invoke-RemotePowerShell {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$Script
|
||||
)
|
||||
|
||||
$bytes = [Text.Encoding]::Unicode.GetBytes($Script)
|
||||
$encoded = [Convert]::ToBase64String($bytes)
|
||||
$sshArgs = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30')
|
||||
$sshArgs += (Get-WorkerSshArgs -RawArgs $Worker.SSHArgs)
|
||||
$sshArgs += "powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -EncodedCommand $encoded"
|
||||
|
||||
& ssh @sshArgs
|
||||
return $LASTEXITCODE
|
||||
}
|
||||
|
||||
function Ensure-ControllerDeployed {
|
||||
param([object]$Worker)
|
||||
$controllerBase64 = Get-ControllerScriptBase64
|
||||
$script = @'
|
||||
`$dataRoot = '{0}'
|
||||
New-Item -ItemType Directory -Path `$dataRoot -Force | Out-Null
|
||||
`$controllerPath = Join-Path `$dataRoot 'controller.ps1'
|
||||
[IO.File]::WriteAllBytes(`$controllerPath, [Convert]::FromBase64String('{1}'))
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $controllerBase64
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
function Ensure-AttachHelperDeployed {
|
||||
param([object]$Worker)
|
||||
$helperBase64 = Get-AttachHelperScriptBase64
|
||||
$script = @'
|
||||
`$dataRoot = '{0}'
|
||||
New-Item -ItemType Directory -Path `$dataRoot -Force | Out-Null
|
||||
`$attachPath = Join-Path `$dataRoot 'attach-helper.ps1'
|
||||
[IO.File]::WriteAllBytes(`$attachPath, [Convert]::FromBase64String('{1}'))
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $helperBase64
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
function Get-RemoteWorkerPaths {
|
||||
param([string]$WorkerType, [string]$WorkerName)
|
||||
$base = Join-Path $global:UnifiedWorkerDataRoot $WorkerType
|
||||
$instance = Join-Path $base $WorkerName
|
||||
[pscustomobject]@{
|
||||
InstanceRoot = $instance
|
||||
MetadataPath = Join-Path $instance 'state\worker-info.json'
|
||||
CommandPath = Join-Path $instance 'state\commands.txt'
|
||||
LogPath = Join-Path $instance 'logs\worker.log'
|
||||
}
|
||||
}
|
||||
|
||||
function Get-EnsureWorkerScript {
|
||||
param(
|
||||
[string]$WorkerName,
|
||||
[string]$WorkerType,
|
||||
[string]$PayloadBase64
|
||||
)
|
||||
|
||||
$paths = Get-RemoteWorkerPaths -WorkerType $WorkerType -WorkerName $WorkerName
|
||||
$controllerPath = Join-Path $global:UnifiedWorkerDataRoot 'controller.ps1'
|
||||
|
||||
@'
|
||||
`$controllerPath = '{0}'
|
||||
`$metaPath = '{1}'
|
||||
`$workerName = '{2}'
|
||||
`$workerType = '{3}'
|
||||
`$payloadBase64 = '{4}'
|
||||
|
||||
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) {
|
||||
`$psExe = (Get-Command pwsh -ErrorAction SilentlyContinue)?.Source
|
||||
if (-not `$psExe) {
|
||||
`$psExe = (Get-Command powershell -ErrorAction Stop).Source
|
||||
}
|
||||
|
||||
`$controllerArgs = @(
|
||||
'-NoLogo','-NoProfile','-ExecutionPolicy','Bypass',
|
||||
'-File',"`$controllerPath",
|
||||
'-WorkerName',"`$workerName",
|
||||
'-WorkerType',"`$workerType",
|
||||
'-PayloadBase64',"`$payloadBase64"
|
||||
)
|
||||
|
||||
Start-Process -FilePath `$psExe -ArgumentList `$controllerArgs -WindowStyle Hidden | Out-Null
|
||||
Write-Host "Worker `$workerName started under controller." -ForegroundColor Green
|
||||
}
|
||||
'@ -f $controllerPath, $paths.MetadataPath, $WorkerName, $WorkerType, $PayloadBase64
|
||||
}
|
||||
|
||||
function Ensure-PersistentWorker {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$WorkerType,
|
||||
[string]$PayloadScript
|
||||
)
|
||||
|
||||
Ensure-ControllerDeployed -Worker $Worker
|
||||
$payloadBase64 = ConvertTo-Base64Unicode -Content $PayloadScript
|
||||
$ensureScript = Get-EnsureWorkerScript -WorkerName $Worker.Name -WorkerType $WorkerType -PayloadBase64 $payloadBase64
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $ensureScript | Out-Null
|
||||
}
|
||||
|
||||
function Invoke-WorkerAttach {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$WorkerType,
|
||||
[switch]$CommandOnly,
|
||||
[string]$Command
|
||||
)
|
||||
|
||||
Ensure-AttachHelperDeployed -Worker $Worker
|
||||
$paramsBlock = @("-WorkerName","$($Worker.Name)","-WorkerType","$WorkerType")
|
||||
if ($CommandOnly) {
|
||||
$paramsBlock += "-CommandOnly"
|
||||
}
|
||||
if ($Command) {
|
||||
$paramsBlock += "-Command"
|
||||
$paramsBlock += $Command
|
||||
}
|
||||
|
||||
$paramsLiteral = ($paramsBlock | ForEach-Object { "'$_'" }) -join ', '
|
||||
$script = @'
|
||||
`$helperPath = Join-Path '{0}' 'attach-helper.ps1'
|
||||
if (-not (Test-Path `$helperPath)) {
|
||||
throw "Attach helper missing at `$helperPath"
|
||||
}
|
||||
`$arguments = @({1})
|
||||
& `$helperPath @arguments
|
||||
'@ -f $global:UnifiedWorkerDataRoot, $paramsLiteral
|
||||
|
||||
Invoke-RemotePowerShell -Worker $Worker -Script $script | Out-Null
|
||||
}
|
||||
|
||||
function Get-RemoteSheepItCommand {
|
||||
param(
|
||||
[string]$RenderKey,
|
||||
@@ -135,6 +325,40 @@ try {
|
||||
throw
|
||||
}
|
||||
}
|
||||
|
||||
function Ensure-SheepItWorkerController {
|
||||
param([object]$Worker)
|
||||
Initialize-SheepItCredentials
|
||||
$payloadScript = Get-RemoteSheepItCommand -RenderKey $script:SheepItRenderKey -Username $script:SheepItUsername
|
||||
Ensure-PersistentWorker -Worker $Worker -WorkerType 'sheepit' -PayloadScript $payloadScript
|
||||
}
|
||||
|
||||
function Start-SheepItWorker {
|
||||
param([object]$Worker)
|
||||
Ensure-SheepItWorkerController -Worker $Worker
|
||||
Invoke-WorkerAttach -Worker $Worker -WorkerType 'sheepit'
|
||||
}
|
||||
|
||||
function Start-AllSheepIt {
|
||||
Initialize-SheepItCredentials
|
||||
foreach ($worker in $workers | Where-Object { $_.Enabled }) {
|
||||
Ensure-SheepItWorkerController -Worker $worker
|
||||
}
|
||||
Write-Host "All enabled SheepIt workers ensured running under controllers." -ForegroundColor Green
|
||||
Write-Host "Use the attach option to monitor any worker." -ForegroundColor Cyan
|
||||
Read-Host "Press Enter to continue" | Out-Null
|
||||
}
|
||||
|
||||
function Send-SheepItCommandAll {
|
||||
param([string]$CommandText)
|
||||
|
||||
foreach ($worker in $workers | Where-Object { $_.Enabled }) {
|
||||
Write-Host "[$($worker.Name)] Sending '$CommandText'..." -ForegroundColor Gray
|
||||
Invoke-WorkerAttach -Worker $worker -WorkerType 'sheepit' -CommandOnly -Command $CommandText
|
||||
}
|
||||
Write-Host "Command '$CommandText' dispatched to all enabled workers." -ForegroundColor Green
|
||||
Read-Host "Press Enter to continue" | Out-Null
|
||||
}
|
||||
catch {
|
||||
Write-Host ('Error: {0}' -f `$_.Exception.Message) -ForegroundColor Red
|
||||
Write-Host ('Stack trace: {0}' -f `$_.ScriptStackTrace) -ForegroundColor DarkRed
|
||||
@@ -142,127 +366,7 @@ catch {
|
||||
"@
|
||||
}
|
||||
|
||||
function New-SheepItSessionScript {
|
||||
param(
|
||||
[object]$Worker,
|
||||
[string]$RemoteCommand
|
||||
)
|
||||
|
||||
$rawArgs = @($Worker.SSHArgs -split '\s+' | Where-Object { $_ -and $_.Trim().Length -gt 0 })
|
||||
$sshArgsJson = if ($rawArgs) { ($rawArgs | ConvertTo-Json -Compress) } else { '[]' }
|
||||
$targetHost = ($rawArgs | Where-Object { $_ -notmatch '^-{1,2}' } | Select-Object -Last 1)
|
||||
$hostLiteral = if ($targetHost) { "'" + ($targetHost -replace "'", "''") + "'" } else { '$null' }
|
||||
$portValue = $null
|
||||
for ($i = 0; $i -lt $rawArgs.Count; $i++) {
|
||||
if ($rawArgs[$i] -eq '-p' -and $i -lt ($rawArgs.Count - 1)) {
|
||||
$portValue = $rawArgs[$i + 1]
|
||||
break
|
||||
}
|
||||
}
|
||||
$portLiteral = if ($portValue) { "'" + ($portValue -replace "'", "''") + "'" } else { '$null' }
|
||||
$remoteScriptBase64 = [Convert]::ToBase64String([System.Text.Encoding]::Unicode.GetBytes($RemoteCommand))
|
||||
|
||||
$scriptContent = @"
|
||||
Write-Host 'Connecting to $($Worker.Name)...' -ForegroundColor Cyan
|
||||
`$sshArgs = ConvertFrom-Json '$sshArgsJson'
|
||||
`$targetHost = $hostLiteral
|
||||
`$port = $portLiteral
|
||||
`$scriptBase64 = '$remoteScriptBase64'
|
||||
|
||||
if (-not `$targetHost) {
|
||||
Write-Host "Unable to determine SSH host for $($Worker.Name)." -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
|
||||
`$localTemp = [System.IO.Path]::GetTempFileName() + '.ps1'
|
||||
`$remoteScript = [System.Text.Encoding]::Unicode.GetString([Convert]::FromBase64String(`$scriptBase64))
|
||||
Set-Content -Path `$localTemp -Value `$remoteScript -Encoding UTF8
|
||||
|
||||
`$remoteFile = "sheepit-$([System.Guid]::NewGuid().ToString()).ps1"
|
||||
|
||||
`$scpArgs = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30')
|
||||
if (`$port) {
|
||||
`$scpArgs += '-P'
|
||||
`$scpArgs += `$port
|
||||
}
|
||||
`$scpArgs += `$localTemp
|
||||
`$scpArgs += ("${targetHost}:`$remoteFile")
|
||||
|
||||
Write-Host "Transferring SheepIt script to `$targetHost..." -ForegroundColor Gray
|
||||
& scp @scpArgs
|
||||
if (`$LASTEXITCODE -ne 0) {
|
||||
Write-Host "Failed to copy script to `$targetHost." -ForegroundColor Red
|
||||
Remove-Item `$localTemp -ErrorAction SilentlyContinue
|
||||
return
|
||||
}
|
||||
|
||||
`$sshBase = @('-o','ServerAliveInterval=60','-o','ServerAliveCountMax=30') + `$sshArgs
|
||||
`$execArgs = `$sshBase + @('powershell','-NoLogo','-NoProfile','-ExecutionPolicy','Bypass','-File',"`$remoteFile")
|
||||
|
||||
Write-Host "Starting SheepIt client on `$targetHost..." -ForegroundColor Cyan
|
||||
& ssh @execArgs
|
||||
|
||||
Write-Host "Cleaning up remote script..." -ForegroundColor Gray
|
||||
`$cleanupArgs = `$sshBase + @('powershell','-NoLogo','-NoProfile','-ExecutionPolicy','Bypass','-Command',"`$ErrorActionPreference='SilentlyContinue'; Remove-Item -LiteralPath ``"$remoteFile``" -ErrorAction SilentlyContinue")
|
||||
& ssh @cleanupArgs | Out-Null
|
||||
|
||||
Remove-Item `$localTemp -ErrorAction SilentlyContinue
|
||||
Write-Host "`nSSH session ended." -ForegroundColor Yellow
|
||||
Read-Host "Press Enter to close"
|
||||
"@
|
||||
|
||||
return $scriptContent
|
||||
}
|
||||
|
||||
function Start-SheepItTab {
|
||||
param(
|
||||
[string]$Title,
|
||||
[string]$Content
|
||||
)
|
||||
|
||||
$tempScript = [System.IO.Path]::GetTempFileName() + '.ps1'
|
||||
Set-Content -Path $tempScript -Value $Content -Encoding UTF8
|
||||
|
||||
if (Get-Command wt.exe -ErrorAction SilentlyContinue) {
|
||||
Start-Process wt.exe -ArgumentList "-w 0 new-tab --title `"$Title`" powershell -NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$tempScript`""
|
||||
} else {
|
||||
Start-Process powershell -ArgumentList "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$tempScript`""
|
||||
}
|
||||
}
|
||||
|
||||
function Start-SheepItWorker {
|
||||
param([object]$Worker)
|
||||
|
||||
if (-not $Worker.Enabled) {
|
||||
Write-Host "$($Worker.Name) is not enabled for SheepIt." -ForegroundColor Yellow
|
||||
return
|
||||
}
|
||||
|
||||
Initialize-SheepItCredentials
|
||||
|
||||
$remoteCommand = Get-RemoteSheepItCommand -RenderKey $script:SheepItRenderKey -Username $script:SheepItUsername
|
||||
$sessionScript = New-SheepItSessionScript -Worker $Worker -RemoteCommand $remoteCommand
|
||||
$title = "$($Worker.Name) - SheepIt"
|
||||
|
||||
Start-SheepItTab -Title $title -Content $sessionScript
|
||||
Write-Host "Opened SheepIt session for $($Worker.Name) in a new terminal tab." -ForegroundColor Green
|
||||
}
|
||||
|
||||
function Start-AllSheepIt {
|
||||
Initialize-SheepItCredentials
|
||||
|
||||
$targets = $workers | Where-Object { $_.Enabled }
|
||||
|
||||
if (-not $targets) {
|
||||
Write-Host "No systems are ready for SheepIt." -ForegroundColor Yellow
|
||||
return
|
||||
}
|
||||
|
||||
foreach ($worker in $targets) {
|
||||
Start-SheepItWorker -Worker $worker
|
||||
Start-Sleep -Milliseconds 200
|
||||
}
|
||||
}
|
||||
|
||||
function Select-SheepItWorker {
|
||||
while ($true) {
|
||||
@@ -308,20 +412,22 @@ Initialize-SheepItCredentials
|
||||
while ($true) {
|
||||
Show-Header
|
||||
Write-Host "Main Menu:" -ForegroundColor Magenta
|
||||
Write-Host "1. Launch SheepIt on a single system" -ForegroundColor Yellow
|
||||
Write-Host "2. Launch SheepIt on all ready systems" -ForegroundColor Yellow
|
||||
Write-Host "3. Exit" -ForegroundColor Yellow
|
||||
Write-Host "1. Launch/Attach SheepIt on a single system" -ForegroundColor Yellow
|
||||
Write-Host "2. Ensure all ready systems are running" -ForegroundColor Yellow
|
||||
Write-Host "3. Pause all workers" -ForegroundColor Yellow
|
||||
Write-Host "4. Resume all workers" -ForegroundColor Yellow
|
||||
Write-Host "5. Quit all workers" -ForegroundColor Yellow
|
||||
Write-Host "6. Exit" -ForegroundColor Yellow
|
||||
|
||||
$choice = Read-Host "Select option (1-3)"
|
||||
|
||||
switch ($choice) {
|
||||
'1' { Select-SheepItWorker }
|
||||
'2' {
|
||||
Start-AllSheepIt
|
||||
Write-Host
|
||||
Read-Host "Press Enter to return to the main menu" | Out-Null
|
||||
}
|
||||
'3' { break }
|
||||
'2' { Start-AllSheepIt }
|
||||
'3' { Send-SheepItCommandAll -CommandText 'pause' }
|
||||
'4' { Send-SheepItCommandAll -CommandText 'resume' }
|
||||
'5' { Send-SheepItCommandAll -CommandText 'quit' }
|
||||
'6' { break }
|
||||
default {
|
||||
Write-Host "Invalid selection." -ForegroundColor Red
|
||||
Start-Sleep -Seconds 1
|
||||
|
||||
Reference in New Issue
Block a user