_etrecord.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# pyre-unsafe

import io
import json
import os
from typing import BinaryIO, Dict, IO, List, Optional, Union
from zipfile import BadZipFile, ZipFile

import torch

from executorch import exir

from executorch.devtools.bundled_program.config import ConfigValue
from executorch.devtools.bundled_program.core import BundledProgram
from executorch.exir import (
    EdgeProgramManager,
    ExecutorchProgram,
    ExecutorchProgramManager,
    ExirExportedProgram,
    ExportedProgram,
)
from executorch.exir.emit._emitter import _DelegateDebugIdentifierMap

from executorch.exir.serde.export_serialize import SerializedArtifact
from executorch.exir.serde.serialize import deserialize, serialize

ProgramInput = ConfigValue
ProgramOutput = torch.Tensor

try:
    # breaking change introduced in python 3.11
    # pyre-ignore
    from enum import StrEnum
except ImportError:
    from enum import Enum

    class StrEnum(str, Enum):
        pass


class ETRecordReservedFileNames(StrEnum):
    ETRECORD_IDENTIFIER = "ETRECORD_V0"
    EXPORTED_PROGRAM = "exported_program"
    EXPORT_GRAPH_ID = "export_graph_id"
    EDGE_DIALECT_EXPORTED_PROGRAM = "edge_dialect_exported_program"
    ET_DIALECT_GRAPH_MODULE = "et_dialect_graph_module"
    DEBUG_HANDLE_MAP_NAME = "debug_handle_map"
    DELEGATE_MAP_NAME = "delegate_map"
    INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME = "instruction_id_to_num_outs_map"
    REFERENCE_OUTPUTS = "reference_outputs"
    REPRESENTATIVE_INPUTS = "representative_inputs"


