Function parse_etrecord

devtools/etrecord/_etrecord.py:744–929

Parses an `ETRecord` file and returns an `ETRecord` object that contains the deserialized graph modules, program buffer, and a debug handle map. In the graph map in the returned `ETRecord` object if a model with multiple entry points was provided originally by the user during `ETRec

(etrecord_path: str)

def parse_etrecord(etrecord_path: str) -> ETRecord:  # noqa: C901
    """
    Parses an `ETRecord` file and returns an `ETRecord` object that contains the deserialized graph
    modules, program buffer, and a debug handle map.
    In the graph map in the returned `ETRecord` object if a model with multiple entry points was provided
    originally by the user during `ETRecord` generation then each entry point will be stored as a separate
    graph module in the `ETRecord` object with the name being `the original module name + "/" + the
    name of the entry point`.

    Args:
        etrecord_path: Path to the `ETRecord` file.

    Returns:
        `ETRecord` object.
    """

    try:
        etrecord_zip = ZipFile(etrecord_path, "r")
    except BadZipFile:
        raise RuntimeError("Invalid etrecord file passed in.")

    file_list = etrecord_zip.namelist()

    if ETRecordReservedFileNames.ETRECORD_IDENTIFIER not in file_list:
        raise RuntimeError(
            "ETRecord identifier missing from etrecord file passed in. Either an invalid file was passed in or the file is corrupt."
        )

    graph_map: Dict[str, ExportedProgram] = {}
    edge_dialect_programs: Dict[str, ExportedProgram] = {}
    debug_handle_map = None
    delegate_map = None
    instruction_id_to_num_outs_map = None
    exported_program = None
    edge_dialect_program: Optional[
        Union[ExportedProgram, Dict[str, ExportedProgram]]
    ] = None
    reference_outputs = None
    representative_inputs = None
    export_graph_id = 0

    serialized_exported_program_files = set()
    serialized_edge_dialect_program_files = set()
    serialized_state_dict_files = set()
    serialized_constants_files = set()
    serialized_example_inputs_files = set()

    edge_dialect_prefix = f"{ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM}/"

    for entry in file_list:
        if entry == ETRecordReservedFileNames.DEBUG_HANDLE_MAP_NAME:
            debug_handle_map = json.loads(
                etrecord_zip.read(ETRecordReservedFileNames.DEBUG_HANDLE_MAP_NAME)
            )
        elif entry == ETRecordReservedFileNames.DELEGATE_MAP_NAME:
            delegate_map = json.loads(
                etrecord_zip.read(ETRecordReservedFileNames.DELEGATE_MAP_NAME)
            )