chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,54 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { ComputeClientError } from "@internal/compute";
|
||||
import { isRetryableCreateError, runnerNameForAttempt } from "./compute.js";
|
||||
|
||||
describe("runnerNameForAttempt", () => {
|
||||
it("keeps the unsuffixed name for the first attempt", () => {
|
||||
expect(runnerNameForAttempt("runner-abc123", 1)).toBe("runner-abc123");
|
||||
});
|
||||
|
||||
it("suffixes retry attempts deterministically", () => {
|
||||
expect(runnerNameForAttempt("runner-abc123", 2)).toBe("runner-abc123-r2");
|
||||
expect(runnerNameForAttempt("runner-abc123", 3)).toBe("runner-abc123-r3");
|
||||
});
|
||||
});
|
||||
|
||||
describe("isRetryableCreateError", () => {
|
||||
it("retries statuses where the create definitely did not commit", () => {
|
||||
expect(isRetryableCreateError(new ComputeClientError(500, "tap busy", "http://gw"))).toBe(true);
|
||||
expect(isRetryableCreateError(new ComputeClientError(503, "no placement", "http://gw"))).toBe(
|
||||
true
|
||||
);
|
||||
});
|
||||
|
||||
it("does not retry lost-response statuses (create may have committed)", () => {
|
||||
expect(isRetryableCreateError(new ComputeClientError(502, "bad gateway", "http://gw"))).toBe(
|
||||
false
|
||||
);
|
||||
expect(
|
||||
isRetryableCreateError(new ComputeClientError(504, "gateway timeout", "http://gw"))
|
||||
).toBe(false);
|
||||
});
|
||||
|
||||
it("does not retry 4xx responses", () => {
|
||||
expect(isRetryableCreateError(new ComputeClientError(400, "bad request", "http://gw"))).toBe(
|
||||
false
|
||||
);
|
||||
expect(isRetryableCreateError(new ComputeClientError(409, "conflict", "http://gw"))).toBe(
|
||||
false
|
||||
);
|
||||
});
|
||||
|
||||
it("does not retry timeouts (instance may still be provisioning)", () => {
|
||||
expect(isRetryableCreateError(new DOMException("timed out", "TimeoutError"))).toBe(false);
|
||||
});
|
||||
|
||||
it("retries network-level fetch failures", () => {
|
||||
expect(isRetryableCreateError(new TypeError("fetch failed"))).toBe(true);
|
||||
});
|
||||
|
||||
it("does not retry unknown errors", () => {
|
||||
expect(isRetryableCreateError(new Error("something else"))).toBe(false);
|
||||
expect(isRetryableCreateError("string error")).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,511 @@
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import { parseTraceparent } from "@trigger.dev/core/v3/isomorphic";
|
||||
import { flattenAttributes } from "@trigger.dev/core/v3/utils/flattenAttributes";
|
||||
import {
|
||||
type WorkloadManager,
|
||||
type WorkloadManagerCreateOptions,
|
||||
type WorkloadManagerOptions,
|
||||
} from "./types.js";
|
||||
import { ComputeClient, ComputeClientError, stripImageDigest } from "@internal/compute";
|
||||
import { setTimeout as sleep } from "node:timers/promises";
|
||||
import { extractTraceparent, getRunnerId } from "../util.js";
|
||||
import type { OtlpTraceService } from "../services/otlpTraceService.js";
|
||||
import { tryCatch } from "@trigger.dev/core";
|
||||
import { encodeBaggage, fromContext } from "../wideEvents/index.js";
|
||||
|
||||
const DEFAULT_CREATE_MAX_ATTEMPTS = 3;
|
||||
const DEFAULT_CREATE_RETRY_BASE_DELAY_MS = 250;
|
||||
|
||||
/**
|
||||
* TEMPORARY (TRI-10293): a failed create can leave its instance name
|
||||
* registered gateway/fcrun-side until async cleanup runs, so a same-name
|
||||
* retry can 409 against our own residue. Until the gateway cleans up
|
||||
* failed-create registrations properly, retry attempts get a deterministic
|
||||
* suffix. Attempt 1 keeps the unsuffixed name so the non-retry path is
|
||||
* unchanged; the suffixed name flows into both the instance name and
|
||||
* TRIGGER_RUNNER_ID, which downstream flows treat as one opaque
|
||||
* self-reported token. Only attempts following a ComputeClientError are
|
||||
* suffixed - network-failure retries keep the same name on purpose, because
|
||||
* the gateway's name-collision 409 is their safety net against
|
||||
* double-creating an instance whose create response was lost.
|
||||
*/
|
||||
export function runnerNameForAttempt(runnerId: string, attempt: number): string {
|
||||
return attempt === 1 ? runnerId : `${runnerId}-r${attempt}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Whether a failed instance create is worth retrying. Only statuses where
|
||||
* the create definitely did NOT commit are retried: 500 means the agent or
|
||||
* fcrun returned a create error (e.g. a netns slot holding the tap busy, a
|
||||
* full node disk - placement may differ on retry), 503 means the gateway
|
||||
* had nowhere to place it. 502/504 are excluded: the gateway emits those
|
||||
* when it fails to reach the node or read its response, which can happen
|
||||
* AFTER the agent committed the create - and the gateway only records the
|
||||
* instance name on a clean 201, so a same-name retry would miss the
|
||||
* collision check and could double-create the VM on another node. 4xx won't
|
||||
* heal on retry, and timeouts may still be provisioning. Network-level
|
||||
* fetch failures are safe: if the gateway processed the create, its name
|
||||
* index is populated and the retry 409s harmlessly.
|
||||
*/
|
||||
export function isRetryableCreateError(error: unknown): boolean {
|
||||
if (error instanceof ComputeClientError) {
|
||||
return error.status === 500 || error.status === 503;
|
||||
}
|
||||
if (error instanceof DOMException && error.name === "TimeoutError") {
|
||||
return false;
|
||||
}
|
||||
// Network-level fetch failures (gateway briefly unreachable)
|
||||
return error instanceof TypeError;
|
||||
}
|
||||
|
||||
type ComputeWorkloadManagerOptions = WorkloadManagerOptions & {
|
||||
gateway: {
|
||||
url: string;
|
||||
authToken?: string;
|
||||
timeoutMs: number;
|
||||
};
|
||||
snapshots: {
|
||||
enabled: boolean;
|
||||
delayMs: number;
|
||||
dispatchLimit: number;
|
||||
callbackUrl: string;
|
||||
};
|
||||
tracing?: OtlpTraceService;
|
||||
runner: {
|
||||
instanceName: string;
|
||||
otelEndpoint: string;
|
||||
prettyLogs: boolean;
|
||||
sendRunDebugLogs: boolean;
|
||||
};
|
||||
createRetry?: {
|
||||
maxAttempts: number;
|
||||
baseDelayMs: number;
|
||||
};
|
||||
};
|
||||
|
||||
export class ComputeWorkloadManager implements WorkloadManager {
|
||||
private readonly logger = new SimpleStructuredLogger("compute-workload-manager");
|
||||
private readonly compute: ComputeClient;
|
||||
private readonly createMaxAttempts: number;
|
||||
private readonly createRetryBaseDelayMs: number;
|
||||
|
||||
constructor(private opts: ComputeWorkloadManagerOptions) {
|
||||
this.createMaxAttempts = opts.createRetry?.maxAttempts ?? DEFAULT_CREATE_MAX_ATTEMPTS;
|
||||
this.createRetryBaseDelayMs =
|
||||
opts.createRetry?.baseDelayMs ?? DEFAULT_CREATE_RETRY_BASE_DELAY_MS;
|
||||
|
||||
if (opts.workloadApiDomain) {
|
||||
this.logger.warn("⚠️ Custom workload API domain", {
|
||||
domain: opts.workloadApiDomain,
|
||||
});
|
||||
}
|
||||
|
||||
this.compute = new ComputeClient({
|
||||
gatewayUrl: opts.gateway.url,
|
||||
authToken: opts.gateway.authToken,
|
||||
timeoutMs: opts.gateway.timeoutMs,
|
||||
// Forward the current wide-event scope's traceparent + request_id so the
|
||||
// downstream service continues the same trace and joins its own wide
|
||||
// events to ours. Additionally serialize caller-supplied meta labels
|
||||
// into the W3C Baggage header so the downstream service auto-stamps
|
||||
// them even on early-error paths that bail before parsing the body.
|
||||
// When called outside a wide-event scope (or when wide events are
|
||||
// disabled), `fromContext` returns undefined and propagation is skipped.
|
||||
getPropagationHeaders: () => {
|
||||
const state = fromContext();
|
||||
if (!state) return {};
|
||||
const headers: Record<string, string> = { "x-request-id": state.requestId };
|
||||
if (state.traceparent) {
|
||||
headers.traceparent = state.traceparent;
|
||||
}
|
||||
const baggage = encodeBaggage(state.meta);
|
||||
if (baggage) {
|
||||
headers.baggage = baggage;
|
||||
}
|
||||
return headers;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
get snapshotsEnabled(): boolean {
|
||||
return this.opts.snapshots.enabled;
|
||||
}
|
||||
|
||||
get snapshotDelayMs(): number {
|
||||
return this.opts.snapshots.delayMs;
|
||||
}
|
||||
|
||||
get snapshotDispatchLimit(): number {
|
||||
return this.opts.snapshots.dispatchLimit;
|
||||
}
|
||||
|
||||
get traceSpansEnabled(): boolean {
|
||||
return !!this.opts.tracing;
|
||||
}
|
||||
|
||||
async create(opts: WorkloadManagerCreateOptions) {
|
||||
const runnerId = getRunnerId(opts.runFriendlyId, opts.nextAttemptNumber);
|
||||
|
||||
const envVars: Record<string, string> = {
|
||||
OTEL_EXPORTER_OTLP_ENDPOINT: this.opts.runner.otelEndpoint,
|
||||
TRIGGER_DEQUEUED_AT_MS: String(opts.dequeuedAt.getTime()),
|
||||
TRIGGER_POD_SCHEDULED_AT_MS: String(Date.now()),
|
||||
TRIGGER_ENV_ID: opts.envId,
|
||||
TRIGGER_DEPLOYMENT_ID: opts.deploymentFriendlyId,
|
||||
TRIGGER_DEPLOYMENT_VERSION: opts.deploymentVersion,
|
||||
TRIGGER_RUN_ID: opts.runFriendlyId,
|
||||
TRIGGER_SNAPSHOT_ID: opts.snapshotFriendlyId,
|
||||
TRIGGER_SUPERVISOR_API_PROTOCOL: this.opts.workloadApiProtocol,
|
||||
TRIGGER_SUPERVISOR_API_PORT: String(this.opts.workloadApiPort),
|
||||
TRIGGER_SUPERVISOR_API_DOMAIN: this.opts.workloadApiDomain ?? "",
|
||||
TRIGGER_WORKER_INSTANCE_NAME: this.opts.runner.instanceName,
|
||||
TRIGGER_RUNNER_ID: runnerId,
|
||||
TRIGGER_MACHINE_CPU: String(opts.machine.cpu),
|
||||
TRIGGER_MACHINE_MEMORY: String(opts.machine.memory),
|
||||
PRETTY_LOGS: String(this.opts.runner.prettyLogs),
|
||||
TRIGGER_SEND_RUN_DEBUG_LOGS: String(this.opts.runner.sendRunDebugLogs),
|
||||
};
|
||||
|
||||
if (this.opts.warmStartUrl) {
|
||||
envVars.TRIGGER_WARM_START_URL = this.opts.warmStartUrl;
|
||||
}
|
||||
|
||||
if (this.snapshotsEnabled && this.opts.metadataUrl) {
|
||||
envVars.TRIGGER_METADATA_URL = this.opts.metadataUrl;
|
||||
}
|
||||
|
||||
if (this.opts.heartbeatIntervalSeconds) {
|
||||
envVars.TRIGGER_HEARTBEAT_INTERVAL_SECONDS = String(this.opts.heartbeatIntervalSeconds);
|
||||
}
|
||||
|
||||
if (this.opts.snapshotPollIntervalSeconds) {
|
||||
envVars.TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS = String(
|
||||
this.opts.snapshotPollIntervalSeconds
|
||||
);
|
||||
}
|
||||
|
||||
if (this.opts.additionalEnvVars) {
|
||||
Object.assign(envVars, this.opts.additionalEnvVars);
|
||||
}
|
||||
|
||||
// Strip image digest - resolve by tag, not digest
|
||||
const imageRef = stripImageDigest(opts.image);
|
||||
|
||||
// Labels forwarded to the compute provider for network-policy selection.
|
||||
// `org` is always set so every run carries its org identity.
|
||||
const labels: Record<string, string> = {
|
||||
org: opts.orgId,
|
||||
};
|
||||
|
||||
// Wide event: single canonical log line emitted in finally
|
||||
const event: Record<string, unknown> = {
|
||||
// High-cardinality identifiers
|
||||
runId: opts.runFriendlyId,
|
||||
runnerId,
|
||||
envId: opts.envId,
|
||||
envType: opts.envType,
|
||||
orgId: opts.orgId,
|
||||
projectId: opts.projectId,
|
||||
deploymentVersion: opts.deploymentVersion,
|
||||
machine: opts.machine.name,
|
||||
// Environment
|
||||
instanceName: this.opts.runner.instanceName,
|
||||
// Supervisor timing
|
||||
dequeueResponseMs: opts.dequeueResponseMs,
|
||||
pollingIntervalMs: opts.pollingIntervalMs,
|
||||
warmStartCheckMs: opts.warmStartCheckMs,
|
||||
// Request
|
||||
image: imageRef,
|
||||
};
|
||||
|
||||
const startMs = performance.now();
|
||||
|
||||
try {
|
||||
const createRequest = {
|
||||
name: runnerId,
|
||||
image: imageRef,
|
||||
env: envVars,
|
||||
cpu: opts.machine.cpu,
|
||||
memory_gb: opts.machine.memory,
|
||||
metadata: {
|
||||
runId: opts.runFriendlyId,
|
||||
envId: opts.envId,
|
||||
envType: opts.envType,
|
||||
orgId: opts.orgId,
|
||||
projectId: opts.projectId,
|
||||
deploymentVersion: opts.deploymentVersion,
|
||||
machine: opts.machine.name,
|
||||
},
|
||||
...(Object.keys(labels).length > 0 ? { labels } : {}),
|
||||
};
|
||||
|
||||
// Retry transient placement failures instead of abandoning the run: a
|
||||
// swallowed create error leaves the run waiting for the run engine's
|
||||
// PENDING_EXECUTING timeout (minutes) before it is redriven, while a
|
||||
// retried create typically succeeds in under a second (TRI-10293).
|
||||
let error: unknown;
|
||||
let data: Awaited<ReturnType<typeof this.compute.instances.create>> | null | undefined;
|
||||
let attempt = 1;
|
||||
// Set after a ComputeClientError: the failed create may have left its
|
||||
// name registered, so subsequent attempts use a suffixed name.
|
||||
let suffixAttempts = false;
|
||||
for (; attempt <= this.createMaxAttempts; attempt++) {
|
||||
const attemptRunnerId = suffixAttempts ? runnerNameForAttempt(runnerId, attempt) : runnerId;
|
||||
[error, data] = await tryCatch(
|
||||
this.compute.instances.create(
|
||||
attemptRunnerId === runnerId
|
||||
? createRequest
|
||||
: {
|
||||
...createRequest,
|
||||
name: attemptRunnerId,
|
||||
env: { ...envVars, TRIGGER_RUNNER_ID: attemptRunnerId },
|
||||
}
|
||||
)
|
||||
);
|
||||
|
||||
if (!error) {
|
||||
event.runnerId = attemptRunnerId;
|
||||
break;
|
||||
}
|
||||
|
||||
if (error instanceof ComputeClientError) {
|
||||
suffixAttempts = true;
|
||||
}
|
||||
|
||||
this.logger.warn("create instance attempt failed", {
|
||||
runnerId: attemptRunnerId,
|
||||
attempt,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
|
||||
if (!isRetryableCreateError(error) || attempt === this.createMaxAttempts) break;
|
||||
await sleep(this.createRetryBaseDelayMs * attempt);
|
||||
}
|
||||
event.createAttempts = attempt;
|
||||
|
||||
if (error || !data) {
|
||||
event.error = error instanceof Error ? error.message : String(error);
|
||||
event.errorType =
|
||||
error instanceof DOMException && error.name === "TimeoutError" ? "timeout" : "fetch";
|
||||
// Intentional: errors are captured in the wide event, not thrown. This matches
|
||||
// the Docker/K8s managers. The run will eventually time out if scheduling fails.
|
||||
return;
|
||||
}
|
||||
|
||||
event.instanceId = data.id;
|
||||
event.ok = true;
|
||||
|
||||
// Parse timing data from compute response (optional - requires gateway timing flag)
|
||||
if (data._timing) {
|
||||
event.timing = data._timing;
|
||||
}
|
||||
|
||||
this.#emitProvisionSpan(opts, startMs, data._timing);
|
||||
} finally {
|
||||
event.durationMs = Math.round(performance.now() - startMs);
|
||||
event.ok ??= false;
|
||||
this.logger.debug("create instance", event);
|
||||
}
|
||||
}
|
||||
|
||||
async snapshot(opts: { runnerId: string; metadata: Record<string, string> }): Promise<boolean> {
|
||||
const [error] = await tryCatch(
|
||||
this.compute.instances.snapshot(opts.runnerId, {
|
||||
callback: {
|
||||
url: this.opts.snapshots.callbackUrl,
|
||||
metadata: opts.metadata,
|
||||
},
|
||||
})
|
||||
);
|
||||
|
||||
if (error) {
|
||||
this.logger.error("snapshot request failed", {
|
||||
runnerId: opts.runnerId,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
this.logger.debug("snapshot request accepted", { runnerId: opts.runnerId });
|
||||
return true;
|
||||
}
|
||||
|
||||
async deleteInstance(runnerId: string): Promise<boolean> {
|
||||
const [error] = await tryCatch(this.compute.instances.delete(runnerId));
|
||||
|
||||
if (error) {
|
||||
this.logger.error("delete instance failed", {
|
||||
runnerId,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
this.logger.debug("delete instance success", { runnerId });
|
||||
return true;
|
||||
}
|
||||
|
||||
#emitProvisionSpan(opts: WorkloadManagerCreateOptions, startMs: number, timing?: unknown) {
|
||||
if (!this.traceSpansEnabled) return;
|
||||
|
||||
const parsed = parseTraceparent(extractTraceparent(opts.traceContext));
|
||||
if (!parsed) return;
|
||||
|
||||
const endMs = performance.now();
|
||||
const now = Date.now();
|
||||
const provisionStartEpochMs = now - (endMs - startMs);
|
||||
const endEpochMs = now;
|
||||
|
||||
// Span starts at dequeue time so events (dequeue) render in the thin-line section
|
||||
// before "Started". The actual provision call time is in provisionStartEpochMs.
|
||||
// Subtract 1ms so compute span always sorts before the attempt span (same dequeue time)
|
||||
const startEpochMs = opts.dequeuedAt.getTime() - 1;
|
||||
|
||||
const spanAttributes: Record<string, string | number | boolean> = {
|
||||
"compute.type": "create",
|
||||
"compute.provision_start_ms": provisionStartEpochMs,
|
||||
...(timing
|
||||
? (flattenAttributes(timing, "compute") as Record<string, string | number | boolean>)
|
||||
: {}),
|
||||
};
|
||||
|
||||
if (opts.dequeueResponseMs !== undefined) {
|
||||
spanAttributes["supervisor.dequeue_response_ms"] = opts.dequeueResponseMs;
|
||||
}
|
||||
if (opts.warmStartCheckMs !== undefined) {
|
||||
spanAttributes["supervisor.warm_start_check_ms"] = opts.warmStartCheckMs;
|
||||
}
|
||||
|
||||
// Use the platform API URL, not the runner OTLP endpoint (which may be a VM gateway IP)
|
||||
this.opts.tracing?.emit({
|
||||
traceId: parsed.traceId,
|
||||
parentSpanId: parsed.spanId,
|
||||
spanName: "compute.provision",
|
||||
startTimeMs: startEpochMs,
|
||||
endTimeMs: endEpochMs,
|
||||
resourceAttributes: {
|
||||
"ctx.environment.id": opts.envId,
|
||||
"ctx.organization.id": opts.orgId,
|
||||
"ctx.project.id": opts.projectId,
|
||||
"ctx.run.id": opts.runFriendlyId,
|
||||
},
|
||||
spanAttributes,
|
||||
});
|
||||
}
|
||||
|
||||
async restore(opts: {
|
||||
snapshotId: string;
|
||||
runnerId: string;
|
||||
runFriendlyId: string;
|
||||
snapshotFriendlyId: string;
|
||||
machine: { cpu: number; memory: number };
|
||||
// Trace context for OTel span emission
|
||||
traceContext?: Record<string, unknown>;
|
||||
envId?: string;
|
||||
orgId?: string;
|
||||
projectId?: string;
|
||||
hasPrivateLink?: boolean;
|
||||
dequeuedAt?: Date;
|
||||
}): Promise<boolean> {
|
||||
const metadata: Record<string, string> = {
|
||||
TRIGGER_RUNNER_ID: opts.runnerId,
|
||||
TRIGGER_RUN_ID: opts.runFriendlyId,
|
||||
TRIGGER_SNAPSHOT_ID: opts.snapshotFriendlyId,
|
||||
TRIGGER_SUPERVISOR_API_PROTOCOL: this.opts.workloadApiProtocol,
|
||||
TRIGGER_SUPERVISOR_API_PORT: String(this.opts.workloadApiPort),
|
||||
TRIGGER_SUPERVISOR_API_DOMAIN: this.opts.workloadApiDomain ?? "",
|
||||
TRIGGER_WORKER_INSTANCE_NAME: this.opts.runner.instanceName,
|
||||
};
|
||||
|
||||
// Resupply labels on restore (the provider doesn't persist them across a
|
||||
// snapshot). orgId is optional on the restore opts type, so guard it.
|
||||
const labels: Record<string, string> = {};
|
||||
if (opts.orgId) {
|
||||
labels.org = opts.orgId;
|
||||
}
|
||||
|
||||
this.logger.verbose("restore request body", {
|
||||
snapshotId: opts.snapshotId,
|
||||
runnerId: opts.runnerId,
|
||||
});
|
||||
|
||||
const startMs = performance.now();
|
||||
|
||||
const [error] = await tryCatch(
|
||||
this.compute.snapshots.restore(opts.snapshotId, {
|
||||
name: opts.runnerId,
|
||||
metadata,
|
||||
cpu: opts.machine.cpu,
|
||||
memory_gb: opts.machine.memory,
|
||||
...(Object.keys(labels).length > 0 ? { labels } : {}),
|
||||
})
|
||||
);
|
||||
|
||||
const durationMs = Math.round(performance.now() - startMs);
|
||||
|
||||
if (error) {
|
||||
this.logger.error("restore request failed", {
|
||||
snapshotId: opts.snapshotId,
|
||||
runnerId: opts.runnerId,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
durationMs,
|
||||
});
|
||||
return false;
|
||||
}
|
||||
|
||||
this.logger.debug("restore request success", {
|
||||
snapshotId: opts.snapshotId,
|
||||
runnerId: opts.runnerId,
|
||||
durationMs,
|
||||
});
|
||||
|
||||
this.#emitRestoreSpan(opts, startMs);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
#emitRestoreSpan(
|
||||
opts: {
|
||||
snapshotId: string;
|
||||
runnerId: string;
|
||||
runFriendlyId: string;
|
||||
traceContext?: Record<string, unknown>;
|
||||
envId?: string;
|
||||
orgId?: string;
|
||||
projectId?: string;
|
||||
dequeuedAt?: Date;
|
||||
},
|
||||
startMs: number
|
||||
) {
|
||||
if (!this.traceSpansEnabled) return;
|
||||
|
||||
const parsed = parseTraceparent(extractTraceparent(opts.traceContext));
|
||||
if (!parsed || !opts.envId || !opts.orgId || !opts.projectId) return;
|
||||
|
||||
const endMs = performance.now();
|
||||
const now = Date.now();
|
||||
const restoreStartEpochMs = now - (endMs - startMs);
|
||||
const endEpochMs = now;
|
||||
|
||||
// Subtract 1ms so restore span always sorts before the attempt span
|
||||
const startEpochMs = (opts.dequeuedAt?.getTime() ?? restoreStartEpochMs) - 1;
|
||||
|
||||
this.opts.tracing?.emit({
|
||||
traceId: parsed.traceId,
|
||||
parentSpanId: parsed.spanId,
|
||||
spanName: "compute.restore",
|
||||
startTimeMs: startEpochMs,
|
||||
endTimeMs: endEpochMs,
|
||||
resourceAttributes: {
|
||||
"ctx.environment.id": opts.envId,
|
||||
"ctx.organization.id": opts.orgId,
|
||||
"ctx.project.id": opts.projectId,
|
||||
"ctx.run.id": opts.runFriendlyId,
|
||||
},
|
||||
spanAttributes: {
|
||||
"compute.type": "restore",
|
||||
"compute.snapshot_id": opts.snapshotId,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import {
|
||||
type WorkloadManager,
|
||||
type WorkloadManagerCreateOptions,
|
||||
type WorkloadManagerOptions,
|
||||
} from "./types.js";
|
||||
import { env } from "../env.js";
|
||||
import { getDockerHostDomain, getRunnerId, normalizeDockerHostUrl } from "../util.js";
|
||||
import Docker from "dockerode";
|
||||
import { tryCatch } from "@trigger.dev/core";
|
||||
import { ECRAuthService } from "./ecrAuth.js";
|
||||
|
||||
export class DockerWorkloadManager implements WorkloadManager {
|
||||
private readonly logger = new SimpleStructuredLogger("docker-workload-manager");
|
||||
private readonly docker: Docker;
|
||||
|
||||
private readonly runnerNetworks: string[];
|
||||
private readonly staticAuth?: Docker.AuthConfig;
|
||||
private readonly platformOverride?: string;
|
||||
private readonly ecrAuthService?: ECRAuthService;
|
||||
|
||||
constructor(private opts: WorkloadManagerOptions) {
|
||||
this.docker = new Docker({
|
||||
version: env.DOCKER_API_VERSION,
|
||||
});
|
||||
|
||||
if (opts.workloadApiDomain) {
|
||||
this.logger.warn("⚠️ Custom workload API domain", {
|
||||
domain: opts.workloadApiDomain,
|
||||
});
|
||||
}
|
||||
|
||||
this.runnerNetworks = env.DOCKER_RUNNER_NETWORKS.split(",");
|
||||
|
||||
this.platformOverride = env.DOCKER_PLATFORM;
|
||||
if (this.platformOverride) {
|
||||
this.logger.info("🖥️ Platform override", {
|
||||
targetPlatform: this.platformOverride,
|
||||
hostPlatform: process.arch,
|
||||
});
|
||||
}
|
||||
|
||||
if (env.DOCKER_REGISTRY_USERNAME && env.DOCKER_REGISTRY_PASSWORD && env.DOCKER_REGISTRY_URL) {
|
||||
this.logger.info("🐋 Using Docker registry credentials", {
|
||||
username: env.DOCKER_REGISTRY_USERNAME,
|
||||
url: env.DOCKER_REGISTRY_URL,
|
||||
});
|
||||
|
||||
this.staticAuth = {
|
||||
username: env.DOCKER_REGISTRY_USERNAME,
|
||||
password: env.DOCKER_REGISTRY_PASSWORD,
|
||||
serveraddress: env.DOCKER_REGISTRY_URL,
|
||||
};
|
||||
} else if (ECRAuthService.hasAWSCredentials()) {
|
||||
this.logger.info("🐋 AWS credentials found, initializing ECR auth service");
|
||||
this.ecrAuthService = new ECRAuthService();
|
||||
} else {
|
||||
this.logger.warn(
|
||||
"🐋 No Docker registry credentials or AWS credentials provided, skipping auth"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async create(opts: WorkloadManagerCreateOptions) {
|
||||
this.logger.verbose("create()", { opts });
|
||||
|
||||
const runnerId = getRunnerId(opts.runFriendlyId, opts.nextAttemptNumber);
|
||||
|
||||
// Build environment variables
|
||||
const envVars: string[] = [
|
||||
`OTEL_EXPORTER_OTLP_ENDPOINT=${env.OTEL_EXPORTER_OTLP_ENDPOINT}`,
|
||||
`TRIGGER_DEQUEUED_AT_MS=${opts.dequeuedAt.getTime()}`,
|
||||
`TRIGGER_POD_SCHEDULED_AT_MS=${Date.now()}`,
|
||||
`TRIGGER_ENV_ID=${opts.envId}`,
|
||||
`TRIGGER_DEPLOYMENT_ID=${opts.deploymentFriendlyId}`,
|
||||
`TRIGGER_DEPLOYMENT_VERSION=${opts.deploymentVersion}`,
|
||||
`TRIGGER_RUN_ID=${opts.runFriendlyId}`,
|
||||
`TRIGGER_SNAPSHOT_ID=${opts.snapshotFriendlyId}`,
|
||||
`TRIGGER_SUPERVISOR_API_PROTOCOL=${this.opts.workloadApiProtocol}`,
|
||||
`TRIGGER_SUPERVISOR_API_PORT=${this.opts.workloadApiPort}`,
|
||||
`TRIGGER_SUPERVISOR_API_DOMAIN=${this.opts.workloadApiDomain ?? getDockerHostDomain()}`,
|
||||
`TRIGGER_WORKER_INSTANCE_NAME=${env.TRIGGER_WORKER_INSTANCE_NAME}`,
|
||||
`TRIGGER_RUNNER_ID=${runnerId}`,
|
||||
`TRIGGER_MACHINE_CPU=${opts.machine.cpu}`,
|
||||
`TRIGGER_MACHINE_MEMORY=${opts.machine.memory}`,
|
||||
`PRETTY_LOGS=${env.RUNNER_PRETTY_LOGS}`,
|
||||
`TRIGGER_SEND_RUN_DEBUG_LOGS=${env.SEND_RUN_DEBUG_LOGS}`,
|
||||
];
|
||||
|
||||
if (this.opts.warmStartUrl) {
|
||||
envVars.push(`TRIGGER_WARM_START_URL=${normalizeDockerHostUrl(this.opts.warmStartUrl)}`);
|
||||
}
|
||||
|
||||
if (this.opts.metadataUrl) {
|
||||
envVars.push(`TRIGGER_METADATA_URL=${this.opts.metadataUrl}`);
|
||||
}
|
||||
|
||||
if (this.opts.heartbeatIntervalSeconds) {
|
||||
envVars.push(`TRIGGER_HEARTBEAT_INTERVAL_SECONDS=${this.opts.heartbeatIntervalSeconds}`);
|
||||
}
|
||||
|
||||
if (this.opts.snapshotPollIntervalSeconds) {
|
||||
envVars.push(
|
||||
`TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS=${this.opts.snapshotPollIntervalSeconds}`
|
||||
);
|
||||
}
|
||||
|
||||
if (this.opts.additionalEnvVars) {
|
||||
Object.entries(this.opts.additionalEnvVars).forEach(([key, value]) => {
|
||||
envVars.push(`${key}=${value}`);
|
||||
});
|
||||
}
|
||||
|
||||
const hostConfig: Docker.HostConfig = {
|
||||
AutoRemove: !!this.opts.dockerAutoremove,
|
||||
};
|
||||
|
||||
const [firstNetwork, ...remainingNetworks] = this.runnerNetworks;
|
||||
|
||||
// Always attach the first network at container creation time. This has the following benefits:
|
||||
// - If there is only a single network to attach, this will prevent having to make a separate request.
|
||||
// - If there are multiple networks to attach, this will ensure the runner won't also be connected to the bridge network
|
||||
hostConfig.NetworkMode = firstNetwork;
|
||||
|
||||
if (env.DOCKER_ENFORCE_MACHINE_PRESETS) {
|
||||
hostConfig.NanoCpus = opts.machine.cpu * 1e9;
|
||||
hostConfig.Memory = opts.machine.memory * 1024 * 1024 * 1024;
|
||||
}
|
||||
|
||||
let imageRef = opts.image;
|
||||
|
||||
if (env.DOCKER_STRIP_IMAGE_DIGEST) {
|
||||
imageRef = opts.image.split("@")[0]!;
|
||||
}
|
||||
|
||||
const containerCreateOpts: Docker.ContainerCreateOptions = {
|
||||
name: runnerId,
|
||||
Hostname: runnerId,
|
||||
HostConfig: hostConfig,
|
||||
Image: imageRef,
|
||||
AttachStdout: false,
|
||||
AttachStderr: false,
|
||||
AttachStdin: false,
|
||||
};
|
||||
|
||||
if (this.platformOverride) {
|
||||
containerCreateOpts.platform = this.platformOverride;
|
||||
}
|
||||
|
||||
const logger = this.logger.child({ opts, containerCreateOpts });
|
||||
|
||||
const [inspectError, inspectResult] = await tryCatch(this.docker.getImage(imageRef).inspect());
|
||||
|
||||
let shouldPull = !!inspectError;
|
||||
if (this.platformOverride) {
|
||||
const imageArchitecture = inspectResult?.Architecture;
|
||||
|
||||
// When the image architecture doesn't match the platform, we need to pull the image
|
||||
if (imageArchitecture && !this.platformOverride.includes(imageArchitecture)) {
|
||||
shouldPull = true;
|
||||
}
|
||||
}
|
||||
|
||||
// If the image is not present, try to pull it
|
||||
if (shouldPull) {
|
||||
logger.info("Pulling image", {
|
||||
error: inspectError,
|
||||
image: opts.image,
|
||||
targetPlatform: this.platformOverride,
|
||||
imageArchitecture: inspectResult?.Architecture,
|
||||
});
|
||||
|
||||
// Get auth config (static or ECR)
|
||||
const authConfig = await this.getAuthConfig();
|
||||
|
||||
// Ensure the image is present
|
||||
const [createImageError, imageResponseReader] = await tryCatch(
|
||||
this.docker.createImage(authConfig, {
|
||||
fromImage: imageRef,
|
||||
...(this.platformOverride ? { platform: this.platformOverride } : {}),
|
||||
})
|
||||
);
|
||||
if (createImageError) {
|
||||
logger.error("Failed to pull image", { error: createImageError });
|
||||
return;
|
||||
}
|
||||
|
||||
const [imageReadError, imageResponse] = await tryCatch(readAllChunks(imageResponseReader));
|
||||
if (imageReadError) {
|
||||
logger.error("failed to read image response", { error: imageReadError });
|
||||
return;
|
||||
}
|
||||
|
||||
logger.debug("pulled image", { image: opts.image, imageResponse });
|
||||
} else {
|
||||
// Image is present, so we can use it to create the container
|
||||
}
|
||||
|
||||
// Create container
|
||||
const [createContainerError, container] = await tryCatch(
|
||||
this.docker.createContainer({
|
||||
...containerCreateOpts,
|
||||
// Add env vars here so they're not logged
|
||||
Env: envVars,
|
||||
})
|
||||
);
|
||||
|
||||
if (createContainerError) {
|
||||
logger.error("Failed to create container", { error: createContainerError });
|
||||
return;
|
||||
}
|
||||
|
||||
// If there are multiple networks to attach to we need to attach the remaining ones after creation
|
||||
if (remainingNetworks.length > 0) {
|
||||
await this.attachContainerToNetworks({
|
||||
containerId: container.id,
|
||||
networkNames: remainingNetworks,
|
||||
});
|
||||
}
|
||||
|
||||
// Start container
|
||||
const [startError, startResult] = await tryCatch(container.start());
|
||||
|
||||
if (startError) {
|
||||
logger.error("Failed to start container", { error: startError, containerId: container.id });
|
||||
return;
|
||||
}
|
||||
|
||||
logger.debug("create succeeded", { startResult, containerId: container.id });
|
||||
}
|
||||
|
||||
/**
|
||||
* Get authentication config for Docker operations
|
||||
* Uses static credentials if available, otherwise attempts ECR auth
|
||||
*/
|
||||
private async getAuthConfig(): Promise<Docker.AuthConfig | undefined> {
|
||||
// Use static credentials if available
|
||||
if (this.staticAuth) {
|
||||
return this.staticAuth;
|
||||
}
|
||||
|
||||
// Use ECR auth if service is available
|
||||
if (this.ecrAuthService) {
|
||||
const ecrAuth = await this.ecrAuthService.getAuthConfig();
|
||||
return ecrAuth || undefined;
|
||||
}
|
||||
|
||||
// No auth available
|
||||
return undefined;
|
||||
}
|
||||
|
||||
private async attachContainerToNetworks({
|
||||
containerId,
|
||||
networkNames,
|
||||
}: {
|
||||
containerId: string;
|
||||
networkNames: string[];
|
||||
}) {
|
||||
this.logger.debug("Attaching container to networks", { containerId, networkNames });
|
||||
|
||||
const [error, networkResults] = await tryCatch(
|
||||
this.docker.listNetworks({
|
||||
filters: {
|
||||
// Full name matches only to prevent unexpected results
|
||||
name: networkNames.map((name) => `^${name}$`),
|
||||
},
|
||||
})
|
||||
);
|
||||
|
||||
if (error) {
|
||||
this.logger.error("Failed to list networks", { networkNames });
|
||||
return;
|
||||
}
|
||||
|
||||
const results = await Promise.allSettled(
|
||||
networkResults.map((networkInfo) => {
|
||||
const network = this.docker.getNetwork(networkInfo.Id);
|
||||
return network.connect({ Container: containerId });
|
||||
})
|
||||
);
|
||||
|
||||
if (results.some((r) => r.status === "rejected")) {
|
||||
this.logger.error("Failed to attach container to some networks", {
|
||||
containerId,
|
||||
networkNames,
|
||||
results,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
this.logger.debug("Attached container to networks", {
|
||||
containerId,
|
||||
networkNames,
|
||||
results,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async function readAllChunks(reader: NodeJS.ReadableStream) {
|
||||
const chunks = [];
|
||||
for await (const chunk of reader) {
|
||||
chunks.push(chunk.toString());
|
||||
}
|
||||
return chunks;
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
import { ECRClient, GetAuthorizationTokenCommand } from "@aws-sdk/client-ecr";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import { tryCatch } from "@trigger.dev/core";
|
||||
import type Docker from "dockerode";
|
||||
|
||||
interface ECRTokenCache {
|
||||
token: string;
|
||||
username: string;
|
||||
serverAddress: string;
|
||||
expiresAt: Date;
|
||||
}
|
||||
|
||||
export class ECRAuthService {
|
||||
private readonly logger = new SimpleStructuredLogger("ecr-auth-service");
|
||||
private readonly ecrClient: ECRClient;
|
||||
private tokenCache: ECRTokenCache | null = null;
|
||||
|
||||
constructor() {
|
||||
this.ecrClient = new ECRClient();
|
||||
|
||||
this.logger.info("🔐 ECR Auth Service initialized", {
|
||||
region: this.ecrClient.config.region,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if we have AWS credentials configured
|
||||
*/
|
||||
static hasAWSCredentials(): boolean {
|
||||
if (process.env.AWS_ACCESS_KEY_ID && process.env.AWS_SECRET_ACCESS_KEY) {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (
|
||||
process.env.AWS_PROFILE ||
|
||||
process.env.AWS_ROLE_ARN ||
|
||||
process.env.AWS_WEB_IDENTITY_TOKEN_FILE
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the current token is still valid with a 10-minute buffer
|
||||
*/
|
||||
private isTokenValid(): boolean {
|
||||
if (!this.tokenCache) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const now = new Date();
|
||||
const bufferMs = 10 * 60 * 1000; // 10 minute buffer before expiration
|
||||
return now < new Date(this.tokenCache.expiresAt.getTime() - bufferMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get a fresh ECR authorization token from AWS
|
||||
*/
|
||||
private async fetchNewToken(): Promise<ECRTokenCache | null> {
|
||||
const [error, response] = await tryCatch(
|
||||
this.ecrClient.send(new GetAuthorizationTokenCommand({}))
|
||||
);
|
||||
|
||||
if (error) {
|
||||
this.logger.error("Failed to get ECR authorization token", { error });
|
||||
return null;
|
||||
}
|
||||
|
||||
const authData = response.authorizationData?.[0];
|
||||
if (!authData?.authorizationToken || !authData.proxyEndpoint) {
|
||||
this.logger.error("Invalid ECR authorization response", { authData });
|
||||
return null;
|
||||
}
|
||||
|
||||
// Decode the base64 token to get username:password
|
||||
const decoded = Buffer.from(authData.authorizationToken, "base64").toString("utf-8");
|
||||
const [username, password] = decoded.split(":", 2);
|
||||
|
||||
if (!username || !password) {
|
||||
this.logger.error("Failed to parse ECR authorization token");
|
||||
return null;
|
||||
}
|
||||
|
||||
const expiresAt = authData.expiresAt || new Date(Date.now() + 12 * 60 * 60 * 1000); // Default 12 hours
|
||||
|
||||
const tokenCache: ECRTokenCache = {
|
||||
token: password,
|
||||
username,
|
||||
serverAddress: authData.proxyEndpoint,
|
||||
expiresAt,
|
||||
};
|
||||
|
||||
this.logger.info("🔐 Successfully fetched ECR token", {
|
||||
username,
|
||||
serverAddress: authData.proxyEndpoint,
|
||||
expiresAt: expiresAt.toISOString(),
|
||||
});
|
||||
|
||||
return tokenCache;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get ECR auth config for Docker operations
|
||||
* Returns cached token if valid, otherwise fetches a new one
|
||||
*/
|
||||
async getAuthConfig(): Promise<Docker.AuthConfig | null> {
|
||||
// Check if cached token is still valid
|
||||
if (this.isTokenValid()) {
|
||||
this.logger.debug("Using cached ECR token");
|
||||
return {
|
||||
username: this.tokenCache!.username,
|
||||
password: this.tokenCache!.token,
|
||||
serveraddress: this.tokenCache!.serverAddress,
|
||||
};
|
||||
}
|
||||
|
||||
// Fetch new token
|
||||
this.logger.info("Fetching new ECR authorization token");
|
||||
const newToken = await this.fetchNewToken();
|
||||
|
||||
if (!newToken) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Cache the new token
|
||||
this.tokenCache = newToken;
|
||||
|
||||
return {
|
||||
username: newToken.username,
|
||||
password: newToken.token,
|
||||
serveraddress: newToken.serverAddress,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear the cached token (useful for testing or forcing refresh)
|
||||
*/
|
||||
clearCache(): void {
|
||||
this.tokenCache = null;
|
||||
this.logger.debug("ECR token cache cleared");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,582 @@
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import {
|
||||
type WorkloadManager,
|
||||
type WorkloadManagerCreateOptions,
|
||||
type WorkloadManagerOptions,
|
||||
} from "./types.js";
|
||||
import type {
|
||||
EnvironmentType,
|
||||
MachinePreset,
|
||||
MachinePresetName,
|
||||
PlacementTag,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import { PlacementTagProcessor } from "@trigger.dev/core/v3/serverOnly";
|
||||
import { env } from "../env.js";
|
||||
import { type K8sApi, createK8sApi, type k8s } from "../clients/kubernetes.js";
|
||||
import { getRunnerId } from "../util.js";
|
||||
|
||||
type ResourceQuantities = {
|
||||
[K in "cpu" | "memory" | "ephemeral-storage"]?: string;
|
||||
};
|
||||
|
||||
const cpuRequestRatioByMachinePreset: Record<MachinePresetName, number | undefined> = {
|
||||
micro: env.KUBERNETES_CPU_REQUEST_RATIO_MICRO,
|
||||
"small-1x": env.KUBERNETES_CPU_REQUEST_RATIO_SMALL_1X,
|
||||
"small-2x": env.KUBERNETES_CPU_REQUEST_RATIO_SMALL_2X,
|
||||
"medium-1x": env.KUBERNETES_CPU_REQUEST_RATIO_MEDIUM_1X,
|
||||
"medium-2x": env.KUBERNETES_CPU_REQUEST_RATIO_MEDIUM_2X,
|
||||
"large-1x": env.KUBERNETES_CPU_REQUEST_RATIO_LARGE_1X,
|
||||
"large-2x": env.KUBERNETES_CPU_REQUEST_RATIO_LARGE_2X,
|
||||
};
|
||||
|
||||
const memoryRequestRatioByMachinePreset: Record<MachinePresetName, number | undefined> = {
|
||||
micro: env.KUBERNETES_MEMORY_REQUEST_RATIO_MICRO,
|
||||
"small-1x": env.KUBERNETES_MEMORY_REQUEST_RATIO_SMALL_1X,
|
||||
"small-2x": env.KUBERNETES_MEMORY_REQUEST_RATIO_SMALL_2X,
|
||||
"medium-1x": env.KUBERNETES_MEMORY_REQUEST_RATIO_MEDIUM_1X,
|
||||
"medium-2x": env.KUBERNETES_MEMORY_REQUEST_RATIO_MEDIUM_2X,
|
||||
"large-1x": env.KUBERNETES_MEMORY_REQUEST_RATIO_LARGE_1X,
|
||||
"large-2x": env.KUBERNETES_MEMORY_REQUEST_RATIO_LARGE_2X,
|
||||
};
|
||||
|
||||
export class KubernetesWorkloadManager implements WorkloadManager {
|
||||
private readonly logger = new SimpleStructuredLogger("kubernetes-workload-provider");
|
||||
private k8s: K8sApi;
|
||||
private namespace = env.KUBERNETES_NAMESPACE;
|
||||
private placementTagProcessor: PlacementTagProcessor;
|
||||
|
||||
// Resource settings
|
||||
private readonly cpuRequestMinCores = env.KUBERNETES_CPU_REQUEST_MIN_CORES;
|
||||
private readonly cpuRequestRatio = env.KUBERNETES_CPU_REQUEST_RATIO;
|
||||
private readonly memoryRequestMinGb = env.KUBERNETES_MEMORY_REQUEST_MIN_GB;
|
||||
private readonly memoryRequestRatio = env.KUBERNETES_MEMORY_REQUEST_RATIO;
|
||||
private readonly memoryOverheadGb = env.KUBERNETES_MEMORY_OVERHEAD_GB;
|
||||
|
||||
constructor(private opts: WorkloadManagerOptions) {
|
||||
this.k8s = createK8sApi();
|
||||
this.placementTagProcessor = new PlacementTagProcessor({
|
||||
enabled: env.PLACEMENT_TAGS_ENABLED,
|
||||
prefix: env.PLACEMENT_TAGS_PREFIX,
|
||||
});
|
||||
|
||||
if (opts.workloadApiDomain) {
|
||||
this.logger.warn("[KubernetesWorkloadManager] ⚠️ Custom workload API domain", {
|
||||
domain: opts.workloadApiDomain,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
private addPlacementTags(
|
||||
podSpec: Omit<k8s.V1PodSpec, "containers">,
|
||||
placementTags?: PlacementTag[]
|
||||
): Omit<k8s.V1PodSpec, "containers"> {
|
||||
const nodeSelector = this.placementTagProcessor.convertToNodeSelector(
|
||||
placementTags,
|
||||
podSpec.nodeSelector
|
||||
);
|
||||
|
||||
return {
|
||||
...podSpec,
|
||||
nodeSelector,
|
||||
};
|
||||
}
|
||||
|
||||
private stripImageDigest(imageRef: string): string {
|
||||
if (!env.KUBERNETES_STRIP_IMAGE_DIGEST) {
|
||||
return imageRef;
|
||||
}
|
||||
|
||||
const atIndex = imageRef.lastIndexOf("@");
|
||||
|
||||
if (atIndex === -1) {
|
||||
return imageRef;
|
||||
}
|
||||
|
||||
return imageRef.substring(0, atIndex);
|
||||
}
|
||||
|
||||
private clamp(value: number, min: number, max: number): number {
|
||||
return Math.min(Math.max(value, min), max);
|
||||
}
|
||||
|
||||
async create(opts: WorkloadManagerCreateOptions) {
|
||||
this.logger.verbose("[KubernetesWorkloadManager] Creating container", { opts });
|
||||
|
||||
const runnerId = getRunnerId(opts.runFriendlyId, opts.nextAttemptNumber);
|
||||
|
||||
try {
|
||||
await this.k8s.core.createNamespacedPod({
|
||||
namespace: this.namespace,
|
||||
body: {
|
||||
metadata: {
|
||||
name: runnerId,
|
||||
namespace: this.namespace,
|
||||
labels: {
|
||||
...this.#getSharedLabels(opts),
|
||||
app: "task-run",
|
||||
"app.kubernetes.io/part-of": "trigger-worker",
|
||||
"app.kubernetes.io/component": "create",
|
||||
},
|
||||
},
|
||||
spec: {
|
||||
...this.addPlacementTags(this.#defaultPodSpec, opts.placementTags),
|
||||
affinity: this.#getAffinity(opts),
|
||||
tolerations: this.#getScheduleTolerations(this.#isScheduledRun(opts)),
|
||||
terminationGracePeriodSeconds: 60 * 60,
|
||||
containers: [
|
||||
{
|
||||
name: "run-controller",
|
||||
image: this.stripImageDigest(opts.image),
|
||||
ports: [
|
||||
{
|
||||
containerPort: 8000,
|
||||
},
|
||||
],
|
||||
resources: this.#getResourcesForMachine(opts.machine),
|
||||
env: [
|
||||
{
|
||||
name: "TRIGGER_DEQUEUED_AT_MS",
|
||||
value: opts.dequeuedAt.getTime().toString(),
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_POD_SCHEDULED_AT_MS",
|
||||
value: Date.now().toString(),
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_RUN_ID",
|
||||
value: opts.runFriendlyId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_ENV_ID",
|
||||
value: opts.envId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_DEPLOYMENT_ID",
|
||||
value: opts.deploymentFriendlyId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_DEPLOYMENT_VERSION",
|
||||
value: opts.deploymentVersion,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SNAPSHOT_ID",
|
||||
value: opts.snapshotFriendlyId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SUPERVISOR_API_PROTOCOL",
|
||||
value: this.opts.workloadApiProtocol,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SUPERVISOR_API_PORT",
|
||||
value: `${this.opts.workloadApiPort}`,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SUPERVISOR_API_DOMAIN",
|
||||
...(this.opts.workloadApiDomain
|
||||
? {
|
||||
value: this.opts.workloadApiDomain,
|
||||
}
|
||||
: {
|
||||
valueFrom: {
|
||||
fieldRef: {
|
||||
fieldPath: "status.hostIP",
|
||||
},
|
||||
},
|
||||
}),
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_WORKER_INSTANCE_NAME",
|
||||
valueFrom: {
|
||||
fieldRef: {
|
||||
fieldPath: "spec.nodeName",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "OTEL_EXPORTER_OTLP_ENDPOINT",
|
||||
value: env.OTEL_EXPORTER_OTLP_ENDPOINT,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_RUNNER_ID",
|
||||
value: runnerId,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_MACHINE_CPU",
|
||||
value: `${opts.machine.cpu}`,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_MACHINE_MEMORY",
|
||||
value: `${opts.machine.memory}`,
|
||||
},
|
||||
{
|
||||
name: "TRIGGER_SEND_RUN_DEBUG_LOGS",
|
||||
value: `${env.SEND_RUN_DEBUG_LOGS}`,
|
||||
},
|
||||
{
|
||||
name: "LIMITS_CPU",
|
||||
valueFrom: {
|
||||
resourceFieldRef: {
|
||||
resource: "limits.cpu",
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "LIMITS_MEMORY",
|
||||
valueFrom: {
|
||||
resourceFieldRef: {
|
||||
resource: "limits.memory",
|
||||
},
|
||||
},
|
||||
},
|
||||
...(this.opts.warmStartUrl
|
||||
? [{ name: "TRIGGER_WARM_START_URL", value: this.opts.warmStartUrl }]
|
||||
: []),
|
||||
...(this.opts.metadataUrl
|
||||
? [{ name: "TRIGGER_METADATA_URL", value: this.opts.metadataUrl }]
|
||||
: []),
|
||||
...(this.opts.heartbeatIntervalSeconds
|
||||
? [
|
||||
{
|
||||
name: "TRIGGER_HEARTBEAT_INTERVAL_SECONDS",
|
||||
value: `${this.opts.heartbeatIntervalSeconds}`,
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...(this.opts.snapshotPollIntervalSeconds
|
||||
? [
|
||||
{
|
||||
name: "TRIGGER_SNAPSHOT_POLL_INTERVAL_SECONDS",
|
||||
value: `${this.opts.snapshotPollIntervalSeconds}`,
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...(this.opts.additionalEnvVars
|
||||
? Object.entries(this.opts.additionalEnvVars).map(([key, value]) => ({
|
||||
name: key,
|
||||
value: value,
|
||||
}))
|
||||
: []),
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
});
|
||||
} catch (err: unknown) {
|
||||
this.#handleK8sError(err);
|
||||
}
|
||||
}
|
||||
|
||||
#throwUnlessRecord(candidate: unknown): asserts candidate is Record<string, unknown> {
|
||||
if (typeof candidate !== "object" || candidate === null) {
|
||||
throw candidate;
|
||||
}
|
||||
}
|
||||
|
||||
#handleK8sError(err: unknown) {
|
||||
this.#throwUnlessRecord(err);
|
||||
|
||||
if ("body" in err && err.body) {
|
||||
this.logger.error("[KubernetesWorkloadManager] Create failed", { rawError: err.body });
|
||||
this.#throwUnlessRecord(err.body);
|
||||
|
||||
if (typeof err.body.message === "string") {
|
||||
throw new Error(err.body?.message);
|
||||
} else {
|
||||
throw err.body;
|
||||
}
|
||||
} else {
|
||||
this.logger.error("[KubernetesWorkloadManager] Create failed", { rawError: err });
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
#envTypeToLabelValue(type: EnvironmentType) {
|
||||
switch (type) {
|
||||
case "PRODUCTION":
|
||||
return "prod";
|
||||
case "STAGING":
|
||||
return "stg";
|
||||
case "DEVELOPMENT":
|
||||
return "dev";
|
||||
case "PREVIEW":
|
||||
return "preview";
|
||||
}
|
||||
}
|
||||
|
||||
private getImagePullSecrets(): k8s.V1LocalObjectReference[] | undefined {
|
||||
return this.opts.imagePullSecrets?.map((name) => ({ name }));
|
||||
}
|
||||
|
||||
get #defaultPodSpec(): Omit<k8s.V1PodSpec, "containers"> {
|
||||
return {
|
||||
restartPolicy: "Never",
|
||||
automountServiceAccountToken: false,
|
||||
imagePullSecrets: this.getImagePullSecrets(),
|
||||
...(env.KUBERNETES_SCHEDULER_NAME
|
||||
? {
|
||||
schedulerName: env.KUBERNETES_SCHEDULER_NAME,
|
||||
}
|
||||
: {}),
|
||||
...(env.KUBERNETES_WORKER_NODETYPE_LABEL
|
||||
? {
|
||||
nodeSelector: {
|
||||
nodetype: env.KUBERNETES_WORKER_NODETYPE_LABEL,
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
...(env.KUBERNETES_POD_DNS_NDOTS_OVERRIDE_ENABLED
|
||||
? {
|
||||
dnsConfig: {
|
||||
options: [{ name: "ndots", value: `${env.KUBERNETES_POD_DNS_NDOTS}` }],
|
||||
},
|
||||
}
|
||||
: {}),
|
||||
};
|
||||
}
|
||||
|
||||
get #defaultResourceRequests(): ResourceQuantities {
|
||||
return {
|
||||
"ephemeral-storage": env.KUBERNETES_EPHEMERAL_STORAGE_SIZE_REQUEST,
|
||||
};
|
||||
}
|
||||
|
||||
get #defaultResourceLimits(): ResourceQuantities {
|
||||
return {
|
||||
"ephemeral-storage": env.KUBERNETES_EPHEMERAL_STORAGE_SIZE_LIMIT,
|
||||
};
|
||||
}
|
||||
|
||||
#isScheduledRun(opts: WorkloadManagerCreateOptions): boolean {
|
||||
return opts.annotations?.rootTriggerSource === "schedule";
|
||||
}
|
||||
|
||||
#getSharedLabels(opts: WorkloadManagerCreateOptions): Record<string, string> {
|
||||
const labels: Record<string, string> = {
|
||||
env: opts.envId,
|
||||
envtype: this.#envTypeToLabelValue(opts.envType),
|
||||
org: opts.orgId,
|
||||
project: opts.projectId,
|
||||
machine: opts.machine.name,
|
||||
// We intentionally use a boolean label rather than exposing the full trigger source
|
||||
// (e.g. sdk, api, cli, mcp, schedule) to keep label cardinality low in metrics.
|
||||
// The schedule vs non-schedule distinction is all we need for the current metrics
|
||||
// and pool-level scheduling decisions; finer-grained source breakdowns live in run annotations.
|
||||
scheduled: String(this.#isScheduledRun(opts)),
|
||||
};
|
||||
|
||||
// Add privatelink label for CiliumNetworkPolicy matching
|
||||
if (opts.hasPrivateLink) {
|
||||
labels.privatelink = opts.orgId;
|
||||
}
|
||||
|
||||
return labels;
|
||||
}
|
||||
|
||||
#getResourceRequestsForMachine(preset: MachinePreset): ResourceQuantities {
|
||||
const cpuRatio = cpuRequestRatioByMachinePreset[preset.name] ?? this.cpuRequestRatio;
|
||||
const memoryRatio = memoryRequestRatioByMachinePreset[preset.name] ?? this.memoryRequestRatio;
|
||||
|
||||
const cpuRequest = preset.cpu * cpuRatio;
|
||||
const memoryRequest = preset.memory * memoryRatio;
|
||||
|
||||
// Clamp between min and max
|
||||
const clampedCpu = this.clamp(cpuRequest, this.cpuRequestMinCores, preset.cpu);
|
||||
const clampedMemory = this.clamp(memoryRequest, this.memoryRequestMinGb, preset.memory);
|
||||
|
||||
return {
|
||||
cpu: `${clampedCpu}`,
|
||||
memory: `${clampedMemory}G`,
|
||||
};
|
||||
}
|
||||
|
||||
#getResourceLimitsForMachine(preset: MachinePreset): ResourceQuantities {
|
||||
const memoryLimit = this.memoryOverheadGb
|
||||
? preset.memory + this.memoryOverheadGb
|
||||
: preset.memory;
|
||||
|
||||
return {
|
||||
cpu: `${preset.cpu}`,
|
||||
memory: `${memoryLimit}G`,
|
||||
};
|
||||
}
|
||||
|
||||
#getResourcesForMachine(preset: MachinePreset): k8s.V1ResourceRequirements {
|
||||
return {
|
||||
requests: {
|
||||
...this.#defaultResourceRequests,
|
||||
...this.#getResourceRequestsForMachine(preset),
|
||||
},
|
||||
limits: {
|
||||
...this.#defaultResourceLimits,
|
||||
...this.#getResourceLimitsForMachine(preset),
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
#isLargeMachine(preset: MachinePreset): boolean {
|
||||
return preset.name.startsWith("large-");
|
||||
}
|
||||
|
||||
#getAffinity(opts: WorkloadManagerCreateOptions): k8s.V1Affinity | undefined {
|
||||
const largeNodeAffinity = this.#getNodeAffinityRules(opts.machine);
|
||||
const scheduleNodeAffinity = this.#getScheduleNodeAffinityRules(this.#isScheduledRun(opts));
|
||||
const podAffinity = this.#getProjectPodAffinity(opts.projectId);
|
||||
|
||||
// Merge node affinity rules from multiple sources
|
||||
const preferred = [
|
||||
...(largeNodeAffinity?.preferredDuringSchedulingIgnoredDuringExecution ?? []),
|
||||
...(scheduleNodeAffinity?.preferredDuringSchedulingIgnoredDuringExecution ?? []),
|
||||
];
|
||||
// Only large machine affinity produces hard requirements (non-large runs must stay off the large pool).
|
||||
// Schedule affinity is soft both ways.
|
||||
const required = [
|
||||
...(largeNodeAffinity?.requiredDuringSchedulingIgnoredDuringExecution?.nodeSelectorTerms ??
|
||||
[]),
|
||||
];
|
||||
|
||||
const hasNodeAffinity = preferred.length > 0 || required.length > 0;
|
||||
|
||||
if (!hasNodeAffinity && !podAffinity) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return {
|
||||
...(hasNodeAffinity && {
|
||||
nodeAffinity: {
|
||||
...(preferred.length > 0 && {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: preferred,
|
||||
}),
|
||||
...(required.length > 0 && {
|
||||
requiredDuringSchedulingIgnoredDuringExecution: { nodeSelectorTerms: required },
|
||||
}),
|
||||
},
|
||||
}),
|
||||
...(podAffinity && { podAffinity }),
|
||||
};
|
||||
}
|
||||
|
||||
#getNodeAffinityRules(preset: MachinePreset): k8s.V1NodeAffinity | undefined {
|
||||
if (!env.KUBERNETES_LARGE_MACHINE_AFFINITY_ENABLED) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
if (this.#isLargeMachine(preset)) {
|
||||
// soft preference for the large-machine pool, falls back to standard if unavailable
|
||||
return {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: [
|
||||
{
|
||||
weight: env.KUBERNETES_LARGE_MACHINE_AFFINITY_WEIGHT,
|
||||
preference: {
|
||||
matchExpressions: [
|
||||
{
|
||||
key: env.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_KEY,
|
||||
operator: "In",
|
||||
values: [env.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_VALUE],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
// not schedulable in the large-machine pool
|
||||
return {
|
||||
requiredDuringSchedulingIgnoredDuringExecution: {
|
||||
nodeSelectorTerms: [
|
||||
{
|
||||
matchExpressions: [
|
||||
{
|
||||
key: env.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_KEY,
|
||||
operator: "NotIn",
|
||||
values: [env.KUBERNETES_LARGE_MACHINE_AFFINITY_POOL_LABEL_VALUE],
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
#getScheduleNodeAffinityRules(isScheduledRun: boolean): k8s.V1NodeAffinity | undefined {
|
||||
if (
|
||||
!env.KUBERNETES_SCHEDULED_RUN_AFFINITY_ENABLED ||
|
||||
!env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE
|
||||
) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
if (isScheduledRun) {
|
||||
// soft preference for the schedule pool
|
||||
return {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: [
|
||||
{
|
||||
weight: env.KUBERNETES_SCHEDULED_RUN_AFFINITY_WEIGHT,
|
||||
preference: {
|
||||
matchExpressions: [
|
||||
{
|
||||
key: env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_KEY,
|
||||
operator: "In",
|
||||
values: [env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
// soft anti-affinity: non-schedule runs prefer to avoid the schedule pool
|
||||
return {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: [
|
||||
{
|
||||
weight: env.KUBERNETES_SCHEDULED_RUN_ANTI_AFFINITY_WEIGHT,
|
||||
preference: {
|
||||
matchExpressions: [
|
||||
{
|
||||
key: env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_KEY,
|
||||
operator: "NotIn",
|
||||
values: [env.KUBERNETES_SCHEDULED_RUN_AFFINITY_POOL_LABEL_VALUE],
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
#getScheduleTolerations(isScheduledRun: boolean): k8s.V1Toleration[] | undefined {
|
||||
if (!isScheduledRun || !env.KUBERNETES_SCHEDULED_RUN_TOLERATIONS?.length) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return env.KUBERNETES_SCHEDULED_RUN_TOLERATIONS;
|
||||
}
|
||||
|
||||
#getProjectPodAffinity(projectId: string): k8s.V1PodAffinity | undefined {
|
||||
if (!env.KUBERNETES_PROJECT_AFFINITY_ENABLED) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return {
|
||||
preferredDuringSchedulingIgnoredDuringExecution: [
|
||||
{
|
||||
weight: env.KUBERNETES_PROJECT_AFFINITY_WEIGHT,
|
||||
podAffinityTerm: {
|
||||
labelSelector: {
|
||||
matchExpressions: [
|
||||
{
|
||||
key: "project",
|
||||
operator: "In",
|
||||
values: [projectId],
|
||||
},
|
||||
],
|
||||
},
|
||||
topologyKey: env.KUBERNETES_PROJECT_AFFINITY_TOPOLOGY_KEY,
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
import type {
|
||||
EnvironmentType,
|
||||
MachinePreset,
|
||||
PlacementTag,
|
||||
RunAnnotations,
|
||||
} from "@trigger.dev/core/v3";
|
||||
|
||||
export interface WorkloadManagerOptions {
|
||||
workloadApiProtocol: "http" | "https";
|
||||
workloadApiDomain?: string; // If unset, will use orchestrator-specific default
|
||||
workloadApiPort: number;
|
||||
warmStartUrl?: string;
|
||||
metadataUrl?: string;
|
||||
imagePullSecrets?: string[];
|
||||
heartbeatIntervalSeconds?: number;
|
||||
snapshotPollIntervalSeconds?: number;
|
||||
additionalEnvVars?: Record<string, string>;
|
||||
dockerAutoremove?: boolean;
|
||||
}
|
||||
|
||||
export interface WorkloadManager {
|
||||
create: (opts: WorkloadManagerCreateOptions) => Promise<unknown>;
|
||||
}
|
||||
|
||||
export interface WorkloadManagerCreateOptions {
|
||||
image: string;
|
||||
machine: MachinePreset;
|
||||
version: string;
|
||||
nextAttemptNumber?: number;
|
||||
dequeuedAt: Date;
|
||||
placementTags?: PlacementTag[];
|
||||
// Timing context (populated by supervisor handler, included in wide event)
|
||||
dequeueResponseMs?: number;
|
||||
pollingIntervalMs?: number;
|
||||
warmStartCheckMs?: number;
|
||||
// identifiers
|
||||
envId: string;
|
||||
envType: EnvironmentType;
|
||||
orgId: string;
|
||||
projectId: string;
|
||||
deploymentFriendlyId: string;
|
||||
deploymentVersion: string;
|
||||
runId: string;
|
||||
runFriendlyId: string;
|
||||
snapshotId: string;
|
||||
snapshotFriendlyId: string;
|
||||
// Trace context for OTel span emission (W3C format: { traceparent: "00-...", tracestate?: "..." })
|
||||
traceContext?: Record<string, unknown>;
|
||||
annotations?: RunAnnotations;
|
||||
// private networking
|
||||
hasPrivateLink?: boolean;
|
||||
}
|
||||
Reference in New Issue
Block a user