feat: extend importBundleArchive for site-packages and fixups

This commit is contained in:
SnapOtter
2026-06-13 16:32:45 +08:00
parent d1e9536cff
commit 5d5240e481
2 changed files with 98 additions and 0 deletions
+39
View File
@@ -1,3 +1,4 @@
import { execFileSync } from "node:child_process";
import { randomUUID } from "node:crypto"; import { randomUUID } from "node:crypto";
import { import {
constants, constants,
@@ -632,6 +633,44 @@ export async function importBundleArchive(
moveTreeRecursive(stagingModels, MODELS_DIR); moveTreeRecursive(stagingModels, MODELS_DIR);
} }
// Move site-packages/* into venv site-packages
const stagingSitePackages = join(stagingDir, "site-packages");
if (existsSync(stagingSitePackages)) {
const venvPath = process.env.PYTHON_VENV_PATH || join(AI_DIR, "venv");
let sitePackagesDir = "";
const libDir = join(venvPath, "lib");
if (existsSync(libDir)) {
const pyDirs = readdirSync(libDir).filter((d) => d.startsWith("python"));
if (pyDirs.length > 0) {
sitePackagesDir = join(libDir, pyDirs[0], "site-packages");
}
}
if (sitePackagesDir && existsSync(sitePackagesDir)) {
moveTreeRecursive(stagingSitePackages, sitePackagesDir);
}
}
// Apply fixups (NCCL wheel) if present
const stagingFixups = join(stagingDir, "fixups");
if (existsSync(stagingFixups)) {
const wheels = readdirSync(stagingFixups).filter((f) => f.endsWith(".whl"));
if (wheels.length > 0) {
const venvPython =
(process.env.PYTHON_VENV_PATH || join(AI_DIR, "venv")) + "/bin/python3";
for (const wheel of wheels) {
try {
execFileSync(venvPython, [
"-m", "pip", "install", "--no-index",
`--find-links=${stagingFixups}`,
wheel.split("-")[0],
], { stdio: "ignore", timeout: 30_000 });
} catch {
// Non-fatal
}
}
}
}
markInstalled(descriptor.bundleId, descriptor.version, descriptor.models); markInstalled(descriptor.bundleId, descriptor.version, descriptor.models);
return { return {
+59
View File
@@ -143,6 +143,41 @@ async function buildSymlinkArchive(): Promise<string> {
return archivePath; return archivePath;
} }
async function createBundleTarWithSitePackages(
bundleId: string,
version: string,
modelFiles: Record<string, string>,
sitePackageFiles: Record<string, string>,
): Promise<string> {
const tarDir = join(testRoot, `tar-src-${randomUUID()}`);
mkdirSync(tarDir, { recursive: true });
writeFileSync(
join(tarDir, "bundle.json"),
JSON.stringify({ bundleId, version, models: Object.keys(modelFiles) }),
);
const modelsSubdir = join(tarDir, "models");
mkdirSync(modelsSubdir, { recursive: true });
for (const [name, content] of Object.entries(modelFiles)) {
const modelPath = join(modelsSubdir, name);
mkdirSync(join(modelPath, ".."), { recursive: true });
writeFileSync(modelPath, content);
}
const spSubdir = join(tarDir, "site-packages");
mkdirSync(spSubdir, { recursive: true });
for (const [name, content] of Object.entries(sitePackageFiles)) {
const spPath = join(spSubdir, name);
mkdirSync(join(spPath, ".."), { recursive: true });
writeFileSync(spPath, content);
}
const tarPath = join(testRoot, `bundle-${randomUUID()}.tar.gz`);
await tar.create({ gzip: true, file: tarPath, cwd: tarDir }, ["."]);
return tarPath;
}
function resetState(): void { function resetState(): void {
writeFileSync(installedPath, JSON.stringify({ bundles: {} }), "utf-8"); writeFileSync(installedPath, JSON.stringify({ bundles: {} }), "utf-8");
invalidateCache(); invalidateCache();
@@ -265,6 +300,30 @@ describe("importBundleArchive", () => {
}); });
}); });
describe("site-packages import", () => {
beforeEach(resetState);
it("extracts site-packages into venv site-packages directory", async () => {
const venvSitePackages = join(aiDir, "venv", "lib", "python3.12", "site-packages");
mkdirSync(venvSitePackages, { recursive: true });
process.env.PYTHON_VENV_PATH = join(aiDir, "venv");
const tarPath = await createBundleTarWithSitePackages(
testBundleId,
testVersion,
{ "mediapipe/face.tflite": "model-data" },
{ "fakepkg/__init__.py": "# fake package" },
);
invalidateCache();
const result = await importBundleArchive(createReadStream(tarPath));
expect(result.bundleId).toBe(testBundleId);
expect(existsSync(join(venvSitePackages, "fakepkg", "__init__.py"))).toBe(true);
delete process.env.PYTHON_VENV_PATH;
});
});
describe("POST /api/v1/admin/features/import", () => { describe("POST /api/v1/admin/features/import", () => {
let app: Awaited<ReturnType<typeof import("fastify")>>["default"] extends ( let app: Awaited<ReturnType<typeof import("fastify")>>["default"] extends (
...args: infer _A ...args: infer _A