fix(cli): a killed remove-background run no longer leaves its CoreML temp files behind (#6582)

* fix(cli): a killed remove-background run no longer leaves its CoreML temp files behind

On Apple Silicon, CoreML writes about 340 MB of compiled models to the system
temp folder for every session and removes them only when the session is
released. A run stopped by a signal never released it, so every cancelled run
left them. A signal now stops the run through the CLI's shutdown, which
releases the session, and the next run removes what a killed run listed in its
own manifest.

* style(cli): wrap the CoreML shutdown comment to the line limit

packages/cli/src/background-removal/coremlTemp.test.ts

@@ -0,0 +1,87 @@
+import { spawnSync } from "node:child_process";
+import {
+  existsSync,
+  mkdirSync,
+  mkdtempSync,
+  readdirSync,
+  readFileSync,
+  rmSync,
+  writeFileSync,
+} from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, describe, expect, it } from "vitest";
+import { sweepDeadCoreMlTemp, trackCoreMlTemp } from "./coremlTemp.js";
+
+const roots: string[] = [];
+afterEach(() => {
+  for (const root of roots.splice(0)) rmSync(root, { recursive: true, force: true });
+});
+
+function scratch() {
+  const root = mkdtempSync(join(tmpdir(), "hf-coreml-temp-"));
+  roots.push(root);
+  const tmp = join(root, "T");
+  const manifests = join(root, "manifests");
+  mkdirSync(tmp);
+  return { tmp, manifests };
+}
+
+const compiled = (tmp: string, id: string) => {
+  mkdirSync(join(tmp, `onnxruntime-${id}.model.mlmodelc`));
+  writeFileSync(join(tmp, `onnxruntime-${id}.model.mlmodel`), "spec");
+};
+
+describe("CoreML temp folders", () => {
+  it("a run killed before it released its session has its compiled models removed by the next run", () => {
+    const { tmp, manifests } = scratch();
+    mkdirSync(join(tmp, "unrelated"));
+    const killed = trackCoreMlTemp({ tmp, manifests, alive: () => true });
+    compiled(tmp, "a");
+    compiled(tmp, "b");
+    killed.record();
+    const path = join(manifests, `${process.pid}.json`);
+    const gone = spawnSync(process.execPath, ["-e", ""]).pid;
+    writeFileSync(path, JSON.stringify({ ...JSON.parse(readFileSync(path, "utf-8")), pid: gone }));
+
+    sweepDeadCoreMlTemp({ tmp, manifests });
+
+    expect(readdirSync(tmp)).toEqual(["unrelated"]);
+    expect(readdirSync(manifests)).toEqual([]);
+  });
+
+  it("leaves a running run's compiled models alone", () => {
+    const { tmp, manifests } = scratch();
+    const running = trackCoreMlTemp({ tmp, manifests });
+    compiled(tmp, "a");
+    running.record();
+    writeFileSync(join(manifests, "999999.json"), JSON.stringify({ pid: 999999, tmp, dirs: [] }));
+
+    sweepDeadCoreMlTemp({ tmp, manifests, alive: (pid) => pid !== 999999 });
+
+    expect(existsSync(join(tmp, "onnxruntime-a.model.mlmodelc"))).toBe(true);
+    expect(readdirSync(manifests)).toEqual([`${process.pid}.json`]);
+  });
+
+  it("removes only compiled models by name, never a path a manifest points elsewhere", () => {
+    const { tmp, manifests } = scratch();
+    mkdirSync(manifests);
+    mkdirSync(join(tmp, "keep"));
+    writeFileSync(
+      join(manifests, "1.json"),
+      JSON.stringify({ pid: 1, tmp, dirs: ["../T/keep", "keep", "onnxruntime-x.model.mlmodelc"] }),
+    );
+
+    sweepDeadCoreMlTemp({ tmp, manifests, alive: () => false });
+
+    expect(existsSync(join(tmp, "keep"))).toBe(true);
+  });
+
+  it("forgets a run that released its session", () => {
+    const { tmp, manifests } = scratch();
+    const run = trackCoreMlTemp({ tmp, manifests });
+    run.record();
+    run.forget();
+    expect(readdirSync(manifests)).toEqual([]);
+  });
+});

packages/cli/src/background-removal/coremlTemp.ts

