Address comments

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: larryliu0820

runtime/TARGETS

@@ -0,0 +1,11 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+runtime.python_library(
+    name = "runtime",
+    srcs = ["__init__.py"],
+    deps = [
+        "//executorch/extension/pybindings:portable_lib",
+    ],
+)

runtime/__init__.py

@@ -12,12 +12,12 @@
     from pathlib import Path
 
     import torch
-    from executorch.runtime import LoadProgramConfig, Runtime
+    from executorch.runtime import Verification, Runtime
 
     et_runtime: Runtime = Runtime.get()
     program: Program = et_runtime.load_program(
         Path("/tmp/program.pte"),
-        config=LoadProgramConfig(verification="internal_consistency"),
+        verification=Verification.Minimal,
     )
     print("Program methods:", program.method_names)
     forward: Method = program.load_method("forward")
@@ -39,26 +39,30 @@
 """
 
 import functools
-from collections import defaultdict
 from pathlib import Path
 from types import ModuleType
-from typing import Any, BinaryIO, Optional, Sequence, Union
+from typing import Any, BinaryIO, Dict, Optional, Sequence, Union
 
-from executorch.extension.pybindings.portable_lib import (
-    _get_operator_names,
-    ExecuTorchModule,
-    MethodMeta,
-    Verification,
-)
+try:
+    from executorch.extension.pybindings.portable_lib import (
+        ExecuTorchModule,
+        MethodMeta,
+        Verification,
+    )
+except ModuleNotFoundError as e:
+    raise ModuleNotFoundError(
+        "Prebuilt <site-packages>/extension/pybindings/_portable_lib.so "
+        "is not found. Please reinstall ExecuTorch from pip."
+    ) from e
 
 
 class Method:
     """An ExecuTorch method, loaded from a Program.
-    TODO: This class should be pybind to the C++ counterpart instead of hosting ExecuTorchModule.
     This can be used to execute the method with inputs.
     """
 
     def __init__(self, method_name: str, module: ExecuTorchModule) -> None:
+        # TODO: This class should be pybind to the C++ counterpart instead of hosting ExecuTorchModule.
         self._method_name = method_name
         self._module = module
 
@@ -89,19 +93,23 @@ class Program:
     This can be used to load the methods/models defined by the program.
     """
 
-    def __init__(self, module: ExecuTorchModule, data: bytes) -> None:
+    def __init__(self, module: ExecuTorchModule, data: Optional[bytes]) -> None:
         # Hold the data so the program is not freed.
         self._data = data
         self._module = module
-        self._methods = defaultdict(str)
+        self._methods: Dict[str, Method] = {}
+        # ExecuTorchModule already pre-loads all Methods when created, so this
+        # doesn't do any extra work. TODO: Don't load a given Method until
+        # load_method() is called. Create a separate Method instance each time,
+        # to allow multiple independent instances of the same model.
         for method_name in self._module.method_names():
             self._methods[method_name] = Method(method_name, self._module)
 
     @property
     def method_names(self) -> Sequence[str]:
         return set(self._methods.keys())
 
-    def load_method(self, name: str) -> Method:
+    def load_method(self, name: str) -> Optional[Method]:
         """Loads a method from the program.
 
         Args:
@@ -110,26 +118,20 @@ def load_method(self, name: str) -> Method:
         Returns:
             The loaded method.
         """
-        return self._methods[name]
+        return self._methods.get(name, None)
 
 
 class OperatorRegistry:
-    """The registry of operators that are available to the runtime.
-
-    Currently only supports printing out all registered operator names.
-    """
+    """The registry of operators that are available to the runtime."""
 
-    def __init__(self) -> None:
-        pass
+    def __init__(self, legacy_module: ModuleType) -> None:
+        # TODO: Expose the kernel callables to Python.
+        self._legacy_module = legacy_module
 
     @property
     def operator_names(self) -> Sequence[str]:
-        """Gets the names of all registered operators.
-
-        Returns:
-            The names of all registered operators.
-        """
-        return _get_operator_names()
+        """The names of all registered operators."""
+        return set(self._legacy_module._get_operator_names())
 
 
 class Runtime:
@@ -142,67 +144,55 @@ class Runtime:
     @staticmethod
     @functools.lru_cache(maxsize=1)
     def get() -> "Runtime":
