refactor: lot 543 — worker index.ts unified on worker-runtime.ts
- index.ts reduced from 583 LOC to 155 LOC (thin orchestrator) - All logic delegated to worker-runtime.ts exports: createShutdownController, createNodeExecutor, runPollCycle, waitForNextPollTick - Removed duplicated executeRun, executeNodeStub, syncQueuedRuns, etc. - Single source of truth for poll/dequeue/execute cycle - 6/6 worker-runtime tests pass Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
+86
-481
@@ -1,8 +1,7 @@
|
||||
/**
|
||||
* KXKM_Clown V2 — Node Engine Worker
|
||||
*
|
||||
* Connects to Postgres, polls for queued runs, executes graph nodes in
|
||||
* topological order using stub executors, and updates run status in the DB.
|
||||
* Thin orchestrator that delegates to worker-runtime.ts for all logic.
|
||||
*/
|
||||
|
||||
import {
|
||||
@@ -15,35 +14,22 @@ import {
|
||||
import {
|
||||
createNodeEngineRegistry,
|
||||
createQueueState,
|
||||
canDequeue,
|
||||
dequeue,
|
||||
enqueue,
|
||||
markComplete,
|
||||
topologicalSort,
|
||||
validateEdgeContracts,
|
||||
collectNodeInputs,
|
||||
resolveFinalStatus,
|
||||
listDefaultRuntimes,
|
||||
createRun,
|
||||
validateJobSpec,
|
||||
buildTrlCommand,
|
||||
DEFAULT_HYPERPARAMS,
|
||||
type TrainingJobSpec,
|
||||
type NodeRun,
|
||||
type RunStep,
|
||||
type GraphNode,
|
||||
type NodeGraph,
|
||||
type StepStatus,
|
||||
type RunStatus,
|
||||
type NodeEngineRegistry,
|
||||
type QueueState,
|
||||
} from "@kxkm/node-engine";
|
||||
import { createIsoTimestamp } from "@kxkm/core";
|
||||
import { formatOverviewLine, ansi } from "@kxkm/tui";
|
||||
import { formatOverviewLine } from "@kxkm/tui";
|
||||
import { execFile } from "node:child_process";
|
||||
import { promisify } from "node:util";
|
||||
import * as path from "node:path";
|
||||
|
||||
import {
|
||||
createShutdownController,
|
||||
createNodeExecutor,
|
||||
runPollCycle,
|
||||
waitForNextPollTick,
|
||||
type WorkerLogger,
|
||||
type WorkerConfig,
|
||||
} from "./worker-runtime.js";
|
||||
|
||||
const execFileAsync = promisify(execFile);
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -52,355 +38,42 @@ const execFileAsync = promisify(execFile);
|
||||
|
||||
const POLL_INTERVAL_MS = 2000;
|
||||
const MAX_CONCURRENCY = Number(process.env.NODE_ENGINE_MAX_CONCURRENCY) || 1;
|
||||
const STEP_DELAY_MS = 100;
|
||||
const DRY_RUN = process.argv.includes("--dry-run") || process.env.DRY_RUN === "1";
|
||||
const PYTHON_BIN = process.env.PYTHON_BIN || "/home/kxkm/venv/bin/python3";
|
||||
const SCRIPTS_DIR = process.env.SCRIPTS_DIR || path.resolve(process.cwd(), "scripts");
|
||||
const TRAINING_TIMEOUT_MS = Number(process.env.TRAINING_TIMEOUT_MS) || 60 * 60 * 1000; // 1h
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Logging helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const DEBUG = process.env.NODE_ENV !== "production" || process.env.DEBUG === "1";
|
||||
|
||||
function log(msg: string): void {
|
||||
if (!DEBUG) return;
|
||||
const ts = new Date().toISOString();
|
||||
console.log(`[${ts}] ${msg}`);
|
||||
}
|
||||
|
||||
function logError(msg: string, err?: unknown): void {
|
||||
const ts = new Date().toISOString();
|
||||
const detail = err instanceof Error ? err.message : String(err ?? "");
|
||||
console.error(`[${ts}] ERROR: ${msg}${detail ? " — " + detail : ""}`);
|
||||
}
|
||||
const config: WorkerConfig = {
|
||||
dryRun: DRY_RUN,
|
||||
stepDelayMs: 100,
|
||||
pythonBin: process.env.PYTHON_BIN || "/home/kxkm/venv/bin/python3",
|
||||
scriptsDir: process.env.SCRIPTS_DIR || path.resolve(process.cwd(), "scripts"),
|
||||
trainingTimeoutMs: Number(process.env.TRAINING_TIMEOUT_MS) || 60 * 60 * 1000,
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Stub executors
|
||||
// Logger
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type StubResult = Record<string, unknown>;
|
||||
|
||||
function delay(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
async function executeNodeStub(
|
||||
nodeType: string,
|
||||
inputs: Record<string, unknown>,
|
||||
params: Record<string, unknown>,
|
||||
): Promise<StubResult> {
|
||||
await delay(STEP_DELAY_MS);
|
||||
|
||||
switch (nodeType) {
|
||||
// Dataset sources — return a mock dataset
|
||||
case "dataset_file":
|
||||
case "dataset_folder":
|
||||
case "huggingface_dataset":
|
||||
case "web_scraper":
|
||||
return { dataset: { items: [], format: "stub" } };
|
||||
|
||||
// Data processing — pass through inputs
|
||||
case "clean_text":
|
||||
case "remove_duplicates":
|
||||
case "split_dataset":
|
||||
return { dataset: inputs.dataset ?? { items: [], format: "stub" } };
|
||||
|
||||
// Dataset builders — pass through as dataset_ready
|
||||
case "format_instruction_dataset":
|
||||
case "chat_dataset":
|
||||
return { dataset_ready: inputs.dataset ?? { items: [], format: "stub" } };
|
||||
|
||||
// Evaluation — execute via eval_model.py
|
||||
case "prompt_test":
|
||||
case "benchmark": {
|
||||
const evalModel = typeof params.model === "string" && params.model
|
||||
? params.model
|
||||
: (inputs.model as Record<string, unknown>)?.modelName as string || "unsloth/llama-3-8b";
|
||||
const adapterPath = (inputs.model as Record<string, unknown>)?.adapterPath as string || undefined;
|
||||
const promptsPath = typeof params.promptsPath === "string" ? params.promptsPath : "";
|
||||
const evalOutputPath = `/tmp/kxkm-eval-${Date.now()}.json`;
|
||||
|
||||
if (DRY_RUN) {
|
||||
log(` [dry-run] would evaluate model=${evalModel} adapter=${adapterPath || "none"}`);
|
||||
return { evaluation: { kind: "dry-run", score: 1 } };
|
||||
}
|
||||
|
||||
if (!promptsPath) {
|
||||
// No prompts file — return stub score
|
||||
log(` [eval] no promptsPath provided — returning stub evaluation`);
|
||||
return { evaluation: { kind: "stub", score: 1 } };
|
||||
}
|
||||
|
||||
const scriptPath = path.join(SCRIPTS_DIR, "eval_model.py");
|
||||
const args = [
|
||||
scriptPath,
|
||||
"--model", evalModel,
|
||||
"--prompts", promptsPath,
|
||||
"--output", evalOutputPath,
|
||||
];
|
||||
if (adapterPath) args.push("--adapter", adapterPath);
|
||||
|
||||
log(` [eval] ${PYTHON_BIN} ${args.join(" ")}`);
|
||||
|
||||
try {
|
||||
const { stdout, stderr } = await execFileAsync(PYTHON_BIN, args, {
|
||||
timeout: TRAINING_TIMEOUT_MS,
|
||||
maxBuffer: 50 * 1024 * 1024,
|
||||
});
|
||||
if (stderr) log(` [eval] stderr: ${stderr.slice(-500)}`);
|
||||
|
||||
const jsonLine = stdout.trim().split("\n").pop() || "{}";
|
||||
let evalResult: Record<string, unknown> = {};
|
||||
try {
|
||||
evalResult = JSON.parse(jsonLine);
|
||||
} catch (parseErr) {
|
||||
logError(` [eval] Failed to parse JSON output`, parseErr);
|
||||
}
|
||||
log(` [eval] result: status=${evalResult.status} score=${evalResult.score}`);
|
||||
|
||||
return {
|
||||
evaluation: {
|
||||
kind: "real",
|
||||
score: evalResult.score,
|
||||
metrics: evalResult.metrics,
|
||||
outputFile: evalOutputPath,
|
||||
},
|
||||
};
|
||||
} catch (err) {
|
||||
logError(` [eval] failed`, err);
|
||||
return { evaluation: { kind: "error", score: 0, error: err instanceof Error ? err.message : String(err) } };
|
||||
}
|
||||
}
|
||||
|
||||
// Training — execute via train_unsloth.py
|
||||
case "sft_training":
|
||||
case "lora_training":
|
||||
case "qlora_training": {
|
||||
const baseModel = typeof params.baseModel === "string" && params.baseModel
|
||||
? params.baseModel
|
||||
: "unsloth/llama-3-8b";
|
||||
const datasetPath = typeof params.datasetPath === "string" ? params.datasetPath : "";
|
||||
const outputDir = typeof params.outputDir === "string"
|
||||
? params.outputDir
|
||||
: `/tmp/kxkm-training-${Date.now()}`;
|
||||
const hp = {
|
||||
...DEFAULT_HYPERPARAMS,
|
||||
...(params.hyperparams && typeof params.hyperparams === "object" ? params.hyperparams : {}),
|
||||
};
|
||||
|
||||
if (DRY_RUN) {
|
||||
const jobSpec = validateJobSpec({
|
||||
type: nodeType as TrainingJobSpec["type"], baseModel,
|
||||
datasetPath: datasetPath || "/data/dataset.jsonl",
|
||||
outputDir, hyperparams: hp,
|
||||
});
|
||||
log(` [dry-run] would execute: ${buildTrlCommand(jobSpec)}`);
|
||||
return { model: { kind: "dry-run", modelName: `${baseModel}-finetuned`, jobSpec } };
|
||||
}
|
||||
|
||||
if (!datasetPath) {
|
||||
return { model: { kind: "error", error: "datasetPath is required for training" } };
|
||||
}
|
||||
|
||||
const scriptPath = path.join(SCRIPTS_DIR, "train_unsloth.py");
|
||||
const args = [
|
||||
scriptPath,
|
||||
"--model", baseModel,
|
||||
"--data", datasetPath,
|
||||
"--output", outputDir,
|
||||
"--method", params.dpo === true ? "dpo" : nodeType === "qlora_training" ? "qlora" : nodeType === "sft_training" ? "sft" : "lora",
|
||||
"--lr", String(hp.learningRate),
|
||||
"--epochs", String(hp.epochs),
|
||||
"--batch-size", String(hp.batchSize),
|
||||
"--lora-rank", String(hp.loraRank),
|
||||
"--lora-alpha", String(hp.loraAlpha),
|
||||
"--warmup-steps", String(hp.warmupSteps),
|
||||
"--max-seq-length", String(hp.maxSeqLength),
|
||||
];
|
||||
if (nodeType === "qlora_training") args.push("--quantize", "4bit");
|
||||
|
||||
log(` [training] ${PYTHON_BIN} ${args.join(" ")}`);
|
||||
|
||||
try {
|
||||
const { stdout, stderr } = await execFileAsync(PYTHON_BIN, args, {
|
||||
timeout: TRAINING_TIMEOUT_MS,
|
||||
maxBuffer: 50 * 1024 * 1024,
|
||||
});
|
||||
if (stderr) log(` [training] stderr: ${stderr.slice(-500)}`);
|
||||
|
||||
const jsonLine = stdout.trim().split("\n").pop() || "{}";
|
||||
let trainResult: Record<string, unknown> = {};
|
||||
try {
|
||||
trainResult = JSON.parse(jsonLine);
|
||||
} catch (parseErr) {
|
||||
logError(` [training] Failed to parse JSON output`, parseErr);
|
||||
}
|
||||
log(` [training] result: status=${trainResult.status} loss=${(trainResult.metrics as Record<string, unknown> | undefined)?.trainLoss}`);
|
||||
|
||||
return {
|
||||
model: {
|
||||
kind: "trained",
|
||||
modelName: `${baseModel}-finetuned`,
|
||||
adapterPath: trainResult.adapterPath || outputDir,
|
||||
metrics: trainResult.metrics,
|
||||
status: trainResult.status,
|
||||
},
|
||||
};
|
||||
} catch (err) {
|
||||
const msg = err instanceof Error ? err.message : String(err);
|
||||
logError(` [training] failed`, err);
|
||||
return { model: { kind: "error", error: msg } };
|
||||
}
|
||||
}
|
||||
|
||||
// Registry
|
||||
case "register_model":
|
||||
return { registered_model: { id: "stub" } };
|
||||
|
||||
// Deployment — import LoRA adapter into Ollama
|
||||
case "deploy_api": {
|
||||
const modelInput = inputs.registered_model as Record<string, unknown> || inputs.model as Record<string, unknown> || {};
|
||||
const adapterPath = modelInput.adapterPath as string || (typeof params.adapterPath === "string" ? params.adapterPath : "");
|
||||
const baseOllamaModel = typeof params.baseOllamaModel === "string" ? params.baseOllamaModel : "llama3.2:1b";
|
||||
const deployName = typeof params.deployName === "string" ? params.deployName : `kxkm-${Date.now()}`;
|
||||
|
||||
if (DRY_RUN || !adapterPath) {
|
||||
log(` [deploy] dry-run or no adapter: base=${baseOllamaModel} name=${deployName}`);
|
||||
return { deployment: { kind: DRY_RUN ? "dry-run" : "stub", id: deployName } };
|
||||
}
|
||||
|
||||
const scriptPath = path.join(SCRIPTS_DIR, "ollama-import-adapter.sh");
|
||||
const args = [scriptPath, "--base-model", baseOllamaModel, "--adapter-path", adapterPath, "--name", deployName];
|
||||
|
||||
log(` [deploy] importing to Ollama: ${deployName} from ${baseOllamaModel} + ${adapterPath}`);
|
||||
|
||||
try {
|
||||
const { stdout, stderr } = await execFileAsync("/bin/bash", args, { timeout: 120_000 });
|
||||
if (stderr) log(` [deploy] stderr: ${stderr.slice(-500)}`);
|
||||
const jsonLine = stdout.trim().split("\n").pop() || "{}";
|
||||
let result;
|
||||
try {
|
||||
result = JSON.parse(jsonLine);
|
||||
} catch {
|
||||
logError(` [deploy] invalid JSON: ${jsonLine.slice(0, 100)}`);
|
||||
return { deployment: { kind: "error", id: deployName, error: "invalid_json_output" } };
|
||||
}
|
||||
log(` [deploy] result: ${JSON.stringify(result)}`);
|
||||
return { deployment: { kind: "ollama", id: deployName, ...result } };
|
||||
} catch (err) {
|
||||
logError(` [deploy] failed`, err);
|
||||
return { deployment: { kind: "error", id: deployName, error: err instanceof Error ? err.message : String(err) } };
|
||||
}
|
||||
}
|
||||
|
||||
default:
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Run executor
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async function executeRun(
|
||||
run: NodeRun,
|
||||
registry: NodeEngineRegistry,
|
||||
options: { shouldCancel?: () => boolean } = {},
|
||||
): Promise<void> {
|
||||
const graph = run.graphSnapshot;
|
||||
const shouldCancel = options.shouldCancel ?? (() => false);
|
||||
|
||||
// Validate edges against registry contracts
|
||||
validateEdgeContracts(graph, registry);
|
||||
|
||||
// Get execution order
|
||||
const sorted = topologicalSort(graph);
|
||||
|
||||
// Track outputs per node for input collection
|
||||
const outputsByNode = new Map<string, Record<string, unknown>>();
|
||||
|
||||
// Mark run as running
|
||||
run.status = "running";
|
||||
run.startedAt = createIsoTimestamp();
|
||||
let cancelled = false;
|
||||
|
||||
log(` Executing ${sorted.length} node(s) in topological order`);
|
||||
|
||||
// Restore already-completed steps (recovery support)
|
||||
for (const node of sorted) {
|
||||
const step = run.steps.find((s) => s.id === node.id);
|
||||
if (step?.status === "completed") {
|
||||
log(` [${node.id}] ${node.type} — already completed (recovered)`);
|
||||
}
|
||||
}
|
||||
|
||||
for (const node of sorted) {
|
||||
const step = run.steps.find((s) => s.id === node.id);
|
||||
if (!step) continue;
|
||||
if (step.status === "completed") continue; // skip recovered steps
|
||||
|
||||
if (shouldCancel() || shutdownRequested) {
|
||||
cancelled = true;
|
||||
log(` [${node.id}] ${node.type} — cancelled`);
|
||||
break;
|
||||
}
|
||||
|
||||
step.status = "running";
|
||||
step.startedAt = createIsoTimestamp();
|
||||
|
||||
log(` [${node.id}] ${node.type} — running`);
|
||||
|
||||
try {
|
||||
const inputs = collectNodeInputs(graph, node.id, outputsByNode);
|
||||
const result = await executeNodeStub(node.type, inputs, node.params);
|
||||
outputsByNode.set(node.id, result);
|
||||
|
||||
step.status = "completed";
|
||||
step.finishedAt = createIsoTimestamp();
|
||||
step.outputs = Object.keys(result);
|
||||
|
||||
log(` [${node.id}] ${node.type} — completed`);
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : String(err);
|
||||
step.status = "failed";
|
||||
step.finishedAt = createIsoTimestamp();
|
||||
step.error = message;
|
||||
|
||||
logError(` [${node.id}] ${node.type} — failed`, err);
|
||||
break; // stop on first failure (like V1)
|
||||
}
|
||||
}
|
||||
|
||||
// Resolve final status from step statuses
|
||||
const stepStatuses = run.steps.map((s) => s.status);
|
||||
run.status = resolveFinalStatus(stepStatuses, cancelled);
|
||||
run.finishedAt = createIsoTimestamp();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shutdown handling
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
let shutdownRequested = false;
|
||||
|
||||
function requestShutdown(): void {
|
||||
if (shutdownRequested) return;
|
||||
shutdownRequested = true;
|
||||
log("Shutdown requested — finishing current work...");
|
||||
}
|
||||
const logger: WorkerLogger = {
|
||||
log(msg: string): void {
|
||||
if (!DEBUG) return;
|
||||
console.log(`[${new Date().toISOString()}] ${msg}`);
|
||||
},
|
||||
error(msg: string, err?: unknown): void {
|
||||
const detail = err instanceof Error ? err.message : String(err ?? "");
|
||||
console.error(`[${new Date().toISOString()}] ERROR: ${msg}${detail ? " — " + detail : ""}`);
|
||||
},
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Global error handlers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
process.on("unhandledRejection", (reason) => {
|
||||
logError("Unhandled promise rejection", reason);
|
||||
logger.error("Unhandled promise rejection", reason);
|
||||
});
|
||||
|
||||
process.on("uncaughtException", (err) => {
|
||||
logError("Uncaught exception — shutting down", err);
|
||||
logger.error("Uncaught exception — shutting down", err);
|
||||
process.exit(1);
|
||||
});
|
||||
|
||||
@@ -409,174 +82,106 @@ process.on("uncaughtException", (err) => {
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async function main(): Promise<void> {
|
||||
// Check DATABASE_URL
|
||||
if (!process.env.DATABASE_URL) {
|
||||
console.warn(
|
||||
"WARNING: DATABASE_URL is not set. " +
|
||||
"The worker requires a Postgres connection. Exiting.",
|
||||
);
|
||||
console.warn("WARNING: DATABASE_URL is not set. The worker requires a Postgres connection. Exiting.");
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
log("Starting KXKM_Clown V2 worker");
|
||||
logger.log("Starting KXKM_Clown V2 worker");
|
||||
|
||||
// 1. Connect to DB
|
||||
// Connect to DB
|
||||
const dbConfig = loadDatabaseConfig();
|
||||
const pool = createPostgresPool(dbConfig);
|
||||
log(`Connected to database: ${dbConfig.connectionString.replace(/\/\/.*@/, "//***@")}`);
|
||||
logger.log(`Connected to database: ${dbConfig.connectionString.replace(/\/\/.*@/, "//***@")}`);
|
||||
|
||||
// 2. Run migrations
|
||||
log("Running database migrations...");
|
||||
// Run migrations
|
||||
logger.log("Running database migrations...");
|
||||
await runMigrations(pool);
|
||||
log("Migrations complete");
|
||||
logger.log("Migrations complete");
|
||||
|
||||
// 3. Create repos
|
||||
// Create repos + registry + queue
|
||||
const graphRepo = createNodeGraphRepo(pool);
|
||||
const runRepo = createNodeRunRepo(pool);
|
||||
|
||||
// 4. Create registry with all built-in node types
|
||||
const registry = createNodeEngineRegistry();
|
||||
const nodeTypes = registry.listNodeTypes();
|
||||
const runtimes = listDefaultRuntimes();
|
||||
log(`Registry loaded: ${nodeTypes.length} node types, ${runtimes.length} runtimes`);
|
||||
logger.log(`Registry loaded: ${registry.listNodeTypes().length} node types, ${runtimes.length} runtimes`);
|
||||
|
||||
// 5. Create queue state
|
||||
const queueState = createQueueState({ maxConcurrency: MAX_CONCURRENCY });
|
||||
log(`Queue state initialised: maxConcurrency=${MAX_CONCURRENCY}`);
|
||||
logger.log(`Queue state initialised: maxConcurrency=${MAX_CONCURRENCY}`);
|
||||
|
||||
// 5b. Recovery — re-queue runs that were running when worker last crashed
|
||||
// Recovery
|
||||
const recovered = await runRepo.recoverStaleRuns();
|
||||
if (recovered.length > 0) {
|
||||
log(`Recovered ${recovered.length} stale run(s): ${recovered.map((r) => r.id).join(", ")}`);
|
||||
logger.log(`Recovered ${recovered.length} stale run(s): ${recovered.map((r) => r.id).join(", ")}`);
|
||||
}
|
||||
|
||||
// 6. Graceful shutdown with forced exit timeout
|
||||
// Shutdown handling
|
||||
const shutdown = createShutdownController();
|
||||
const SHUTDOWN_TIMEOUT_MS = 30_000;
|
||||
function handleShutdownSignal(signal: string) {
|
||||
log(`${signal} received`);
|
||||
requestShutdown();
|
||||
logger.log(`${signal} received`);
|
||||
shutdown.requestShutdown();
|
||||
setTimeout(() => {
|
||||
logError(`Forced exit after ${SHUTDOWN_TIMEOUT_MS}ms timeout`);
|
||||
logger.error(`Forced exit after ${SHUTDOWN_TIMEOUT_MS}ms timeout`);
|
||||
process.exit(1);
|
||||
}, SHUTDOWN_TIMEOUT_MS).unref();
|
||||
}
|
||||
process.on("SIGTERM", () => handleShutdownSignal("SIGTERM"));
|
||||
process.on("SIGINT", () => handleShutdownSignal("SIGINT"));
|
||||
|
||||
// 7. Poll loop
|
||||
log(`Entering poll loop (interval=${POLL_INTERVAL_MS}ms)`);
|
||||
// Create node executor
|
||||
const executeNode = createNodeExecutor(config, execFileAsync, logger);
|
||||
|
||||
while (!shutdownRequested) {
|
||||
// Poll loop
|
||||
logger.log(`Entering poll loop (interval=${POLL_INTERVAL_MS}ms)`);
|
||||
|
||||
while (!shutdown.isShutdownRequested()) {
|
||||
try {
|
||||
// Fetch queued runs from the DB
|
||||
const queuedDbRuns = await runRepo.listByStatus("queued", 20);
|
||||
|
||||
// Sync DB queued runs into in-memory queue state
|
||||
for (const dbRun of queuedDbRuns) {
|
||||
enqueue(queueState, dbRun.id);
|
||||
}
|
||||
|
||||
// Process runs while we have capacity
|
||||
while (canDequeue(queueState) && !shutdownRequested) {
|
||||
const runId = dequeue(queueState);
|
||||
if (!runId) break;
|
||||
|
||||
log(`Dequeued run: ${runId}`);
|
||||
|
||||
// Update status to running in DB
|
||||
await runRepo.updateStatus(runId, "running");
|
||||
|
||||
// Look up the run from the DB
|
||||
const dbRun = await runRepo.findById(runId);
|
||||
if (!dbRun) {
|
||||
logError(`Run ${runId} not found in DB — skipping`);
|
||||
markComplete(queueState, runId);
|
||||
continue;
|
||||
}
|
||||
|
||||
const graphRecord = await graphRepo.findById(dbRun.graphId);
|
||||
if (!graphRecord) {
|
||||
logError(`Graph ${dbRun.graphId} for run ${runId} not found — marking failed`);
|
||||
await runRepo.updateStatus(runId, "failed");
|
||||
markComplete(queueState, runId);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Build a NodeGraph from the record (minimal — no edges/nodes stored in DB yet,
|
||||
// so we create a stub graph from the record for now)
|
||||
const graph: NodeGraph = {
|
||||
id: graphRecord.id,
|
||||
name: graphRecord.name,
|
||||
description: graphRecord.description,
|
||||
nodes: [],
|
||||
edges: [],
|
||||
createdAt: createIsoTimestamp(),
|
||||
updatedAt: createIsoTimestamp(),
|
||||
};
|
||||
|
||||
// Create an in-memory run with steps from the graph
|
||||
const nodeRun = createRun(graph, "worker");
|
||||
// Override the id to match the DB run
|
||||
(nodeRun as { id: string }).id = runId;
|
||||
|
||||
try {
|
||||
await executeRun(nodeRun, registry, {
|
||||
shouldCancel: () => {
|
||||
// Check if run was cancelled in DB (async check would be better
|
||||
// but keeping it simple — the cancel is also checked via shutdownRequested)
|
||||
return shutdownRequested;
|
||||
},
|
||||
});
|
||||
|
||||
// Persist final status
|
||||
await runRepo.updateStatus(runId, nodeRun.status);
|
||||
log(`Run ${runId} finished with status: ${nodeRun.status}`);
|
||||
} catch (err) {
|
||||
logError(`Run ${runId} failed unexpectedly`, err);
|
||||
await runRepo.updateStatus(runId, "failed");
|
||||
} finally {
|
||||
markComplete(queueState, runId);
|
||||
}
|
||||
}
|
||||
|
||||
// Log overview periodically
|
||||
const overview = formatOverviewLine({
|
||||
queue: {
|
||||
desiredWorkers: MAX_CONCURRENCY,
|
||||
activeWorkers: queueState.running.length,
|
||||
queuedRuns: queueState.queued.length,
|
||||
runningRuns: queueState.running.length,
|
||||
},
|
||||
registry: {
|
||||
graphs: (await graphRepo.list()).length,
|
||||
models: 0,
|
||||
},
|
||||
storage: {
|
||||
backend: "postgres",
|
||||
artifacts: "filesystem",
|
||||
},
|
||||
const result = await runPollCycle({
|
||||
queueState,
|
||||
runRepo,
|
||||
graphRepo,
|
||||
registry,
|
||||
executeNode,
|
||||
shutdown,
|
||||
logger,
|
||||
});
|
||||
|
||||
if (queuedDbRuns.length > 0) {
|
||||
log(`Poll status: ${overview}`);
|
||||
if (result.queuedDbRuns > 0) {
|
||||
const overview = formatOverviewLine({
|
||||
queue: {
|
||||
desiredWorkers: MAX_CONCURRENCY,
|
||||
activeWorkers: queueState.running.length,
|
||||
queuedRuns: queueState.queued.length,
|
||||
runningRuns: queueState.running.length,
|
||||
},
|
||||
registry: {
|
||||
graphs: (await graphRepo.list()).length,
|
||||
models: 0,
|
||||
},
|
||||
storage: {
|
||||
backend: "postgres",
|
||||
artifacts: "filesystem",
|
||||
},
|
||||
});
|
||||
logger.log(`Poll status: ${overview}`);
|
||||
}
|
||||
} catch (err) {
|
||||
logError("Poll loop error", err);
|
||||
logger.error("Poll loop error", err);
|
||||
}
|
||||
|
||||
// Wait before next poll (check shutdown frequently)
|
||||
if (!shutdownRequested) {
|
||||
await delay(POLL_INTERVAL_MS);
|
||||
if (!shutdown.isShutdownRequested()) {
|
||||
await waitForNextPollTick(POLL_INTERVAL_MS);
|
||||
}
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
log("Shutting down...");
|
||||
logger.log("Shutting down...");
|
||||
await pool.end();
|
||||
log("Worker stopped");
|
||||
logger.log("Worker stopped");
|
||||
}
|
||||
|
||||
main().catch((err) => {
|
||||
logError("Fatal error in worker", err);
|
||||
logger.error("Fatal error in worker", err);
|
||||
process.exit(1);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user