294 lines
14 KiB
JavaScript
294 lines
14 KiB
JavaScript
import crypto from "node:crypto";
|
|
import fs from "node:fs";
|
|
import fsp from "node:fs/promises";
|
|
import http from "node:http";
|
|
import path from "node:path";
|
|
import { spawn } from "node:child_process";
|
|
import { fileURLToPath } from "node:url";
|
|
|
|
const moduleRoot = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "../../..");
|
|
const SHA256 = /^[a-f0-9]{64}$/;
|
|
const PROJECT_ID = /^[A-Za-z0-9_-]{1,64}$/;
|
|
const GRID_NAME = /^[A-Za-z0-9_.:-]{1,255}$/;
|
|
const MAX_SOURCE_BYTES = 512 * 1024 * 1024;
|
|
const MAX_LOG_BYTES = 1024 * 1024;
|
|
|
|
function stableJson(value) {
|
|
if (Array.isArray(value)) return `[${value.map(stableJson).join(",")}]`;
|
|
if (value && typeof value === "object") return `{${Object.keys(value).sort().map((key) => `${JSON.stringify(key)}:${stableJson(value[key])}`).join(",")}}`;
|
|
return JSON.stringify(value);
|
|
}
|
|
|
|
function fileSha256(file) {
|
|
return crypto.createHash("sha256").update(fs.readFileSync(file)).digest("hex");
|
|
}
|
|
|
|
async function runProcess(command, args, options = {}) {
|
|
return new Promise((resolve, reject) => {
|
|
const child = spawn(command, args, { stdio: ["ignore", "pipe", "pipe"], ...options });
|
|
let output = "";
|
|
const append = (chunk) => { output = `${output}${chunk}`.slice(-MAX_LOG_BYTES); options.onLog?.(chunk.toString()); };
|
|
child.stdout.on("data", append);
|
|
child.stderr.on("data", append);
|
|
child.once("error", reject);
|
|
child.once("exit", (code, signal) => resolve({ code, signal, output }));
|
|
options.onChild?.(child);
|
|
});
|
|
}
|
|
|
|
function parseHeaders(request) {
|
|
const projectId = request.headers["x-vdb-project-id"] ?? "vdb-server";
|
|
const sourcePath = request.headers["x-vdb-source-path"] ?? "//volumes/upload.vdb";
|
|
const expectedSha256 = request.headers["x-vdb-source-sha256"];
|
|
const quantization = request.headers["x-vdb-quantization"] ?? "LOSSLESS";
|
|
const chunkByteLength = Number(request.headers["x-vdb-chunk-bytes"] ?? 4 * 1024 * 1024);
|
|
const selectedGrids = String(request.headers["x-vdb-grids"] ?? "").split(",").filter(Boolean);
|
|
if (typeof projectId !== "string" || !PROJECT_ID.test(projectId)) throw new Error("VDB_CONVERSION_INVALID: project id");
|
|
if (typeof sourcePath !== "string" || !sourcePath.startsWith("//") || !sourcePath.toLowerCase().endsWith(".vdb") || sourcePath.includes("..")) throw new Error("VDB_CONVERSION_INVALID: source path");
|
|
if (expectedSha256 !== undefined && (typeof expectedSha256 !== "string" || !SHA256.test(expectedSha256))) throw new Error("VDB_CONVERSION_INVALID: source SHA-256");
|
|
if (quantization !== "LOSSLESS" && quantization !== "FP16") throw new Error("VDB_CONVERSION_INVALID: quantization");
|
|
if (!Number.isSafeInteger(chunkByteLength) || chunkByteLength < 64 * 1024 || chunkByteLength > 16 * 1024 * 1024 || chunkByteLength % 32 !== 0) throw new Error("VDB_CONVERSION_INVALID: chunk size");
|
|
if (selectedGrids.length > 64 || new Set(selectedGrids).size !== selectedGrids.length || selectedGrids.some((name) => !GRID_NAME.test(name))) throw new Error("VDB_CONVERSION_INVALID: grid allowlist");
|
|
return { projectId, sourcePath, expectedSha256, quantization, chunkByteLength, selectedGrids };
|
|
}
|
|
|
|
async function receiveSource(request, file) {
|
|
const hash = crypto.createHash("sha256");
|
|
let bytes = 0;
|
|
const output = fs.createWriteStream(file, { flags: "wx", mode: 0o600 });
|
|
try {
|
|
for await (const chunk of request) {
|
|
bytes += chunk.length;
|
|
if (bytes > MAX_SOURCE_BYTES) throw new Error("NON_MESH_VDB_BUDGET_EXCEEDED: source upload");
|
|
hash.update(chunk);
|
|
if (!output.write(chunk)) await new Promise((resolve) => output.once("drain", resolve));
|
|
}
|
|
await new Promise((resolve, reject) => output.end((error) => error ? reject(error) : resolve()));
|
|
}
|
|
catch (error) {
|
|
output.destroy();
|
|
await fsp.rm(file, { force: true });
|
|
throw error;
|
|
}
|
|
if (bytes === 0) throw new Error("VDB_CONVERSION_INVALID: empty source upload");
|
|
return { bytes, sha256: hash.digest("hex") };
|
|
}
|
|
|
|
function json(response, status, value) {
|
|
const body = Buffer.from(`${JSON.stringify(value)}\n`);
|
|
response.writeHead(status, { "Content-Type": "application/json", "Content-Length": body.length, "Cache-Control": "no-store" });
|
|
response.end(body);
|
|
}
|
|
|
|
export class VDBJobService {
|
|
constructor(options = {}) {
|
|
this.converter = path.resolve(options.converter ?? path.join(moduleRoot, "build_vdb_tools/vdb_to_nanovdb"));
|
|
this.manifestBuilder = path.resolve(options.manifestBuilder ?? path.join(moduleRoot, "tools/vdb/build-nanovdb-manifest.mjs"));
|
|
this.root = path.resolve(options.root ?? path.join(moduleRoot, "build_vdb_server"));
|
|
this.secret = options.secret ?? process.env.VDB_SERVER_SIGNING_KEY;
|
|
this.timeoutMs = options.timeoutMs ?? 120_000;
|
|
this.jobs = new Map();
|
|
this.byKey = new Map();
|
|
if (!this.secret || Buffer.byteLength(this.secret) < 32) throw new Error("VDB_SERVER_CONFIG_INVALID: signing key must be at least 32 bytes");
|
|
if (!fs.existsSync(this.converter) || !fs.existsSync(this.manifestBuilder)) throw new Error("VDB_SERVER_CONFIG_INVALID: converter or manifest builder is missing");
|
|
fs.mkdirSync(path.join(this.root, "jobs"), { recursive: true, mode: 0o700 });
|
|
fs.mkdirSync(path.join(this.root, "incoming"), { recursive: true, mode: 0o700 });
|
|
this.converterSha256 = fileSha256(this.converter);
|
|
}
|
|
|
|
summary(job) {
|
|
return {
|
|
id: job.id,
|
|
key: job.key,
|
|
state: job.state,
|
|
progress: job.progress,
|
|
sourceSha256: job.sourceSha256,
|
|
sourceBytes: job.sourceBytes,
|
|
createdAt: job.createdAt,
|
|
updatedAt: job.updatedAt,
|
|
error: job.error,
|
|
artifacts: job.state === "SUCCEEDED" ? {
|
|
manifest: `/v1/vdb/jobs/${job.id}/manifest`,
|
|
bundle: `/v1/vdb/jobs/${job.id}/bundle`,
|
|
report: `/v1/vdb/jobs/${job.id}/report`,
|
|
signature: job.signature,
|
|
} : undefined,
|
|
sandbox: "bwrap-unshare-all+readonly-root+prlimit",
|
|
};
|
|
}
|
|
|
|
persist(job) {
|
|
job.updatedAt = new Date().toISOString();
|
|
fs.writeFileSync(path.join(job.directory, "job.json"), `${JSON.stringify(this.summary(job), null, 2)}\n`, { mode: 0o600 });
|
|
}
|
|
|
|
appendLog(job, value) {
|
|
job.logs = `${job.logs}${value}`.slice(-MAX_LOG_BYTES);
|
|
}
|
|
|
|
async create(request) {
|
|
const metadata = parseHeaders(request);
|
|
const incoming = path.join(this.root, "incoming", `${crypto.randomUUID()}.vdb`);
|
|
const source = await receiveSource(request, incoming);
|
|
if (metadata.expectedSha256 && metadata.expectedSha256 !== source.sha256) {
|
|
await fsp.rm(incoming, { force: true });
|
|
throw new Error("NANOVDB_HASH_MISMATCH: uploaded source");
|
|
}
|
|
const key = crypto.createHash("sha256").update(stableJson({
|
|
schemaVersion: 1,
|
|
sourceSha256: source.sha256,
|
|
quantization: metadata.quantization,
|
|
chunkByteLength: metadata.chunkByteLength,
|
|
selectedGrids: metadata.selectedGrids,
|
|
converterSha256: this.converterSha256,
|
|
})).digest("hex");
|
|
const existing = this.byKey.get(key);
|
|
if (existing && existing.state !== "FAILED" && existing.state !== "CANCELLED") {
|
|
await fsp.rm(incoming, { force: true });
|
|
return { job: existing, deduplicated: true };
|
|
}
|
|
const id = `vdb-${key.slice(0, 24)}`;
|
|
const directory = path.join(this.root, "jobs", id);
|
|
await fsp.rm(directory, { recursive: true, force: true });
|
|
await fsp.mkdir(directory, { recursive: true, mode: 0o700 });
|
|
await fsp.rename(incoming, path.join(directory, "source.vdb"));
|
|
const job = {
|
|
id, key, directory, metadata, sourceSha256: source.sha256, sourceBytes: source.bytes,
|
|
state: "QUEUED", progress: 0, createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(),
|
|
logs: "", error: undefined, child: undefined, timer: undefined, signature: undefined,
|
|
};
|
|
this.jobs.set(id, job);
|
|
this.byKey.set(key, job);
|
|
this.persist(job);
|
|
setImmediate(() => void this.run(job));
|
|
return { job, deduplicated: false };
|
|
}
|
|
|
|
async run(job) {
|
|
if (job.state === "CANCELLED") return;
|
|
job.state = "RUNNING";
|
|
job.progress = 0.1;
|
|
this.persist(job);
|
|
const output = path.join(job.directory, "bundle.nvdb");
|
|
const report = path.join(job.directory, "report.json");
|
|
const cancelFile = path.join(job.directory, "cancel");
|
|
const converterArgs = [
|
|
"--input", path.join(job.directory, "source.vdb"), "--output", output, "--report", report,
|
|
"--quantization", job.metadata.quantization, "--cancel-file", cancelFile, "--timeout-ms", String(this.timeoutMs),
|
|
...job.metadata.selectedGrids.flatMap((name) => ["--grid", name]),
|
|
];
|
|
const args = [
|
|
"--die-with-parent", "--new-session", "--unshare-all", "--ro-bind", "/", "/", "--dev", "/dev", "--proc", "/proc",
|
|
"--bind", job.directory, job.directory, "--chdir", job.directory,
|
|
"/usr/bin/prlimit", "--as=2147483648", "--cpu=120", "--fsize=1200000000", "--nproc=128", "--", this.converter, ...converterArgs,
|
|
];
|
|
job.timer = setTimeout(() => {
|
|
this.appendLog(job, "VDB_SERVER_TIMEOUT: terminating sandbox\n");
|
|
this.kill(job);
|
|
}, this.timeoutMs + 2_000);
|
|
try {
|
|
const converted = await runProcess("/usr/bin/bwrap", args, {
|
|
detached: true,
|
|
onLog: (value) => this.appendLog(job, value),
|
|
onChild: (child) => { job.child = child; },
|
|
});
|
|
job.child = undefined;
|
|
if (job.state === "CANCELLED") return;
|
|
if (converted.code !== 0) throw new Error(`VDB_CONVERSION_FAILED: sandbox exited ${converted.code ?? converted.signal}`);
|
|
job.progress = 0.8;
|
|
this.persist(job);
|
|
const manifest = path.join(job.directory, "manifest.json");
|
|
const built = await runProcess(process.execPath, [this.manifestBuilder,
|
|
"--source", path.join(job.directory, "source.vdb"), "--bundle", output, "--report", report,
|
|
"--output", manifest, "--converter", this.converter, "--converter-target", "SERVER",
|
|
"--project-id", job.metadata.projectId, "--source-path", job.metadata.sourcePath,
|
|
"--bundle-path", `//volumes/${job.key}.nvdb`, "--chunk-bytes", String(job.metadata.chunkByteLength),
|
|
], { onLog: (value) => this.appendLog(job, value) });
|
|
if (built.code !== 0) throw new Error(`VDB_MANIFEST_FAILED: builder exited ${built.code ?? built.signal}`);
|
|
const manifestSha256 = fileSha256(manifest);
|
|
const bundleSha256 = fileSha256(output);
|
|
job.signature = crypto.createHmac("sha256", this.secret).update(`${job.id}:${manifestSha256}:${bundleSha256}`).digest("hex");
|
|
job.state = "SUCCEEDED";
|
|
job.progress = 1;
|
|
this.persist(job);
|
|
}
|
|
catch (error) {
|
|
if (job.state !== "CANCELLED") {
|
|
job.state = "FAILED";
|
|
job.error = error instanceof Error ? error.message : String(error);
|
|
this.persist(job);
|
|
}
|
|
await Promise.all([fsp.rm(output, { force: true }), fsp.rm(report, { force: true }), fsp.rm(path.join(job.directory, "manifest.json"), { force: true })]);
|
|
}
|
|
finally {
|
|
if (job.timer) clearTimeout(job.timer);
|
|
job.timer = undefined;
|
|
job.child = undefined;
|
|
fs.writeFileSync(path.join(job.directory, "job.log"), job.logs, { mode: 0o600 });
|
|
}
|
|
}
|
|
|
|
kill(job) {
|
|
if (!job.child?.pid) return;
|
|
try { process.kill(-job.child.pid, "SIGTERM"); } catch { /* already exited */ }
|
|
const pid = job.child.pid;
|
|
setTimeout(() => { try { process.kill(-pid, "SIGKILL"); } catch { /* already exited */ } }, 1_000).unref();
|
|
}
|
|
|
|
cancel(id) {
|
|
const job = this.jobs.get(id);
|
|
if (!job) return undefined;
|
|
if (["SUCCEEDED", "FAILED", "CANCELLED"].includes(job.state)) return job;
|
|
job.state = "CANCELLED";
|
|
job.progress = 0;
|
|
fs.writeFileSync(path.join(job.directory, "cancel"), "cancelled\n", { mode: 0o600 });
|
|
this.kill(job);
|
|
this.persist(job);
|
|
return job;
|
|
}
|
|
|
|
artifact(job, name) {
|
|
const allowed = { manifest: "manifest.json", bundle: "bundle.nvdb", report: "report.json", logs: "job.log" };
|
|
if (job.state !== "SUCCEEDED" && name !== "logs") return undefined;
|
|
const filename = allowed[name];
|
|
if (!filename) return undefined;
|
|
const file = path.join(job.directory, filename);
|
|
return fs.existsSync(file) ? file : undefined;
|
|
}
|
|
}
|
|
|
|
export function createVDBJobHttpServer(service) {
|
|
return http.createServer(async (request, response) => {
|
|
try {
|
|
const url = new URL(request.url ?? "/", "http://127.0.0.1");
|
|
if (request.method === "GET" && url.pathname === "/healthz") return json(response, 200, { ok: true, converterSha256: service.converterSha256, sandbox: true });
|
|
if (request.method === "POST" && url.pathname === "/v1/vdb/jobs") {
|
|
const created = await service.create(request);
|
|
return json(response, created.deduplicated ? 200 : 202, { ...service.summary(created.job), deduplicated: created.deduplicated });
|
|
}
|
|
const match = url.pathname.match(/^\/v1\/vdb\/jobs\/([A-Za-z0-9-]+)(?:\/(manifest|bundle|report|logs))?$/);
|
|
if (!match) return json(response, 404, { error: "NOT_FOUND" });
|
|
const job = service.jobs.get(match[1]);
|
|
if (!job) return json(response, 404, { error: "VDB_JOB_NOT_FOUND" });
|
|
if (request.method === "DELETE" && !match[2]) return json(response, 200, service.summary(service.cancel(job.id)));
|
|
if (request.method !== "GET") return json(response, 405, { error: "METHOD_NOT_ALLOWED" });
|
|
if (!match[2]) return json(response, 200, service.summary(job));
|
|
const artifact = service.artifact(job, match[2]);
|
|
if (!artifact) return json(response, 409, { error: "VDB_ARTIFACT_NOT_READY", state: job.state });
|
|
const stat = fs.statSync(artifact);
|
|
response.writeHead(200, {
|
|
"Content-Type": match[2] === "bundle" ? "application/x-nanovdb" : match[2] === "logs" ? "text/plain" : "application/json",
|
|
"Content-Length": stat.size,
|
|
"X-Content-SHA256": fileSha256(artifact),
|
|
"X-VDB-Signature": job.signature ?? "",
|
|
"Cache-Control": "private, immutable",
|
|
});
|
|
fs.createReadStream(artifact).pipe(response);
|
|
}
|
|
catch (error) {
|
|
if (!response.headersSent) json(response, 400, { error: error instanceof Error ? error.message : String(error) });
|
|
else response.destroy(error instanceof Error ? error : undefined);
|
|
}
|
|
});
|
|
}
|