feat: classify sensitive columns locally
This commit is contained in:
Codex
@@ -0,0 +1,254 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
|
||||
import { tmpdir } from "node:os";
|
||||
import { z } from "zod";
|
||||
import type {
|
||||
LocalNerCandidate,
|
||||
LocalNerDetector,
|
||||
LocalNerEvidence,
|
||||
} from "./sensitivity-classifier.js";
|
||||
|
||||
const MAX_LINE_BYTES = 64 * 1024;
|
||||
const candidateSchema = z.object({
|
||||
columnId: z.uuid(),
|
||||
text: z.string().min(1).max(500),
|
||||
}).strict();
|
||||
const workerMessageSchema = z.union([
|
||||
z.object({ ready: z.literal(true) }).strict(),
|
||||
z.object({
|
||||
id: z.uuid(),
|
||||
ok: z.literal(true),
|
||||
evidence: z.array(z.object({
|
||||
columnId: z.uuid(),
|
||||
label: z.string().min(1).max(80),
|
||||
confidence: z.number().min(0).max(1),
|
||||
}).strict()).max(1_000),
|
||||
}).strict(),
|
||||
z.object({ id: z.uuid(), ok: z.literal(false), error: z.string().min(1).max(80) }).strict(),
|
||||
]);
|
||||
|
||||
export class LocalNerUnavailableError extends Error {
|
||||
constructor() {
|
||||
super("local NER is unavailable");
|
||||
this.name = "LocalNerUnavailableError";
|
||||
}
|
||||
}
|
||||
|
||||
interface PendingRequest {
|
||||
resolve: (value: readonly LocalNerEvidence[]) => void;
|
||||
reject: (error: Error) => void;
|
||||
timer: ReturnType<typeof setTimeout>;
|
||||
signal: AbortSignal;
|
||||
cancel: () => void;
|
||||
}
|
||||
|
||||
/** Persistent JSONL adapter for the optional, CPU-only Python NER worker. */
|
||||
export class PythonLocalNerDetector implements LocalNerDetector {
|
||||
private child?: ChildProcessWithoutNullStreams;
|
||||
private ready?: Promise<void>;
|
||||
private readyResolve?: () => void;
|
||||
private readyReject?: (error: Error) => void;
|
||||
private workerReady = false;
|
||||
private stdout = "";
|
||||
private readonly pending = new Map<string, PendingRequest>();
|
||||
|
||||
constructor(private readonly options: {
|
||||
pythonExecutable: string;
|
||||
workerScript: string;
|
||||
modelPath: string;
|
||||
cwd: string;
|
||||
threads?: number;
|
||||
startupTimeoutMs?: number;
|
||||
}) {}
|
||||
|
||||
async warmup(): Promise<void> {
|
||||
await this.ensureStarted();
|
||||
}
|
||||
|
||||
isReady(): boolean {
|
||||
return this.workerReady
|
||||
&& this.child !== undefined
|
||||
&& this.child.exitCode === null
|
||||
&& this.child.signalCode === null;
|
||||
}
|
||||
|
||||
async detect(
|
||||
candidates: readonly LocalNerCandidate[],
|
||||
signal: AbortSignal,
|
||||
deadline: number,
|
||||
): Promise<readonly LocalNerEvidence[]> {
|
||||
const parsed = z.array(candidateSchema).min(1).max(128).parse(candidates);
|
||||
if (signal.aborted || deadline <= Date.now()) throw new LocalNerUnavailableError();
|
||||
await this.ensureStartedWithin(signal, deadline);
|
||||
if (!this.child || this.child.exitCode !== null || this.child.signalCode !== null) {
|
||||
throw new LocalNerUnavailableError();
|
||||
}
|
||||
const id = randomUUID();
|
||||
return await new Promise<readonly LocalNerEvidence[]>((resolve, reject) => {
|
||||
const fail = () => {
|
||||
this.finishPending(id);
|
||||
reject(new LocalNerUnavailableError());
|
||||
this.stopWorker();
|
||||
};
|
||||
const timer = setTimeout(fail, Math.max(1, Math.floor(deadline - Date.now())));
|
||||
const cancel = fail;
|
||||
const pending: PendingRequest = { resolve, reject, timer, signal, cancel };
|
||||
this.pending.set(id, pending);
|
||||
signal.addEventListener("abort", cancel, { once: true });
|
||||
this.child!.stdin.write(`${JSON.stringify({ id, candidates: parsed })}\n`, (error) => {
|
||||
if (error) fail();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async close(): Promise<void> {
|
||||
const child = this.child;
|
||||
if (!child || child.exitCode !== null || child.signalCode !== null) return;
|
||||
await new Promise<void>((resolve) => {
|
||||
child.once("close", () => resolve());
|
||||
child.kill("SIGTERM");
|
||||
setTimeout(() => {
|
||||
if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL");
|
||||
}, 250).unref();
|
||||
});
|
||||
}
|
||||
|
||||
private async ensureStarted(): Promise<void> {
|
||||
if (this.ready) return await this.ready;
|
||||
this.ready = new Promise<void>((resolve, reject) => {
|
||||
this.readyResolve = resolve;
|
||||
this.readyReject = reject;
|
||||
});
|
||||
const threads = String(this.options.threads ?? 2);
|
||||
const inheritedRuntimeEnvironment = Object.fromEntries([
|
||||
"PATH", "SystemRoot", "WINDIR", "PATHEXT", "TMPDIR", "TEMP", "TMP", "LANG", "LC_ALL",
|
||||
].flatMap((name) => process.env[name] === undefined ? [] : [[name, process.env[name]!]]));
|
||||
const child = spawn(this.options.pythonExecutable, [
|
||||
"-I",
|
||||
"-B",
|
||||
this.options.workerScript,
|
||||
"--model",
|
||||
this.options.modelPath,
|
||||
"--threads",
|
||||
threads,
|
||||
], {
|
||||
cwd: this.options.cwd,
|
||||
stdio: ["pipe", "pipe", "pipe"],
|
||||
env: {
|
||||
...inheritedRuntimeEnvironment,
|
||||
HOME: process.env.HOME ?? tmpdir(),
|
||||
CUDA_VISIBLE_DEVICES: "",
|
||||
HIP_VISIBLE_DEVICES: "",
|
||||
HF_HUB_OFFLINE: "1",
|
||||
HF_HUB_DISABLE_TELEMETRY: "1",
|
||||
TRANSFORMERS_OFFLINE: "1",
|
||||
TOKENIZERS_PARALLELISM: "false",
|
||||
PYTHONNOUSERSITE: "1",
|
||||
OMP_NUM_THREADS: threads,
|
||||
MKL_NUM_THREADS: threads,
|
||||
OPENBLAS_NUM_THREADS: threads,
|
||||
HTTP_PROXY: "",
|
||||
HTTPS_PROXY: "",
|
||||
ALL_PROXY: "",
|
||||
NO_PROXY: "*",
|
||||
},
|
||||
});
|
||||
this.child = child;
|
||||
child.stdout.setEncoding("utf8");
|
||||
child.stdout.on("data", (chunk: string) => this.receive(chunk));
|
||||
child.stderr.resume();
|
||||
child.once("error", () => this.failWorker());
|
||||
child.once("close", () => this.failWorker());
|
||||
const startupTimer = setTimeout(() => this.failWorker(), this.options.startupTimeoutMs ?? 120_000);
|
||||
startupTimer.unref();
|
||||
try {
|
||||
await this.ready;
|
||||
} finally {
|
||||
clearTimeout(startupTimer);
|
||||
}
|
||||
}
|
||||
|
||||
private async ensureStartedWithin(signal: AbortSignal, deadline: number): Promise<void> {
|
||||
const started = this.ensureStarted();
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
let settled = false;
|
||||
const finish = (error?: Error, stopWorker = false) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
clearTimeout(timer);
|
||||
signal.removeEventListener("abort", cancel);
|
||||
if (stopWorker) this.failWorker();
|
||||
if (error) reject(error);
|
||||
else resolve();
|
||||
};
|
||||
const cancel = () => finish(new LocalNerUnavailableError(), true);
|
||||
const timer = setTimeout(cancel, Math.max(1, Math.floor(deadline - Date.now())));
|
||||
signal.addEventListener("abort", cancel, { once: true });
|
||||
void started.then(
|
||||
() => finish(),
|
||||
() => finish(new LocalNerUnavailableError()),
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
private receive(chunk: string): void {
|
||||
this.stdout += chunk;
|
||||
if (Buffer.byteLength(this.stdout, "utf8") > MAX_LINE_BYTES) {
|
||||
this.failWorker();
|
||||
return;
|
||||
}
|
||||
let newline: number;
|
||||
while ((newline = this.stdout.indexOf("\n")) >= 0) {
|
||||
const line = this.stdout.slice(0, newline);
|
||||
this.stdout = this.stdout.slice(newline + 1);
|
||||
if (!line) continue;
|
||||
try {
|
||||
const message = workerMessageSchema.parse(JSON.parse(line));
|
||||
if ("ready" in message) {
|
||||
this.workerReady = true;
|
||||
this.readyResolve?.();
|
||||
this.readyResolve = undefined;
|
||||
this.readyReject = undefined;
|
||||
continue;
|
||||
}
|
||||
const pending = this.pending.get(message.id);
|
||||
if (!pending) continue;
|
||||
this.finishPending(message.id);
|
||||
if (message.ok) pending.resolve(message.evidence);
|
||||
else pending.reject(new LocalNerUnavailableError());
|
||||
} catch {
|
||||
this.failWorker();
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private finishPending(id: string): void {
|
||||
const pending = this.pending.get(id);
|
||||
if (!pending) return;
|
||||
clearTimeout(pending.timer);
|
||||
pending.signal.removeEventListener("abort", pending.cancel);
|
||||
this.pending.delete(id);
|
||||
}
|
||||
|
||||
private stopWorker(): void {
|
||||
const child = this.child;
|
||||
if (child && child.exitCode === null && child.signalCode === null) child.kill("SIGTERM");
|
||||
}
|
||||
|
||||
private failWorker(): void {
|
||||
const error = new LocalNerUnavailableError();
|
||||
this.readyReject?.(error);
|
||||
this.readyResolve = undefined;
|
||||
this.readyReject = undefined;
|
||||
for (const [id, pending] of this.pending) {
|
||||
this.finishPending(id);
|
||||
pending.reject(error);
|
||||
}
|
||||
this.stopWorker();
|
||||
this.child = undefined;
|
||||
this.ready = undefined;
|
||||
this.workerReady = false;
|
||||
this.stdout = "";
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user