class ETRecord:
    def __init__(
        self,
        exported_program: Optional[ExportedProgram] = None,
        export_graph_id: Optional[int] = None,
        edge_dialect_program: Optional[
            Union[ExportedProgram, Dict[str, ExportedProgram]]
        ] = None,
        graph_map: Optional[Dict[str, ExportedProgram]] = None,
        _debug_handle_map: Optional[Dict[int, Union[int, List[int]]]] = None,
        _delegate_map: Optional[
            Dict[str, Dict[int, Dict[str, Union[str, _DelegateDebugIdentifierMap]]]]
        ] = None,
        _instruction_id_to_num_outs_map: Optional[
            Dict[str, Dict[int, Union[int, List[int]]]]
        ] = None,
        _reference_outputs: Optional[Dict[str, List[ProgramOutput]]] = None,
        _representative_inputs: Optional[List[ProgramInput]] = None,
    ):
        """
        Please do not construct an ETRecord object directly.

        If you want to create an ETRecord for logging AOT information to further analysis, please mark `generate_etrecord`
        as True in your export api, and get the ETRecord object from the `ExecutorchProgramManager`.
        For exmaple:
        ```python
            exported_program = torch.export.export(model, inputs)
            edge_program = to_edge_transform_and_lower(exported_program, generate_etrecord=True)
            executorch_program = edge_program.to_executorch()
            etrecord = executorch_program.get_etrecord()
        ```

        If user need to create an ETRecord manually, please use the `create_etrecord` function.

        **EXPERIMENTAL**: This API supports multiple methods. For example:
        ```python
        lowered_and_edge = to_edge_transform_and_lower(
            {
                "vision_encoder": vision_encoder_ep,
                "token_embedding": token_embedding_ep,
                "text_decoder": causal_llm_ep,
            },
            partitioner={
                "vision_encoder": [XnnpackPartitioner()],
                "token_embedding": [XnnpackPartitioner()],
                "text_decoder": [
                    XnnpackPartitioner(
                        config_precisions=ConfigPrecisionType.DYNAMIC_QUANT,
                        per_op_mode=True,
                    ),
                    XnnpackPartitioner(),
                ],
            },
            compile_config=EdgeCompileConfig(_check_ir_validity=False),
            constant_methods=manager.metadata,
            generate_etrecord=True,  # Enable ETRecord generation for all 3 methods
        )
        ```
        """

        self.exported_program = exported_program
        self.export_graph_id = export_graph_id
        self.edge_dialect_program = edge_dialect_program
        self.graph_map = graph_map
        self.__debug_handle_map = _debug_handle_map  # Use private attribute
        self._delegate_map = _delegate_map
        self._instruction_id_to_num_outs_map = _instruction_id_to_num_outs_map
        self._reference_outputs = _reference_outputs
        self._representative_inputs = _representative_inputs

    @property
    def _debug_handle_map(self):
        return self.__debug_handle_map

    @_debug_handle_map.setter
    def _debug_handle_map(self, value):
        self.__debug_handle_map = value

    def save(self, path: Union[str, os.PathLike, BinaryIO, IO[bytes]]) -> None:
        """
        Serialize and save the ETRecord to the specified path for use in Inspector. The ETRecord
        should contains at least edge dialect program and executorch program information for further
        analysis, otherwise it will raise an exception.

        Args:
            path: Path where the ETRecord file will be saved to.

        Raises:
            RuntimeError: If the ETRecord does not contain essential information for Inpector.
        """
        if isinstance(path, (str, os.PathLike)):
            # pyre-ignore[6]: In call `os.fspath`, for 1st positional argument, expected `str` but got `Union[PathLike[typing.Any], str]`
            path = os.fspath(path)

        if not (self.edge_dialect_program and self._debug_handle_map):
            raise RuntimeError(
                "ETRecord must contain edge dialect program and executorch program to be saved"
            )

        # Normalize edge_dialect_program to dict format for consistent handling
        if isinstance(self.edge_dialect_program, ExportedProgram):
            self._edge_dialect_programs_dict: Dict[str, ExportedProgram] = {
                "forward": self.edge_dialect_program
            }
        else:
            self._edge_dialect_programs_dict = self.edge_dialect_program

        etrecord_zip = ZipFile(path, "w")

        try:
            self._write_identifier(etrecord_zip)
            self._save_programs(etrecord_zip)
            self._save_graph_map(etrecord_zip)
            self._save_metadata(etrecord_zip)
        finally:
            etrecord_zip.close()

    def _write_identifier(self, etrecord_zip: ZipFile) -> None:
        """Write the magic file identifier."""
        etrecord_zip.writestr(ETRecordReservedFileNames.ETRECORD_IDENTIFIER, "")

    def _save_programs(self, etrecord_zip: ZipFile) -> None:
        """Save exported program and edge dialect program(s)."""
        if self.exported_program is not None:
            self._save_exported_program(
                etrecord_zip,
                ETRecordReservedFileNames.EXPORTED_PROGRAM,
                "",
                self.exported_program,
            )

        # Save all edge dialect programs (supports multiple methods)
        for method_name, edge_program in self._edge_dialect_programs_dict.items():
            self._save_edge_dialect_program(etrecord_zip, method_name, edge_program)

    def _save_graph_map(self, etrecord_zip: ZipFile) -> None:
        """Save graph map if present."""
        if self.graph_map is not None:
            # pyre-ignore[16]: Undefined attribute [16]: `Optional` has no attribute `items`.
            for module_name, export_module in self.graph_map.items():
                if "/" in module_name:
                    base_name, method_name = module_name.rsplit("/", 1)
                    self._save_exported_program(
                        etrecord_zip, base_name, method_name, export_module
                    )
                else:
                    self._save_exported_program(
                        etrecord_zip, module_name, "forward", export_module
                    )

    def _save_metadata(self, etrecord_zip: ZipFile) -> None:
        """Save debug maps, reference outputs, and other metadata."""
        if self._debug_handle_map is not None:
            etrecord_zip.writestr(
                ETRecordReservedFileNames.DEBUG_HANDLE_MAP_NAME,
                json.dumps(self._debug_handle_map),
            )

        if self._delegate_map is not None:
            etrecord_zip.writestr(
                ETRecordReservedFileNames.DELEGATE_MAP_NAME,
                json.dumps(self._delegate_map),
            )

        if self._instruction_id_to_num_outs_map is not None:
            etrecord_zip.writestr(
                ETRecordReservedFileNames.INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME,
                json.dumps(self._instruction_id_to_num_outs_map),
            )

        if self._reference_outputs is not None:
            buf = io.BytesIO()
            torch.save(self._reference_outputs, buf)
            etrecord_zip.writestr(
                ETRecordReservedFileNames.REFERENCE_OUTPUTS,
                buf.getvalue(),
            )

        if self._representative_inputs is not None:
            buf = io.BytesIO()
            torch.save(self._representative_inputs, buf)
            etrecord_zip.writestr(
                ETRecordReservedFileNames.REPRESENTATIVE_INPUTS,
                buf.getvalue(),
            )

        if self.export_graph_id is not None:
            etrecord_zip.writestr(
                ETRecordReservedFileNames.EXPORT_GRAPH_ID,
                json.dumps(self.export_graph_id),
            )

    def copy(self) -> "ETRecord":
        """
        Create a shallow copy of this ETRecord suitable for transformation stages.

        This creates a new ETRecord instance with the same edge-stage data but without
        executorch-specific data. Useful when transitioning between transformation stages
        while preserving immutable snapshot semantics.

        Returns:
            ETRecord: A new ETRecord with edge-stage data copied over.
        """
        return ETRecord(
            exported_program=self.exported_program,
            export_graph_id=self.export_graph_id,
            edge_dialect_program=self.edge_dialect_program,
            graph_map=self.graph_map,
            # Explicitly exclude executorch-specific fields for clean separation
            _debug_handle_map=None,
            _delegate_map=None,
            _instruction_id_to_num_outs_map=None,
            _reference_outputs=None,
            _representative_inputs=None,
        )

    def _save_exported_program(
        self,
        etrecord_zip: ZipFile,
        module_name: str,
        method_name: str,
        ep: ExportedProgram,
    ) -> None:
        """Save an exported program to the ETRecord zip file."""
        serialized_artifact = serialize(ep)
        assert isinstance(serialized_artifact.exported_program, bytes)

        method_name = f"/{method_name}" if method_name != "" else ""
        base_name = f"{module_name}{method_name}"

        etrecord_zip.writestr(base_name, serialized_artifact.exported_program)
        etrecord_zip.writestr(f"{base_name}_state_dict", serialized_artifact.state_dict)
        etrecord_zip.writestr(f"{base_name}_constants", serialized_artifact.constants)
        etrecord_zip.writestr(
            f"{base_name}_example_inputs", serialized_artifact.example_inputs
        )

    def _save_edge_dialect_program(
        self,
        etrecord_zip: ZipFile,
        method_name: str,
        edge_dialect_program: ExportedProgram,
    ) -> None:
        """Save the edge dialect program to the ETRecord zip file."""
        serialized_artifact = serialize(edge_dialect_program)
        assert isinstance(serialized_artifact.exported_program, bytes)

        # Use format: edge_dialect_exported_program/method_name for multi-method support
        base_name = (
            f"{ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM}/{method_name}"
        )
        etrecord_zip.writestr(base_name, serialized_artifact.exported_program)
        etrecord_zip.writestr(f"{base_name}_state_dict", serialized_artifact.state_dict)
        etrecord_zip.writestr(f"{base_name}_constants", serialized_artifact.constants)
        etrecord_zip.writestr(
            f"{base_name}_example_inputs", serialized_artifact.example_inputs
        )

    def add_extra_export_modules(
        self,
        extra_recorded_export_modules: Dict[
            str,
            Union[
                ExportedProgram,
                ExirExportedProgram,
                EdgeProgramManager,
            ],
        ],
    ) -> None:
        """
        Add extra export modules to the ETRecord after it has been created.

        This method allows users to add more export modules they want to record
        to an existing ETRecord instance. The modules will be added to the graph_map
        and will be included when the ETRecord is saved.

        Args:
            extra_recorded_export_modules: A dictionary of graph modules with the key being
                the user provided name and the value being the corresponding exported module.
                The exported graph modules can be either the output of `torch.export()` or `exir.to_edge()`.
        """
        if self.graph_map is None:
            self.graph_map = {}

        # Now self.graph_map is guaranteed to be non-None
        graph_map = self.graph_map
        for module_name, export_module in extra_recorded_export_modules.items():
            _add_module_to_graph_map(graph_map, module_name, export_module)

    def add_executorch_program(
        self,
        executorch_program: Union[
            ExecutorchProgram,
            ExecutorchProgramManager,
            BundledProgram,
        ],
    ) -> None:
        """
        Add executorch program data to the ETRecord after it has been created.

        This method allows users to add executorch program data they want to record
        to an existing ETRecord instance. The executorch program data includes debug handle map,
        delegate map, reference outputs, and representative inputs that will be included
        when the ETRecord is saved.

        Args:
            executorch_program: The ExecuTorch program for this model returned by the call to
                `to_executorch()` or the `BundledProgram` of this model.

        Raises:
            RuntimeError: If executorch program data already exists in the ETRecord.
        """
        # Check if executorch program data already exists
        if (
            self._debug_handle_map is not None
            or self._delegate_map is not None
            or self._instruction_id_to_num_outs_map is not None
            or self._reference_outputs is not None
            or self._representative_inputs is not None
        ):
            raise RuntimeError(
                "Executorch program data already exists in the ETRecord. "
                "Cannot add executorch program data when it already exists."
            )

        # Process executorch program and extract data
        (
            debug_handle_map,
            delegate_map,
            instruction_id_to_num_outs_map,
            reference_outputs,
            representative_inputs,
        ) = _process_executorch_program(executorch_program)

        # Set the extracted data
        self._debug_handle_map = debug_handle_map
        self._delegate_map = delegate_map
        self._instruction_id_to_num_outs_map = instruction_id_to_num_outs_map
        self._reference_outputs = reference_outputs
        self._representative_inputs = representative_inputs

    def add_exported_program(
        self,
        exported_program: Optional[Union[ExportedProgram, Dict[str, ExportedProgram]]],
    ) -> None:
        """
        Add exported program to the ETRecord after it has been created.

        This method allows users to add an exported program they want to record
        to an existing ETRecord instance. The exported program will be included
        when the ETRecord is saved.

        Args:
            exported_program: The exported program for this model returned by the call to
                `torch.export()` or a dictionary with method names as keys and exported programs as values.
                Can be None, in which case no exported program data will be added.

        Raises:
            RuntimeError: If exported program already exists in the ETRecord.
        """
        # Check if exported program already exists
        if self.exported_program is not None or self.export_graph_id is not None:
            raise RuntimeError(
                "Exported program already exists in the ETRecord. "
                "Cannot add exported program when it already exists."
            )

        # Process exported program and extract data
        processed_exported_program, export_graph_id = _process_exported_program(
            exported_program
        )

        # Set the extracted data
        self.exported_program = processed_exported_program
        self.export_graph_id = export_graph_id

    def add_edge_dialect_program(
        self,
        edge_dialect_program: Union[EdgeProgramManager, ExirExportedProgram],
    ) -> None:
        """
        Add edge dialect program to the ETRecord after it has been created.

        This method allows users to add an edge dialect program they want to record
        to an existing ETRecord instance. The edge dialect program will be included
        when the ETRecord is saved.

        Args:
            edge_dialect_program: The edge dialect program for this model returned by the call to
                `to_edge()` or `EdgeProgramManager` for this model.

        Raises:
            RuntimeError: If edge dialect program already exists in the ETRecord.
        """
        # Check if edge dialect program already exists
        if self.edge_dialect_program is not None:
            raise RuntimeError(
                "Edge dialect program already exists in the ETRecord. "
                "Cannot add edge dialect program when it already exists."
            )

        # Process edge dialect program and extract data

        processed_edge_dialect_program = _process_edge_dialect_program(
            edge_dialect_program
        )

        # Set the extracted data
        self.edge_dialect_program = processed_edge_dialect_program

    def update_representative_inputs(
        self,
        representative_inputs: Union[List[ProgramInput], BundledProgram],
    ) -> None:
        """
        Update the representative inputs in the ETRecord.

        This method allows users to customize the representative inputs that will be
        included when the ETRecord is saved. The representative inputs can be provided
        directly as a list or extracted from a BundledProgram.

        Args:
            representative_inputs: Either a list of ProgramInput objects or a BundledProgram
                from which representative inputs will be extracted.
        """
        if isinstance(representative_inputs, BundledProgram):
            self._representative_inputs = _get_representative_inputs(
                representative_inputs
            )
        else:
            self._representative_inputs = representative_inputs

    def update_reference_outputs(
        self,
        reference_outputs: Union[
            Dict[str, List[ProgramOutput]], List[ProgramOutput], BundledProgram
        ],
    ) -> None:
        """
        Update the reference outputs in the ETRecord.

        This method allows users to customize the reference outputs that will be
        included when the ETRecord is saved. The reference outputs can be provided
        directly as a dictionary mapping method names to lists of outputs, as a
        single list of outputs (which will be treated as {"forward": List[ProgramOutput]}),
        or extracted from a BundledProgram.

        Args:
            reference_outputs: Either a dictionary mapping method names to lists of
                ProgramOutput objects, a single list of ProgramOutput objects (treated
                as outputs for the "forward" method), or a BundledProgram from which
                reference outputs will be extracted.
        """
        if isinstance(reference_outputs, BundledProgram):
            self._reference_outputs = _get_reference_outputs(reference_outputs)
        elif isinstance(reference_outputs, list):
            self._reference_outputs = {"forward": reference_outputs}
        else:
            self._reference_outputs = reference_outputs