-        """Gets a Runtime singleton.
-
-        Raises:
-            ValueError: The requested config is not known.
-            ModuleNotFoundError: The prebuilt _portable_lib.so is not found.
-        """
-        try:
-            import executorch.extension.pybindings.portable_lib as legacy_module
-        except ModuleNotFoundError as e:
-            raise ModuleNotFoundError(
-                "Prebuilt <site-packages>/extension/pybindings/_portable_lib.so is not found. Please reinstall ExecuTorch from pip."
-            ) from e
+        """Gets the Runtime singleton."""
+        import executorch.extension.pybindings.portable_lib as legacy_module
 
         return Runtime(legacy_module=legacy_module)
 
     def __init__(self, *, legacy_module: ModuleType) -> None:
-        # TODO: Expose the kernel callables to Python.
         # Public attributes.
-        self.operator_registry = OperatorRegistry()
+        self.operator_registry = OperatorRegistry(legacy_module)
         # Private attributes.
         self._legacy_module = legacy_module
 
     def load_program(
         self,
         data: Union[bytes, bytearray, BinaryIO, Path, str],
         *,
-        verification_config: Optional[Verification] = Verification.InternalConsistency,
+        verification: Verification = Verification.InternalConsistency,
     ) -> Program:
         """Loads an ExecuTorch program from a PTE binary.
 
         Args:
-            data: The binary program data to load; typically PTE data. Note that
-                this can also load PTE data that is wrapped inside a bundled
-                program, but it will not provide access to the bundled program's
-                test/validation data.
-            verification_config: The configuration for program verification.
+            data: The binary program data to load; typically PTE data.
+            verification: level of program verification to perform.
 
         Returns:
             The loaded program.
         """
-        if isinstance(data, Path):
-            with data.open("rb") as f:
-                data = f.read()
+        if isinstance(data, (Path, str)):
+            m = self._legacy_module._load_for_executorch(
+                str(data),
+                enable_etdump=False,
+                debug_buffer_size=0,
+                program_verification=verification,
+            )
+            return Program(m, data=None)
         elif isinstance(data, BinaryIO):
-            data = data.read()
+            data_bytes = data.read()
         elif isinstance(data, bytearray):
-            data = bytes(data)
-        elif isinstance(data, str):
-            with open(data, "rb") as f:
-                data = f.read()
+            data_bytes = bytes(data)
         elif isinstance(data, bytes):
-            pass
+            data_bytes = data
         else:
             raise TypeError(
-                f"Expected data to be bytes, bytearray, a string to a valid .pte path, or a file-like object, but got {type(data).__name__}."
+                f"Expected data to be bytes, bytearray, a path to a .pte file, or a file-like object, but got {type(data).__name__}."
             )
         m = self._legacy_module._load_for_executorch_from_buffer(
-            data,
+            data_bytes,
             enable_etdump=False,
             debug_buffer_size=0,
-            program_verification=verification_config,
+            program_verification=verification,
         )
 
-        return Program(m, data=data)
+        return Program(m, data=data_bytes)

runtime/test/TARGETS

@@ -7,6 +7,6 @@ runtime.python_test(
     srcs = ["test_runtime.py"],
     deps = [
         "//executorch/extension/pybindings/test:make_test",
-        "//executorch/extension/pybindings:portable_lib",
+        "//executorch/runtime:runtime",
     ],
 )

runtime/test/test_runtime.py

@@ -9,24 +9,23 @@
 from pathlib import Path
 
 import torch
-from executorch.extension.pybindings.portable_lib import Verification
 
 from executorch.extension.pybindings.test.make_test import (
     create_program,
     ModuleAdd,
     ModuleMulti,
 )
-from executorch.runtime import Runtime
+from executorch.runtime import Runtime, Verification
 
 
 class RuntimeTest(unittest.TestCase):
     def test_smoke(self):
         ep, inputs = create_program(ModuleAdd())
         runtime = Runtime.get()
-
-        program = runtime.load_program(
-            ep.buffer, verification_config=Verification.Minimal
-        )
+        # Demonstrate that get() returns a singleton.
+        runtime2 = Runtime.get()
+        self.assertTrue(runtime is runtime2)
+        program = runtime.load_program(ep.buffer, verification=Verification.Minimal)
         method = program.load_method("forward")
         outputs = method.execute(inputs)
         self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1]))
@@ -35,10 +34,8 @@ def test_module_with_multiple_method_names(self):
         ep, inputs = create_program(ModuleMulti())
         runtime = Runtime.get()
 
-        program = runtime.load_program(
-            ep.buffer, verification_config=Verification.Minimal
-        )
-        self.assertEqual(program.method_names, set("forward", "forward2"))
+        program = runtime.load_program(ep.buffer, verification=Verification.Minimal)
+        self.assertEqual(program.method_names, set({"forward", "forward2"}))
         method = program.load_method("forward")
         outputs = method.execute(inputs)
         self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1]))