pytorch/torch/utils/_zip.py
Sahan Paliskara 3921ed31f3 Use multipy.package in multipy/runtime (#111) (#82690)
Summary:
This change points multipy::runtime to use multipy.package instead of torch.package by copying `_deploy.py` (which is used in order to pass objects in and out of interpreters) into multipy as well as making the neccessary changes to allow `multipy::runtime` to access `multipy.package` and `_deploy.py`.

X-link: https://github.com/pytorch/multipy/pull/111

Reviewed By: d4l3k

Differential Revision: D38337551

Pulled By: PaliC

Pull Request resolved: https://github.com/pytorch/pytorch/pull/82690
Approved by: https://github.com/d4l3k
2022-08-03 19:11:12 +00:00

63 lines
2.1 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, prepend_str=""):
stripped_file_path = prepend_str + 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("--prepend_str", help="A string to prepend onto all paths of a file in the zip", default="")
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
prepend_str = args.prepend_str
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, prepend_str=prepend_str)
else:
write_to_zip(p, strip_file_dir + "/", zf, prepend_str=prepend_str)