Extend data file support for other pybindings (pytorch#14131)

Summary:

This diff extends the support for external data_file to other executorch runtime pybindings.

Differential Revision: D82046648

Author: Yi Li

extension/pybindings/pybindings.cpp

@@ -161,10 +161,31 @@ void setup_output_storage(
 inline std::unique_ptr<Module> load_module_from_buffer(
     const void* ptr,
     size_t ptr_len,
+    std::optional<const std::string>& data_map_path,
     std::unique_ptr<runtime::EventTracer> event_tracer,
     Program::Verification program_verification) {
   EXECUTORCH_SCOPE_PROF("load_module_from_buffer");
   auto loader = std::make_unique<BufferDataLoader>(ptr, ptr_len);
+
+  if (data_map_path.has_value()) {
+    Result<MmapDataLoader> data_map_loader_res = MmapDataLoader::from(
+        data_map_path->c_str(),
+        MmapDataLoader::MlockConfig::UseMlockIgnoreErrors);
+    THROW_IF_ERROR(
+        data_map_loader_res.error(),
+        "Failed to create MmapDataLoader from file %s, error: 0x:%" PRIx32,
+        data_map_path->c_str(),
+        static_cast<uint32_t>(data_map_loader_res.error()));
+    auto data_map_loader =
+        std::make_unique<MmapDataLoader>(std::move(data_map_loader_res.get()));
+    return std::make_unique<Module>(
+        std::move(loader),
+        nullptr, // memory_allocator
+        nullptr, // temp_allocator
+        std::move(event_tracer), // event_tracer
+        std::move(data_map_loader)); // data_map_loader
+  }
+
   return std::make_unique<Module>(
       std::move(loader),
       nullptr, // memory_allocator
@@ -504,6 +525,7 @@ struct PyMethodMeta final {
 struct PyModule final {
   explicit PyModule(
       const py::bytes& buffer,
+      std::optional<const std::string>& data_path,
       bool enable_etdump,
       size_t debug_buffer_size = 0,
       Program::Verification program_verification =
@@ -512,12 +534,14 @@ struct PyModule final {
         module_(load_module_from_buffer(
             buffer.cast<std::string_view>().data(),
             py::len(buffer),
+            data_path,
             setup_event_tracer(enable_etdump, debug_buffer_size),
             program_verification)) {}
 
   explicit PyModule(
       const void* ptr,
       size_t ptr_len,
+      std::optional<const std::string>& data_path,
       bool enable_etdump,
       size_t debug_buffer_size = 0,
       Program::Verification program_verification =
@@ -526,6 +550,7 @@ struct PyModule final {
         module_(load_module_from_buffer(
             ptr,
             ptr_len,
+            data_path,
             setup_event_tracer(enable_etdump, debug_buffer_size),
             program_verification)) {}
 
@@ -551,12 +576,13 @@ struct PyModule final {
   // Module is only valid as long as the python buffer is alive.
   static std::unique_ptr<PyModule> load_from_buffer(
       const py::bytes& buffer,
+      std::optional<const std::string>& data_path,
       bool enable_etdump,
       size_t debug_buffer_size = 0,
       Program::Verification program_verification =
           Program::Verification::InternalConsistency) {
     return std::make_unique<PyModule>(
-        buffer, enable_etdump, debug_buffer_size, program_verification);
+        buffer, data_path, enable_etdump, debug_buffer_size, program_verification);
   }
 
   static std::unique_ptr<PyModule> load_from_file(
@@ -576,11 +602,13 @@ struct PyModule final {
 
   static std::unique_ptr<PyModule> load_from_bundled_program(
       PyBundledModule& m,
+      std::optional<const std::string>& data_path,
       bool enable_etdump,
       size_t debug_buffer_size = 0) {
     return std::make_unique<PyModule>(
         m.get_program_ptr(),
         m.get_program_len(),
+        data_path,
         enable_etdump,
         debug_buffer_size);
   }
@@ -1390,6 +1418,7 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
       "_load_for_executorch_from_buffer",
       &PyModule::load_from_buffer,
       py::arg("buffer"),
+      py::arg("data_path") = std::nullopt,
       py::arg("enable_etdump") = false,
       py::arg("debug_buffer_size") = 0,
       py::arg("program_verification") =
@@ -1399,6 +1428,7 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
       "_load_for_executorch_from_bundled_program",
       &PyModule::load_from_bundled_program,
       py::arg("ptr"),
+      py::arg("data_path") = std::nullopt,
       py::arg("enable_etdump") = false,
       py::arg("debug_buffer_size") = 0,
       call_guard);

extension/pybindings/test/TARGETS

@@ -24,6 +24,9 @@ runtime.python_library(
         "//executorch/exir/emit:lib",
         "//executorch/exir/passes:lib",
         "//executorch/runtime/core:core",
+        "//executorch/devtools/bundled_program:core",
+        "//executorch/devtools/bundled_program:config",
+        "//executorch/devtools/bundled_program/serialize:lib",
     ],
 )
 

extension/pybindings/test/test_pybindings.py

@@ -651,8 +651,65 @@ def test_program_data_separation(self) -> None:
                 )
                 ptd.write(tensor_data)
 
-            executorch_program = self.runtime._load_for_executorch(pte_file, ptd_file)
-
             expected = eager_module(inputs[0])
-            executorch_output = executorch_program.forward(inputs)[0]
-            self.assertTrue(torch.allclose(expected, executorch_output))
+
+            # Test 1: File-based loading with external data file
+            executorch_module_file = self.runtime._load_for_executorch(pte_file, ptd_file)
+            executorch_output_file = executorch_module_file.forward(inputs)[0]
+            self.assertTrue(torch.allclose(expected, executorch_output_file))
+
+            # Test 2: Buffer-based loading with external data file
+            with open(pte_file, "rb") as f:
+                program_buffer = f.read()
+            executorch_module_buffer = self.load_fn(program_buffer, ptd_file)
+            executorch_output_buffer = executorch_module_buffer.forward(inputs)[0]
+            self.assertTrue(torch.allclose(expected, executorch_output_buffer))
+
+            # Test 3: Buffer-based loading without external data file (should fail or work differently)
+            # This should fail because the program expects external data
+            with self.assertRaises(RuntimeError):
+                executorch_module_no_data = self.load_fn(program_buffer)
+                executorch_module_no_data.forward(inputs)
+
+            # Test 4: Test with invalid data file path (should fail)
+            invalid_ptd_file = os.path.join(tmpdir, "nonexistent.ptd")
+            with self.assertRaises(RuntimeError):
+                self.load_fn(program_buffer, invalid_ptd_file)
+
+            # Test 5: Test bundled program loading with external data
+            # First create a bundled program with external constants
+            from executorch.devtools.bundled_program.core import BundledProgram
+            from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite
+            from executorch.devtools.bundled_program.serialize import (
+                serialize_from_bundled_program_to_flatbuffer,
+            )
+
+            method_test_suites = [
+                MethodTestSuite(
+                    method_name="forward",
+                    test_cases=[
+                        MethodTestCase(
+                            inputs=input,
+                            expected_outputs=expected,
+                        )
+                        for input in inputs
+                    ],
+                ),
+            ]
+            bundled_program = BundledProgram(exec_program, method_test_suites)
+            bundled_buffer = serialize_from_bundled_program_to_flatbuffer(bundled_program)
+            bundled_module = self.runtime._load_bundled_program_from_buffer(bundled_buffer)
+
+            # Load module from bundled program with external data
+            executorch_module_bundled = self.runtime._load_for_executorch_from_bundled_program(
+                bundled_module, ptd_file
+            )
+            executorch_output_bundled = executorch_module_bundled.forward(inputs)[0]
+            self.assertTrue(torch.allclose(expected, executorch_output_bundled))
+
+            # Test 6: Bundled program without external data should fail
+            with self.assertRaises(RuntimeError):
+                executorch_module_bundled_no_data = self.runtime._load_for_executorch_from_bundled_program(
+                    bundled_module
+                )
+                executorch_module_bundled_no_data.forward(inputs)