Reduce binary size, limit asynchronous/backgroud thread stuff to training only.

onnxruntime/core/session/inference_session.cc

@@ -56,8 +56,10 @@
 #include "core/session/custom_ops.h"
 #endif
 
+#ifdef ENABLE_TRAINING
 #include "core/providers/cpu/controlflow/event_pool.h"
 #include "core/providers/cpu/controlflow/message_queue.h"
+#endif
 
 using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::experimental;
@@ -378,6 +380,7 @@ InferenceSession::~InferenceSession() {
     }
   }
 
+#ifdef ENABLE_TRAINING
   // TODO: find a better way to terminate the background thread
   // backward is not completed yet, set terminate_flag to True
   if (task_.bg_thread_future_.valid()) {
@@ -385,6 +388,7 @@ InferenceSession::~InferenceSession() {
     Status s = ContinueRunInBackground({});
     ORT_UNUSED_PARAMETER(s);
   }
+#endif
 
 #ifdef ONNXRUNTIME_ENABLE_INSTRUMENT
   if (session_activity_started_)
@@ -1728,6 +1732,7 @@ common::Status InferenceSession::Run(IOBinding& io_binding) {
   return Run(run_options, io_binding);
 }
 
+#ifdef ENABLE_TRAINING
 common::Status InferenceSession::RunInBackgroundAndWaitForYield(RunOptions& run_options, IOBinding& io_binding,
                                                                 std::vector<OrtValue>& user_outputs) {
   const int64_t main_thread_event_id = 0;
@@ -1781,6 +1786,7 @@ common::Status InferenceSession::ContinueRunInBackground(const std::vector<OrtVa
 
   return bg_thread_status;
 }
+#endif
 
 template <typename T>
 void InferenceSession::StartProfiling(const std::basic_string<T>& file_prefix) {

onnxruntime/core/session/inference_session.h

@@ -5,8 +5,11 @@
 
 #include <string>
 #include <unordered_map>
+
+#ifdef ENABLE_TRAINING
 #include <thread>
 #include <future>
+#endif
 
 #include "core/common/common.h"
 #include "core/common/logging/logging.h"
@@ -299,12 +302,14 @@ class InferenceSession {
   virtual common::Status Run(const RunOptions& run_options, IOBinding& io_binding) ORT_MUST_USE_RESULT;
   common::Status Run(IOBinding& io_binding) ORT_MUST_USE_RESULT;
 
+#ifdef ENABLE_TRAINING
   // For ORTModule.forward()
   virtual common::Status RunInBackgroundAndWaitForYield(RunOptions& run_options, IOBinding& io_binding,
                                                         std::vector<OrtValue>& user_outputs) ORT_MUST_USE_RESULT;
 
   // For ORTModule.backward()
   common::Status ContinueRunInBackground(const std::vector<OrtValue>& backward_output_grads) ORT_MUST_USE_RESULT;
+#endif
 
   /**
     * @return pair.first = OK; FAIL otherwise. pair.second is non-NULL when pair.first = OK.
@@ -670,6 +675,7 @@ class InferenceSession {
   // those into new OrtValue instances, at which point we won't free them until the InferenceSession goes away.
   std::vector<uint8_t> ort_format_model_bytes_;
 
+#ifdef ENABLE_TRAINING
   // background thread for RunInBackgroundAndWaitForYield
   struct Task {
     std::thread bg_thread_;
@@ -677,10 +683,11 @@ class InferenceSession {
     std::future<Status> bg_thread_future_;
     bool* terminate_flag_ = nullptr;
   } task_;
-
+#endif
   std::shared_ptr<onnxruntime::AllocatorManager> allocator_manager_;
 };
 
+
 struct SessionIOBinding {
  public:
   SessionIOBinding(InferenceSession* session);

onnxruntime/python/dlpack_convertor.cc

@@ -1,6 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
-
+#ifdef ENABLE_TRAINING
 #include "python/dlpack_convertor.h"
 
 namespace onnxruntime {
@@ -122,3 +122,4 @@ DLManagedTensor* ort_value_to_dlpack(const OrtValue& ort_value) {
 
 }  // namespace python
 }  // namespace onnxruntime
+#endif

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -25,7 +25,10 @@
 #include "core/platform/env.h"
 #include "core/session/IOBinding.h"
 #include "core/session/abi_session_options_impl.h"
+
+#ifdef ENABLE_TRAINING
 #include "python/dlpack_convertor.h"
+#endif
 
 // execution provider factory creator headers
 #include "core/providers/cpu/cpu_provider_factory_creator.h"
@@ -1057,6 +1060,7 @@ void addOpSchemaSubmodule(py::module& m) {
 
 #endif  //onnxruntime_PYBIND_EXPORT_OPSCHEMA
 
+#ifdef ENABLE_TRAINING
 void dlpack_capsule_destructor(PyObject* data) {
   DLManagedTensor* dlmanged_tensor = (DLManagedTensor*)PyCapsule_GetPointer(data, "dltensor");
   if (dlmanged_tensor) {
@@ -1068,6 +1072,7 @@ void dlpack_capsule_destructor(PyObject* data) {
     PyErr_Clear();
   }
 }
+#endif
 
 void addObjectMethods(py::module& m, Environment& env) {
   py::enum_<GraphOptimizationLevel>(m, "GraphOptimizationLevel")
@@ -1221,6 +1226,7 @@ void addObjectMethods(py::module& m, Environment& env) {
 
         return ml_value;
       })
+#ifdef ENABLE_TRAINING
       .def_static("ortvalue_from_data_ptr", [](std::vector<int64_t>& shape, py::object& element_type,
                                                OrtDevice& device, int64_t data_ptr) {
         ORT_ENFORCE(data_ptr != 0, "Pointer to data memory is invalid");
@@ -1242,6 +1248,7 @@ void addObjectMethods(py::module& m, Environment& env) {
 
         return ort_value;
       })
+#endif
       .def("data_ptr", [](OrtValue* ml_value) -> int64_t {
         // TODO: Assumes that the OrtValue is a Tensor, make this generic to handle non-Tensors
         ORT_ENFORCE(ml_value->IsTensor(), "Only OrtValues that are Tensors are currently supported");
@@ -1304,11 +1311,14 @@ void addObjectMethods(py::module& m, Environment& env) {
 #endif
         return obj;
       })
+#ifdef ENABLE_TRAINING
       .def("to_dlpack", [](OrtValue* ort_value) -> py::object {
         DLManagedTensor* dlmanaged_tensor = ort_value_to_dlpack(*ort_value);
         return py::reinterpret_steal<py::object>(
             PyCapsule_New(dlmanaged_tensor, "dltensor", dlpack_capsule_destructor));
-      });
+      })
+#endif
+;
 
   py::class_<SessionIOBinding> session_io_binding(m, "SessionIOBinding");
   session_io_binding
@@ -1831,6 +1841,7 @@ including arg name, arg type (contains both type and shape).)pbdoc")
         if (!status.IsOK())
           throw std::runtime_error("Error in execution: " + status.ErrorMessage());
       })
+#ifdef ENABLE_TRAINING
       .def("run_forward", [](PyInferenceSession* sess, SessionIOBinding& io_binding, RunOptions& run_options) -> std::vector<OrtValue> {
         std::vector<OrtValue> module_outputs;
         Status status = sess->GetSessionHandle()->RunInBackgroundAndWaitForYield(run_options, *io_binding.Get(), module_outputs);
@@ -1844,7 +1855,9 @@ including arg name, arg type (contains both type and shape).)pbdoc")
         Status status = sess->GetSessionHandle()->ContinueRunInBackground(backward_output_grads);
         if (!status.IsOK())
           throw std::runtime_error("Error in execution: " + status.ErrorMessage());
-      });
+      })
+#endif
+;
 
   py::enum_<onnxruntime::ArenaExtendStrategy>(m, "ArenaExtendStrategy", py::arithmetic())
       .value("kNextPowerOfTwo", onnxruntime::ArenaExtendStrategy::kNextPowerOfTwo)