def _get_reference_outputs(
    bundled_program: BundledProgram,
) -> Dict[str, List[ProgramOutput]]:
    """
    Extracts out the expected outputs from the bundled program, keyed by the method names.
    """
    reference_outputs = {}
    for method_test_suite in bundled_program.method_test_suites:
        reference_outputs[method_test_suite.method_name] = []
        for test_case in method_test_suite.test_cases:
            if not test_case.expected_outputs:
                raise ValueError(
                    f"Missing at least one set of expected outputs for method {method_test_suite.method_name}."
                )
            reference_outputs[method_test_suite.method_name].append(
                test_case.expected_outputs
            )
    return reference_outputs


def _get_representative_inputs(
    bundled_program: BundledProgram,
) -> Optional[List[ProgramInput]]:
    """
    Extracts out the inputs from the bundled program, keyed by the method names.
    """
    for method_test_suite in bundled_program.method_test_suites:
        if method_test_suite.method_name == "forward":
            if not method_test_suite.test_cases:
                raise ValueError(
                    "The 'forward' method is defined, but no corresponding input test cases are provided."
                )
            # Get first example input from the forward method
            test_case = method_test_suite.test_cases[0]
            return test_case.inputs

    # If the forward method is not defined, return None to indicate that there are no representative inputs for the model.
    return None


