Major ui speed improvements. Moved file server out of next js app and made it multithreadded. Significantly faster downloads for files, images, and videos.
This commit is contained in:
Jaret Burkett
441
ui/cron/fileServer.ts
Normal file
441
ui/cron/fileServer.ts
Normal file
@@ -0,0 +1,441 @@
|
||||
/**
|
||||
* Front server for the UI: fast file downloads + reverse proxy + Next.js supervisor.
|
||||
*
|
||||
* Owns the single public port. GET/HEAD /api/files/* are served directly with
|
||||
* raw Node streams (full HTTP Range support, no Next.js per-request overhead);
|
||||
* everything else — including the dev-mode HMR WebSocket — is piped unmodified
|
||||
* to a Next.js server this process spawns on an ephemeral, loopback-only port.
|
||||
* The internal port is never fixed (no collisions) and is not reachable from
|
||||
* the network, so the externally visible surface is identical to running
|
||||
* `next start` alone: one port.
|
||||
*
|
||||
* With AI_TOOLKIT_FILE_SERVER_WORKERS > 1 (default: min(4, cores) in start
|
||||
* mode, 1 in dev mode), the front server runs as that many cluster workers
|
||||
* sharing the public socket, so file serving is parallel across CPU cores and
|
||||
* a burst of thumbnail requests never queues behind a multi-GB download.
|
||||
*
|
||||
* This process is optional. The equivalent (slower) route at
|
||||
* src/app/api/files/[...filePath]/route.ts still exists, so deployments that
|
||||
* run the worker and `next start` directly keep working exactly as before.
|
||||
*
|
||||
* Usage:
|
||||
* node dist/cron/fileServer.js start --port 8675
|
||||
* ts-node-dev cron/fileServer.ts dev --port 3000
|
||||
*/
|
||||
import cluster from 'cluster';
|
||||
import { spawn } from 'child_process';
|
||||
import http from 'http';
|
||||
import net from 'net';
|
||||
import fs from 'fs';
|
||||
import os from 'os';
|
||||
import path from 'path';
|
||||
import { pipeline } from 'stream';
|
||||
import prisma from './prisma';
|
||||
import { defaultDatasetsFolder, defaultTrainFolder, defaultDataRoot } from './paths';
|
||||
|
||||
const isDev = process.argv.includes('dev');
|
||||
|
||||
function argValue(name: string, fallback: number): number {
|
||||
const i = process.argv.indexOf(name);
|
||||
if (i !== -1 && process.argv[i + 1]) {
|
||||
const v = parseInt(process.argv[i + 1], 10);
|
||||
if (!isNaN(v)) return v;
|
||||
}
|
||||
return fallback;
|
||||
}
|
||||
|
||||
const PUBLIC_PORT = argValue('--port', isDev ? 3000 : 8675);
|
||||
const UPSTREAM_HOST = '127.0.0.1';
|
||||
|
||||
const numWorkers = (() => {
|
||||
const env = parseInt(process.env.AI_TOOLKIT_FILE_SERVER_WORKERS || '', 10);
|
||||
if (!isNaN(env) && env > 0) return env;
|
||||
return isDev ? 1 : Math.min(4, os.cpus().length);
|
||||
})();
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// File serving (same security model as the Next.js route). Settings can be
|
||||
// changed at runtime from the UI in the Next.js process, so cache briefly
|
||||
// instead of forever.
|
||||
// ---------------------------------------------------------------------------
|
||||
type Roots = { datasets: string; training: string; data: string };
|
||||
let rootsCache: { roots: Roots; ts: number } | null = null;
|
||||
|
||||
async function getRoots(): Promise<Roots> {
|
||||
if (rootsCache && Date.now() - rootsCache.ts < 10_000) {
|
||||
return rootsCache.roots;
|
||||
}
|
||||
const rows = await prisma.settings.findMany({
|
||||
where: { key: { in: ['DATASETS_FOLDER', 'TRAINING_FOLDER', 'DATA_ROOT'] } },
|
||||
});
|
||||
const fromRow = (key: string, fallback: string) => {
|
||||
const row = rows.find(r => r.key === key);
|
||||
return row?.value && row.value !== '' ? row.value : fallback;
|
||||
};
|
||||
const roots: Roots = {
|
||||
datasets: fromRow('DATASETS_FOLDER', defaultDatasetsFolder),
|
||||
training: fromRow('TRAINING_FOLDER', defaultTrainFolder),
|
||||
data: fromRow('DATA_ROOT', defaultDataRoot),
|
||||
};
|
||||
rootsCache = { roots, ts: Date.now() };
|
||||
return roots;
|
||||
}
|
||||
|
||||
const contentTypeMap: { [key: string]: string } = {
|
||||
'.jpg': 'image/jpeg',
|
||||
'.jpeg': 'image/jpeg',
|
||||
'.png': 'image/png',
|
||||
'.gif': 'image/gif',
|
||||
'.webp': 'image/webp',
|
||||
'.svg': 'image/svg+xml',
|
||||
'.bmp': 'image/bmp',
|
||||
'.safetensors': 'application/octet-stream',
|
||||
'.zip': 'application/zip',
|
||||
// Videos
|
||||
'.mp4': 'video/mp4',
|
||||
'.avi': 'video/x-msvideo',
|
||||
'.mov': 'video/quicktime',
|
||||
'.mkv': 'video/x-matroska',
|
||||
'.wmv': 'video/x-ms-wmv',
|
||||
'.m4v': 'video/x-m4v',
|
||||
'.flv': 'video/x-flv',
|
||||
// Audio
|
||||
'.mp3': 'audio/mpeg',
|
||||
'.wav': 'audio/wav',
|
||||
'.flac': 'audio/flac',
|
||||
'.ogg': 'audio/ogg',
|
||||
};
|
||||
|
||||
async function serveFile(req: http.IncomingMessage, res: http.ServerResponse, prefix: string): Promise<void> {
|
||||
try {
|
||||
// /api/img/ serves media inline with ETag revalidation; /api/files/ is a
|
||||
// forced download. Both mirror their Next.js route's exact semantics.
|
||||
const isImg = prefix === '/api/img/';
|
||||
const urlPath = (req.url || '').split('?')[0];
|
||||
const rest = urlPath.slice(prefix.length);
|
||||
const decodedFilePath = rest
|
||||
.split('/')
|
||||
.map(decodeURIComponent)
|
||||
.join('/');
|
||||
|
||||
const resolvedFilePath = path.resolve(decodedFilePath);
|
||||
const roots = await getRoots();
|
||||
const allowedDirs = isImg ? [roots.datasets, roots.training, roots.data] : [roots.datasets, roots.training];
|
||||
const isAllowed = allowedDirs.some(
|
||||
allowedDir => resolvedFilePath === allowedDir || resolvedFilePath.startsWith(allowedDir + path.sep),
|
||||
);
|
||||
if (!isAllowed) {
|
||||
console.warn(`Access denied: ${resolvedFilePath} not in ${allowedDirs.join(', ')}`);
|
||||
res.writeHead(403);
|
||||
res.end('Access denied');
|
||||
return;
|
||||
}
|
||||
|
||||
let stat: fs.Stats;
|
||||
try {
|
||||
stat = await fs.promises.stat(resolvedFilePath);
|
||||
} catch {
|
||||
res.writeHead(404);
|
||||
res.end('File not found');
|
||||
return;
|
||||
}
|
||||
if (!stat.isFile()) {
|
||||
res.writeHead(400);
|
||||
res.end('Not a file');
|
||||
return;
|
||||
}
|
||||
|
||||
const filename = path.basename(resolvedFilePath);
|
||||
const ext = path.extname(resolvedFilePath).toLowerCase();
|
||||
const contentType = contentTypeMap[ext] || 'application/octet-stream';
|
||||
|
||||
const headers: http.OutgoingHttpHeaders = {
|
||||
'Content-Type': contentType,
|
||||
'Accept-Ranges': 'bytes',
|
||||
};
|
||||
if (isImg) {
|
||||
// Weak ETag from inode/size/mtime — cheap and stable enough for
|
||||
// revalidation of thumbnails/videos the UI requests over and over.
|
||||
const etag = `W/"${stat.ino.toString(36)}-${stat.size.toString(36)}-${stat.mtimeMs.toString(36)}"`;
|
||||
headers['Cache-Control'] = 'public, max-age=86400, immutable';
|
||||
headers['ETag'] = etag;
|
||||
if (req.headers['if-none-match'] === etag) {
|
||||
res.writeHead(304, { ETag: etag, 'Cache-Control': headers['Cache-Control'] });
|
||||
res.end();
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
headers['Cache-Control'] = 'public, max-age=86400';
|
||||
headers['Content-Disposition'] = `attachment; filename="${encodeURIComponent(filename)}"`;
|
||||
headers['X-Content-Type-Options'] = 'nosniff';
|
||||
}
|
||||
|
||||
const range = req.headers.range;
|
||||
let start = 0;
|
||||
let end = stat.size - 1;
|
||||
let status = 200;
|
||||
|
||||
if (range) {
|
||||
// Open-ended ranges (`bytes=0-`) serve to EOF — capping them forces
|
||||
// clients into serial re-requests and RTT-bound throughput.
|
||||
const parts = range.replace(/bytes=/, '').split('-');
|
||||
start = parseInt(parts[0], 10);
|
||||
end = parts[1] ? parseInt(parts[1], 10) : stat.size - 1;
|
||||
if (isNaN(start) || isNaN(end) || start > end || start >= stat.size) {
|
||||
res.writeHead(416, { 'Content-Range': `bytes */${stat.size}` });
|
||||
res.end();
|
||||
return;
|
||||
}
|
||||
status = 206;
|
||||
headers['Content-Range'] = `bytes ${start}-${end}/${stat.size}`;
|
||||
}
|
||||
|
||||
headers['Content-Length'] = String(end - start + 1);
|
||||
res.writeHead(status, headers);
|
||||
if (req.method === 'HEAD') {
|
||||
res.end();
|
||||
return;
|
||||
}
|
||||
if (res.destroyed) {
|
||||
// Client disconnected while we were stat-ing; piping into a destroyed
|
||||
// stream throws.
|
||||
return;
|
||||
}
|
||||
|
||||
const fileStream = fs.createReadStream(resolvedFilePath, {
|
||||
start,
|
||||
end,
|
||||
highWaterMark: 4 * 1024 * 1024, // large buffer to minimize per-chunk overhead
|
||||
});
|
||||
pipeline(fileStream, res, err => {
|
||||
// Client disconnects mid-download land here; just make sure the socket
|
||||
// is gone.
|
||||
if (err) res.destroy();
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Error serving file:', error);
|
||||
if (!res.headersSent) res.writeHead(500);
|
||||
res.end();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Passthrough proxy to the Next.js server
|
||||
// ---------------------------------------------------------------------------
|
||||
const upstreamAgent = new http.Agent({ keepAlive: true });
|
||||
|
||||
// Hop-by-hop headers must not be forwarded: Node de-chunks bodies on read and
|
||||
// re-frames them on write, so forwarding e.g. `transfer-encoding` would
|
||||
// corrupt the stream.
|
||||
const HOP_BY_HOP = [
|
||||
'connection',
|
||||
'keep-alive',
|
||||
'proxy-authenticate',
|
||||
'proxy-authorization',
|
||||
'te',
|
||||
'trailer',
|
||||
'transfer-encoding',
|
||||
'upgrade',
|
||||
];
|
||||
|
||||
function stripHopByHop(headers: http.IncomingHttpHeaders): http.IncomingHttpHeaders {
|
||||
const out = { ...headers };
|
||||
for (const h of HOP_BY_HOP) delete out[h];
|
||||
return out;
|
||||
}
|
||||
|
||||
function proxy(req: http.IncomingMessage, res: http.ServerResponse, upstreamPort: number, attempt = 0): void {
|
||||
const bodyless = req.method === 'GET' || req.method === 'HEAD';
|
||||
const upstreamReq = http.request(
|
||||
{
|
||||
host: UPSTREAM_HOST,
|
||||
port: upstreamPort,
|
||||
path: req.url,
|
||||
method: req.method,
|
||||
headers: stripHopByHop(req.headers),
|
||||
agent: upstreamAgent,
|
||||
},
|
||||
upstreamRes => {
|
||||
if (res.destroyed) {
|
||||
// Client disconnected while Next.js was working; drain and drop.
|
||||
upstreamRes.resume();
|
||||
return;
|
||||
}
|
||||
res.writeHead(upstreamRes.statusCode || 502, stripHopByHop(upstreamRes.headers));
|
||||
pipeline(upstreamRes, res, () => { });
|
||||
},
|
||||
);
|
||||
// If the client goes away, stop the upstream work too. No-op after normal
|
||||
// completion.
|
||||
const abortUpstream = () => upstreamReq.destroy();
|
||||
res.once('close', abortUpstream);
|
||||
upstreamReq.on('error', (err: NodeJS.ErrnoException) => {
|
||||
res.removeListener('close', abortUpstream);
|
||||
// Next.js is still booting. Nothing was sent on a refused connection, so
|
||||
// bodyless requests are safe to retry (request bodies can't be replayed).
|
||||
if (err.code === 'ECONNREFUSED' && bodyless && attempt < 120 && !res.destroyed) {
|
||||
setTimeout(() => proxy(req, res, upstreamPort, attempt + 1), 250);
|
||||
return;
|
||||
}
|
||||
if (res.destroyed) return;
|
||||
if (!res.headersSent) res.writeHead(502);
|
||||
res.end('UI server unavailable');
|
||||
});
|
||||
if (bodyless) {
|
||||
upstreamReq.end();
|
||||
} else {
|
||||
pipeline(req, upstreamReq, () => { });
|
||||
}
|
||||
}
|
||||
|
||||
function startServer(publicPort: number, upstreamPort: number): void {
|
||||
const server = http.createServer((req, res) => {
|
||||
const urlPath = (req.url || '').split('?')[0];
|
||||
if (req.method === 'GET' || req.method === 'HEAD') {
|
||||
const prefix = ['/api/files/', '/api/img/'].find(p => urlPath.startsWith(p));
|
||||
if (prefix) {
|
||||
serveFile(req, res, prefix);
|
||||
return;
|
||||
}
|
||||
}
|
||||
proxy(req, res, upstreamPort);
|
||||
});
|
||||
|
||||
// WebSocket upgrades (dev HMR) are forwarded as raw TCP with the original
|
||||
// request head replayed, so the hop-by-hop upgrade headers pass through
|
||||
// intact.
|
||||
server.on('upgrade', (req, socket, head) => {
|
||||
const upstream = net.connect(upstreamPort, UPSTREAM_HOST, () => {
|
||||
if (socket.destroyed) {
|
||||
upstream.destroy();
|
||||
return;
|
||||
}
|
||||
let raw = `${req.method} ${req.url} HTTP/1.1\r\n`;
|
||||
for (let i = 0; i < req.rawHeaders.length; i += 2) {
|
||||
raw += `${req.rawHeaders[i]}: ${req.rawHeaders[i + 1]}\r\n`;
|
||||
}
|
||||
raw += '\r\n';
|
||||
upstream.write(raw);
|
||||
if (head && head.length) upstream.write(head);
|
||||
upstream.pipe(socket);
|
||||
socket.pipe(upstream);
|
||||
});
|
||||
upstream.on('error', () => socket.destroy());
|
||||
socket.on('error', () => upstream.destroy());
|
||||
});
|
||||
|
||||
// Uploads can be huge (100gb body limit upstream); the default 5-minute
|
||||
// request timeout would kill them mid-transfer.
|
||||
server.requestTimeout = 0;
|
||||
|
||||
server.listen(publicPort);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Supervisor: spawn Next.js on an ephemeral loopback port, then serve
|
||||
// ---------------------------------------------------------------------------
|
||||
function getFreePort(): Promise<number> {
|
||||
return new Promise((resolve, reject) => {
|
||||
const probe = net.createServer();
|
||||
probe.on('error', reject);
|
||||
probe.listen(0, UPSTREAM_HOST, () => {
|
||||
const port = (probe.address() as net.AddressInfo).port;
|
||||
probe.close(err => (err ? reject(err) : resolve(port)));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async function primaryMain(): Promise<void> {
|
||||
const upstreamPort = await getFreePort();
|
||||
|
||||
const nextBin = require.resolve('next/dist/bin/next');
|
||||
const nextArgs = isDev ? ['dev', '--turbopack'] : ['start'];
|
||||
nextArgs.push('--port', String(upstreamPort), '--hostname', UPSTREAM_HOST);
|
||||
|
||||
const nextChild = spawn(process.execPath, [nextBin, ...nextArgs], {
|
||||
stdio: ['inherit', 'pipe', 'pipe'],
|
||||
cwd: process.cwd(),
|
||||
});
|
||||
|
||||
// Next.js prints its banner with the internal loopback port ("- Local:
|
||||
// http://127.0.0.1:39865"). Users who follow that link get the
|
||||
// non-accelerated server (or nothing, remotely). Rewrite every mention of
|
||||
// the internal address to the real public port before it reaches the
|
||||
// console.
|
||||
const internalAddr = new RegExp(`(https?://)?(127\\.0\\.0\\.1|localhost):${upstreamPort}`, 'g');
|
||||
const rewritePipe = (from: NodeJS.ReadableStream, to: NodeJS.WritableStream) => {
|
||||
let buffered = '';
|
||||
from.on('data', (chunk: Buffer) => {
|
||||
buffered += chunk.toString();
|
||||
const lines = buffered.split('\n');
|
||||
buffered = lines.pop() || '';
|
||||
for (const line of lines) {
|
||||
to.write(line.replace(internalAddr, `http://localhost:${PUBLIC_PORT}`) + '\n');
|
||||
}
|
||||
});
|
||||
from.on('end', () => {
|
||||
if (buffered) to.write(buffered.replace(internalAddr, `http://localhost:${PUBLIC_PORT}`) + '\n');
|
||||
});
|
||||
};
|
||||
rewritePipe(nextChild.stdout!, process.stdout);
|
||||
rewritePipe(nextChild.stderr!, process.stderr);
|
||||
|
||||
let shuttingDown = false;
|
||||
const shutdown = (code: number) => {
|
||||
if (shuttingDown) return;
|
||||
shuttingDown = true;
|
||||
nextChild.kill('SIGTERM');
|
||||
for (const worker of Object.values(cluster.workers || {})) {
|
||||
worker?.kill();
|
||||
}
|
||||
process.exit(code);
|
||||
};
|
||||
|
||||
nextChild.on('exit', code => {
|
||||
if (!shuttingDown) {
|
||||
console.error(`Next.js exited with code ${code}, shutting down file server`);
|
||||
shutdown(code ?? 1);
|
||||
}
|
||||
});
|
||||
process.on('SIGINT', () => shutdown(0));
|
||||
process.on('SIGTERM', () => shutdown(0));
|
||||
// Never leave an orphaned Next.js child behind if we die unexpectedly.
|
||||
process.on('uncaughtException', err => {
|
||||
console.error('File server uncaught exception:', err);
|
||||
shutdown(1);
|
||||
});
|
||||
|
||||
console.log(`AI Toolkit UI: http://localhost:${PUBLIC_PORT} (${numWorkers} file server worker${numWorkers === 1 ? '' : 's'})`);
|
||||
|
||||
if (numWorkers <= 1) {
|
||||
startServer(PUBLIC_PORT, upstreamPort);
|
||||
return;
|
||||
}
|
||||
|
||||
// AI_TOOLKIT_QUIET_PATHS keeps each forked worker from re-printing the
|
||||
// TOOLKIT_ROOT line from paths.ts.
|
||||
const workerEnv = {
|
||||
UPSTREAM_PORT: String(upstreamPort),
|
||||
PUBLIC_PORT: String(PUBLIC_PORT),
|
||||
AI_TOOLKIT_QUIET_PATHS: '1',
|
||||
};
|
||||
for (let i = 0; i < numWorkers; i++) {
|
||||
cluster.fork(workerEnv);
|
||||
}
|
||||
cluster.on('exit', worker => {
|
||||
if (!shuttingDown) {
|
||||
console.warn(`File server worker ${worker.id} died, restarting`);
|
||||
cluster.fork(workerEnv);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
if (cluster.isPrimary) {
|
||||
primaryMain().catch(err => {
|
||||
console.error('File server failed to start:', err);
|
||||
process.exit(1);
|
||||
});
|
||||
} else {
|
||||
process.on('disconnect', () => process.exit(0));
|
||||
startServer(parseInt(process.env.PUBLIC_PORT!, 10), parseInt(process.env.UPSTREAM_PORT!, 10));
|
||||
}
|
||||
@@ -6,7 +6,11 @@ export const defaultTrainFolder = path.join(TOOLKIT_ROOT, 'output');
|
||||
export const defaultDatasetsFolder = path.join(TOOLKIT_ROOT, 'datasets');
|
||||
export const defaultDataRoot = path.join(TOOLKIT_ROOT, 'data');
|
||||
|
||||
console.log('TOOLKIT_ROOT:', TOOLKIT_ROOT);
|
||||
// Forked file-server workers set AI_TOOLKIT_QUIET_PATHS so this line prints
|
||||
// once per launched process group, not once per worker.
|
||||
if (!process.env.AI_TOOLKIT_QUIET_PATHS) {
|
||||
console.log('TOOLKIT_ROOT:', TOOLKIT_ROOT);
|
||||
}
|
||||
|
||||
export const getTrainingFolder = async () => {
|
||||
const key = 'TRAINING_FOLDER';
|
||||
|
||||
@@ -3,9 +3,9 @@
|
||||
"version": "0.1.0",
|
||||
"private": true,
|
||||
"scripts": {
|
||||
"dev": "concurrently -k -n WORKER,UI \"ts-node-dev --project tsconfig.worker.json --respawn --watch cron --transpile-only cron/worker.ts\" \"next dev --turbopack\"",
|
||||
"dev": "concurrently -k -n WORKER,UI \"ts-node-dev --project tsconfig.worker.json --respawn --watch cron --transpile-only cron/worker.ts\" \"ts-node-dev --project tsconfig.worker.json --respawn --watch cron --transpile-only cron/fileServer.ts dev --port 3000\"",
|
||||
"build": "tsc -p tsconfig.worker.json && next build",
|
||||
"start": "concurrently --restart-tries -1 --restart-after 1000 -n WORKER,UI \"node dist/cron/worker.js\" \"next start --port 8675\"",
|
||||
"start": "concurrently --restart-tries -1 --restart-after 1000 -n WORKER,UI \"node dist/cron/worker.js\" \"node dist/cron/fileServer.js start --port 8675\"",
|
||||
"build_and_start": "npm install && npm run update_db && npm run build && npm run start",
|
||||
"lint": "next lint",
|
||||
"update_db": "npx prisma generate && npx prisma db push",
|
||||
|
||||
@@ -1 +1 @@
|
||||
VERSION = "0.11.0"
|
||||
VERSION = "0.11.1"
|
||||
|
||||
Reference in New Issue
Block a user