Let minibenchmark dup the model file descriptor if the model is passed in as a file descriptor.

There is a chance that fd is opened with O_CLOEXEC, and will be closed when starting child process. We dup the input fd, and pass the new fd to the child process.

PiperOrigin-RevId: 510979253

Author: hehhjiang

tensorflow/lite/experimental/acceleration/mini_benchmark/BUILD

@@ -462,6 +462,7 @@ cc_library(
         ":runner",
         ":status_codes",
         ":validator",
+        "@com_google_absl//absl/strings",
         "@flatbuffers",
         "//tensorflow/lite/experimental/acceleration/configuration/c:delegate_plugin",
         "//tensorflow/lite:minimal_logging",

tensorflow/lite/experimental/acceleration/mini_benchmark/blocking_validator_runner_test.cc

@@ -106,6 +106,40 @@ TEST_F(BlockingValidatorRunnerTest, SucceedWithEmbeddedValidation) {
   }
 }
 
+TEST_F(BlockingValidatorRunnerTest, SucceedWithFdCloexecEmbeddedValidation) {
+  if (!should_perform_test_) {
+    std::cerr << "Skipping test";
+    return;
+  }
+
+  options_.model_fd = open(options_.model_path.c_str(), O_RDONLY | O_CLOEXEC);
+  ASSERT_GE(options_.model_fd, 0);
+  struct stat stat_buf = {0};
+  ASSERT_EQ(fstat(options_.model_fd, &stat_buf), 0);
+  options_.model_size = stat_buf.st_size;
+  options_.model_offset = 0;
+  options_.model_path.clear();
+
+  BlockingValidatorRunner runner(options_);
+  ASSERT_EQ(runner.Init(), kMinibenchmarkSuccess);
+  FlatBufferBuilder fbb;
+#ifdef __ANDROID__
+  fbb.Finish(CreateTFLiteSettings(fbb, Delegate_GPU));
+#else
+  fbb.Finish(CreateTFLiteSettings(fbb));
+#endif  // __ANDROID__
+
+  std::vector<FlatBufferBuilder> results = runner.TriggerValidation(
+      {flatbuffers::GetRoot<TFLiteSettings>(fbb.GetBufferPointer())});
+  EXPECT_THAT(results, testing::Not(testing::IsEmpty()));
+  for (auto& result : results) {
+    const BenchmarkEvent* event =
+        GetRoot<BenchmarkEvent>(result.GetBufferPointer());
+    EXPECT_EQ(event->event_type(), BenchmarkEventType_END);
+    EXPECT_TRUE(event->result()->ok());
+  }
+}
+
 TEST_F(BlockingValidatorRunnerTest, SucceedWithFdModelCustomValidation) {
   if (!should_perform_test_) {
     std::cerr << "Skipping test";

tensorflow/lite/experimental/acceleration/mini_benchmark/validator_runner_impl.cc

@@ -23,6 +23,10 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/strings/match.h"
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/str_split.h"
 #include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
 #include "tensorflow/lite/experimental/acceleration/configuration/configuration_generated.h"
 #include "tensorflow/lite/experimental/acceleration/mini_benchmark/benchmark_result_evaluator.h"
@@ -54,6 +58,56 @@ std::unique_ptr<FlatBufferBuilder> CopyModel(
   return copy;
 }
 
+// A simple holder for file descriptor that will close the file descriptor at
+// destruction time.
+class FdHolder {
+ public:
+  explicit FdHolder(int fd) : fd_(fd) {}
+
+  // Move only.
+  FdHolder(FdHolder&& other) = default;
+  FdHolder& operator=(FdHolder&& other) = default;
+
+  ~FdHolder() {
+    if (fd_ > 0) {
+      close(fd_);
+    }
+  }
+
+ private:
+  int fd_;
+};
+
+// Returns a FdHolder that will close the duped file descriptor when going out
+// of scope. If the model is passed in as a file descriptor, update the
+// model_path with a duped file descriptor. The original file descriptor may be
+// opened with FD_CLOEXEC, and cannot be read from the child process.
+std::unique_ptr<FdHolder> UpdateModelPathIfUsingFd(std::string& model_path) {
+  if (!absl::StartsWith(model_path, "fd:")) {
+    return nullptr;
+  }
+  std::vector<std::string> parts = absl::StrSplit(model_path, ':');
+  int model_fd;
+  if (!absl::SimpleAtoi(parts[1], &model_fd)) {
+    TFLITE_LOG_PROD(TFLITE_LOG_ERROR,
+                    "Failed to parse file descriptor %s from model_path %s",
+                    parts[1].c_str(), model_path.c_str());
+    return nullptr;
+  }
+  int new_fd = dup(model_fd);
+  if (new_fd < 0) {
+    TFLITE_LOG_PROD(
+        TFLITE_LOG_ERROR,
+        "Failed to dup() file descriptor. Original fd: %d errno: %d", model_fd,
+        errno);
+    return nullptr;
+  }
+
+  parts[1] = std::to_string(new_fd);
+  model_path = absl::StrJoin(parts, ":");
+  return std::make_unique<FdHolder>(new_fd);
+}
+
 }  // namespace
 
 MinibenchmarkStatus ValidatorRunnerImpl::Init() {
@@ -143,7 +197,7 @@ void ValidatorRunnerImpl::TriggerValidationAsync(
   // error_reporter is not passed in because the ownership cannot be passed to
   // the thread.
   std::thread detached_thread(
-      [model_path = fd_or_model_path_, storage_path = storage_path_,
+      [original_model_path = fd_or_model_path_, storage_path = storage_path_,
        data_directory_path = data_directory_path_,
        tflite_settings = std::move(tflite_settings),
        validation_entrypoint_name =
@@ -157,6 +211,10 @@ void ValidatorRunnerImpl::TriggerValidationAsync(
         if (!lock.TryLock()) {
           return;
         }
+
+        std::string model_path = original_model_path;
+        std::unique_ptr<FdHolder> fd_holder =
+            UpdateModelPathIfUsingFd(model_path);
         for (auto& one_setting : *tflite_settings) {
           FlatbufferStorage<BenchmarkEvent> storage(storage_path);
           TFLiteSettingsT tflite_settings_obj;

tensorflow/lite/tools/model_loader.cc

@@ -56,14 +56,21 @@ bool PathModelLoader::InitInternal() {
 
 bool MmapModelLoader::InitInternal() {
   if (model_fd_ < 0 || model_offset_ < 0 || model_size_ < 0) {
+    TFLITE_LOG_PROD(
+        TFLITE_LOG_ERROR,
+        "Invalid model file descriptor. file descriptor: %d model_offset: "
+        "%d model_size: %d",
+        model_fd_, model_offset_, model_size_);
     return false;
   }
   if (!MMAPAllocation::IsSupported()) {
+    TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "MMAPAllocation is not supported.");
     return false;
   }
   auto allocation = std::make_unique<MMAPAllocation>(
       model_fd_, model_offset_, model_size_, tflite::DefaultErrorReporter());
   if (!allocation->valid()) {
+    TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "MMAPAllocation is not valid.");
     return false;
   }
   model_ = FlatBufferModel::VerifyAndBuildFromAllocation(std::move(allocation));
@@ -72,6 +79,8 @@ bool MmapModelLoader::InitInternal() {
 
 bool PipeModelLoader::InitInternal() {
   if (pipe_fd_ < 0) {
+    TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Invalid pipe file descriptor %d",
+                    pipe_fd_);
     return false;
   }
 
@@ -89,7 +98,7 @@ bool PipeModelLoader::InitInternal() {
   // Close the read pipe.
   close(pipe_fd_);
   if (read_bytes < 0 || remaining_bytes != 0) {
-    TFLITE_LOG_PROD(TFLITE_LOG_INFO,
+    TFLITE_LOG_PROD(TFLITE_LOG_ERROR,
                     "Read Model from pipe failed: %s. Expect to read %d bytes, "
                     "%d bytes missing.",
                     std::strerror(errno), model_size_, remaining_bytes);
@@ -112,6 +121,7 @@ std::unique_ptr<ModelLoader> CreateModelLoaderFromPath(absl::string_view path) {
     if (parts.size() != 4 || !absl::SimpleAtoi(parts[1], &model_fd) ||
         !absl::SimpleAtoi(parts[2], &model_offset) ||
         !absl::SimpleAtoi(parts[3], &model_size)) {
+      TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to parse model path: %s", path);
       return nullptr;
     }
     return std::make_unique<MmapModelLoader>(model_fd, model_offset,
@@ -124,6 +134,7 @@ std::unique_ptr<ModelLoader> CreateModelLoaderFromPath(absl::string_view path) {
     if (parts.size() != 4 || !absl::SimpleAtoi(parts[1], &read_fd) ||
         !absl::SimpleAtoi(parts[2], &write_fd) ||
         !absl::SimpleAtoi(parts[3], &model_size)) {
+      TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "Failed to parse model path: %s", path);
       return nullptr;
     }
     // If set, close the write pipe for the read process / thread.

tensorflow/lite/tools/model_loader.h

@@ -34,7 +34,7 @@ namespace tools {
 // Class to load the Model.
 class ModelLoader {
  public:
-  virtual ~ModelLoader() {}
+  virtual ~ModelLoader() = default;
 
   // Return whether the model is loaded successfully.
   virtual bool Init();