def generate_etrecord(
    et_record: Union[str, os.PathLike, BinaryIO, IO[bytes]],
    edge_dialect_program: Union[EdgeProgramManager, ExirExportedProgram],
    executorch_program: Union[
        ExecutorchProgram,
        ExecutorchProgramManager,
        BundledProgram,
    ],
    exported_program: Optional[
        Union[ExportedProgram, Dict[str, ExportedProgram]]
    ] = None,
    extra_recorded_export_modules: Optional[
        Dict[
            str,
            Union[
                ExportedProgram,
                ExirExportedProgram,
                EdgeProgramManager,
            ],
        ]
    ] = None,
) -> None:
    """
    Generates an `ETRecord` from the given objects, serializes it and saves it to the given path.
    The objects that will be serialized to an `ETRecord` are all the graph modules present
    in the `extra_recorded_export_modules` dict, the graph module present in the edge dialect program object,
    and also the graph module present in the ExecuTorch program object, which
    is the closest graph module representation of what is eventually run on the device.
    In addition to all the graph modules, we also serialize the program buffer, which the users
    can provide to the ExecuTorch runtime to run the model, and the debug handle map
    for Developer Tools usage.

    Args:
        et_record: Path to where the `ETRecord` file will be saved to.
        edge_dialect_program: `EdgeProgramManager` for this model returned by the call to to_edge()
        executorch_program: The ExecuTorch program for this model returned by the call to `to_executorch()` or the `BundledProgram` of this model
        exported_program: Optional graph module for this model returned by the call to `torch.export` from nn.Module.
        extra_recorded_export_modules [Optional]: **Should be ignored by OSS users**. A dictionary of graph modules with the key being the user provided name and the
            value being the corresponding exported module. The exported graph modules can be either the
            output of `torch.export()` or `exir.to_edge()`.

    Returns:
        None
    """
    etrecord = ETRecord()
    etrecord.add_exported_program(exported_program)
    etrecord.add_edge_dialect_program(edge_dialect_program)
    etrecord.add_executorch_program(executorch_program)

    # Add extra export modules if user provided
    if extra_recorded_export_modules is not None:
        etrecord.add_extra_export_modules(extra_recorded_export_modules)

    etrecord.save(et_record)


