File size: 9,993 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 |
# mypy: allow-untyped-decorators
# mypy: allow-untyped-defs
"""
Freeze Python packages.
Freezing makes it possible to ship arbitrary Python modules as part of a C++
library. The Python source of the module is compiled to bytecode and written
to `.c` files, to be imported by Python's built-in FrozenImporter.
In a normal Python installation, FrozenImporter is only used to bootstrap the
initialization of the import machinery. Python's importers are defined in
Python (see `_bootstrap.py` and `_bootstrap_external.py`) but need to be
retrieved before any importers are available. Freezing the module bytecode
resolves this circular dependency.
This script will freeze the Python standard library. It produces two things:
- Bytecode files: A set of `.c` that define C variables containing Python bytecode.
- Main file: A `main.c` file listing all of these modules in the right form to be
consumed by FrozenImporter.
The library that wishes to these modules make them available to the local
Python instance by extending `PyImport_FrozenModules` appropriately (see
https://docs.python.org/3/c-api/import.html#c.PyImport_FrozenModules).
"""
import argparse
import functools
import itertools
import marshal
import os
import types
from dataclasses import dataclass
from pathlib import Path
PATH_MARKER = "<Generated by torch::deploy>"
MAIN_INCLUDES = """#include <Python.h>
"""
MAIN_PREFIX_TEMPLATE = """
// Compiled standard library modules. These should be appended to the existing
// `PyImport_FrozenModules` that ships with CPython.
struct _frozen {}[] = {{
"""
FAKE_PREFIX = MAIN_PREFIX_TEMPLATE.format("_PyImport_FrozenModules")
MAIN_SUFFIX = """\
{0, 0, 0} /* sentinel */
};
"""
# Exclude some standard library modules to:
# 1. Slim down the final frozen lib.
# 2. Remove functionality we don't want to support.
DENY_LIST = [
# Interface to unix databases
"dbm",
# ncurses bindings (terminal interfaces)
"curses",
# Tcl/Tk GUI
"tkinter",
"tkinter",
# Tests for the standard library
"test",
"tests",
"idle_test",
"__phello__.foo.py",
# importlib frozen modules. These are already baked into CPython.
"_bootstrap.py",
"_bootstrap_external.py",
]
NUM_BYTECODE_FILES = 5
def indent_msg(fn):
@functools.wraps(fn)
def wrapper(*args, **kwargs):
args[0].indent += 1
ret = fn(*args, **kwargs)
args[0].indent -= 1
return ret
return wrapper
@dataclass
class FrozenModule:
# The fully qualified module name, e.g. 'foo.bar.baz'
module_name: str
# The name of the C variable that holds the bytecode, e.g. 'M_foo__bar__baz'
c_name: str
# The size of the C variable. Negative if this module is a package.
size: int
# The frozen bytecode
bytecode: bytes
class Freezer:
def __init__(self, verbose: bool):
self.frozen_modules: list[FrozenModule] = []
self.indent: int = 0
self.verbose: bool = verbose
def msg(self, path: Path, code: str):
if not self.verbose:
return
# P: package dir
# F: python file
# S: skipped (not a package dir)
# X: skipped (deny-listed)
# N: skipped (not a python file)
print(" " * self.indent, end="")
print(f"{code} {path}")
def write_bytecode(self, install_root):
"""
Write the `.c` files containing the frozen bytecode.
Shared frozen modules evenly across the files.
"""
bytecode_file_names = [f"bytecode_{i}.c" for i in range(NUM_BYTECODE_FILES)]
bytecode_files = [
open(os.path.join(install_root, name), "w") for name in bytecode_file_names
]
it = itertools.cycle(bytecode_files)
for m in self.frozen_modules:
self.write_frozen(m, next(it))
for f in bytecode_files:
f.close()
def write_main(self, install_root, oss, symbol_name):
"""Write the `main.c` file containing a table enumerating all the frozen modules."""
with open(os.path.join(install_root, "main.c"), "w") as outfp:
outfp.write(MAIN_INCLUDES)
for m in self.frozen_modules:
outfp.write(f"extern unsigned char {m.c_name}[];\n")
outfp.write(MAIN_PREFIX_TEMPLATE.format(symbol_name))
for m in self.frozen_modules:
outfp.write(f'\t{{"{m.module_name}", {m.c_name}, {m.size}}},\n')
outfp.write(MAIN_SUFFIX)
if oss:
outfp.write(FAKE_PREFIX)
outfp.write(MAIN_SUFFIX)
def write_frozen(self, m: FrozenModule, outfp):
"""Write a single frozen module's bytecode out to a C variable."""
outfp.write(f"unsigned char {m.c_name}[] = {{")
for i in range(0, len(m.bytecode), 16):
outfp.write("\n\t")
for c in bytes(m.bytecode[i : i + 16]):
outfp.write(f"{c:d},")
outfp.write("\n};\n")
def compile_path(self, path: Path, top_package_path: Path):
"""Entry point for compiling a Path object."""
if path.is_dir():
self.compile_package(path, top_package_path)
else:
self.compile_file(path, top_package_path)
@indent_msg
def compile_package(self, path: Path, top_package_path: Path):
"""Compile all the files within a Python package dir."""
assert path.is_dir()
if path.name in DENY_LIST:
self.msg(path, "X")
return
# Python packages are directories that have __init__.py in them.
is_package_dir = any(child.name == "__init__.py" for child in path.iterdir())
if not is_package_dir:
self.msg(path, "S")
return
self.msg(path, "P")
# Recursively compile all children in this dir
for child in path.iterdir():
self.compile_path(child, top_package_path)
def get_module_qualname(self, file_path: Path, top_package_path: Path) -> list[str]:
# `path` looks like 'Lib/foo/bar/baz.py'
# chop off 'Lib/' to get something that represents a Python module hierarchy.
# e.g. 'foo/bar/baz.py', which maps to 'foo.bar.baz'
normalized_path = file_path.relative_to(top_package_path.parent)
if normalized_path.name == "__init__.py":
# Special handling for `__init__.py`. In this case, this file
# specifies that the containing directory should be treated as a package.
# For 'foo/bar/baz/__init__.py':
# - The module name is 'baz'
module_basename = normalized_path.parent.name
# - The parent is foo.bar (need to shave off the 'baz')
module_parent = normalized_path.parent.parent.parts
else:
module_basename = normalized_path.stem
module_parent = normalized_path.parent.parts
return list(module_parent) + [module_basename]
def compile_string(self, file_content: str) -> types.CodeType:
# instead of passing in the real build time path to 'compile', we
# pass in a marker instead. This prevents the build time path being
# leaked to runtime. That path may not be available at runtime.
# Setting the path to a mark make sure it's a hard error rather
# than a flaky error when inspect module tries to retrieve python source
# code during torchscripting.
path_marker = PATH_MARKER
return compile(file_content, path_marker, "exec")
@indent_msg
def compile_file(self, path: Path, top_package_path: Path):
"""
Compile a Python source file to frozen bytecode.
Append the result to `self.frozen_modules`.
"""
assert path.is_file()
if path.suffix != ".py":
self.msg(path, "N")
return
if path.name in DENY_LIST:
self.msg(path, "X")
return
self.msg(path, "F")
module_qualname = self.get_module_qualname(path, top_package_path)
module_mangled_name = "__".join(module_qualname)
c_name = "M_" + module_mangled_name
with open(path) as src_file:
co = self.compile_string(src_file.read())
bytecode = marshal.dumps(co)
size = len(bytecode)
if path.name == "__init__.py":
# Python packages are signified by negative size.
size = -size
self.frozen_modules.append(
FrozenModule(".".join(module_qualname), c_name, size, bytecode)
)
def main() -> None:
parser = argparse.ArgumentParser(description="Compile py source")
parser.add_argument("paths", nargs="*", help="Paths to freeze.")
parser.add_argument("--verbose", action="store_true", help="Print debug logs")
parser.add_argument(
"--install-dir", "--install_dir", help="Root directory for all output files"
)
parser.add_argument(
"--oss",
action="store_true",
help="If it's OSS build, add a fake _PyImport_FrozenModules",
)
parser.add_argument(
"--symbol-name",
"--symbol_name",
help="The name of the frozen module array symbol to generate",
default="_PyImport_FrozenModules_torch",
)
args = parser.parse_args()
f = Freezer(args.verbose)
for p in args.paths:
path = Path(p)
if path.is_dir() and not Path.exists(path / "__init__.py"):
# this 'top level path p' is a standard directory containing modules,
# not a module itself
# each 'mod' could be a dir containing __init__.py or .py file
# NB: sorted to make sure this is deterministic
for mod in sorted(path.glob("*")):
f.compile_path(mod, mod)
else:
f.compile_path(path, path)
f.write_bytecode(args.install_dir)
f.write_main(args.install_dir, args.oss, args.symbol_name)
if __name__ == "__main__":
main() # pragma: no cover
|