Advance WebGPU volume and bounded workflows

2026-08-14 18:08:29 -04:00
parent 3da1dfc804
commit 68d50f810f
119 changed files with 9028 additions and 430 deletions
--- a/tools/vdb/server/vdb-job-server.mjs
+++ b/tools/vdb/server/vdb-job-server.mjs
@@ -0,0 +1,18 @@
+#!/usr/bin/env node
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+import { VDBJobService, createVDBJobHttpServer } from "./vdb-job-service.mjs";
+
+const root = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "../../..");
+const host = process.env.VDB_SERVER_HOST ?? "127.0.0.1";
+const port = Number(process.env.VDB_SERVER_PORT ?? 8787);
+const service = new VDBJobService({
+  converter: process.env.VDB_CONVERTER ?? path.join(root, "build_vdb_tools/vdb_to_nanovdb"),
+  root: process.env.VDB_SERVER_DATA ?? path.join(root, "build_vdb_server"),
+});
+const server = createVDBJobHttpServer(service);
+server.listen(port, host, () => process.stdout.write(`vdb-job-server-ready http://${host}:${port}\n`));
+
+for (const signal of ["SIGINT", "SIGTERM"]) {
+  process.on(signal, () => server.close(() => process.exit(0)));
+}
--- a/tools/vdb/server/vdb-job-service.mjs
+++ b/tools/vdb/server/vdb-job-service.mjs
@@ -0,0 +1,293 @@
+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);
+    }
+  });
+}