def _process_exported_program(
    exported_program: Optional[Union[ExportedProgram, Dict[str, ExportedProgram]]]
) -> tuple[Optional[ExportedProgram], Optional[int]]:
    """Process exported program and return the processed program and export graph id."""
    processed_exported_program = None
    export_graph_id = None

    if exported_program is not None:
        if isinstance(exported_program, dict) and "forward" in exported_program:
            processed_exported_program = exported_program["forward"]
        elif isinstance(exported_program, ExportedProgram):
            processed_exported_program = exported_program

        if processed_exported_program is not None:
            export_graph_id = id(processed_exported_program.graph)

    return processed_exported_program, export_graph_id


def _validate_module_name(module_name: str) -> None:
    """Validate that module name is not a reserved name."""
    contains_reserved_name = any(
        reserved_name in module_name for reserved_name in ETRecordReservedFileNames
    )
    if contains_reserved_name:
        raise RuntimeError(
            f"The name {module_name} provided in the extra_recorded_export_modules dict is a reserved name in the ETRecord namespace."
        )


def _add_module_to_graph_map(
    graph_map: Dict[str, ExportedProgram],
    module_name: str,
    export_module: Union[ExportedProgram, ExirExportedProgram, EdgeProgramManager],
) -> None:
    """Add export module to graph map based on its type."""
    _validate_module_name(module_name)

    if isinstance(export_module, ExirExportedProgram):
        graph_map[f"{module_name}/forward"] = export_module.exported_program
    elif isinstance(export_module, ExportedProgram):
        graph_map[f"{module_name}/forward"] = export_module
    elif isinstance(
        export_module,
        (EdgeProgramManager, exir.program._program.EdgeProgramManager),
    ):
        for method in export_module.methods:
            graph_map[f"{module_name}/{method}"] = export_module.exported_program(
                method
            )
    else:
        raise RuntimeError(f"Unsupported graph module type. {type(export_module)}")


