From 11ba3c24bad139eec6a34aade6fbc8ac786da5fa Mon Sep 17 00:00:00 2001
From: Vivek Panyam <vip@uber.com>
Date: Thu, 28 May 2020 12:45:00 -0400
Subject: [PATCH] [OPE] Pass RuntimeOptions to the worker process (#363)

This PR passes user-specified `RuntimeOptions` to the OPE worker process when loading a model.

This enables us to easily add options without having to write special logic for OPE.

See #361 and #364 for examples.
---
 source/neuropod/multiprocess/multiprocess.cc        | 8 ++++++++
 source/neuropod/multiprocess/multiprocess_worker.cc | 7 ++++++-
 source/neuropod/multiprocess/ope_load_config.hh     | 4 ++++
 source/neuropod/tests/test_ipc_serialization.cc     | 3 +++
 4 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/source/neuropod/multiprocess/multiprocess.cc b/source/neuropod/multiprocess/multiprocess.cc
index b5f68107..b0420629 100644
--- a/source/neuropod/multiprocess/multiprocess.cc
+++ b/source/neuropod/multiprocess/multiprocess.cc
@@ -190,6 +190,14 @@ class MultiprocessNeuropodBackend : public NeuropodBackendWithDefaultAllocator<S
         load_config_.neuropod_path             = neuropod_path_;
         load_config_.default_backend_overrides = default_backend_overrides;
 
+        // Copy options into the load configuration
+        // Note: some of these options will be overridden in the worker process
+        load_config_.opts = options_;
+
+        // Since we're using CUDA_VISIBLE_DEVICES to set the appropriate device above,
+        // we'll just tell the worker to use GPU0
+        load_config_.opts.visible_device = Device::GPU0;
+
         if (options.load_model_at_construction)
         {
             load_model();
diff --git a/source/neuropod/multiprocess/multiprocess_worker.cc b/source/neuropod/multiprocess/multiprocess_worker.cc
index a80ac627..eed3fff5 100644
--- a/source/neuropod/multiprocess/multiprocess_worker.cc
+++ b/source/neuropod/multiprocess/multiprocess_worker.cc
@@ -57,8 +57,13 @@ void multiprocess_worker_loop(const std::string &control_queue_name)
                 ope_load_config config;
                 received.get(config);
 
+                // Override some options
+                auto &opts                      = config.opts;
+                opts.load_model_at_construction = true;
+                opts.use_ope                    = false;
+
                 // Load a neuropod
-                neuropod  = stdx::make_unique<Neuropod>(config.neuropod_path, config.default_backend_overrides);
+                neuropod  = stdx::make_unique<Neuropod>(config.neuropod_path, config.default_backend_overrides, opts);
                 allocator = neuropod->get_tensor_allocator();
                 inputs.clear();
                 control_channel.send_message(LOAD_SUCCESS);
diff --git a/source/neuropod/multiprocess/ope_load_config.hh b/source/neuropod/multiprocess/ope_load_config.hh
index c402d726..ac591815 100644
--- a/source/neuropod/multiprocess/ope_load_config.hh
+++ b/source/neuropod/multiprocess/ope_load_config.hh
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "neuropod/internal/backend_registration.hh"
 #include "neuropod/multiprocess/serialization/ipc_serialization.hh"
+#include "neuropod/options.hh"
 
 namespace neuropod
 {
@@ -29,6 +30,9 @@ struct ope_load_config
 
     // See the docs in `neuropod.hh`
     std::vector<BackendLoadSpec> default_backend_overrides;
+
+    // Options to pass to the worker process
+    RuntimeOptions opts;
 };
 
 } // namespace neuropod
diff --git a/source/neuropod/tests/test_ipc_serialization.cc b/source/neuropod/tests/test_ipc_serialization.cc
index 6fd24991..915b9e6f 100644
--- a/source/neuropod/tests/test_ipc_serialization.cc
+++ b/source/neuropod/tests/test_ipc_serialization.cc
@@ -76,10 +76,13 @@ TEST(test_ipc_serialization, ope_load_config)
         {"torchscript", "1.12.0", "/some/path/to/neuropod_torchscrtipt_backend.so"},
     };
 
+    expected.opts.visible_device = neuropod::Device::GPU2;
+
     const auto actual = serialize_deserialize(expected);
 
     EXPECT_EQ(expected.neuropod_path, actual.neuropod_path);
     EXPECT_EQ(expected.default_backend_overrides, actual.default_backend_overrides);
+    EXPECT_EQ(expected.opts.visible_device, actual.opts.visible_device);
 }
 
 TEST(test_ipc_serialization, neuropod_value_map)