diff --git a/ui/cron/fileServer.ts b/ui/cron/fileServer.ts new file mode 100644 index 0000000..52b0f50 --- /dev/null +++ b/ui/cron/fileServer.ts @@ -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 { + 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 { + 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 { + 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 { + 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)); +} diff --git a/ui/cron/paths.ts b/ui/cron/paths.ts index ef28b97..a2a4eab 100644 --- a/ui/cron/paths.ts +++ b/ui/cron/paths.ts @@ -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'; diff --git a/ui/package.json b/ui/package.json index 86d990a..bea992b 100644 --- a/ui/package.json +++ b/ui/package.json @@ -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", diff --git a/version.py b/version.py index 056f0f4..b5ce99b 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.11.0" +VERSION = "0.11.1"