def _process_edge_dialect_program(
    edge_dialect_program: Union[EdgeProgramManager, ExirExportedProgram]
) -> Union[ExportedProgram, Dict[str, ExportedProgram]]:
    """Process edge dialect program and return the exported program(s).

    For EdgeProgramManager with multiple methods, returns a Dict[str, ExportedProgram]
    mapping method names to their exported programs. For single-method cases or
    ExirExportedProgram, returns a single ExportedProgram.
    """
    if isinstance(
        edge_dialect_program,
        (EdgeProgramManager, exir.program._program.EdgeProgramManager),
    ):
        methods = edge_dialect_program.methods
        if len(methods) == 1:
            # Single method case - return the ExportedProgram directly
            method_name = next(iter(methods))
            return edge_dialect_program.exported_program(method_name)
        else:
            # Multiple methods - return a dict of all methods
            return {
                method: edge_dialect_program.exported_program(method)
                for method in methods
            }
    elif isinstance(edge_dialect_program, ExirExportedProgram):
        return edge_dialect_program.exported_program
    else:
        raise RuntimeError(
            f"Unsupported type of edge_dialect_program passed in {type(edge_dialect_program)}."
        )


def _process_executorch_program(
    executorch_program: Union[
        ExecutorchProgram, ExecutorchProgramManager, BundledProgram
    ]
) -> tuple[
    Optional[Dict], Optional[Dict], Optional[Dict], Optional[Dict], Optional[List]
]:
    """Process executorch program and return debug maps and bundled program data."""
    if isinstance(executorch_program, BundledProgram):
        reference_outputs = _get_reference_outputs(executorch_program)
        representative_inputs = _get_representative_inputs(executorch_program)
        # pyre-ignore[16]: Item `None` of `typing.Union[None, exir.program._program.ExecutorchProgram, exir.program._program.ExecutorchProgramManager]` has no attribute `debug_handle_map`
        debug_handle_map = executorch_program.executorch_program.debug_handle_map
        # pyre-ignore[16]: Item `None` of `typing.Union[None, exir.program._program.ExecutorchProgram, exir.program._program.ExecutorchProgramManager]` has no attribute `debug_handle_map`
        delegate_map = executorch_program.executorch_program.delegate_map
        # pyre-ignore[16]: Item `None` of `typing.Union[None, exir.program._program.ExecutorchProgram, exir.program._program.ExecutorchProgramManager]` has no attribute `instruction_id_to_num_outs_map`
        instruction_id_to_num_outs_map = (
            executorch_program.executorch_program.instruction_id_to_num_outs_map
        )
        return (
            debug_handle_map,
            delegate_map,
            instruction_id_to_num_outs_map,
            reference_outputs,
            representative_inputs,
        )
    else:
        debug_handle_map = executorch_program.debug_handle_map
        delegate_map = executorch_program.delegate_map
        instruction_id_to_num_outs_map = (
            executorch_program.instruction_id_to_num_outs_map
        )
        return (
            debug_handle_map,
            delegate_map,
            instruction_id_to_num_outs_map,
            None,
            None,
        )


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)
            )
        elif entry == ETRecordReservedFileNames.INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME:
            instruction_id_to_num_outs_map = json.loads(
                etrecord_zip.read(
                    ETRecordReservedFileNames.INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME
                )
            )
        elif entry == ETRecordReservedFileNames.ETRECORD_IDENTIFIER:
            continue
        elif entry == ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM:
            # Old format: single edge dialect program (backward compatibility)
            serialized_artifact = SerializedArtifact(
                etrecord_zip.read(
                    ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM
                ),
                etrecord_zip.read(f"{entry}_state_dict"),
                etrecord_zip.read(f"{entry}_constants"),
                etrecord_zip.read(f"{entry}_example_inputs"),
            )
            edge_dialect_program = deserialize(serialized_artifact)
        elif entry.startswith(edge_dialect_prefix) and not entry.endswith(
            ("_state_dict", "_constants", "_example_inputs")
        ):
            # New format: edge_dialect_exported_program/method_name
            serialized_edge_dialect_program_files.add(entry)
        elif entry == ETRecordReservedFileNames.EXPORTED_PROGRAM:
            serialized_artifact = SerializedArtifact(
                etrecord_zip.read(ETRecordReservedFileNames.EXPORTED_PROGRAM),
                etrecord_zip.read(f"{entry}_state_dict"),
                etrecord_zip.read(f"{entry}_constants"),
                etrecord_zip.read(f"{entry}_example_inputs"),
            )
            exported_program = deserialize(serialized_artifact)
        elif entry == ETRecordReservedFileNames.REFERENCE_OUTPUTS:
            try:
                reference_outputs = torch.load(
                    io.BytesIO(
                        etrecord_zip.read(ETRecordReservedFileNames.REFERENCE_OUTPUTS)
                    ),
                    weights_only=True,
                )
            except Exception as e:
                raise RuntimeError(
                    "Failed to load reference_outputs from ETRecord. "
                    "This ETRecord file may have been created with an older "
                    "version that used pickle serialization. Please regenerate "
                    "the ETRecord file with the current version of ExecuTorch."
                ) from e
        elif entry == ETRecordReservedFileNames.REPRESENTATIVE_INPUTS:
            try:
                representative_inputs = torch.load(
                    io.BytesIO(
                        etrecord_zip.read(
                            ETRecordReservedFileNames.REPRESENTATIVE_INPUTS
                        )
                    ),
                    weights_only=True,
                )
            except Exception as e:
                raise RuntimeError(
                    "Failed to load representative_inputs from ETRecord. "
                    "This ETRecord file may have been created with an older "
                    "version that used pickle serialization. Please regenerate "
                    "the ETRecord file with the current version of ExecuTorch."
                ) from e
        elif entry == ETRecordReservedFileNames.EXPORT_GRAPH_ID:
            export_graph_id = json.loads(
                etrecord_zip.read(ETRecordReservedFileNames.EXPORT_GRAPH_ID)
            )
        else:
            if entry.endswith("state_dict"):
                serialized_state_dict_files.add(entry)
            elif entry.endswith("constants"):
                serialized_constants_files.add(entry)
            elif entry.endswith("example_inputs"):
                serialized_example_inputs_files.add(entry)
            else:
                serialized_exported_program_files.add(entry)

    # Parse new format edge dialect programs (multi-method support)
    for serialized_file in serialized_edge_dialect_program_files:
        serialized_state_dict_file = f"{serialized_file}_state_dict"
        serialized_constants_file = f"{serialized_file}_constants"
        serialized_example_inputs_file = f"{serialized_file}_example_inputs"
        serialized_artifact = SerializedArtifact(
            etrecord_zip.read(serialized_file),
            etrecord_zip.read(serialized_state_dict_file),
            etrecord_zip.read(serialized_constants_file),
            etrecord_zip.read(serialized_example_inputs_file),
        )
        # Extract method name from path: edge_dialect_exported_program/method_name -> method_name
        method_name = serialized_file[len(edge_dialect_prefix) :]
        edge_dialect_programs[method_name] = deserialize(serialized_artifact)

    # If we found multi-method edge dialect programs, use them
    if edge_dialect_programs:
        if len(edge_dialect_programs) == 1:
            # Single method - store as ExportedProgram for backward compatibility
            edge_dialect_program = next(iter(edge_dialect_programs.values()))
        else:
            # Multiple methods - store as dict
            edge_dialect_program = edge_dialect_programs

    for serialized_file in serialized_exported_program_files:
        serialized_state_dict_file = f"{serialized_file}_state_dict"
        serialized_constants_file = f"{serialized_file}_constants"
        serialized_example_inputs_file = f"{serialized_file}_example_inputs"
        assert (
            serialized_state_dict_file in serialized_state_dict_files
        ), f"Could not find corresponding state dict file for {serialized_file}."
        serialized_artifact = SerializedArtifact(
            etrecord_zip.read(serialized_file),
            etrecord_zip.read(serialized_state_dict_file),
            etrecord_zip.read(serialized_constants_file),
            etrecord_zip.read(serialized_example_inputs_file),
        )
        graph_map[serialized_file] = deserialize(serialized_artifact)

    return ETRecord(
        exported_program=exported_program,
        edge_dialect_program=edge_dialect_program,
        graph_map=graph_map,
        _debug_handle_map=debug_handle_map,
        _delegate_map=delegate_map,
        _instruction_id_to_num_outs_map=instruction_id_to_num_outs_map,
        _reference_outputs=reference_outputs,
        _representative_inputs=representative_inputs,
        export_graph_id=export_graph_id,
    )