From 990c52132b59971084b58822482e4693391698e9 Mon Sep 17 00:00:00 2001 From: ivanmilevtues Date: Wed, 19 Aug 2026 18:55:41 +0200 Subject: [PATCH] [python] Improve bytecode knapsack selection with static analysis --- .changeset/python-bytecode-packing.md | 6 + .../vercel_python_analysis/src/imports.rs | 240 +++++++++ .../crates/vercel_python_analysis/src/lib.rs | 11 +- .../vercel_python_analysis/wit/world.wit | 32 ++ packages/python-analysis/src/index.ts | 17 +- .../src/semantic/import-graph.ts | 360 ++++++++++++++ .../python-analysis/test/import-graph.test.ts | 264 ++++++++++ packages/python/src/bytecode-packing.ts | 105 ++++ packages/python/src/compileall.ts | 121 ++++- packages/python/src/import-closure.ts | 97 ++++ packages/python/src/index.ts | 467 ++++++++++-------- .../python/src/installed-distributions.ts | 74 ++- packages/python/templates/vc_compileall.py | 24 +- packages/python/test/pycache-prefix.test.ts | 40 +- .../python/test/unit.bytecode-fill.test.ts | 297 ----------- .../python/test/unit.bytecode-packing.test.ts | 252 ++++++++++ packages/python/test/unit.compileall.test.ts | 33 +- .../python/test/unit.import-closure.test.ts | 87 ++++ .../test/unit.installed-distributions.test.ts | 14 +- packages/python/test/unit.test.ts | 195 ++++++++ 20 files changed, 2166 insertions(+), 570 deletions(-) create mode 100644 .changeset/python-bytecode-packing.md create mode 100644 packages/python-analysis/crates/vercel_python_analysis/src/imports.rs create mode 100644 packages/python-analysis/src/semantic/import-graph.ts create mode 100644 packages/python-analysis/test/import-graph.test.ts create mode 100644 packages/python/src/bytecode-packing.ts create mode 100644 packages/python/src/import-closure.ts delete mode 100644 packages/python/test/unit.bytecode-fill.test.ts create mode 100644 packages/python/test/unit.bytecode-packing.test.ts create mode 100644 packages/python/test/unit.import-closure.test.ts diff --git a/.changeset/python-bytecode-packing.md b/.changeset/python-bytecode-packing.md new file mode 100644 index 000000000000..8e37090ac37d --- /dev/null +++ b/.changeset/python-bytecode-packing.md @@ -0,0 +1,6 @@ +--- +'@vercel/python-analysis': minor +'@vercel/python': minor +--- + +Rank precompiled bytecode by cold-start value instead of size. On every bytecode fill path (standard, runtime-install knapsack, bytecode-first, large functions), when `.pyc` files overflow the zip's fill capacity, selection is now per file: bytecode for modules the app actually imports at startup first (from a new static AST import closure in `@vercel/python-analysis` — no user code runs at build time), ranked by measured compile cost per byte, then everything else by the same ranking. Compile timings are now measured per file during `compileall`. Falls back to density-only, then per-file size ordering, so no build is worse than the old per-package size knapsack. When all bytecode fits, everything ships with no analysis. The import closure is bounded by a 30s timeout, and setting `VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS=1` disables the closure and timing-based ranking entirely, reverting selection to per-file size ordering. diff --git a/packages/python-analysis/crates/vercel_python_analysis/src/imports.rs b/packages/python-analysis/crates/vercel_python_analysis/src/imports.rs new file mode 100644 index 000000000000..08e654889f2f --- /dev/null +++ b/packages/python-analysis/crates/vercel_python_analysis/src/imports.rs @@ -0,0 +1,240 @@ +//! Import statement extraction for Python source code. +//! +//! Walks the module's statement tree and records every `import` / +//! `from ... import ...` with the syntactic context needed by the caller to +//! decide whether the import runs at module-import time: +//! +//! - Function and method bodies are lazy (recorded with +//! `is_module_level = false`). +//! - `if TYPE_CHECKING:` blocks never run at runtime (`in_type_checking`). +//! - Everything else at module level — including If/Try/With/loop/class +//! bodies — executes at import time and is descended into. + +use ruff_python_ast::{Expr, Stmt}; +use ruff_python_parser::parse_module; + +use crate::bindings::ImportStmt; + +/// Extract all import statements from Python source with their context. +/// Returns an empty vec for invalid syntax. +pub(crate) fn extract_imports_impl(source: &str) -> Vec { + let parsed = match parse_module(source) { + Ok(parsed) => parsed, + Err(_) => return Vec::new(), + }; + + let mut out = Vec::new(); + walk_suite(parsed.suite(), false, false, &mut out); + out +} + +fn walk_suite(suite: &[Stmt], in_function: bool, in_type_checking: bool, out: &mut Vec) { + for stmt in suite { + match stmt { + Stmt::Import(import) => { + for alias in &import.names { + out.push(ImportStmt { + module: Some(alias.name.to_string()), + level: 0, + names: Vec::new(), + is_module_level: !in_function, + in_type_checking, + }); + } + } + Stmt::ImportFrom(import_from) => { + out.push(ImportStmt { + module: import_from.module.as_ref().map(|m| m.to_string()), + level: u8::try_from(import_from.level).unwrap_or(u8::MAX), + names: import_from + .names + .iter() + .map(|alias| alias.name.to_string()) + .collect(), + is_module_level: !in_function, + in_type_checking, + }); + } + // Function bodies are lazy: recorded but flagged not-module-level. + // (Async functions are Stmt::FunctionDef with is_async = true.) + Stmt::FunctionDef(_) => { + walk_suite(stmt_body(stmt), true, in_type_checking, out); + } + // Class bodies execute at import time; methods within are lazy. + Stmt::ClassDef(class_def) => { + for sub in &class_def.body { + match sub { + Stmt::FunctionDef(_) => { + walk_suite(stmt_body(sub), true, in_type_checking, out); + } + _ => walk_suite( + std::slice::from_ref(sub), + in_function, + in_type_checking, + out, + ), + } + } + } + Stmt::If(if_stmt) => { + let branch_tc = in_type_checking || is_type_checking_test(&if_stmt.test); + walk_suite(&if_stmt.body, in_function, branch_tc, out); + for elif in &if_stmt.elif_else_clauses { + let clause_tc = + elif.test.as_ref().map_or(in_type_checking, |test| { + in_type_checking || is_type_checking_test(test) + }); + walk_suite(&elif.body, in_function, clause_tc, out); + } + } + Stmt::Try(try_stmt) => { + walk_suite(&try_stmt.body, in_function, in_type_checking, out); + for handler in &try_stmt.handlers { + let ruff_python_ast::ExceptHandler::ExceptHandler(handler) = handler; + walk_suite(&handler.body, in_function, in_type_checking, out); + } + walk_suite(&try_stmt.orelse, in_function, in_type_checking, out); + walk_suite(&try_stmt.finalbody, in_function, in_type_checking, out); + } + Stmt::With(with_stmt) => { + walk_suite(&with_stmt.body, in_function, in_type_checking, out); + } + Stmt::While(while_stmt) => { + walk_suite(&while_stmt.body, in_function, in_type_checking, out); + walk_suite(&while_stmt.orelse, in_function, in_type_checking, out); + } + Stmt::For(for_stmt) => { + walk_suite(&for_stmt.body, in_function, in_type_checking, out); + walk_suite(&for_stmt.orelse, in_function, in_type_checking, out); + } + Stmt::Match(match_stmt) => { + for case in &match_stmt.cases { + walk_suite(&case.body, in_function, in_type_checking, out); + } + } + _ => {} + } + } +} + +fn stmt_body(stmt: &Stmt) -> &[Stmt] { + match stmt { + Stmt::FunctionDef(def) => &def.body, + _ => &[], + } +} + +/// Matches `if TYPE_CHECKING:` and `if typing.TYPE_CHECKING:`. +fn is_type_checking_test(test: &Expr) -> bool { + match test { + Expr::Name(name) => name.id.as_str() == "TYPE_CHECKING", + Expr::Attribute(attr) => attr.attr.as_str() == "TYPE_CHECKING", + _ => false, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn one(source: &str) -> Vec { + extract_imports_impl(source) + } + + #[test] + fn test_plain_import() { + let stmts = one("import a.b.c\nimport d"); + assert_eq!(stmts.len(), 2); + assert_eq!(stmts[0].module.as_deref(), Some("a.b.c")); + assert_eq!(stmts[0].level, 0); + assert!(stmts[0].names.is_empty()); + assert!(stmts[0].is_module_level); + assert!(!stmts[0].in_type_checking); + assert_eq!(stmts[1].module.as_deref(), Some("d")); + } + + #[test] + fn test_from_import() { + let stmts = one("from a.b import c, d as e"); + assert_eq!(stmts.len(), 1); + assert_eq!(stmts[0].module.as_deref(), Some("a.b")); + assert_eq!(stmts[0].level, 0); + assert_eq!(stmts[0].names, vec!["c".to_string(), "d".to_string()]); + } + + #[test] + fn test_relative_imports() { + let stmts = one("from . import x\nfrom ..pkg import y\nfrom .sub import z"); + assert_eq!(stmts.len(), 3); + assert_eq!(stmts[0].module, None); + assert_eq!(stmts[0].level, 1); + assert_eq!(stmts[0].names, vec!["x".to_string()]); + assert_eq!(stmts[1].module.as_deref(), Some("pkg")); + assert_eq!(stmts[1].level, 2); + assert_eq!(stmts[2].module.as_deref(), Some("sub")); + assert_eq!(stmts[2].level, 1); + } + + #[test] + fn test_function_level_imports_flagged() { + let stmts = one("import top\ndef f():\n import inner\n async def g():\n import deep"); + assert_eq!(stmts.len(), 3); + assert!(stmts[0].is_module_level); + assert!(!stmts[1].is_module_level); + assert!(!stmts[2].is_module_level); + } + + #[test] + fn test_type_checking_block() { + let stmts = one( + "from typing import TYPE_CHECKING\nif TYPE_CHECKING:\n import typeshed_only\nimport runtime", + ); + assert_eq!(stmts.len(), 3); + assert!(!stmts[0].in_type_checking); + assert!(stmts[1].in_type_checking); + assert!(!stmts[2].in_type_checking); + } + + #[test] + fn test_typing_attribute_type_checking() { + let stmts = one("import typing\nif typing.TYPE_CHECKING:\n import t\nelse:\n import r"); + assert_eq!(stmts.len(), 3); + assert!(stmts[1].in_type_checking); + assert!(!stmts[2].in_type_checking); + } + + #[test] + fn test_try_except_both_branches() { + let stmts = one("try:\n import orjson\nexcept ImportError:\n import json\nfinally:\n import cleanup_mod"); + assert_eq!(stmts.len(), 3); + assert!(stmts.iter().all(|s| s.is_module_level)); + } + + #[test] + fn test_class_body_executes_methods_do_not() { + let stmts = one("class A:\n import class_level\n def m(self):\n import method_level"); + assert_eq!(stmts.len(), 2); + assert!(stmts[0].is_module_level); + assert!(!stmts[1].is_module_level); + } + + #[test] + fn test_module_level_if_both_branches() { + let stmts = one("import sys\nif sys.platform == 'win32':\n import win\nelse:\n import posix"); + assert_eq!(stmts.len(), 3); + assert!(stmts.iter().all(|s| s.is_module_level)); + assert!(stmts.iter().all(|s| !s.in_type_checking)); + } + + #[test] + fn test_invalid_syntax() { + assert!(one("def broken(").is_empty()); + } + + #[test] + fn test_match_statement() { + let stmts = one("match x:\n case 1:\n import one_mod"); + assert_eq!(stmts.len(), 1); + assert!(stmts[0].is_module_level); + } +} diff --git a/packages/python-analysis/crates/vercel_python_analysis/src/lib.rs b/packages/python-analysis/crates/vercel_python_analysis/src/lib.rs index 1d82bcfbf203..4e3cacbe2809 100644 --- a/packages/python-analysis/crates/vercel_python_analysis/src/lib.rs +++ b/packages/python-analysis/crates/vercel_python_analysis/src/lib.rs @@ -5,6 +5,7 @@ mod bindings { } mod dist_metadata; mod entrypoint; +mod imports; mod pep508; mod requirements_txt; @@ -21,10 +22,11 @@ use uv_distribution_filename::WheelFilename; use uv_pep508::{MarkerEnvironmentBuilder, MarkerTree}; use uv_platform_tags::{Arch, Os, Platform, Tags}; -use crate::bindings::{DirectUrlInfo, DistMetadata, ParsedReqEntry, ParsedRequirementsTxt, RecordEntry}; +use crate::bindings::{DirectUrlInfo, DistMetadata, ImportStmt, ParsedReqEntry, ParsedRequirementsTxt, RecordEntry}; use crate::entrypoint::{ find_app_or_handler_impl, contains_top_level_callable_impl, get_string_constant_impl, }; +use crate::imports::extract_imports_impl; /// Single-poll executor for WASM: all stub I/O resolves synchronously via host-bridge, /// so the future is guaranteed to be ready on the first poll. @@ -66,6 +68,13 @@ impl crate::bindings::Guest for PythonAnalyzer { get_string_constant_impl(&source, &name) } + /// Extract every import statement in Python source with its syntactic + /// context (module-level / TYPE_CHECKING). Returns an empty list for + /// invalid Python syntax. + fn extract_imports(source: String) -> Vec { + extract_imports_impl(&source) + } + fn parse_dist_metadata(content: Vec) -> Result { dist_metadata::metadata::parse(&content) } diff --git a/packages/python-analysis/crates/vercel_python_analysis/wit/world.wit b/packages/python-analysis/crates/vercel_python_analysis/wit/world.wit index 6495b0fdf2d3..101d6ce2f882 100644 --- a/packages/python-analysis/crates/vercel_python_analysis/wit/world.wit +++ b/packages/python-analysis/crates/vercel_python_analysis/wit/world.wit @@ -55,6 +55,38 @@ world python-analysis { /// Returns false for invalid Python syntax. export contains-top-level-callable: func(source: string, name: string) -> bool; + /// A single import statement extracted from Python source, with the + /// syntactic context needed to decide whether it runs at import time. + /// + /// Policy (which contexts count as "runs at import") is deliberately left + /// to the caller so heuristics can evolve without rebuilding the WASM. + record import-stmt { + /// Dotted module path for from-imports (`from a.b import c` -> "a.b"). + /// None for `from . import x` (module resolved via level only) and + /// always set for plain imports (`import a.b` -> "a.b"). + module: option, + /// Relative-import level (number of leading dots). 0 = absolute. + level: u8, + /// Imported names for from-imports (empty for plain imports). + /// Each name may itself be a submodule of `module`. + names: list, + /// True when the statement is not nested inside any function body + /// (methods included). Function-level imports are lazy: they run on + /// first call, not at module import. + is-module-level: bool, + /// True when nested under `if TYPE_CHECKING:` (or + /// `typing.TYPE_CHECKING`). Such imports never run at runtime. + in-type-checking: bool, + } + + /// Extract every import statement in Python source with its syntactic + /// context. Descends through module-level If/Try/With/loop/class bodies + /// (which execute at import time) but records function-body imports with + /// is-module-level = false. + /// + /// Returns an empty list for invalid Python syntax. + export extract-imports: func(source: string) -> list; + // ========================================================================= // Installed package distribution metadata parsing // ========================================================================= diff --git a/packages/python-analysis/src/index.ts b/packages/python-analysis/src/index.ts index 4c4346f67516..206e55b253d5 100644 --- a/packages/python-analysis/src/index.ts +++ b/packages/python-analysis/src/index.ts @@ -17,6 +17,18 @@ export { getStringConstant, } from './semantic/entrypoints'; +// ============================================================================= +// Static import-graph analysis (WASM-based extraction + module resolution) +// ============================================================================= + +export { collectImportClosure, extractImports } from './semantic/import-graph'; + +export type { + ImportClosureOptions, + ImportClosureResult, + ImportStmt, +} from './semantic/import-graph'; + // ============================================================================= // Installed package analysis (WASM-based .dist-info parsing) // ============================================================================= @@ -86,10 +98,7 @@ export { // Wheel compatibility checking // ============================================================================= -export { - evaluateMarker, - isWheelCompatible, -} from './manifest/wheel-compat'; +export { evaluateMarker, isWheelCompatible } from './manifest/wheel-compat'; // ============================================================================= // Python selection (runtime + types) diff --git a/packages/python-analysis/src/semantic/import-graph.ts b/packages/python-analysis/src/semantic/import-graph.ts new file mode 100644 index 000000000000..b739314f1fee --- /dev/null +++ b/packages/python-analysis/src/semantic/import-graph.ts @@ -0,0 +1,360 @@ +/** + * Static import-graph analysis: compute the transitive closure of modules a + * Python application imports at startup, without executing any user code. + * + * The closure feeds bytecode packing. Both error directions are safe: + * over-included modules waste capacity, and missed modules (plugin loaders, + * imports made by compiled extensions) still ship from residual capacity. + * + * Resolution semantics mirror CPython: + * - dotted name -> `a/b/__init__.py` | `a/b.py` | PEP 420 namespace dir + * - importing `a.b.c` executes every parent `__init__.py` (emitted + recursed) + * - `from a import b` probes `b` as a submodule of `a` + * - relative imports resolve against the importer's package by dot level + * - first search root wins, so app modules shadow vendor (runtime sys.path) + * + * Startup heuristics (applied on top of the raw WASM extraction): + * - imports inside function bodies are lazy and excluded + * - `if TYPE_CHECKING:` blocks never run at runtime and are excluded + * - both branches of module-level `if` / `try` are included + */ + +import fs from 'fs'; +import { isAbsolute, join, dirname, resolve, sep } from 'path'; +import { importWasmModule } from '../wasm/load'; + +export interface ImportStmt { + /** Dotted module path for from-imports; undefined for `from . import x`. */ + module?: string; + /** Relative-import level (number of leading dots). 0 = absolute. */ + level: number; + /** Imported names for from-imports (each may be a submodule). */ + names: string[]; + /** False when nested in a function body (lazy import). */ + isModuleLevel: boolean; + /** True when under `if TYPE_CHECKING:` (never runs at runtime). */ + inTypeChecking: boolean; +} + +/** + * Extract every import statement in Python source with its syntactic + * context. Returns an empty list for invalid syntax. + */ +export async function extractImports(source: string): Promise { + // Cheap pre-filter: the WASM parse is wasted on files without imports. + if (!source.includes('import')) { + return []; + } + const mod = await importWasmModule(); + return mod.extractImports(source); +} + +export interface ImportClosureOptions { + /** + * Entry points of the closure: absolute `.py` file paths and/or dotted + * module or object names (e.g. Django's ROOT_URLCONF / INSTALLED_APPS + * strings), resolved to their longest importable prefix against searchRoots. + */ + seeds: string[]; + /** + * Ordered module search roots (e.g. [workPath, ...sitePackageDirs]). + * First match wins, so earlier roots shadow later ones. + */ + searchRoots: string[]; + /** + * Approximate safety bound on parsed files, checked per frontier batch. + * On overflow the partial closure is returned with `truncated: true` + * (remaining modules fall to the residual bytecode tier). + */ + maxFiles?: number; +} + +export interface ImportClosureResult { + /** Absolute paths of every `.py` file imported at startup. */ + files: Set; + truncated: boolean; +} + +const DEFAULT_MAX_FILES = 30000; + +type FsKind = 'file' | 'dir' | null; + +interface ResolutionCache { + statKinds: Map>; + names: Map>; +} + +async function readStatKind(path: string): Promise { + try { + const stats = await fs.promises.stat(path); + if (stats.isFile()) return 'file'; + if (stats.isDirectory()) return 'dir'; + } catch { + // unreadable / missing + } + return null; +} + +function statKind(path: string, cache: ResolutionCache): Promise { + const cached = cache.statKinds.get(path); + if (cached) return cached; + + const pending = readStatKind(path); + cache.statKinds.set(path, pending); + return pending; +} + +interface ResolvedName { + /** Resolved module file, or a directory for a namespace package. */ + target: string; + targetKind: FsKind; + /** `__init__.py` of every parent package (executed on import). */ + parentChain: string[]; +} + +/** + * Resolve dotted parts under each search dir. Importing `a.b.c` executes the + * `__init__.py` of `a` and `a.b`, so the parent chain is emitted alongside. + */ +async function resolveNamePartsUncached( + parts: string[], + searchDirs: string[], + cache: ResolutionCache +): Promise { + const rel = join(...parts); + for (const root of searchDirs) { + let target: string | null = null; + let targetKind: FsKind = null; + const candidateInit = join(root, rel, '__init__.py'); + if ((await statKind(candidateInit, cache)) === 'file') { + target = candidateInit; + targetKind = 'file'; + } else { + const candidateModule = join(root, `${rel}.py`); + if ((await statKind(candidateModule, cache)) === 'file') { + target = candidateModule; + targetKind = 'file'; + } else if ((await statKind(join(root, rel), cache)) === 'dir') { + target = join(root, rel); + targetKind = 'dir'; // PEP 420 namespace package + } + } + if (target !== null) { + const parentChain: string[] = []; + let acc = root; + for (const part of parts.slice(0, -1)) { + acc = join(acc, part); + const init = join(acc, '__init__.py'); + if ((await statKind(init, cache)) === 'file') { + parentChain.push(init); + } + } + return { target, targetKind, parentChain }; + } + } + return null; +} + +/** + * `searchKey` must uniquely identify `searchDirs` (precomputed once per + * dir set — building keys per call dominated resolution on large apps). + */ +function resolveNameParts( + parts: string[], + searchDirs: string[], + searchKey: string, + cache: ResolutionCache +): Promise { + if (parts.length === 0) return Promise.resolve(null); + + const key = `${searchKey}\0${parts.join('.')}`; + const cached = cache.names.get(key); + if (cached) return cached; + + const pending = resolveNamePartsUncached(parts, searchDirs, cache); + cache.names.set(key, pending); + return pending; +} + +/** Resolve a seed that may end in an object name rather than a module. */ +async function resolveSeedName( + parts: string[], + searchDirs: string[], + searchKey: string, + cache: ResolutionCache +): Promise { + for (let length = parts.length; length > 0; length--) { + const resolved = await resolveNameParts( + parts.slice(0, length), + searchDirs, + searchKey, + cache + ); + if (resolved) return resolved; + } + return null; +} + +/** Directory to search for submodules of the module at `target`. */ +function submoduleDir(target: string, targetKind: FsKind): string | null { + if (targetKind === 'dir') return target; + if (targetKind === 'file' && target.endsWith(`${sep}__init__.py`)) { + return dirname(target); + } + return null; +} + +/** + * Resolve one import statement to the files it executes (module itself plus + * parent packages). Unresolvable names (stdlib, missing) are skipped — they + * are never part of the bundle. + */ +async function resolveImport( + stmt: ImportStmt, + importerPath: string, + roots: string[], + rootsKey: string, + cache: ResolutionCache +): Promise { + const resolved = new Set(); + + let search = roots; + let searchKey = rootsKey; + if (stmt.level > 0) { + // level=1 -> the importer's own package directory. + let base = dirname(importerPath); + for (let i = 1; i < stmt.level; i++) { + base = dirname(base); + } + search = [base]; + searchKey = base; + } + + const parts = stmt.module ? stmt.module.split('.') : []; + let target: string; + let targetKind: FsKind; + if (parts.length > 0) { + const name = await resolveNameParts(parts, search, searchKey, cache); + if (!name) return []; + target = name.target; + targetKind = name.targetKind; + if (targetKind === 'file') resolved.add(target); + for (const init of name.parentChain) { + resolved.add(init); + } + } else if (stmt.level > 0) { + // `from . import x` -> names resolved in the importer's package dir. + target = search[0]; + targetKind = 'dir'; + } else { + return []; + } + + // For from-imports, each imported name may itself be a submodule. + const subDir = submoduleDir(target, targetKind); + if (subDir) { + for (const name of stmt.names) { + if (name === '*') continue; + const sub = await resolveNameParts([name], [subDir], subDir, cache); + if (sub) { + if (sub.targetKind === 'file') resolved.add(sub.target); + for (const init of sub.parentChain) { + resolved.add(init); + } + } + } + } + + return [...resolved]; +} + +/** True when a seed looks like a dotted module name rather than a path. */ +function isModuleNameSeed(seed: string): boolean { + return !isAbsolute(seed) && /^[A-Za-z_][\w]*(\.[A-Za-z_][\w]*)*$/.test(seed); +} + +/** + * Compute the transitive closure of modules imported at startup. + * + * Only module-level imports outside `if TYPE_CHECKING:` blocks are followed + * (see module docstring). Files are read and parsed frontier-by-frontier in + * parallel; cycles are cut by the visited set. + */ +export async function collectImportClosure({ + seeds, + searchRoots, + maxFiles = DEFAULT_MAX_FILES, +}: ImportClosureOptions): Promise { + const roots = searchRoots.map(r => resolve(r)); + const rootsKey = roots.join('\0'); + const cache: ResolutionCache = { + statKinds: new Map(), + names: new Map(), + }; + const visited = new Set(); + let frontier: string[] = []; + let truncated = false; + + for (const seed of seeds) { + if (isModuleNameSeed(seed)) { + const name = await resolveSeedName( + seed.split('.'), + roots, + rootsKey, + cache + ); + if (name) { + if (name.targetKind === 'file') frontier.push(name.target); + frontier.push(...name.parentChain); + } + } else { + const seedPath = resolve(seed); + if ((await statKind(seedPath, cache)) === 'file') { + frontier.push(seedPath); + } + } + } + + while (frontier.length > 0) { + const batch = [...new Set(frontier)].filter(p => !visited.has(p)); + frontier = []; + if (batch.length === 0) continue; + + for (const path of batch) { + visited.add(path); + } + // Checked per batch: visited may overshoot maxFiles by one frontier, + // and the final batch joins the closure unparsed. Both are safe — + // truncation only shrinks the ranking input. + if (visited.size > maxFiles) { + truncated = true; + break; + } + + const dependencies = await Promise.all( + batch.map(async importerPath => { + let source: string; + try { + source = await fs.promises.readFile(importerPath, 'utf8'); + } catch { + return []; + } + const stmts = await extractImports(source); + const files = await Promise.all( + stmts + .filter(s => s.isModuleLevel && !s.inTypeChecking) + .map(s => resolveImport(s, importerPath, roots, rootsKey, cache)) + ); + return files.flat(); + }) + ); + + for (const dep of dependencies.flat()) { + if (!visited.has(dep)) { + frontier.push(dep); + } + } + } + + return { files: visited, truncated }; +} diff --git a/packages/python-analysis/test/import-graph.test.ts b/packages/python-analysis/test/import-graph.test.ts new file mode 100644 index 000000000000..5eadb468a508 --- /dev/null +++ b/packages/python-analysis/test/import-graph.test.ts @@ -0,0 +1,264 @@ +import { describe, it, expect, beforeAll, afterAll, vi } from 'vitest'; +import fs from 'fs'; +import os from 'os'; +import path from 'path'; +import { collectImportClosure, extractImports } from '../src'; + +/** + * Fixture tree built per test run in a temp dir. Layout: + * + * app/ <- search root 0 (shadows vendor) + * main.py entrypoint + * pkg/__init__.py `from . import core`, lazy import in function + * pkg/core.py `from vendor_pkg import thing` + * pkg/util.py only reachable via TYPE_CHECKING + * requests.py shadows the vendored `requests` package + * venv/site-packages/ <- search root 1 + * vendor_pkg/__init__.py namespace-relative imports, sub-package chain + * vendor_pkg/deep/mod.py reached via `import vendor_pkg.deep.mod` + * vendor_pkg/lazy_mod.py only imported inside a function -> excluded + * requests/__init__.py shadowed by app/requests.py + * unseen/__init__.py never imported anywhere + */ + +let tmpDir: string; +let appRoot: string; +let sitePackages: string; + +function write(rel: string, content: string) { + const abs = path.join(tmpDir, rel); + fs.mkdirSync(path.dirname(abs), { recursive: true }); + fs.writeFileSync(abs, content); +} + +beforeAll(() => { + tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'import-graph-test-')); + appRoot = path.join(tmpDir, 'app'); + sitePackages = path.join(tmpDir, 'venv', 'site-packages'); + + write('app/main.py', 'from pkg import core\nimport requests\n'); + write( + 'app/pkg/__init__.py', + 'from . import core\ndef f():\n import vendor_pkg.lazy_mod\n' + ); + write( + 'app/pkg/core.py', + 'from vendor_pkg import thing\nfrom typing import TYPE_CHECKING\n' + + 'if TYPE_CHECKING:\n from . import util\n' + ); + write('app/pkg/util.py', '# only reachable via TYPE_CHECKING\n'); + write('app/requests.py', '# app-local module shadows the vendored package\n'); + + write( + 'venv/site-packages/vendor_pkg/__init__.py', + 'from . import thing\nimport vendor_pkg.deep.mod\n' + ); + write( + 'venv/site-packages/vendor_pkg/thing.py', + 'try:\n import ujson\nexcept ImportError:\n import json\n' + ); + write('venv/site-packages/vendor_pkg/deep/__init__.py', ''); + write('venv/site-packages/vendor_pkg/deep/mod.py', ''); + write('venv/site-packages/vendor_pkg/lazy_mod.py', ''); + write('venv/site-packages/requests/__init__.py', '# vendored, shadowed\n'); + write('venv/site-packages/unseen/__init__.py', ''); +}); + +afterAll(() => { + fs.rmSync(tmpDir, { recursive: true, force: true }); +}); + +function rel(paths: Set): string[] { + return [...paths] + .map(p => path.relative(tmpDir, p).split(path.sep).join('/')) + .sort(); +} + +describe('extractImports', () => { + it('flags module-level vs lazy and TYPE_CHECKING imports', async () => { + const stmts = await extractImports( + 'import a\ndef f():\n import b\nif TYPE_CHECKING:\n import c\n' + ); + expect(stmts).toHaveLength(3); + expect(stmts[0]).toMatchObject({ module: 'a', isModuleLevel: true }); + expect(stmts[1]).toMatchObject({ module: 'b', isModuleLevel: false }); + expect(stmts[2]).toMatchObject({ module: 'c', inTypeChecking: true }); + }); + + it('captures relative level and from-import names', async () => { + const stmts = await extractImports('from ..pkg import x, y\n'); + expect(stmts).toHaveLength(1); + expect(stmts[0]).toMatchObject({ + module: 'pkg', + level: 2, + names: ['x', 'y'], + }); + }); + + it('returns an empty list for invalid syntax', async () => { + expect(await extractImports('def broken(')).toEqual([]); + }); +}); + +describe('collectImportClosure', () => { + it('computes the transitive closure with shadowing and parent chains', async () => { + const { files, truncated } = await collectImportClosure({ + seeds: [path.join(appRoot, 'main.py')], + searchRoots: [appRoot, sitePackages], + }); + + expect(truncated).toBe(false); + expect(rel(files)).toEqual([ + 'app/main.py', + 'app/pkg/__init__.py', + 'app/pkg/core.py', + // app module shadows the vendored `requests` package + 'app/requests.py', + 'venv/site-packages/vendor_pkg/__init__.py', + 'venv/site-packages/vendor_pkg/deep/__init__.py', + 'venv/site-packages/vendor_pkg/deep/mod.py', + 'venv/site-packages/vendor_pkg/thing.py', + ]); + }); + + it('resolves dotted-name seeds against the search roots', async () => { + const { files } = await collectImportClosure({ + seeds: ['vendor_pkg.deep.mod'], + searchRoots: [appRoot, sitePackages], + }); + + // Parent packages execute on import: both __init__ files are included. + expect(rel(files)).toEqual([ + 'venv/site-packages/vendor_pkg/__init__.py', + 'venv/site-packages/vendor_pkg/deep/__init__.py', + 'venv/site-packages/vendor_pkg/deep/mod.py', + 'venv/site-packages/vendor_pkg/thing.py', + ]); + }); + + it('resolves the longest importable prefix of object-path seeds', async () => { + write('app/myapp/__init__.py', ''); + write('app/myapp/apps.py', ''); + + const { files } = await collectImportClosure({ + seeds: ['myapp.apps.MyAppConfig'], + searchRoots: [appRoot, sitePackages], + }); + + expect(rel(files)).toEqual(['app/myapp/__init__.py', 'app/myapp/apps.py']); + }); + + it('prefers a package over a same-named module file', async () => { + write('app/collision.py', 'import module_winner\n'); + write('app/collision/__init__.py', 'import package_winner\n'); + write('app/module_winner.py', ''); + write('app/package_winner.py', ''); + + const { files } = await collectImportClosure({ + seeds: ['collision'], + searchRoots: [appRoot], + }); + + expect(rel(files)).toEqual([ + 'app/collision/__init__.py', + 'app/package_winner.py', + ]); + }); + + it('includes both branches of module-level try/except', async () => { + const { files } = await collectImportClosure({ + seeds: ['vendor_pkg.thing'], + searchRoots: [sitePackages], + }); + // ujson is missing (skipped silently); json is stdlib (skipped silently). + expect(rel(files)).toEqual([ + 'venv/site-packages/vendor_pkg/__init__.py', + 'venv/site-packages/vendor_pkg/deep/__init__.py', + 'venv/site-packages/vendor_pkg/deep/mod.py', + 'venv/site-packages/vendor_pkg/thing.py', + ]); + }); + + it('marks the result truncated when maxFiles is exceeded', async () => { + const { truncated } = await collectImportClosure({ + seeds: [path.join(appRoot, 'main.py')], + searchRoots: [appRoot, sitePackages], + maxFiles: 2, + }); + expect(truncated).toBe(true); + }); + + it('survives cycles', async () => { + write('app/cycle_a.py', 'import cycle_b\n'); + write('app/cycle_b.py', 'import cycle_a\n'); + const { files } = await collectImportClosure({ + seeds: [path.join(appRoot, 'cycle_a.py')], + searchRoots: [appRoot, sitePackages], + }); + expect(rel(files)).toEqual(['app/cycle_a.py', 'app/cycle_b.py']); + }); + + it('shares filesystem probes for repeated imports and misses', async () => { + write('app/cache_seed.py', 'import cache_first\nimport cache_second\n'); + write('app/cache_first.py', 'import cache_shared\nimport cache_missing\n'); + write('app/cache_second.py', 'import cache_shared\nimport cache_missing\n'); + write('app/cache_shared.py', ''); + + const statSpy = vi.spyOn(fs.promises, 'stat'); + try { + const { files } = await collectImportClosure({ + seeds: [path.join(appRoot, 'cache_seed.py')], + searchRoots: [appRoot, sitePackages], + }); + + expect(rel(files)).toEqual([ + 'app/cache_first.py', + 'app/cache_second.py', + 'app/cache_seed.py', + 'app/cache_shared.py', + ]); + + const statCounts = new Map(); + for (const [candidate] of statSpy.mock.calls) { + const candidatePath = String(candidate); + statCounts.set(candidatePath, (statCounts.get(candidatePath) ?? 0) + 1); + } + + const sharedCandidates = [ + path.join(appRoot, 'cache_shared', '__init__.py'), + path.join(appRoot, 'cache_shared.py'), + ]; + const missingCandidates = [appRoot, sitePackages].flatMap(root => [ + path.join(root, 'cache_missing', '__init__.py'), + path.join(root, 'cache_missing.py'), + path.join(root, 'cache_missing'), + ]); + + for (const candidate of [...sharedCandidates, ...missingCandidates]) { + expect(statCounts.get(candidate)).toBe(1); + } + } finally { + statSpy.mockRestore(); + } + }); + + it('does not retain resolution misses between closure runs', async () => { + write('app/cache_scope_seed.py', 'import cache_scope_late\n'); + + const first = await collectImportClosure({ + seeds: [path.join(appRoot, 'cache_scope_seed.py')], + searchRoots: [appRoot], + }); + expect(rel(first.files)).toEqual(['app/cache_scope_seed.py']); + + write('app/cache_scope_late.py', ''); + const second = await collectImportClosure({ + seeds: [path.join(appRoot, 'cache_scope_seed.py')], + searchRoots: [appRoot], + }); + expect(rel(second.files)).toEqual([ + 'app/cache_scope_late.py', + 'app/cache_scope_seed.py', + ]); + }); +}); diff --git a/packages/python/src/bytecode-packing.ts b/packages/python/src/bytecode-packing.ts new file mode 100644 index 000000000000..1ee5c23a415d --- /dev/null +++ b/packages/python/src/bytecode-packing.ts @@ -0,0 +1,105 @@ +/** + * Value-ranked bytecode packing: select `.pyc` files to maximise compile + * time avoided at cold start. + * + * Modules imported at startup (static import closure) rank first, then + * everything else, both ordered by compile seconds per byte. + * + * Fallbacks: no closure -> compile density only; no timings -> size desc. + * Every fallback is at least as good as the old per-package size knapsack. + */ + +import type { Files } from '@vercel/build-utils'; +import type { BytecodeItem } from './compileall'; + +/** + * Kill switch for the import closure and timing-based ranking, set to + * `1`/`true`. Selection falls back to per-file size ordering; bytecode + * still ships. + */ +export function isBytecodeAnalysisDisabled(): boolean { + const val = process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS; + if (val === undefined || val === '') return false; + + const lower = val.toLowerCase(); + return lower === '1' || lower === 'true'; +} + +export interface RankedBytecodeItem extends BytecodeItem { + imported: boolean; + /** Compile seconds for the source file; undefined when timings are missing. */ + compileSeconds?: number; +} + +/** + * Attach import-closure membership and compile timings to collected items. + */ +export function annotateBytecodeItems( + items: BytecodeItem[], + importedModules: ReadonlySet | undefined, + timings: ReadonlyMap | undefined +): RankedBytecodeItem[] { + return items.map(item => ({ + ...item, + imported: importedModules?.has(item.moduleKey) ?? false, + compileSeconds: timings?.get(item.sourceAbsPath), + })); +} + +/** + * Order items best-first for greedy filling. + * + * Sort keys, in order: + * - imported first (only when the closure ran — otherwise all items are + * unimported and the key is inert) + * - compile density (seconds per byte) desc; items without a timing take + * the median density so they rank by size among equals + * - size desc + */ +export function rankBytecodeItems( + items: RankedBytecodeItem[] +): RankedBytecodeItem[] { + const densities = items + .filter(i => i.compileSeconds != null && i.size > 0) + .map(i => (i.compileSeconds as number) / i.size) + .sort((a, b) => a - b); + const medianDensity = + densities.length > 0 ? densities[Math.floor(densities.length / 2)] : null; + + return [...items].sort((a, b) => { + if (a.imported !== b.imported) return a.imported ? -1 : 1; + if (medianDensity !== null) { + const da = + a.compileSeconds != null && a.size > 0 + ? a.compileSeconds / a.size + : medianDensity; + const db = + b.compileSeconds != null && b.size > 0 + ? b.compileSeconds / b.size + : medianDensity; + if (da !== db) return db - da; + } + return b.size - a.size; + }); +} + +/** + * Greedy fill: add ranked items to `files` while they fit the capacity, + * skipping items that don't (a small file may still fit after a large one + * is skipped). Returns the remaining capacity. + */ +export function fillBytecodeWithinCapacity( + files: Files, + rankedItems: RankedBytecodeItem[], + capacity: number +): number { + let remaining = capacity; + for (const item of rankedItems) { + if (remaining <= 0) break; + if (item.size <= remaining) { + files[item.bundlePath] = item.file; + remaining -= item.size; + } + } + return remaining; +} diff --git a/packages/python/src/compileall.ts b/packages/python/src/compileall.ts index ecd1b4ad28ad..2423fbc75e5f 100644 --- a/packages/python/src/compileall.ts +++ b/packages/python/src/compileall.ts @@ -57,6 +57,16 @@ export function shouldCompileAll({ return isCompileAllFlagEnabled(); } +export interface CompileAllResult { + success: boolean; + /** + * Per-file compile seconds keyed by the exact source paths passed in. + * Undefined when the coordinator did not emit a table (old template, + * mocked subprocess, write failure) — callers fall back to size ranking. + */ + timings?: Map; +} + /** * Runs the Python compile coordinator to precompile `.py` files into `.pyc` * bytecode. @@ -77,10 +87,10 @@ export async function runCompileAll({ sourceFiles: string[]; env?: NodeJS.ProcessEnv; pycachePrefix?: string; -}): Promise { +}): Promise { const uniqueSourceFiles = [...new Set(sourceFiles)]; if (uniqueSourceFiles.length === 0) { - return false; + return { success: false }; } let tempDir: string | undefined; @@ -91,6 +101,7 @@ export async function runCompileAll({ ); const listPath = join(tempDir, 'pysources.json'); await fs.promises.writeFile(listPath, JSON.stringify(uniqueSourceFiles)); + const timingsPath = join(tempDir, 'timings.json'); const scriptPath = join(__dirname, '..', 'templates', 'vc_compileall.py'); const baseEnv = env || process.env; @@ -98,14 +109,22 @@ export async function runCompileAll({ ? { ...baseEnv, PYTHONPYCACHEPREFIX: pycachePrefix } : baseEnv; - await execa(pythonBin, [scriptPath, listPath], { + await execa(pythonBin, [scriptPath, listPath, timingsPath], { env: subprocessEnv, timeout: COMPILEALL_TIMEOUT_MS, }); - return true; + + let timings: Map | undefined; + try { + const raw = await fs.promises.readFile(timingsPath, 'utf8'); + timings = new Map(Object.entries(JSON.parse(raw))); + } catch (err) { + debug(`compileall timings unavailable: ${String(err)}`); + } + return { success: true, timings }; } catch (err) { debug(`compileall error details: ${JSON.stringify(err)}`); - return false; + return { success: false }; } finally { if (tempDir) { try { @@ -210,6 +229,24 @@ export function derivePrefixPycBundlePath( return `${PYCACHE_PREFIX_DIR}/${rel}`; } +/** + * A single `.pyc` candidate for bytecode packing, at per-file granularity. + */ +export interface BytecodeItem { + /** Bundle-relative path of the `.pyc` file. */ + bundlePath: string; + file: FileFsRef; + /** Uncompressed `.pyc` size in bytes (the knapsack weight). */ + size: number; + /** + * Module identity for import-closure membership: workPath-relative for + * app files, site-packages-relative for vendor files (forward slashes). + */ + moduleKey: string; + /** Absolute fs path of the `.py` source (join key for compile timings). */ + sourceAbsPath: string; +} + export interface BytecodeCollectionResult { /** FileFsRef entries for .pyc files, keyed by bundle-relative path. */ files: Files; @@ -217,6 +254,8 @@ export interface BytecodeCollectionResult { totalSize: number; /** Per-item bytecode sizes for knapsack packing (keyed by package name or bundle path). */ perItemSizes: Map; + /** Per-file candidates for import-aware, value-ranked packing. */ + items: BytecodeItem[]; } /** @@ -239,14 +278,20 @@ export async function collectAppPrefixBytecodeFiles({ pythonMajor: number; pythonMinor: number; }): Promise { - const pending: { bundlePath: string; srcFsPath: string }[] = []; + const pending: { + bundlePath: string; + srcFsPath: string; + moduleKey: string; + sourceAbsPath: string; + }[] = []; for (const bundlePath of Object.keys(appFiles)) { if (!bundlePath.endsWith('.py')) continue; + const sourceAbsPath = join(workPath, bundlePath.replaceAll('/', sep)); const stagedFsPath = deriveStagedPycFsPath( stagingDir, - join(workPath, bundlePath.replaceAll('/', sep)), + sourceAbsPath, pythonMajor, pythonMinor ); @@ -257,14 +302,25 @@ export async function collectAppPrefixBytecodeFiles({ ); if (!stagedFsPath || !pycBundlePath) continue; - pending.push({ bundlePath: pycBundlePath, srcFsPath: stagedFsPath }); + pending.push({ + bundlePath: pycBundlePath, + srcFsPath: stagedFsPath, + moduleKey: bundlePath, + sourceAbsPath, + }); } const results = await Promise.all( - pending.map(async ({ bundlePath, srcFsPath }) => { + pending.map(async ({ bundlePath, srcFsPath, moduleKey, sourceAbsPath }) => { try { const stats = await fs.promises.stat(srcFsPath); - return { bundlePath, srcFsPath, size: stats.size }; + return { + bundlePath, + srcFsPath, + moduleKey, + sourceAbsPath, + size: stats.size, + }; } catch { return null; } @@ -273,19 +329,28 @@ export async function collectAppPrefixBytecodeFiles({ const files: Files = {}; const perItemSizes = new Map(); + const items: BytecodeItem[] = []; let totalSize = 0; for (const result of results) { if (!result) continue; - files[result.bundlePath] = new FileFsRef({ + const file = new FileFsRef({ fsPath: result.srcFsPath, size: result.size, }); + files[result.bundlePath] = file; perItemSizes.set(result.bundlePath, result.size); + items.push({ + bundlePath: result.bundlePath, + file, + size: result.size, + moduleKey: result.moduleKey, + sourceAbsPath: result.sourceAbsPath, + }); totalSize += result.size; } - return { files, totalSize, perItemSizes }; + return { files, totalSize, perItemSizes, items }; } export async function collectAppBytecodeFiles({ @@ -299,7 +364,12 @@ export async function collectAppBytecodeFiles({ pythonMajor: number; pythonMinor: number; }): Promise { - const pending: { bundlePath: string; srcFsPath: string }[] = []; + const pending: { + bundlePath: string; + srcFsPath: string; + moduleKey: string; + sourceAbsPath: string; + }[] = []; for (const bundlePath of Object.keys(appFiles)) { const pycRel = derivePycPath(bundlePath, pythonMajor, pythonMinor); @@ -308,14 +378,22 @@ export async function collectAppBytecodeFiles({ pending.push({ bundlePath: pycRel, srcFsPath: join(workPath, pycRel.replaceAll('/', sep)), + moduleKey: bundlePath, + sourceAbsPath: join(workPath, bundlePath.replaceAll('/', sep)), }); } const results = await Promise.all( - pending.map(async ({ bundlePath, srcFsPath }) => { + pending.map(async ({ bundlePath, srcFsPath, moduleKey, sourceAbsPath }) => { try { const stats = await fs.promises.stat(srcFsPath); - return { bundlePath, srcFsPath, size: stats.size }; + return { + bundlePath, + srcFsPath, + moduleKey, + sourceAbsPath, + size: stats.size, + }; } catch { return null; } @@ -324,17 +402,26 @@ export async function collectAppBytecodeFiles({ const files: Files = {}; const perItemSizes = new Map(); + const items: BytecodeItem[] = []; let totalSize = 0; for (const result of results) { if (!result) continue; - files[result.bundlePath] = new FileFsRef({ + const file = new FileFsRef({ fsPath: result.srcFsPath, size: result.size, }); + files[result.bundlePath] = file; perItemSizes.set(result.bundlePath, result.size); + items.push({ + bundlePath: result.bundlePath, + file, + size: result.size, + moduleKey: result.moduleKey, + sourceAbsPath: result.sourceAbsPath, + }); totalSize += result.size; } - return { files, totalSize, perItemSizes }; + return { files, totalSize, perItemSizes, items }; } diff --git a/packages/python/src/import-closure.ts b/packages/python/src/import-closure.ts new file mode 100644 index 000000000000..073047ab7c8e --- /dev/null +++ b/packages/python/src/import-closure.ts @@ -0,0 +1,97 @@ +import { join } from 'path'; +import { debug } from '@vercel/build-utils'; +import type { ImportClosureOptions } from '@vercel/python-analysis'; +import { + generatedPythonPathToModule, + getGeneratedQueueHandlerPath, + getSubscriberOutputPath, + type Subscriber, + type SubscriberDeclaration, +} from './subscribers'; +import { getWorkflowOutputPath, type PyprojectWorkflow } from './workflows'; +import type { WorkflowServingMode } from './sdk-detection'; + +const RUNTIME_BOOTSTRAP_MODULE = 'vercel_runtime.vc_init'; + +/** Converts a stalled import closure into a skipped ranking input. */ +export const IMPORT_CLOSURE_TIMEOUT_MS = 30_000; + +/** Resolve to undefined when `promise` exceeds `ms`. Never rejects on time. */ +export async function withTimeout( + promise: Promise, + ms: number, + label: string +): Promise { + let timer: NodeJS.Timeout | undefined; + const timeout = new Promise(resolvePromise => { + timer = setTimeout(() => { + debug(`${label} timed out after ${ms}ms`); + resolvePromise(undefined); + }, ms); + }); + try { + return await Promise.race([promise, timeout]); + } finally { + clearTimeout(timer); + } +} + +export function getImportClosureOptions({ + workPath, + entrypoint, + frameworkSeeds, + extraPythonPath, + subscriberDeclarations, + subscribers, + workflows, + workflowMode, + sitePackageDirs, +}: { + workPath: string; + entrypoint?: string; + frameworkSeeds: string[]; + extraPythonPath?: string; + subscriberDeclarations: Pick[]; + subscribers: Pick[]; + workflows: Pick[]; + workflowMode: WorkflowServingMode; + sitePackageDirs: string[]; +}): ImportClosureOptions { + const workerModules = [ + ...subscriberDeclarations.map(declaration => declaration.moduleName), + ...workflows.map(workflow => workflow.moduleName), + ]; + const generatedWorkerModules = [ + ...subscribers.map(subscriber => + generatedPythonPathToModule( + getGeneratedQueueHandlerPath(getSubscriberOutputPath(subscriber.name)) + ) + ), + ...(workflowMode === 'queue' + ? workflows.map(workflow => + generatedPythonPathToModule( + getGeneratedQueueHandlerPath(getWorkflowOutputPath(workflow.name)) + ) + ) + : []), + ]; + + return { + seeds: [ + ...new Set([ + RUNTIME_BOOTSTRAP_MODULE, + ...(entrypoint ? [join(workPath, entrypoint)] : []), + ...frameworkSeeds, + ...workerModules, + ...generatedWorkerModules, + ]), + ], + searchRoots: [ + ...new Set([ + ...(extraPythonPath ? [extraPythonPath] : []), + workPath, + ...sitePackageDirs, + ]), + ], + }; +} diff --git a/packages/python/src/index.ts b/packages/python/src/index.ts index 86a8a2e314fc..46ee21618991 100644 --- a/packages/python/src/index.ts +++ b/packages/python/src/index.ts @@ -1,6 +1,15 @@ import assert from 'assert'; import fs from 'fs'; -import { join, dirname, basename, parse } from 'path'; +import { + join, + dirname, + basename, + parse, + relative, + isAbsolute, + sep, + resolve, +} from 'path'; import { VERCEL_RUNTIME_VERSION, VERCEL_WORKERS_VERSION, @@ -57,7 +66,6 @@ import { BYTECODE_FILL_CEILING_BYTES, LARGE_FUNCTION_FILL_CEILING_BYTES, LAMBDA_SIZE_THRESHOLD_BYTES, - lambdaKnapsack, calculateBundleSize, RUNTIME_DEPS_DIR, type GenerateBundleResult, @@ -93,9 +101,16 @@ import { type FastAPICollectStaticResult, } from './fastapi'; import { + collectImportClosure, containsTopLevelCallable, type PyProjectToml, } from '@vercel/python-analysis'; +import { + annotateBytecodeItems, + fillBytecodeWithinCapacity, + isBytecodeAnalysisDisabled, + rankBytecodeItems, +} from './bytecode-packing'; import { collectAppBytecodeFiles, collectAppPrefixBytecodeFiles, @@ -123,6 +138,11 @@ import { WORKFLOW_TOPIC_PATTERN, type PyprojectWorkflow, } from './workflows'; +import { + getImportClosureOptions, + IMPORT_CLOSURE_TIMEOUT_MS, + withTimeout, +} from './import-closure'; const writeFile = fs.promises.writeFile; const PYTHON_ENTRYPOINT_DOCS_URL = @@ -228,134 +248,32 @@ function addFiles(target: Files, source: Files) { } } -function addBytecodeWithinCapacity( - files: Files, - bytecodeInfo: BytecodeCollectionResult | undefined, - capacity: number -): number { - if (!bytecodeInfo || bytecodeInfo.totalSize <= 0 || capacity <= 0) { - return capacity; - } - - if (bytecodeInfo.totalSize <= capacity) { - addFiles(files, bytecodeInfo.files); - return capacity - bytecodeInfo.totalSize; - } - - const selected = lambdaKnapsack(bytecodeInfo.perItemSizes, capacity); - let remainingCapacity = capacity; - for (const p of selected) { - const file = bytecodeInfo.files[p]; - if (!file) continue; - files[p] = file; - remainingCapacity -= bytecodeInfo.perItemSizes.get(p) ?? 0; - } - - return remainingCapacity; -} - -async function addVendorBytecodeWithinCapacity({ - files, - installedDistributions, - vendorDir, - bytecodeInfo, - capacity, -}: { - files: Files; - installedDistributions: Pick< - InstalledPythonDistributions, - 'collectBytecodeFiles' - >; - vendorDir: string; - bytecodeInfo: BytecodeCollectionResult | undefined; - capacity: number; -}): Promise { - if (!bytecodeInfo || bytecodeInfo.totalSize <= 0 || capacity <= 0) { - return capacity; - } - - if (bytecodeInfo.totalSize <= capacity) { - addFiles(files, bytecodeInfo.files); - return capacity - bytecodeInfo.totalSize; - } - - const selectedPkgs = lambdaKnapsack(bytecodeInfo.perItemSizes, capacity); - if (selectedPkgs.length === 0) return capacity; - - const selectedBytecode = await installedDistributions.collectBytecodeFiles({ - vendorDirName: vendorDir, - includePackages: selectedPkgs, - }); - addFiles(files, selectedBytecode.files); - return capacity - selectedBytecode.totalSize; -} - -/** - * Add vendor bytecode within `capacity`, in tiers: earlier tiers get - * capacity first; packages outside every tier are never collected. A tier - * of `undefined` collects everything. Returns the remaining capacity. - */ -export async function addVendorBytecodeInTiers({ - files, - installedDistributions, - vendorDir, - capacity, - vendorPackageTiers, -}: { - files: Files; - installedDistributions: Pick< - InstalledPythonDistributions, - 'collectBytecodeFiles' - >; - vendorDir: string; - capacity: number; - vendorPackageTiers: (string[] | undefined)[]; -}): Promise { - let remainingCapacity = capacity; - for (const tier of vendorPackageTiers) { - if (remainingCapacity <= 0) break; - if (tier && tier.length === 0) continue; - const bytecodeInfo = await installedDistributions.collectBytecodeFiles({ - vendorDirName: vendorDir, - includePackages: tier, - }); - remainingCapacity = await addVendorBytecodeWithinCapacity({ - files, - installedDistributions, - vendorDir, - bytecodeInfo, - capacity: remainingCapacity, - }); - } - return remainingCapacity; -} - /** - * Add vendor bytecode produced by a collector within `capacity`. When the - * full collection doesn't fit, knapsacks per-package sizes and re-collects - * only the selected packages. Returns the remaining capacity. + * Map absolute `.py` paths from the import closure to the module keys used + * by bytecode items: workPath-relative for app files, site-packages-relative + * for vendor files (forward slashes). Files outside every root (stdlib, + * venv internals) are dropped — they are never part of the bundle. */ -export async function addCollectedVendorBytecode({ - files, - capacity, - collect, -}: { - files: Files; - capacity: number; - collect: (includePackages?: string[]) => Promise; -}): Promise { - if (capacity <= 0) return capacity; - const info = await collect(undefined); - if (!info || info.totalSize <= 0) return capacity; - if (info.totalSize <= capacity) { - addFiles(files, info.files); - return capacity - info.totalSize; +export function moduleKeysForClosurePaths( + paths: Iterable, + workPath: string, + sitePackageDirs: string[] +): Set { + const keys = new Set(); + // Most specific roots first: the venv lives inside workPath + // (.vercel/python/.venv), so vendor files must match site-packages + // before the app root claims them. + const roots = [...sitePackageDirs, workPath].map(r => resolve(r)); + for (const p of paths) { + for (const root of roots) { + const rel = relative(root, p); + if (rel && !rel.startsWith('..') && !isAbsolute(rel)) { + keys.add(rel.split(sep).join('/')); + break; + } + } } - const selected = lambdaKnapsack(info.perItemSizes, capacity); - if (selected.length === 0) return capacity; - const selectedInfo = await collect(selected); - addFiles(files, selectedInfo.files); - return capacity - selectedInfo.totalSize; + return keys; } interface FrameworkHookContext { @@ -374,6 +292,11 @@ interface FrameworkHookResult { interface DjangoFrameworkHookResult extends FrameworkHookResult { djangoStatic: DjangoCollectStaticResult | null; + /** + * Dotted module names Django loads via settings strings (settings module, + * ROOT_URLCONF, INSTALLED_APPS, MIDDLEWARE); seeds the import closure. + */ + importSeeds?: string[]; } interface FastAPIFrameworkHookResult extends FrameworkHookResult { @@ -468,9 +391,21 @@ const frameworkHooks: Partial> = { djangoVersion ); } + + // Django wires apps together via settings strings rather than imports. + // Entries may name a class (e.g. MIDDLEWARE); seed resolution trims + // trailing components until a module resolves. + const importSeeds = [ + settingsModule, + djangoSettings['ROOT_URLCONF'], + ...((djangoSettings['INSTALLED_APPS'] as string[] | undefined) ?? []), + ...((djangoSettings['MIDDLEWARE'] as string[] | undefined) ?? []), + ].filter((s): s is string => typeof s === 'string'); + return { entrypoint: resolvedEntrypoint, djangoStatic, + importSeeds, extraPythonPath: baseDir ? join(workPath, baseDir) : undefined, }; }, @@ -1174,6 +1109,8 @@ export const build: BuildVX = async ({ const fastapiStatic: FastAPICollectStaticResult | null = (hookResult as FastAPIFrameworkHookResult | undefined)?.fastapiStatic ?? null; + const importSeeds: string[] = + (hookResult as DjangoFrameworkHookResult | undefined)?.importSeeds ?? []; const cdnOutputDir = djangoStatic?.cdnOutputDir ?? fastapiStatic?.cdnOutputDir ?? null; @@ -1568,21 +1505,23 @@ export const build: BuildVX = async ({ }: { includePackages?: string[]; pycachePrefix?: string; - }) => { - if (!compileAllOptions) return; + }): Promise | undefined> => { + if (!compileAllOptions) return undefined; const vendorSourceFiles = installedDistributions.getPythonSourceFiles(includePackages); + let timings: Map | undefined; await builderSpan .child('vc.builder.python.compileall') .trace(async compileSpan => { console.log('Compiling Python bytecode...'); - await runCompileAll({ + const result = await runCompileAll({ ...compileAllOptions, sourceFiles: [...appPythonSourceFiles, ...vendorSourceFiles], pycachePrefix, }); + timings = result.timings; compileSpan.setAttributes({ 'python.compileall.enabled': 'true', @@ -1594,46 +1533,153 @@ export const build: BuildVX = async ({ ), }); }); + return timings; + }; + + // Static import closure (no user code runs), computed at most once per + // build and only when a bytecode fill overflows. Undefined on failure + // or timeout, degrading ranking to compile density only. + let importClosurePromise: Promise | undefined> | undefined; + const getImportClosureKeys = (): Promise | undefined> => { + importClosurePromise ??= (async () => { + try { + const sitePackageDirs = installedDistributions.getSitePackageDirs(); + const closure = await withTimeout( + collectImportClosure( + getImportClosureOptions({ + workPath, + entrypoint, + frameworkSeeds: importSeeds, + extraPythonPath: hookResult?.extraPythonPath, + subscriberDeclarations, + subscribers, + workflows, + workflowMode, + sitePackageDirs, + }) + ), + IMPORT_CLOSURE_TIMEOUT_MS, + 'import closure' + ); + if (!closure) return undefined; + const keys = moduleKeysForClosurePaths( + closure.files, + workPath, + sitePackageDirs + ); + debug( + `import closure: ${closure.files.size} files, ` + + `${keys.size} bundled modules` + + (closure.truncated ? ' (truncated)' : '') + ); + return keys; + } catch (err) { + debug( + `import closure unavailable, ranking by compile density only: ${err}` + ); + return undefined; + } + })(); + return importClosurePromise; }; - // Precompile bytecode and fill remaining capacity up to capacityBytes. - // Only .pyc for .py files already in the bundle are collected, so - // excluded source can't re-enter as .pyc. Bytecode is a pure + // Value-ranked bytecode fill shared by every packing path. When all + // `.pyc` fit, ship them all with no analysis. On overflow, prefer + // modules in the import closure, ranked by compile seconds per byte. + // Returns bytes added. + const fillBytecodeWithValueRanking = async ({ + items, + totalSize, + capacity, + timings, + }: { + items: BytecodeCollectionResult['items']; + totalSize: number; + capacity: number; + timings: Map | undefined; + }): Promise => { + if (totalSize <= 0 || capacity <= 0) return 0; + + if (totalSize <= capacity) { + for (const item of items) { + files[item.bundlePath] = item.file; + } + return totalSize; + } + + return bundleSpan + .child('vc.builder.python.bundle.optimize') + .trace(async optimizeSpan => { + console.log('Optimizing Python bundle...'); + + // Kill switch: revert to per-file size ordering. + const analysisDisabled = isBytecodeAnalysisDisabled(); + if (analysisDisabled) { + debug( + 'bytecode analysis disabled via ' + + 'VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS; ranking by size' + ); + } + const importedModules = analysisDisabled + ? undefined + : await getImportClosureKeys(); + const ranked = rankBytecodeItems( + annotateBytecodeItems( + items, + importedModules, + analysisDisabled ? undefined : timings + ) + ); + const selectedSize = + capacity - fillBytecodeWithinCapacity(files, ranked, capacity); + + optimizeSpan.setAttributes({ + 'python.bundle.optimize.bytecodeCoveragePercent': ( + (selectedSize / totalSize) * + 100 + ).toFixed(2), + }); + + return selectedSize; + }); + }; + + // Precompile bytecode and fill remaining capacity up to capacityBytes + // using value-ranked selection. Only `.pyc` for `.py` files already in + // the bundle (plus vendor packages in `includePackages`) is collected, + // so excluded source can't re-enter as `.pyc`. Bytecode is a pure // optimization: failures are logged and the build continues. - // `vendorPackageTiers` restricts/prioritizes vendor collection; - // omitted = one unrestricted pass. - const runCompileAllAndFillBytecode = async ( + const runAdjacentCompileAndFill = async ( capacityBytes: number, - vendorPackageTiers?: string[][] + includePackages?: string[] ) => { try { - await compileSources({ - includePackages: vendorPackageTiers?.flat(), - }); + const pyMajor = pythonVersion.major; + const pyMinor = pythonVersion.minor; + if (pyMajor == null || pyMinor == null) return; - const currentSize = await calculateBundleSize(files); - let remainingCapacity = capacityBytes - currentSize; + const timings = await compileSources({ includePackages }); - if (pythonVersion.major != null && pythonVersion.minor != null) { - const appBytecodeInfo = await collectAppBytecodeFiles({ - workPath, - files, - pythonMajor: pythonVersion.major, - pythonMinor: pythonVersion.minor, - }); - remainingCapacity = addBytecodeWithinCapacity( - files, - appBytecodeInfo, - remainingCapacity - ); - } + const currentSize = await calculateBundleSize(files); + const remaining = capacityBytes - currentSize; + if (remaining <= 0) return; - await addVendorBytecodeInTiers({ + const appInfo = await collectAppBytecodeFiles({ + workPath, files, - installedDistributions, - vendorDir, - capacity: remainingCapacity, - vendorPackageTiers: vendorPackageTiers ?? [undefined], + pythonMajor: pyMajor, + pythonMinor: pyMinor, + }); + const vendorInfo = await installedDistributions.collectBytecodeFiles({ + vendorDirName: vendorDir, + includePackages, + }); + + await fillBytecodeWithValueRanking({ + items: [...appInfo.items, ...vendorInfo.items], + totalSize: appInfo.totalSize + vendorInfo.totalSize, + capacity: remaining, + timings, }); } catch (err) { console.log( @@ -1646,7 +1692,8 @@ export const build: BuildVX = async ({ // Bytecode-first fill: ship a pycache-prefix tree covering the app, // bundled vendor packages, and the packages installed into /tmp at // cold start (safe: `uv sync --frozen` installs the exact versions - // the bytecode was compiled from). Failures degrade to no bytecode. + // the bytecode was compiled from). Selection is value-ranked like + // every other path. Failures degrade to no bytecode. const runPrefixCompileAndFill = async ( bundleResult: GenerateBundleResult ) => { @@ -1657,7 +1704,7 @@ export const build: BuildVX = async ({ // Skip the compile entirely when the zip has no slack for bytecode // (e.g. very large always-bundled private packages). const currentSize = await calculateBundleSize(files); - let remainingCapacity = BYTECODE_FILL_CEILING_BYTES - currentSize; + const remainingCapacity = BYTECODE_FILL_CEILING_BYTES - currentSize; if (remainingCapacity <= 0) { debug( `skipping bytecode precompilation: no zip capacity remaining ` + @@ -1671,18 +1718,23 @@ export const build: BuildVX = async ({ await fs.promises.rm(stagingDir, { recursive: true, force: true }); await fs.promises.mkdir(stagingDir, { recursive: true }); - await compileSources({ + const alwaysBundled = bundleResult.alwaysBundledPackages ?? []; + const bundledPublic = bundleResult.bundledPublicPackages ?? []; + const externalized = bundleResult.externalizedPublicPackages ?? []; + + const timings = await compileSources({ includePackages: [ - ...(bundleResult.alwaysBundledPackages ?? []), - ...(bundleResult.bundledPublicPackages ?? []), - ...(bundleResult.externalizedPublicPackages ?? []), + ...alwaysBundled, + ...bundledPublic, + ...externalized, ], pycachePrefix: stagingDir, }); - const beforeCount = Object.keys(files).length; - - // Tier 1: app source (always imported at cold start). + // Candidates: app source, bundled vendor (/var/task/_vendor), and + // externalized packages (installed into /tmp at cold start). All + // carry module keys the closure can match, so one ranking covers + // the union. const appInfo = await collectAppPrefixBytecodeFiles({ stagingDir, workPath, @@ -1691,40 +1743,35 @@ export const build: BuildVX = async ({ pythonMajor: pyMajor, pythonMinor: pyMinor, }); - remainingCapacity = addBytecodeWithinCapacity( - files, - appInfo, - remainingCapacity - ); - - // Tier 2: bundled vendor packages, imported from /var/task/_vendor. - const alwaysBundled = bundleResult.alwaysBundledPackages ?? []; - remainingCapacity = await addCollectedVendorBytecode({ - files, - capacity: remainingCapacity, - collect: include => - installedDistributions.collectPrefixBytecodeFiles({ - stagingDir, - runtimeRoot: `/var/task/${vendorDir}`, - includePackages: include ?? alwaysBundled, - }), - }); + const bundledVendorInfo = + await installedDistributions.collectPrefixBytecodeFiles({ + stagingDir, + runtimeRoot: `/var/task/${vendorDir}`, + includePackages: [...alwaysBundled, ...bundledPublic], + }); + const externalizedInfo = + await installedDistributions.collectPrefixBytecodeFiles({ + stagingDir, + runtimeRoot: `${RUNTIME_DEPS_DIR}/lib/python${pyMajor}.${pyMinor}/site-packages`, + includePackages: externalized, + }); - // Tier 3: externalized packages, installed into /tmp at cold start. - const externalized = bundleResult.externalizedPublicPackages ?? []; - await addCollectedVendorBytecode({ - files, + const bytesAdded = await fillBytecodeWithValueRanking({ + items: [ + ...appInfo.items, + ...bundledVendorInfo.items, + ...externalizedInfo.items, + ], + totalSize: + appInfo.totalSize + + bundledVendorInfo.totalSize + + externalizedInfo.totalSize, capacity: remainingCapacity, - collect: include => - installedDistributions.collectPrefixBytecodeFiles({ - stagingDir, - runtimeRoot: `${RUNTIME_DEPS_DIR}/lib/python${pyMajor}.${pyMinor}/site-packages`, - includePackages: include ?? externalized, - }), + timings, }); // Point the runtime at the tree only when bytecode shipped. - if (Object.keys(files).length > beforeCount) { + if (bytesAdded > 0) { lambdaEnv.PYTHONPYCACHEPREFIX = RUNTIME_PYCACHE_PREFIX; } } catch (err) { @@ -1762,24 +1809,22 @@ export const build: BuildVX = async ({ packingMode = 'hive'; announceLargeFunction(); if (compileAllEnabled) { - await runCompileAllAndFillBytecode( - LARGE_FUNCTION_FILL_CEILING_BYTES - ); + await runAdjacentCompileAndFill(LARGE_FUNCTION_FILL_CEILING_BYTES); } } else if (bundleResult.packingMode === 'bytecode-first') { await runPrefixCompileAndFill(bundleResult); } else if (compileAllEnabled) { // Knapsack packing (bytecode-first skipped or fell back): fill - // the slack under the ceiling with bytecode for in-zip packages. - // Always-bundled packages get capacity first. Skip only when the - // bundle already exceeds the fill ceiling, since nothing could - // ship. + // the slack under the ceiling with bytecode for in-zip packages, + // selected by import closure and compile density. Skip only when + // the bundle already exceeds the fill ceiling, since nothing + // could ship. const currentSize = await calculateBundleSize(files); const capacity = BYTECODE_FILL_CEILING_BYTES - currentSize; if (capacity > 0) { - await runCompileAllAndFillBytecode(BYTECODE_FILL_CEILING_BYTES, [ - bundleResult.alwaysBundledPackages ?? [], - bundleResult.bundledPublicPackages ?? [], + await runAdjacentCompileAndFill(BYTECODE_FILL_CEILING_BYTES, [ + ...(bundleResult.alwaysBundledPackages ?? []), + ...(bundleResult.bundledPublicPackages ?? []), ]); } else { debug( @@ -1798,9 +1843,7 @@ export const build: BuildVX = async ({ announceLargeFunction(); } if (compileAllEnabled) { - await runCompileAllAndFillBytecode( - LARGE_FUNCTION_FILL_CEILING_BYTES - ); + await runAdjacentCompileAndFill(LARGE_FUNCTION_FILL_CEILING_BYTES); } } else { packingMode = 'standard'; @@ -1811,7 +1854,7 @@ export const build: BuildVX = async ({ const capacity = BYTECODE_FILL_CEILING_BYTES - depAnalysis.totalBundleSize; if (capacity > 0) { - await runCompileAllAndFillBytecode(BYTECODE_FILL_CEILING_BYTES); + await runAdjacentCompileAndFill(BYTECODE_FILL_CEILING_BYTES); } else { debug( `skipping bytecode precompilation: no zip capacity remaining ` + diff --git a/packages/python/src/installed-distributions.ts b/packages/python/src/installed-distributions.ts index aa0de0950247..be7dac7a6db4 100644 --- a/packages/python/src/installed-distributions.ts +++ b/packages/python/src/installed-distributions.ts @@ -11,6 +11,7 @@ import { derivePycPath, deriveStagedPycFsPath, type BytecodeCollectionResult, + type BytecodeItem, } from './compileall'; import { getVenvSitePackagesDirs } from './install'; @@ -145,6 +146,14 @@ export class InstalledPythonDistributions { this.pythonMinor = options.pythonMinor; } + /** + * Site-packages roots of the build venv (resolved). Used as import-closure + * search roots and to map traced module files to vendor module keys. + */ + getSitePackageDirs(): string[] { + return this.sitePackageDirs; + } + async mirrorPackagesIntoVendor({ vendorDirName, includePackages, @@ -271,13 +280,15 @@ export class InstalledPythonDistributions { includePackages?: string[]; }): Promise { if (this.pythonMajor == null || this.pythonMinor == null) { - return { files: {}, totalSize: 0, perItemSizes: new Map() }; + return { files: {}, totalSize: 0, perItemSizes: new Map(), items: [] }; } interface PendingEntry { bundlePath: string; srcFsPath: string; packageName: string; + moduleKey: string; + sourceAbsPath: string; } const pending: PendingEntry[] = []; @@ -289,8 +300,9 @@ export class InstalledPythonDistributions { for (const { packageName, sitePackagesDir, files } of distributionGroups) { for (const { relativePath } of files) { + const moduleKey = relativePath.replaceAll(sep, '/'); const pycRelativePath = derivePycPath( - relativePath.replaceAll(sep, '/'), + moduleKey, this.pythonMajor, this.pythonMinor ); @@ -301,6 +313,8 @@ export class InstalledPythonDistributions { bundlePath: join(vendorDirName, pycFilePath).replace(/\\/g, '/'), srcFsPath: join(sitePackagesDir, pycFilePath), packageName, + moduleKey, + sourceAbsPath: join(sitePackagesDir, relativePath), }); } } @@ -324,13 +338,15 @@ export class InstalledPythonDistributions { includePackages?: string[]; }): Promise { if (this.pythonMajor == null || this.pythonMinor == null) { - return { files: {}, totalSize: 0, perItemSizes: new Map() }; + return { files: {}, totalSize: 0, perItemSizes: new Map(), items: [] }; } interface PendingEntry { bundlePath: string; srcFsPath: string; packageName: string; + moduleKey: string; + sourceAbsPath: string; } const pending: PendingEntry[] = []; @@ -358,7 +374,13 @@ export class InstalledPythonDistributions { ); if (!srcFsPath || !bundlePath) continue; - pending.push({ bundlePath, srcFsPath, packageName }); + pending.push({ + bundlePath, + srcFsPath, + packageName, + moduleKey: recordPath, + sourceAbsPath: absolutePath, + }); } } @@ -377,36 +399,62 @@ export class InstalledPythonDistributions { bundlePath: string; srcFsPath: string; packageName: string; + moduleKey: string; + sourceAbsPath: string; }[] ): Promise { const results = await Promise.all( - pending.map(async ({ bundlePath, srcFsPath, packageName }) => { - try { - const stats = await fs.promises.stat(srcFsPath); - return { bundlePath, srcFsPath, size: stats.size, packageName }; - } catch { - return null; + pending.map( + async ({ + bundlePath, + srcFsPath, + packageName, + moduleKey, + sourceAbsPath, + }) => { + try { + const stats = await fs.promises.stat(srcFsPath); + return { + bundlePath, + srcFsPath, + size: stats.size, + packageName, + moduleKey, + sourceAbsPath, + }; + } catch { + return null; + } } - }) + ) ); const files: Files = {}; let totalSize = 0; const perItemSizes = new Map(); + const items: BytecodeItem[] = []; for (const result of results) { if (!result) continue; - files[result.bundlePath] = new FileFsRef({ + const file = new FileFsRef({ fsPath: result.srcFsPath, size: result.size, }); + files[result.bundlePath] = file; totalSize += result.size; perItemSizes.set( result.packageName, (perItemSizes.get(result.packageName) ?? 0) + result.size ); + items.push({ + bundlePath: result.bundlePath, + file, + size: result.size, + moduleKey: result.moduleKey, + sourceAbsPath: result.sourceAbsPath, + }); } - return { files, totalSize, perItemSizes }; + return { files, totalSize, perItemSizes, items }; } } diff --git a/packages/python/templates/vc_compileall.py b/packages/python/templates/vc_compileall.py index c85d07f2ee24..83324b81be96 100644 --- a/packages/python/templates/vc_compileall.py +++ b/packages/python/templates/vc_compileall.py @@ -2,21 +2,24 @@ import json import signal import sys +import time from multiprocessing import Pool from py_compile import PycInvalidationMode def compile_source(source_file): try: - return compileall.compile_file( + start = time.perf_counter() + ok = compileall.compile_file( source_file, force=True, quiet=1, invalidation_mode=PycInvalidationMode.UNCHECKED_HASH, ) + return (source_file, ok, time.perf_counter() - start) except Exception as error: print(f"Failed to compile {source_file}: {error}", file=sys.stderr) - return False + return (source_file, False, 0.0) def main(): @@ -26,6 +29,10 @@ def main(): if not source_files: return 0 + # Optional argv[2]: write per-file compile timings (seconds) as JSON, + # used by bytecode packing to rank files by compile cost per byte. + timings_path = sys.argv[2] if len(sys.argv) > 2 else None + try: with Pool() as pool: def abort_pool(signum, _): @@ -35,11 +42,22 @@ def abort_pool(signum, _): for sig in signal.SIGINT, signal.SIGTERM: signal.signal(sig, abort_pool) - pool.map(compile_source, source_files) + results = pool.map(compile_source, source_files) except Exception as error: print(f"Bytecode compilation unavailable: {error}", file=sys.stderr) return 1 + if timings_path: + try: + timings = { + path: elapsed for path, ok, elapsed in results if ok and elapsed > 0 + } + with open(timings_path, "w", encoding="utf-8") as out: + json.dump(timings, out) + except Exception as error: + # Timings are advisory; never fail the compile over them. + print(f"Could not write compile timings: {error}", file=sys.stderr) + return 0 diff --git a/packages/python/test/pycache-prefix.test.ts b/packages/python/test/pycache-prefix.test.ts index 3ebd7fd9d339..a8f4162636a9 100644 --- a/packages/python/test/pycache-prefix.test.ts +++ b/packages/python/test/pycache-prefix.test.ts @@ -74,9 +74,11 @@ describe('explicit-list compilation layout (real CPython)', () => { fs.writeFileSync(srcPath, 'X = 1\n'); } - await expect( - runCompileAll({ pythonBin, sourceFiles: sourcePaths }) - ).resolves.toBe(true); + const result = await runCompileAll({ pythonBin, sourceFiles: sourcePaths }); + expect(result.success).toBe(true); + for (const srcPath of sourcePaths) { + expect(result.timings?.get(srcPath)).toBeGreaterThan(0); + } for (const srcPath of sourcePaths) { expect( @@ -105,13 +107,15 @@ describe('explicit-list compilation layout (real CPython)', () => { fs.writeFileSync(srcPath, 'X = 1\n'); } - await expect( - runCompileAll({ - pythonBin, - sourceFiles: sourcePaths, - pycachePrefix: stagingDir, - }) - ).resolves.toBe(true); + const result = await runCompileAll({ + pythonBin, + sourceFiles: sourcePaths, + pycachePrefix: stagingDir, + }); + expect(result.success).toBe(true); + for (const srcPath of sourcePaths) { + expect(result.timings?.get(srcPath)).toBeGreaterThan(0); + } for (const srcPath of sourcePaths) { const derived = deriveStagedPycFsPath(stagingDir, srcPath, major, minor); @@ -144,13 +148,15 @@ describe('explicit-list compilation layout (real CPython)', () => { fs.writeFileSync(validSource, 'X = 1\n'); fs.writeFileSync(invalidSource, 'def invalid syntax\n'); - await expect( - runCompileAll({ - pythonBin, - sourceFiles: [validSource, invalidSource], - pycachePrefix: stagingDir, - }) - ).resolves.toBe(true); + const result = await runCompileAll({ + pythonBin, + sourceFiles: [validSource, invalidSource], + pycachePrefix: stagingDir, + }); + expect(result.success).toBe(true); + // The failed source contributes no timing entry. + expect(result.timings?.get(validSource)).toBeGreaterThan(0); + expect(result.timings?.has(invalidSource)).toBe(false); const validBytecode = deriveStagedPycFsPath( stagingDir, diff --git a/packages/python/test/unit.bytecode-fill.test.ts b/packages/python/test/unit.bytecode-fill.test.ts deleted file mode 100644 index 7b2b5f05b475..000000000000 --- a/packages/python/test/unit.bytecode-fill.test.ts +++ /dev/null @@ -1,297 +0,0 @@ -import { describe, it, expect } from 'vitest'; -import { FileBlob, type Files } from '@vercel/build-utils'; -import { - addCollectedVendorBytecode, - addVendorBytecodeInTiers, -} from '../src/index'; -import type { BytecodeCollectionResult } from '../src/compileall'; - -const MB = 1024 * 1024; - -/** - * Stub externalizer whose collectBytecodeFiles() serves a fixed set of - * per-package bytecode entries, honoring the includePackages restriction - * the same way the real implementation does (missing names are skipped). - */ -function makeStubExternalizer( - packages: Record -) { - const collectCalls: (string[] | undefined)[] = []; - const stub = { - collectCalls, - async collectBytecodeFiles({ - includePackages, - }: { - vendorDirName: string; - includePackages?: string[]; - }) { - collectCalls.push(includePackages); - const names = includePackages ?? Object.keys(packages); - const files: Files = {}; - const perItemSizes = new Map(); - let totalSize = 0; - for (const name of names) { - const pkg = packages[name]; - if (!pkg) continue; - files[pkg.bundlePath] = new FileBlob({ data: 'pyc' }); - perItemSizes.set(name, pkg.size); - totalSize += pkg.size; - } - return { files, totalSize, perItemSizes }; - }, - }; - return stub; -} - -describe('addVendorBytecodeInTiers', () => { - const packages = { - 'private-pkg': { - bundlePath: '_vendor/private/__pycache__/a.pyc', - size: 2 * MB, - }, - 'wheelless-pkg': { - bundlePath: '_vendor/wheelless/__pycache__/b.pyc', - size: 3 * MB, - }, - 'public-big': { - bundlePath: '_vendor/big/__pycache__/c.pyc', - size: 10 * MB, - }, - 'public-small': { - bundlePath: '_vendor/small/__pycache__/d.pyc', - size: 1 * MB, - }, - 'externalized-pkg': { - bundlePath: '_vendor/ext/__pycache__/e.pyc', - size: 4 * MB, - }, - }; - - it('adds bytecode for every tier when capacity fits and returns the remainder', async () => { - const stub = makeStubExternalizer(packages); - const files: Files = {}; - - const remaining = await addVendorBytecodeInTiers({ - files, - installedDistributions: stub, - vendorDir: '_vendor', - capacity: 20 * MB, - vendorPackageTiers: [ - ['private-pkg', 'wheelless-pkg'], - ['public-big', 'public-small'], - ], - }); - - expect(Object.keys(files).sort()).toEqual( - [ - packages['private-pkg'].bundlePath, - packages['wheelless-pkg'].bundlePath, - packages['public-big'].bundlePath, - packages['public-small'].bundlePath, - ].sort() - ); - expect(remaining).toBe(4 * MB); - }); - - it('never collects packages outside the tiers (externalized deps)', async () => { - const stub = makeStubExternalizer(packages); - const files: Files = {}; - - await addVendorBytecodeInTiers({ - files, - installedDistributions: stub, - vendorDir: '_vendor', - capacity: 100 * MB, - vendorPackageTiers: [['private-pkg'], ['public-big']], - }); - - expect(files[packages['externalized-pkg'].bundlePath]).toBeUndefined(); - for (const call of stub.collectCalls) { - expect(call).toBeDefined(); - expect(call).not.toContain('externalized-pkg'); - } - }); - - it('gives earlier tiers capacity first', async () => { - const stub = makeStubExternalizer(packages); - const files: Files = {}; - - // Capacity fits tier 1 (5MB) plus only the small public package. - const remaining = await addVendorBytecodeInTiers({ - files, - installedDistributions: stub, - vendorDir: '_vendor', - capacity: 6 * MB, - vendorPackageTiers: [ - ['private-pkg', 'wheelless-pkg'], - ['public-big', 'public-small'], - ], - }); - - expect(files[packages['private-pkg'].bundlePath]).toBeDefined(); - expect(files[packages['wheelless-pkg'].bundlePath]).toBeDefined(); - expect(files[packages['public-small'].bundlePath]).toBeDefined(); - expect(files[packages['public-big'].bundlePath]).toBeUndefined(); - expect(remaining).toBe(0); - }); - - it('stops once capacity is exhausted', async () => { - const stub = makeStubExternalizer(packages); - const files: Files = {}; - - const remaining = await addVendorBytecodeInTiers({ - files, - installedDistributions: stub, - vendorDir: '_vendor', - capacity: 5 * MB, - vendorPackageTiers: [ - ['private-pkg', 'wheelless-pkg'], - ['public-big', 'public-small'], - ], - }); - - expect(remaining).toBe(0); - // Tier 2 is never collected: capacity hit zero after tier 1. - expect(stub.collectCalls).toHaveLength(1); - expect(files[packages['public-small'].bundlePath]).toBeUndefined(); - }); - - it('skips empty tiers without collecting', async () => { - const stub = makeStubExternalizer(packages); - const files: Files = {}; - - await addVendorBytecodeInTiers({ - files, - installedDistributions: stub, - vendorDir: '_vendor', - capacity: 100 * MB, - vendorPackageTiers: [[], ['public-small']], - }); - - expect(stub.collectCalls).toHaveLength(1); - expect(stub.collectCalls[0]).toEqual(['public-small']); - }); - - it('collects every vendor package for an undefined tier', async () => { - const stub = makeStubExternalizer(packages); - const files: Files = {}; - - await addVendorBytecodeInTiers({ - files, - installedDistributions: stub, - vendorDir: '_vendor', - capacity: 100 * MB, - vendorPackageTiers: [undefined], - }); - - expect(Object.keys(files)).toHaveLength(5); - expect(stub.collectCalls).toEqual([undefined]); - }); - - it('adds nothing when capacity is zero or negative', async () => { - const stub = makeStubExternalizer(packages); - const files: Files = {}; - - const remaining = await addVendorBytecodeInTiers({ - files, - installedDistributions: stub, - vendorDir: '_vendor', - capacity: -1 * MB, - vendorPackageTiers: [['private-pkg']], - }); - - expect(Object.keys(files)).toHaveLength(0); - expect(stub.collectCalls).toHaveLength(0); - expect(remaining).toBe(-1 * MB); - }); -}); - -describe('addCollectedVendorBytecode', () => { - function makeCollector( - packages: Record - ) { - const calls: (string[] | undefined)[] = []; - const collect = async ( - includePackages?: string[] - ): Promise => { - calls.push(includePackages); - const names = includePackages ?? Object.keys(packages); - const files: Files = {}; - const perItemSizes = new Map(); - let totalSize = 0; - for (const name of names) { - const pkg = packages[name]; - if (!pkg) continue; - files[pkg.bundlePath] = new FileBlob({ data: 'pyc' }); - perItemSizes.set(name, pkg.size); - totalSize += pkg.size; - } - return { files, totalSize, perItemSizes }; - }; - return { collect, calls }; - } - - const packages = { - big: { bundlePath: '_vc_pycache/tmp/x/big.pyc', size: 10 * MB }, - small: { bundlePath: '_vc_pycache/tmp/x/small.pyc', size: 1 * MB }, - }; - - it('adds the full collection when it fits', async () => { - const { collect, calls } = makeCollector(packages); - const files: Files = {}; - - const remaining = await addCollectedVendorBytecode({ - files, - capacity: 20 * MB, - collect, - }); - - expect(Object.keys(files)).toHaveLength(2); - expect(remaining).toBe(9 * MB); - expect(calls).toEqual([undefined]); - }); - - it('knapsacks and re-collects when the collection exceeds capacity', async () => { - const { collect, calls } = makeCollector(packages); - const files: Files = {}; - - const remaining = await addCollectedVendorBytecode({ - files, - capacity: 5 * MB, - collect, - }); - - expect(Object.keys(files)).toEqual([packages.small.bundlePath]); - expect(remaining).toBe(4 * MB); - expect(calls).toEqual([undefined, ['small']]); - }); - - it('does nothing with zero or negative capacity', async () => { - const { collect, calls } = makeCollector(packages); - const files: Files = {}; - - const remaining = await addCollectedVendorBytecode({ - files, - capacity: 0, - collect, - }); - - expect(Object.keys(files)).toHaveLength(0); - expect(remaining).toBe(0); - expect(calls).toHaveLength(0); - }); - - it('does nothing when the collection is empty', async () => { - const { collect } = makeCollector({}); - const files: Files = {}; - - const remaining = await addCollectedVendorBytecode({ - files, - capacity: 5 * MB, - collect, - }); - - expect(Object.keys(files)).toHaveLength(0); - expect(remaining).toBe(5 * MB); - }); -}); diff --git a/packages/python/test/unit.bytecode-packing.test.ts b/packages/python/test/unit.bytecode-packing.test.ts new file mode 100644 index 000000000000..6eed3eea9741 --- /dev/null +++ b/packages/python/test/unit.bytecode-packing.test.ts @@ -0,0 +1,252 @@ +import { describe, it, expect, afterEach } from 'vitest'; +import { FileBlob, type Files } from '@vercel/build-utils'; +import { + annotateBytecodeItems, + fillBytecodeWithinCapacity, + isBytecodeAnalysisDisabled, + rankBytecodeItems, +} from '../src/bytecode-packing'; +import { moduleKeysForClosurePaths } from '../src/index'; +import type { BytecodeItem } from '../src/compileall'; +import { join } from 'path'; + +const MB = 1024 * 1024; + +function makeItem( + bundlePath: string, + size: number, + moduleKey?: string, + sourceAbsPath?: string +): BytecodeItem { + return { + bundlePath, + file: new FileBlob({ data: 'pyc' }) as unknown as BytecodeItem['file'], + size, + moduleKey: moduleKey ?? bundlePath, + sourceAbsPath: sourceAbsPath ?? `/src/${moduleKey ?? bundlePath}`, + }; +} + +describe('isBytecodeAnalysisDisabled', () => { + const original = process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS; + + afterEach(() => { + if (original === undefined) { + delete process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS; + } else { + process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS = original; + } + }); + + it.each(['1', 'true', 'TRUE', 'True'])('is disabled for %j', value => { + process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS = value; + expect(isBytecodeAnalysisDisabled()).toBe(true); + }); + + it.each(['', '0', 'false', 'no'])('is enabled for %j', value => { + process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS = value; + expect(isBytecodeAnalysisDisabled()).toBe(false); + }); + + it('is enabled when unset', () => { + delete process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS; + expect(isBytecodeAnalysisDisabled()).toBe(false); + }); +}); + +describe('rankBytecodeItems', () => { + it('ranks imported modules ahead of unimported ones', () => { + const items = annotateBytecodeItems( + [ + makeItem('_vendor/huge/__init__.pyc', 25 * MB, 'huge/__init__.py'), + makeItem('_vendor/small/__init__.pyc', 1 * MB, 'small/__init__.py'), + ], + new Set(['small/__init__.py']), + undefined + ); + + const ranked = rankBytecodeItems(items); + expect(ranked.map(i => i.bundlePath)).toEqual([ + '_vendor/small/__init__.pyc', + '_vendor/huge/__init__.pyc', + ]); + }); + + it('ranks by compile density within a tier, not raw size', () => { + // twilio-like: 25 MB, cheap per byte. pydantic-like: 5 MB, expensive. + const items = annotateBytecodeItems( + [ + makeItem('a.pyc', 25 * MB, 'a.py', '/src/a.py'), + makeItem('b.pyc', 5 * MB, 'b.py', '/src/b.py'), + ], + new Set(['a.py', 'b.py']), + new Map([ + ['/src/a.py', 0.25], // 0.01 s/MB + ['/src/b.py', 0.5], // 0.1 s/MB + ]) + ); + + const ranked = rankBytecodeItems(items); + expect(ranked.map(i => i.bundlePath)).toEqual(['b.pyc', 'a.pyc']); + }); + + it('falls back to per-file size ordering without timings', () => { + const items = annotateBytecodeItems( + [makeItem('small.pyc', 1 * MB), makeItem('big.pyc', 10 * MB)], + undefined, + undefined + ); + + const ranked = rankBytecodeItems(items); + expect(ranked.map(i => i.bundlePath)).toEqual(['big.pyc', 'small.pyc']); + }); + + it('ranks items with missing timings at the median density', () => { + const items = annotateBytecodeItems( + [ + makeItem('timed-dense.pyc', 10 * MB, 'a.py', '/src/a.py'), + makeItem('timed-sparse.pyc', 1 * MB, 'b.py', '/src/b.py'), + makeItem('untimed.pyc', 5 * MB, 'c.py', '/src/c.py'), + ], + undefined, + new Map([ + ['/src/a.py', 10], // 1 s/MB + ['/src/b.py', 0.001], // 0.001 s/MB + ]) + ); + + // median of [1, 0.001] -> lower-middle element after sort = index 1 of 2? + // densities sorted: [0.001, 1]; median index = floor(2/2) = 1 -> 1. + // untimed takes density 1, tying with timed-dense; tie broken by size. + const ranked = rankBytecodeItems(items); + expect(ranked.map(i => i.bundlePath)).toEqual([ + 'timed-dense.pyc', + 'untimed.pyc', + 'timed-sparse.pyc', + ]); + }); +}); + +describe('fillBytecodeWithinCapacity', () => { + it('fills greedily in rank order and returns the remainder', () => { + const files: Files = {}; + const items = annotateBytecodeItems( + [makeItem('a.pyc', 10 * MB), makeItem('b.pyc', 5 * MB)], + undefined, + undefined + ); + + const remaining = fillBytecodeWithinCapacity( + files, + rankBytecodeItems(items), + 20 * MB + ); + expect(Object.keys(files)).toEqual(['a.pyc', 'b.pyc']); + expect(remaining).toBe(5 * MB); + }); + + it('skips items that do not fit and continues with smaller ones', () => { + const files: Files = {}; + const items = annotateBytecodeItems( + [makeItem('big.pyc', 10 * MB), makeItem('small.pyc', 4 * MB)], + undefined, + undefined + ); + + const remaining = fillBytecodeWithinCapacity( + files, + rankBytecodeItems(items), + 5 * MB + ); + expect(Object.keys(files)).toEqual(['small.pyc']); + expect(remaining).toBe(1 * MB); + }); + + it('stops at zero remaining capacity', () => { + const files: Files = {}; + const items = annotateBytecodeItems( + [makeItem('a.pyc', 3 * MB), makeItem('b.pyc', 3 * MB)], + undefined, + undefined + ); + + const remaining = fillBytecodeWithinCapacity( + files, + rankBytecodeItems(items), + 3 * MB + ); + expect(Object.keys(files)).toEqual(['a.pyc']); + expect(remaining).toBe(0); + }); + + it('ships imported packages partially at per-file granularity', () => { + // The hubspot argument: a partially-imported package contributes its + // imported slice, not all-or-nothing. + const files: Files = {}; + const items = annotateBytecodeItems( + [ + makeItem('_vendor/sdk/a.pyc', 2 * MB, 'sdk/a.py'), + makeItem('_vendor/sdk/b.pyc', 2 * MB, 'sdk/b.py'), + makeItem('_vendor/sdk/generated1.pyc', 8 * MB, 'sdk/generated1.py'), + makeItem('_vendor/sdk/generated2.pyc', 8 * MB, 'sdk/generated2.py'), + ], + new Set(['sdk/a.py', 'sdk/b.py']), + undefined + ); + + const remaining = fillBytecodeWithinCapacity( + files, + rankBytecodeItems(items), + 13 * MB + ); + expect(Object.keys(files).sort()).toEqual([ + '_vendor/sdk/a.pyc', + '_vendor/sdk/b.pyc', + // one generated file fits in the leftover capacity + '_vendor/sdk/generated1.pyc', + ]); + expect(remaining).toBe(1 * MB); + }); +}); + +describe('moduleKeysForClosurePaths', () => { + const workPath = join('/work'); + const sitePackages = join('/work/.venv/lib/python3.12/site-packages'); + + it('maps app files to workPath-relative keys and vendor files to site-packages-relative keys', () => { + const keys = moduleKeysForClosurePaths( + [ + join(workPath, 'main.py'), + join(workPath, 'routes/users.py'), + join(sitePackages, 'fastapi/applications.py'), + ], + workPath, + [sitePackages] + ); + + expect([...keys].sort()).toEqual([ + 'fastapi/applications.py', + 'main.py', + 'routes/users.py', + ]); + }); + + it('drops files outside every root (stdlib is never bundled)', () => { + const keys = moduleKeysForClosurePaths( + [join('/usr/lib/python3.12/json/__init__.py')], + workPath, + [sitePackages] + ); + expect(keys.size).toBe(0); + }); + + it('maps vendor files to site-packages keys even when the venv is nested under workPath', () => { + // Production layout: the build venv lives at workPath/.vercel/python/.venv. + const keys = moduleKeysForClosurePaths( + [join(sitePackages, 'fastapi/__init__.py'), join(workPath, 'main.py')], + workPath, + [sitePackages] + ); + expect([...keys].sort()).toEqual(['fastapi/__init__.py', 'main.py']); + }); +}); diff --git a/packages/python/test/unit.compileall.test.ts b/packages/python/test/unit.compileall.test.ts index 2421f6d38b15..f7a02429de2d 100644 --- a/packages/python/test/unit.compileall.test.ts +++ b/packages/python/test/unit.compileall.test.ts @@ -149,12 +149,13 @@ describe('runCompileAll', () => { sourceFiles: ['/work/app.py', '/work/pkg/mod.py', '/work/app.py'], env, }) - ).resolves.toBe(true); + ).resolves.toEqual({ success: true, timings: undefined }); const args = mockedExeca.mock.calls[0][1]; expect(args).toEqual([ expect.stringMatching(/templates[/\\\\]vc_compileall\.py$/), listPath, + expect.stringMatching(/timings\.json$/), ]); expect(mockedExeca).toHaveBeenCalledWith( '/work/.vercel/python/.venv/bin/python', @@ -169,7 +170,7 @@ describe('runCompileAll', () => { it('does not invoke the coordinator when there are no source files', async () => { await expect( runCompileAll({ pythonBin: 'python3', sourceFiles: [] }) - ).resolves.toBe(false); + ).resolves.toEqual({ success: false }); expect(mockedExeca).not.toHaveBeenCalled(); }); @@ -191,7 +192,7 @@ describe('runCompileAll', () => { pythonBin: 'python3', sourceFiles: ['/work/app.py'], }) - ).resolves.toBe(false); + ).resolves.toEqual({ success: false }); expect(fs.existsSync(listPath)).toBe(false); expect(fs.existsSync(path.dirname(listPath))).toBe(false); @@ -211,12 +212,35 @@ describe('runCompileAll', () => { pythonBin: 'python3', sourceFiles: ['/work/app.py'], }) - ).resolves.toBe(false); + ).resolves.toEqual({ success: false }); expect(fs.existsSync(listPath)).toBe(false); expect(fs.existsSync(path.dirname(listPath))).toBe(false); }); + it('parses the per-file timings table when the coordinator writes one', async () => { + mockedExeca.mockImplementation(((_file, args: string[]) => { + fs.writeFileSync( + args[2], + JSON.stringify({ '/work/app.py': 0.012, '/work/pkg/mod.py': 0.004 }) + ); + return Promise.resolve({}); + }) as any); + + await expect( + runCompileAll({ + pythonBin: 'python3', + sourceFiles: ['/work/app.py', '/work/pkg/mod.py'], + }) + ).resolves.toEqual({ + success: true, + timings: new Map([ + ['/work/app.py', 0.012], + ['/work/pkg/mod.py', 0.004], + ]), + }); + }); + it('sets PYTHONPYCACHEPREFIX on the subprocess when provided', async () => { mockedExeca.mockResolvedValue({} as any); const env = { VIRTUAL_ENV: '/work/.vercel/python/.venv' }; @@ -232,6 +256,7 @@ describe('runCompileAll', () => { [ expect.stringMatching(/templates[/\\\\]vc_compileall\.py$/), expect.stringMatching(/pysources\.json$/), + expect.stringMatching(/timings\.json$/), ], { env: { ...env, PYTHONPYCACHEPREFIX: '/work/.vercel/python/pycache' }, diff --git a/packages/python/test/unit.import-closure.test.ts b/packages/python/test/unit.import-closure.test.ts new file mode 100644 index 000000000000..cfa48cb7acc0 --- /dev/null +++ b/packages/python/test/unit.import-closure.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, it } from 'vitest'; +import path from 'path'; +import { getImportClosureOptions, withTimeout } from '../src/import-closure'; + +describe('getImportClosureOptions', () => { + it('includes runtime, framework, and worker startup modules', () => { + const workPath = path.join('/work'); + + const options = getImportClosureOptions({ + workPath, + entrypoint: 'config/wsgi.py', + frameworkSeeds: ['config.settings', 'myapp.apps.MyAppConfig'], + subscriberDeclarations: [{ moduleName: 'workers.subscriber' }], + subscribers: [{ name: 'workers-subscriber' }], + workflows: [{ name: 'flows_workflows', moduleName: 'flows.workflows' }], + workflowMode: 'queue', + sitePackageDirs: [path.join('/venv/site-packages')], + }); + + expect(options.seeds).toEqual([ + 'vercel_runtime.vc_init', + path.join(workPath, 'config/wsgi.py'), + 'config.settings', + 'myapp.apps.MyAppConfig', + 'workers.subscriber', + 'flows.workflows', + '_vc_queue_handlers._py_subscribers_workers_subscriber', + '_vc_queue_handlers._py_workflows_flows__workflows', + ]); + }); + + it('searches the nested Django application root before the project root', () => { + const workPath = path.join('/work'); + const djangoPath = path.join(workPath, 'mysite'); + const sitePackages = path.join('/venv/site-packages'); + + const options = getImportClosureOptions({ + workPath, + frameworkSeeds: [], + extraPythonPath: djangoPath, + subscriberDeclarations: [], + subscribers: [], + workflows: [], + workflowMode: 'workers', + sitePackageDirs: [sitePackages], + }); + + expect(options.searchRoots).toEqual([djangoPath, workPath, sitePackages]); + }); + + it('deduplicates seeds and search roots', () => { + const workPath = path.join('/work'); + + const options = getImportClosureOptions({ + workPath, + frameworkSeeds: ['vercel_runtime.vc_init'], + extraPythonPath: workPath, + subscriberDeclarations: [{ moduleName: 'worker' }], + subscribers: [], + workflows: [{ name: 'worker', moduleName: 'worker' }], + workflowMode: 'workers', + sitePackageDirs: [workPath], + }); + + expect(options.seeds).toEqual(['vercel_runtime.vc_init', 'worker']); + expect(options.searchRoots).toEqual([workPath]); + }); +}); + +describe('withTimeout', () => { + it('resolves the value when the promise settles in time', async () => { + await expect(withTimeout(Promise.resolve(42), 1000, 'test')).resolves.toBe( + 42 + ); + }); + + it('resolves undefined when the promise exceeds the timeout', async () => { + const never = new Promise(() => {}); + await expect(withTimeout(never, 10, 'test')).resolves.toBeUndefined(); + }); + + it('rejects when the promise rejects in time', async () => { + await expect( + withTimeout(Promise.reject(new Error('boom')), 1000, 'test') + ).rejects.toThrow('boom'); + }); +}); diff --git a/packages/python/test/unit.installed-distributions.test.ts b/packages/python/test/unit.installed-distributions.test.ts index 68dda375eaae..3f1dd906e7c7 100644 --- a/packages/python/test/unit.installed-distributions.test.ts +++ b/packages/python/test/unit.installed-distributions.test.ts @@ -379,12 +379,22 @@ describe('InstalledPythonDistributions', () => { expect( await installed.collectBytecodeFiles({ vendorDirName: '_vendor' }) - ).toEqual({ files: {}, totalSize: 0, perItemSizes: new Map() }); + ).toEqual({ + files: {}, + totalSize: 0, + perItemSizes: new Map(), + items: [], + }); expect( await installed.collectPrefixBytecodeFiles({ stagingDir: '/tmp/staging', runtimeRoot: '/var/task/_vendor', }) - ).toEqual({ files: {}, totalSize: 0, perItemSizes: new Map() }); + ).toEqual({ + files: {}, + totalSize: 0, + perItemSizes: new Map(), + items: [], + }); }); }); diff --git a/packages/python/test/unit.test.ts b/packages/python/test/unit.test.ts index 2bd2415fb665..3a931a6e219c 100644 --- a/packages/python/test/unit.test.ts +++ b/packages/python/test/unit.test.ts @@ -54,6 +54,17 @@ vi.mock('execa', () => ({ default: vi.fn(), })); +// Pass-through wrapper so tests can assert whether the import closure ran. +vi.mock('@vercel/python-analysis', async () => { + const real = await vi.importActual( + '@vercel/python-analysis' + ); + return { + ...real, + collectImportClosure: vi.fn(real.collectImportClosure), + }; +}); + // Imports after mocks are set up (vitest hoists vi.mock calls) import { resolvePythonVersion, @@ -62,6 +73,7 @@ import { getInstalledPythonsFromFilesystem, } from '../src/version'; import type { PythonConstraint, PythonPackage } from '@vercel/python-analysis'; +import { collectImportClosure } from '@vercel/python-analysis'; import { build, getDevSidecars, prepareCache } from '../src/index'; import type { BuildResultV3, BuildResultV2 } from '@vercel/build-utils'; import { createVenvEnv, getVenvBinDir } from '../src/utils'; @@ -95,10 +107,12 @@ import { import { getWorkflowOutputPath } from '../src/workflows'; import { FileBlob, + FileFsRef, Span, download, sanitizeConsumerName, } from '@vercel/build-utils'; +import { derivePycPath } from '../src/compileall'; import { getServiceCrons } from '../src/crons'; import { entrypointToModule, @@ -1609,6 +1623,187 @@ describe('file exclusions', () => { }); }); +describe('bundle optimization telemetry', () => { + const originalCompileAllEnv = process.env.VERCEL_PYTHON_COMPILEALL; + const originalDisableAnalysisEnv = + process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS; + const MB = 1024 * 1024; + + afterEach(() => { + vi.mocked(execa).mockReset(); + vi.mocked(collectImportClosure).mockClear(); + if (originalCompileAllEnv === undefined) { + delete process.env.VERCEL_PYTHON_COMPILEALL; + } else { + process.env.VERCEL_PYTHON_COMPILEALL = originalCompileAllEnv; + } + if (originalDisableAnalysisEnv === undefined) { + delete process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS; + } else { + process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS = + originalDisableAnalysisEnv; + } + }); + + async function buildWithBytecode({ + payloadSize, + pycSize, + }: { + payloadSize: number; + pycSize: number; + }) { + const workPath = path.join( + tmpdir(), + `python-bundle-optimize-${Date.now()}-${Math.random().toString(36).slice(2)}` + ); + fs.mkdirSync(workPath, { recursive: true }); + + const handlerPath = path.join(workPath, 'handler.py'); + const secondaryPath = path.join(workPath, 'secondary.py'); + const payloadPath = path.join(workPath, 'payload.bin'); + fs.writeFileSync(handlerPath, 'def app(environ, start_response): pass\n'); + fs.writeFileSync(secondaryPath, 'SECONDARY = True\n'); + const payloadFd = fs.openSync(payloadPath, 'w'); + fs.ftruncateSync(payloadFd, payloadSize); + fs.closeSync(payloadFd); + + const events: any[] = []; + const span = new Span({ + name: 'vc.builder', + reporter: { report: event => events.push(event) }, + }); + const logSpy = vi.spyOn(console, 'log').mockImplementation(() => {}); + + makeMockPython('3.9'); + process.env.VERCEL_PYTHON_COMPILEALL = '1'; + vi.mocked(execa).mockImplementation((async (_file, args: string[]) => { + if (args[0]?.endsWith('vc_compileall.py')) { + const sourceFiles = JSON.parse( + fs.readFileSync(args[1], 'utf8') + ) as string[]; + for (const sourceFile of sourceFiles) { + const sourceRelPath = path + .relative(workPath, sourceFile) + .split(path.sep) + .join('/'); + const pycRelPath = derivePycPath(sourceRelPath, 3, 9); + if (!pycRelPath) continue; + const pycPath = path.join(workPath, pycRelPath); + fs.mkdirSync(path.dirname(pycPath), { recursive: true }); + const pycFd = fs.openSync(pycPath, 'w'); + fs.ftruncateSync(pycFd, pycSize); + fs.closeSync(pycFd); + } + fs.writeFileSync( + args[2], + JSON.stringify( + Object.fromEntries(sourceFiles.map(sourceFile => [sourceFile, 0.1])) + ) + ); + } + return { stdout: '', stderr: '' } as any; + }) as any); + + try { + await build({ + workPath, + files: { + 'handler.py': new FileFsRef({ fsPath: handlerPath }), + 'secondary.py': new FileFsRef({ fsPath: secondaryPath }), + 'payload.bin': new FileFsRef({ fsPath: payloadPath }), + }, + entrypoint: 'handler.py', + meta: { isDev: false, skipDownload: true }, + config: {}, + repoRootPath: workPath, + span, + }); + return { events, logSpy }; + } catch (error) { + logSpy.mockRestore(); + throw error; + } finally { + fs.removeSync(workPath); + } + } + + it('reports bytecode coverage on overflow', async () => { + const { events, logSpy } = await buildWithBytecode({ + payloadSize: 218.5 * MB, + pycSize: MB, + }); + + try { + expect(logSpy).toHaveBeenCalledWith('Optimizing Python bundle...'); + expect( + logSpy.mock.calls.filter( + ([message]) => message === 'Optimizing Python bundle...' + ) + ).toHaveLength(1); + + const optimizeSpans = events.filter( + event => event.name === 'vc.builder.python.bundle.optimize' + ); + expect(optimizeSpans).toHaveLength(1); + const bundleSpan = events.find( + event => event.name === 'vc.builder.python.bundle' + ); + expect(optimizeSpans[0].parentId).toBe(bundleSpan.id); + expect(optimizeSpans[0].tags).toEqual({ + 'python.bundle.optimize.bytecodeCoveragePercent': '50.00', + }); + } finally { + logSpy.mockRestore(); + } + }); + + it('does not report or trace optimization when bytecode fits', async () => { + const { events, logSpy } = await buildWithBytecode({ + payloadSize: 0, + pycSize: MB, + }); + + try { + expect(logSpy).not.toHaveBeenCalledWith('Optimizing Python bundle...'); + expect( + events.some(event => event.name === 'vc.builder.python.bundle.optimize') + ).toBe(false); + expect(collectImportClosure).not.toHaveBeenCalled(); + } finally { + logSpy.mockRestore(); + } + }); + + it('runs the import closure on overflow by default', async () => { + const { logSpy } = await buildWithBytecode({ + payloadSize: 218.5 * MB, + pycSize: MB, + }); + logSpy.mockRestore(); + + expect(collectImportClosure).toHaveBeenCalledTimes(1); + }); + + it('skips the import closure when VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS is set', async () => { + process.env.VERCEL_PYTHON_DISABLE_BYTECODE_ANALYSIS = '1'; + const { events, logSpy } = await buildWithBytecode({ + payloadSize: 218.5 * MB, + pycSize: MB, + }); + logSpy.mockRestore(); + + expect(collectImportClosure).not.toHaveBeenCalled(); + // Bytecode still ships, selected by size only. + const optimizeSpans = events.filter( + event => event.name === 'vc.builder.python.bundle.optimize' + ); + expect(optimizeSpans).toHaveLength(1); + expect(optimizeSpans[0].tags).toEqual({ + 'python.bundle.optimize.bytecodeCoveragePercent': '50.00', + }); + }); +}); + describe('python version selection from uv.lock and pyproject.toml', () => { let mockWorkPath: string;