feat: rewrite install_feature.py for pre-built tar bundles

2026-08-03 07:46:42 +02:00 · 2026-06-13 16:37:00 +08:00
parent 5d5240e481
commit a4fa3ce2a7
2 changed files with 460 additions and 523 deletions
@@ -0,0 +1,167 @@
+import { spawnSync, execFileSync } from "node:child_process";
+import {
+  existsSync,
+  mkdirSync,
+  mkdtempSync,
+  readFileSync,
+  rmSync,
+  writeFileSync,
+} from "node:fs";
+import { createHash } from "node:crypto";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+
+const scriptPath = join(process.cwd(), "packages/ai/python/install_feature.py");
+
+let tempDir: string;
+let aiDir: string;
+let modelsDir: string;
+let venvDir: string;
+let sitePackagesDir: string;
+let manifestPath: string;
+
+beforeEach(() => {
+  tempDir = mkdtempSync(join(tmpdir(), "snapotter-install-test-"));
+  aiDir = join(tempDir, "ai");
+  modelsDir = join(aiDir, "models");
+  venvDir = join(aiDir, "venv");
+  sitePackagesDir = join(venvDir, "lib", "python3.12", "site-packages");
+  manifestPath = join(tempDir, "feature-manifest.json");
+
+  mkdirSync(sitePackagesDir, { recursive: true });
+  mkdirSync(modelsDir, { recursive: true });
+  mkdirSync(join(aiDir, "staging"), { recursive: true });
+  writeFileSync(join(aiDir, "installed.json"), JSON.stringify({ bundles: {} }));
+});
+
+afterEach(() => {
+  rmSync(tempDir, { recursive: true, force: true });
+});
+
+function createTestTar(bundleId: string): { tarPath: string; sha256: string } {
+  const buildDir = join(tempDir, "build");
+  mkdirSync(join(buildDir, "models", "testmodel"), { recursive: true });
+  mkdirSync(join(buildDir, "site-packages", "testpkg"), { recursive: true });
+  writeFileSync(join(buildDir, "models", "testmodel", "weights.bin"), "model-weights");
+  writeFileSync(join(buildDir, "site-packages", "testpkg", "__init__.py"), "# test");
+  writeFileSync(
+    join(buildDir, "bundle.json"),
+    JSON.stringify({
+      bundleId,
+      version: "1.0.0-test",
+      arch: "amd64-gpu",
+      imageVersion: "2.0.0",
+      pythonVersion: "3.12",
+      models: ["testmodel"],
+    }),
+  );
+
+  const tarPath = join(tempDir, `${bundleId}-test.tar.gz`);
+  execFileSync("tar", ["czf", tarPath, "-C", buildDir, "."]);
+  rmSync(buildDir, { recursive: true });
+
+  const hash = createHash("sha256").update(readFileSync(tarPath)).digest("hex");
+  return { tarPath, sha256: hash };
+}
+
+function writeManifest(bundleId: string, tarPath: string, sha256: string) {
+  const size = readFileSync(tarPath).length;
+  const manifest = {
+    manifestVersion: 2,
+    imageVersion: "2.0.0",
+    pythonVersion: "3.12",
+    basePackages: [],
+    bundleRepo: "snapotter/feature-bundles",
+    bundles: {
+      [bundleId]: {
+        name: "Test Bundle",
+        archives: {
+          "amd64-gpu": { file: tarPath, sha256, compressedSize: size, extractedSize: size * 2 },
+          "arm64-cpu": { file: tarPath, sha256, compressedSize: size, extractedSize: size * 2 },
+        },
+        models: [{ id: "testmodel", path: "testmodel/weights.bin", minSize: 0 }],
+        enablesTools: [],
+      },
+    },
+  };
+  writeFileSync(manifestPath, JSON.stringify(manifest));
+}
+
+describe("install_feature.py prebuilt mode", () => {
+  it("extracts models and site-packages from a local tar", () => {
+    const { tarPath, sha256 } = createTestTar("face-detection");
+    writeManifest("face-detection", tarPath, sha256);
+
+    const result = spawnSync(
+      "python3",
+      [scriptPath, "face-detection", manifestPath, modelsDir],
+      {
+        env: {
+          ...process.env,
+          DATA_DIR: tempDir,
+          PYTHON_VENV_PATH: venvDir,
+          SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath,
+        },
+        timeout: 30_000,
+      },
+    );
+
+    expect(result.status, `stderr: ${result.stderr?.toString()}`).toBe(0);
+    expect(existsSync(join(modelsDir, "testmodel", "weights.bin"))).toBe(true);
+    expect(existsSync(join(sitePackagesDir, "testpkg", "__init__.py"))).toBe(true);
+
+    const installed = JSON.parse(readFileSync(join(aiDir, "installed.json"), "utf-8"));
+    expect(installed.bundles["face-detection"]).toBeDefined();
+    expect(installed.bundles["face-detection"].version).toBe("1.0.0-test");
+  });
+
+  it("exits non-zero when checksum mismatches", () => {
+    const { tarPath } = createTestTar("face-detection");
+    writeManifest("face-detection", tarPath, "badhash".padEnd(64, "0"));
+
+    const result = spawnSync(
+      "python3",
+      [scriptPath, "face-detection", manifestPath, modelsDir],
+      {
+        env: {
+          ...process.env,
+          DATA_DIR: tempDir,
+          PYTHON_VENV_PATH: venvDir,
+          SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath,
+        },
+        timeout: 30_000,
+      },
+    );
+
+    expect(result.status).not.toBe(0);
+  });
+
+  it("writes progress JSON to stderr", () => {
+    const { tarPath, sha256 } = createTestTar("face-detection");
+    writeManifest("face-detection", tarPath, sha256);
+
+    const result = spawnSync(
+      "python3",
+      [scriptPath, "face-detection", manifestPath, modelsDir],
+      {
+        env: {
+          ...process.env,
+          DATA_DIR: tempDir,
+          PYTHON_VENV_PATH: venvDir,
+          SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath,
+        },
+        timeout: 30_000,
+      },
+    );
+
+    const stderr = result.stderr?.toString() ?? "";
+    const progressLines = stderr.split("\n").filter((l) => {
+      try { const p = JSON.parse(l); return typeof p.progress === "number"; } catch { return false; }
+    });
+    expect(progressLines.length).toBeGreaterThan(0);
+
+    const last = JSON.parse(progressLines[progressLines.length - 1]);
+    expect(last.progress).toBe(100);
+  });
+});