Files
TradeBot/tools/install_windows_training_agent.ps1
T

169 lines
6.1 KiB
PowerShell

[CmdletBinding()]
param(
[string]$TaskName = "TradeBot Windows Training Agent",
[string]$ApiBaseUrl = "https://tb.kusoft.xyz",
[string]$ApiAuth = "",
[int]$PollSeconds = 10,
[int]$WatchdogMinutes = 5,
[string]$RepoRoot = "",
[string]$CredentialPath = "",
[switch]$StartNow,
[switch]$KeepLegacyRetrainer
)
$ErrorActionPreference = "Stop"
if (-not $RepoRoot) {
$RepoRoot = (Resolve-Path (Join-Path $PSScriptRoot "..")).Path
}
if (-not $CredentialPath) {
$CredentialPath = Join-Path $env:LOCALAPPDATA "TradeBot\training-agent.token"
}
$runner = Join-Path $RepoRoot "tools\run_windows_training_agent.ps1"
if (-not (Test-Path -LiteralPath $runner)) {
throw "Windows training agent runner not found: $runner"
}
$credentialDirectory = Split-Path -Parent $CredentialPath
New-Item -ItemType Directory -Path $credentialDirectory -Force | Out-Null
if ($ApiAuth) {
$ApiAuth.Trim() |
ConvertTo-SecureString -AsPlainText -Force |
ConvertFrom-SecureString |
Set-Content -LiteralPath $CredentialPath -Encoding UTF8
}
if (-not (Test-Path -LiteralPath $CredentialPath)) {
throw "ApiAuth is required for the first installation."
}
# Remove the legacy plaintext secret from the user environment. The new runner
# decrypts the DPAPI-protected credential only inside the agent process tree.
[Environment]::SetEnvironmentVariable("TRADEBOT_API_AUTH", $null, "User")
Remove-Item Env:TRADEBOT_API_AUTH -ErrorAction SilentlyContinue
[Environment]::SetEnvironmentVariable("TRADEBOT_API_BASE_URL", $ApiBaseUrl, "User")
[Environment]::SetEnvironmentVariable("TRADEBOT_TRAINING_WORKER_NAME", $env:COMPUTERNAME, "User")
if (-not $KeepLegacyRetrainer) {
foreach ($legacyName in @("TradeBot PyTorch Forecaster Retrainer", "TradeBot LSTM Retrainer")) {
try {
$legacyTask = Get-ScheduledTask -TaskName $legacyName -ErrorAction SilentlyContinue
if ($legacyTask) {
Unregister-ScheduledTask -TaskName $legacyName -Confirm:$false
Write-Host "Removed legacy scheduled task '$legacyName'."
}
}
catch {
Write-Warning "Could not remove legacy scheduled task '$legacyName': $($_.Exception.Message)"
}
}
}
$currentUser = [System.Security.Principal.WindowsIdentity]::GetCurrent().Name
$principal = New-Object System.Security.Principal.WindowsPrincipal(
[System.Security.Principal.WindowsIdentity]::GetCurrent()
)
$isAdministrator = $principal.IsInRole([System.Security.Principal.WindowsBuiltInRole]::Administrator)
$powershell = (Get-Command powershell.exe -ErrorAction Stop).Source
$runnerArguments = @(
"-NoProfile",
"-WindowStyle", "Hidden",
"-ExecutionPolicy", "Bypass",
"-File", "`"$runner`"",
"-RepoRoot", "`"$RepoRoot`"",
"-ApiBaseUrl", "`"$ApiBaseUrl`"",
"-CredentialPath", "`"$CredentialPath`"",
"-WorkerName", "`"$env:COMPUTERNAME`"",
"-PollSeconds", $PollSeconds.ToString()
) -join " "
$startupShortcut = Join-Path ([Environment]::GetFolderPath("Startup")) "$TaskName.lnk"
$runKey = "HKCU:\Software\Microsoft\Windows\CurrentVersion\Run"
Remove-ItemProperty -Path $runKey -Name "TradeBotWindowsTrainingAgent" -ErrorAction SilentlyContinue
$installMode = "startup shortcut"
if ($isAdministrator) {
if (Test-Path -LiteralPath $startupShortcut) {
Remove-Item -LiteralPath $startupShortcut -Force
}
$action = New-ScheduledTaskAction -Execute $powershell -Argument $runnerArguments -WorkingDirectory $RepoRoot
$trigger = @(
New-ScheduledTaskTrigger -AtLogOn -User $currentUser
New-ScheduledTaskTrigger -AtStartup
New-ScheduledTaskTrigger `
-Once `
-At (Get-Date).AddMinutes(1) `
-RepetitionInterval (New-TimeSpan -Minutes $WatchdogMinutes) `
-RepetitionDuration (New-TimeSpan -Days 3650)
)
$taskPrincipal = New-ScheduledTaskPrincipal `
-UserId $currentUser `
-LogonType Interactive `
-RunLevel Limited
$settings = New-ScheduledTaskSettingsSet `
-StartWhenAvailable `
-MultipleInstances IgnoreNew `
-AllowStartIfOnBatteries `
-DontStopIfGoingOnBatteries `
-RestartCount 999 `
-RestartInterval (New-TimeSpan -Minutes 1) `
-ExecutionTimeLimit (New-TimeSpan -Days 30)
Register-ScheduledTask `
-TaskName $TaskName `
-Action $action `
-Trigger $trigger `
-Principal $taskPrincipal `
-Settings $settings `
-Description "Keeps the TradeBot Windows training agent online and polls the bot API for retrain jobs." `
-Force | Out-Null
$installMode = "scheduled task"
}
else {
try {
$existingTask = Get-ScheduledTask -TaskName $TaskName -ErrorAction SilentlyContinue
if ($existingTask) {
Unregister-ScheduledTask -TaskName $TaskName -Confirm:$false
}
}
catch {
Write-Warning "Could not remove an existing elevated task: $($_.Exception.Message)"
}
$shell = New-Object -ComObject WScript.Shell
$shortcut = $shell.CreateShortcut($startupShortcut)
$shortcut.TargetPath = $powershell
$shortcut.Arguments = $runnerArguments
$shortcut.WorkingDirectory = $RepoRoot
$shortcut.WindowStyle = 7
$shortcut.Description = "TradeBot Windows Training Agent"
$shortcut.Save()
}
Get-CimInstance Win32_Process |
Where-Object {
$_.ProcessId -ne $PID -and
$_.CommandLine -and
($_.CommandLine -match [regex]::Escape("windows_training_agent.py") -or
$_.CommandLine -match [regex]::Escape("run_windows_training_agent.ps1"))
} |
ForEach-Object { Stop-Process -Id $_.ProcessId -Force -ErrorAction SilentlyContinue }
if ($StartNow) {
if ($installMode -eq "scheduled task") {
Start-ScheduledTask -TaskName $TaskName
}
else {
Start-Process `
-FilePath $powershell `
-ArgumentList $runnerArguments `
-WorkingDirectory $RepoRoot `
-WindowStyle Hidden | Out-Null
}
}
Write-Host "Installed '$TaskName' using $installMode."
Write-Host "Agent API: $ApiBaseUrl"
Write-Host "Encrypted credential: $CredentialPath"
Write-Host "Agent runner: $runner"