mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
1ecfa1d61a
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/71072 This PR replaces the old logic of loading frozen torch through cpython by directly loading zipped torch modules directly onto deploy interpreter. We use elf file to load the zip file as its' section and load it back in the interpreter executable. Then, we directly insert the zip file into sys.path of the each initialized interpreter. Python has implicit ZipImporter module that can load modules from zip file as long as they are inside sys.path. Test Plan: buck test //caffe2/torch/csrc/deploy:test_deploy Reviewed By: shunting314 Differential Revision: D32442552 fbshipit-source-id: 627f0e91e40e72217f3ceac79002e1d8308735d5
60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
import argparse
|
|
import glob
|
|
import os
|
|
from pathlib import Path
|
|
from zipfile import ZipFile
|
|
|
|
# Exclude some standard library modules to:
|
|
# 1. Slim down the final zipped file size
|
|
# 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",
|
|
]
|
|
|
|
def remove_prefix(text, prefix):
|
|
if text.startswith(prefix):
|
|
return text[len(prefix):]
|
|
return text
|
|
|
|
def write_to_zip(file_path, strip_file_path, zf):
|
|
stripped_file_path = remove_prefix(file_path, strip_file_dir + "/")
|
|
path = Path(stripped_file_path)
|
|
if path.name in DENY_LIST:
|
|
return
|
|
zf.write(file_path, stripped_file_path)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Zip py source")
|
|
parser.add_argument("paths", nargs="*", help="Paths to zip.")
|
|
parser.add_argument("--install_dir", help="Root directory for all output files")
|
|
parser.add_argument("--strip_dir", help="The absolute directory we want to remove from zip")
|
|
parser.add_argument("--zip_name", help="Output zip name")
|
|
args = parser.parse_args()
|
|
|
|
zip_file_name = args.install_dir + '/' + args.zip_name
|
|
strip_file_dir = args.strip_dir
|
|
zf = ZipFile(zip_file_name, mode='w')
|
|
|
|
for p in args.paths:
|
|
if os.path.isdir(p):
|
|
files = glob.glob(p + "/**/*.py", recursive=True)
|
|
for file_path in files:
|
|
# strip the absolute path
|
|
write_to_zip(file_path, strip_file_dir + "/", zf)
|
|
else:
|
|
write_to_zip(p, strip_file_dir + "/", zf)
|