@@ -0,0 +1,82 @@
+import { mkdirSync, readdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { homedir, tmpdir } from "node:os";
+import { basename, join } from "node:path";
+import { processIsAlive } from "../utils/processTree.js";
+
+const COMPILED_MODEL = /^onnxruntime-.*\.mlmodelc?$/;
+
+const COREML_TEMP_MANIFESTS = join(
+  homedir(),
+  ".cache",
+  "hyperframes",
+  "background-removal",
+  "coreml-temp",
+);
+
+interface Manifest {
+  pid: number;
+  tmp: string;
+  dirs: string[];
+}
+
+interface Options {
+  manifests?: string;
+  tmp?: string;
+  alive?: (pid: number) => boolean;
+}
+
+const compiledModels = (tmp: string): string[] =>
+  readdirSync(tmp).filter((name) => COMPILED_MODEL.test(name));
+
+function readManifest(path: string): Manifest | undefined {
+  try {
+    const parsed = JSON.parse(readFileSync(path, "utf-8")) as Manifest;
+    return Number.isInteger(parsed.pid) && Array.isArray(parsed.dirs) ? parsed : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+/** Removes the compiled models a run left in the system temp folder when it was killed before it could release them. */
+export function sweepDeadCoreMlTemp(options: Options = {}): void {
+  const { manifests = COREML_TEMP_MANIFESTS, alive = processIsAlive } = options;
+  let names: string[];
+  try {
+    names = readdirSync(manifests);
+  } catch {
+    return;
+  }
+  for (const name of names) {
+    const path = join(manifests, name);
+    const manifest = readManifest(path);
+    if (manifest && (manifest.pid === process.pid || alive(manifest.pid))) continue;
+    for (const dir of manifest?.dirs ?? []) {
+      if (basename(dir) === dir && COMPILED_MODEL.test(dir)) {
+        rmSync(join(manifest!.tmp, dir), { recursive: true, force: true });
+      }
+    }
+    rmSync(path, { force: true });
+  }
+}
+
+/**
+ * CoreML compiles the model into one temp folder per partition and only removes them when the session is released;
+ * a killed run leaves them behind. This names the ones this run made so the next run can remove them.
+ */
+export function trackCoreMlTemp(options: Options = {}) {
+  const { manifests = COREML_TEMP_MANIFESTS, tmp = tmpdir() } = options;
+  sweepDeadCoreMlTemp(options);
+  const before = new Set(compiledModels(tmp));
+  const path = join(manifests, `${process.pid}.json`);
+  return {
+    record(): void {
+      const dirs = compiledModels(tmp).filter((name) => !before.has(name));
+      mkdirSync(manifests, { recursive: true });
+      const manifest: Manifest = { pid: process.pid, tmp, dirs };
+      writeFileSync(path, JSON.stringify(manifest));
+    },
+    forget(): void {
+      rmSync(path, { force: true });
+    },
+  };
+}

packages/cli/src/background-removal/inference.coremlTemp.test.ts

@@ -0,0 +1,95 @@
+import { spawnSync } from "node:child_process";
+import {
+  existsSync,
+  mkdirSync,
+  mkdtempSync,
+  readdirSync,
+  readFileSync,
+  rmSync,
+  writeFileSync,
+} from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const scratch = vi.hoisted(() => ({ tmp: "", manifests: "", compiled: 0 }));
+const cleanups = vi.hoisted(() => new Set<() => Promise<void>>());
+
+vi.mock("../browser/launch.js", () => ({
+  registerBrowserSessionCleanup: (cleanup: () => Promise<void>) => {
+    cleanups.add(cleanup);
+    return () => cleanups.delete(cleanup);
+  },
+}));
+vi.mock("./manager.js", () => ({
+  ensureModel: async () => "model.onnx",
+  selectProviders: () => ({ providers: ["coreml", "cpu"], label: "CoreML" }),
+}));
+vi.mock("../utils/optionalPackages.js", () => ({
+  loadSharp: async () => ({}),
+  loadOptionalPackage: async () => ({
+    InferenceSession: {
+      create: async () => {
+        for (let i = 0; i < 3; i++) {
+          mkdirSync(join(scratch.tmp, `onnxruntime-${scratch.compiled++}.model.mlmodelc`));
+        }
+        return { inputNames: ["in"], outputNames: ["out"], release: vi.fn(async () => undefined) };
+      },
+    },
+  }),
+}));
+vi.mock("./coremlTemp.js", async (importActual) => {
+  const actual = await importActual<typeof import("./coremlTemp.js")>();
+  return {
+    ...actual,
+    trackCoreMlTemp: () =>
+      actual.trackCoreMlTemp({ tmp: scratch.tmp, manifests: scratch.manifests }),
+  };
+});
+
+import { createSession } from "./inference.js";
+
+const compiledDirs = () => readdirSync(scratch.tmp).filter((n) => n.startsWith("onnxruntime-"));
+
+beforeEach(() => {
+  const root = mkdtempSync(join(tmpdir(), "hf-coreml-session-"));
+  scratch.tmp = join(root, "T");
+  scratch.manifests = join(root, "manifests");
+  scratch.compiled = 0;
+  mkdirSync(scratch.tmp);
+});
+afterEach(() => {
+  cleanups.clear();
+  rmSync(join(scratch.tmp, ".."), { recursive: true, force: true });
+});
+
+describe("CoreML sessions and the system temp folder", () => {
+  it("the next run removes the compiled models a killed run left behind", async () => {
+    await createSession();
+    expect(compiledDirs()).toHaveLength(3);
+    const manifestPath = join(scratch.manifests, `${process.pid}.json`);
+    const gone = spawnSync(process.execPath, ["-e", ""]).pid!;
+    writeFileSync(
+      manifestPath,
+      JSON.stringify({ ...JSON.parse(readFileSync(manifestPath, "utf-8")), pid: gone }),
+    );
+
+    await createSession();
+
+    expect(compiledDirs()).toEqual([
+      "onnxruntime-3.model.mlmodelc",
+      "onnxruntime-4.model.mlmodelc",
+      "onnxruntime-5.model.mlmodelc",
+    ]);
+  });
+
+  it("a signal releases the session through the CLI's shutdown, and a closed session forgets it", async () => {
+    const session = await createSession();
+    expect(cleanups.size).toBe(1);
+    const [stop] = [...cleanups];
+    await stop!();
+    await session.close();
+    expect(cleanups.size).toBe(0);
+    expect(existsSync(join(scratch.manifests, `${process.pid}.json`))).toBe(false);
+  });
+});

packages/cli/src/background-removal/inference.ts

@@ -8,6 +8,8 @@
 import type { InferenceSession, Tensor } from "onnxruntime-node";
 import type sharpType from "sharp";
 import { loadOptionalPackage, loadSharp } from "../utils/optionalPackages.js";
+import { registerBrowserSessionCleanup } from "../browser/launch.js";
+import { trackCoreMlTemp } from "./coremlTemp.js";
 import { ensureModel, selectProviders, type Device, type ModelId } from "./manager.js";
 
 const INPUT_SIZE = 320;
@@ -71,6 +73,9 @@ export async function createSession(options: CreateSessionOptions = {}): Promise
       graphOptimizationLevel: "all",
     });
 
+  const coreMlTemp = choice.label === "CoreML" ? trackCoreMlTemp() : undefined;
+  // A signal stops the run through the CLI's shutdown, which releases the session so CoreML removes its temp files.
+  const unregister = coreMlTemp && registerBrowserSessionCleanup(async () => session?.release());
   let session: InferenceSession;
   let providerUsed = choice.label;
   try {
@@ -80,8 +85,13 @@ export async function createSession(options: CreateSessionOptions = {}): Promise
     options.onProgress?.(
       `${choice.label} provider failed (${(err as Error).message}); falling back to CPU.`,
     );
-    session = await tryCreate(["cpu"]);
+    session = await tryCreate(["cpu"]).catch((error: unknown) => {
+      unregister?.();
+      throw error;
+    });
     providerUsed = "CPU";
+  } finally {
+    coreMlTemp?.record();
   }
 
   const inputName = session.inputNames[0];
@@ -145,7 +155,9 @@ export async function createSession(options: CreateSessionOptions = {}): Promise
       );
     },
     async close() {
+      unregister?.();
       await session.release();
+      coreMlTemp?.forget();
     },
   };
 }