import io import json import logging import os import shlex import subprocess import tempfile import zipfile from pathlib import Path from typing import Dict, List, Optional, Union import torch import torch._inductor import torch.utils._pytree as pytree from torch._inductor import exc from torch._inductor.cpp_builder import BuildOptionsBase, CppBuilder from torch.export._tree_utils import reorder_kwargs from .pt2_archive_constants import AOTINDUCTOR_DIR, ARCHIVE_VERSION log = logging.getLogger(__name__) class PT2ArchiveWriter: def __init__(self, archive_path: Union[str, io.BytesIO]) -> None: self.archive_path: Union[str, io.BytesIO] = archive_path self.archive_file: Optional[zipfile.ZipFile] = None def __enter__(self) -> "PT2ArchiveWriter": assert self.archive_file is None self.archive_file = zipfile.ZipFile( self.archive_path, "w", compression=zipfile.ZIP_STORED ) self.writestr("version", str(ARCHIVE_VERSION)) self.writestr("archive_format", "pt2") return self def __exit__(self, *args) -> None: # type: ignore[no-untyped-def] assert self.archive_file is not None self.archive_file.close() self.archive_file = None return None def writestr(self, name: str, data: Union[bytes, str]) -> None: assert self.archive_file is not None self.archive_file.writestr(name, data) def write_file(self, name: str, file_path: str) -> None: """ Copy a file into the archive. name: The destination file inside the archive. file_path: The source file on disk. """ assert Path(file_path).is_file(), f"{file_path} is not a valid file path" assert self.archive_file is not None self.archive_file.write(file_path, arcname=name) class PT2ArchiveReader: def __init__(self, archive_path: str) -> None: self.archive_path: str = archive_path self.archive_file: Optional[zipfile.ZipFile] = None def __enter__(self) -> "PT2ArchiveReader": self.archive_file = zipfile.ZipFile( self.archive_path, "r", compression=zipfile.ZIP_STORED ) return self def __exit__(self, *args) -> None: # type: ignore[no-untyped-def] if self.archive_file is not None: self.archive_file.close() return None def read(self, name: str) -> bytes: assert self.archive_file is not None return self.archive_file.read(name) def extract_to_path(self, member: str, path: str) -> str: assert self.archive_file is not None return self.archive_file.extract(member, path) def extractall(self, path: str) -> None: assert self.archive_file is not None self.archive_file.extractall(path) def get_file_names(self) -> List[str]: assert self.archive_file is not None return self.archive_file.namelist() def _run_command_and_check(cmd: str) -> None: cmd = shlex.split(cmd) try: subprocess.run(cmd, check=True) except subprocess.CalledProcessError as e: raise exc.CppCompileError(cmd, e.output) from e def compile_so(aoti_dir: str, aoti_files: List[str], so_path: str) -> str: def get_aoti_file_with_suffix(suffix: str) -> str: for file in aoti_files: if file.endswith(suffix): return file raise RuntimeError(f"Unable to find file with suffix {suffix}") # Compile all the files into a .so cpp_file = os.path.join(aoti_dir, get_aoti_file_with_suffix(".cpp")) consts_o = os.path.join(aoti_dir, get_aoti_file_with_suffix(".o")) file_name = os.path.splitext(cpp_file)[0] # Parse compile flags and build the .o file with open(file_name + "_compile_flags.json") as f: compile_flags = json.load(f) compile_options = BuildOptionsBase(**compile_flags) object_builder = CppBuilder( name=file_name, sources=cpp_file, BuildOption=compile_options, ) compile_cmd = object_builder.get_command_line() output_o = object_builder.get_target_file_path() _run_command_and_check(compile_cmd) # Parse linker flags and build the .so file with open(file_name + "_linker_flags.json") as f: linker_flags = json.load(f) linker_options = BuildOptionsBase(**linker_flags) so_builder = CppBuilder( name=os.path.split(so_path)[-1], sources=[output_o, consts_o], BuildOption=linker_options, output_dir=so_path, ) link_cmd = so_builder.get_command_line() output_so = so_builder.get_target_file_path() _run_command_and_check(link_cmd) # mmapped weights serialized_weights_filename = file_name + "_serialized_weights.bin" if serialized_weights_filename in aoti_files: with open(serialized_weights_filename, "rb") as f_weights: serialized_weights = f_weights.read() with open(output_so, "a+b") as f_so: so_size = f_so.tell() # Page align the weights f_so.write(b" " * (16384 - so_size % 16384)) f_so.write(serialized_weights) return output_so def package_aoti( archive_file: Union[str, io.BytesIO], aoti_files: Union[List[str], Dict[str, List[str]]], ) -> Union[str, io.BytesIO]: """ Saves the AOTInductor generated files to the PT2Archive format. Args: archive_file: The file name to save the package to. aoti_files: This can either be a singular path to a directory containing the AOTInductor files, or a dictionary mapping the model name to the path to its AOTInductor generated files. """ if isinstance(aoti_files, list): aoti_files = {"model": aoti_files} assert isinstance(aoti_files, dict), ( "Please pass a list of AOTI generated files to be packaged or " "a dictionary mapping model names to their list of AOTI generated " "files. You can get this list of files through calling " "`torch._inductor.aot_compile(..., options={aot_inductor.package=True})`" ) assert isinstance(archive_file, io.BytesIO) or ( isinstance(archive_file, str) and archive_file.endswith(".pt2") ), f"Expect archive file to be a file ending in .pt2, or is a buffer. Instead got {archive_file}" # Save using the PT2 packaging format # (https://docs.google.com/document/d/1jLPp8MN8Whs0-VW9PmJ93Yg02W85tpujvHrTa1pc5x8/edit#heading=h.v2y2jgnwc56a) with PT2ArchiveWriter(archive_file) as archive_writer: for model_name, files in aoti_files.items(): num_so_files = 0 num_cpp_files = 0 for file in files: if file == "": continue if file.endswith(".so"): num_so_files += 1 if num_so_files > 1: raise RuntimeError( f"Multiple .so files found in {files}. " "You might need to clear your cache " "directory before calling aoti_compile again." ) if file.endswith(".cpp"): num_cpp_files += 1 if num_so_files > 1: raise RuntimeError( f"Multiple .cpp files found in {files}. " "You might need to clear your cache " "directory before calling aoti_compile again." ) filename = os.path.basename(file) new_filepath = os.path.join(AOTINDUCTOR_DIR, model_name, filename) log.debug( "Saving AOTI generated file %s to archive in %s", file, new_filepath ) archive_writer.write_file( str(new_filepath), file, ) if isinstance(archive_file, io.BytesIO): archive_file.seek(0) return archive_file class AOTICompiledModel: """ Callable AOT Inductor loaded model from a .pt2 """ def __init__(self, loader: torch._C._aoti.AOTIModelPackageLoader) -> None: self.loader = loader def __call__(self, *args, **kwargs): # type: ignore[no-untyped-def] call_spec = self.loader.get_call_spec() # type: ignore[attr-defined] in_spec = pytree.treespec_loads(call_spec[0]) out_spec = pytree.treespec_loads(call_spec[1]) flat_inputs = pytree.tree_flatten((args, reorder_kwargs(kwargs, in_spec)))[0] flat_inputs = [x for x in flat_inputs if isinstance(x, torch.Tensor)] flat_outputs = self.loader.run(flat_inputs) # type: ignore[attr-defined] return pytree.tree_unflatten(flat_outputs, out_spec) def get_metadata(self) -> Dict[str, str]: return self.loader.get_metadata() # type: ignore[attr-defined] def load_constants( self, constants_map: Dict[str, torch.Tensor], *, check_full_update: bool, ) -> None: """ Given a mapping of constant fqns to tensors, load the constants into the model. You can use ``get_constant_fqns`` to get the list of constant fqns that are needed in the compiled model. Args: constants_map: A mapping of constant fqns to tensors. check_full_update: Whether to add check to see if all the constants are updated and have values. """ self.loader.load_constants(constants_map, False, check_full_update) # type: ignore[attr-defined] def get_constant_fqns(self) -> List[str]: return self.loader.get_constant_fqns() # type: ignore[attr-defined] def load_package(path: Union[str, io.BytesIO], model_name: str = "model") -> AOTICompiledModel: # type: ignore[type-arg] assert isinstance(path, io.BytesIO) or ( isinstance(path, str) and path.endswith(".pt2") ), f"Unable to load package. Path must be a buffer or a file ending in .pt2. Instead got {path}" if isinstance(path, io.BytesIO): with tempfile.NamedTemporaryFile(suffix=".pt2") as f: # TODO(angelayi): We shouldn't need to do this -- miniz should # handle reading the buffer. This is just a temporary workaround f.write(path.read()) path.seek(0) log.debug("Writing buffer to tmp file located at %s.", f.name) loader = torch._C._aoti.AOTIModelPackageLoader(f.name, model_name) # type: ignore[call-arg] return AOTICompiledModel(loader) loader = torch._C._aoti.AOTIModelPackageLoader(path, model_name) # type: ignore[call-arg] return AOTICompiledModel(loader)