Size default thread pools from ORT_*_NUM_THREADS environment variables

.github/workflows/android.yml

@@ -79,8 +79,8 @@ jobs:
       run: |
         set -e -x
         BINARY_SIZE_THRESHOLD_ARGS=""
-        echo "Binary size threshold in bytes: 1440768"
-        BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes 1440768"
+        echo "Binary size threshold in bytes: 1451008"
+        BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes 1451008"
 
         # Ensure ANDROID_NDK_HOME is available and get its real path
         if [ -z "$ANDROID_NDK_HOME" ]; then

onnxruntime/core/util/thread_utils.cc

@@ -13,6 +13,7 @@
 #include "core/session/ort_apis.h"
 #include "core/common/string_utils.h"
 #include "core/common/logging/logging.h"
+#include "core/platform/env_var_utils.h"
 
 std::ostream& operator<<(std::ostream& os, const OrtThreadPoolParams& params) {
   os << "OrtThreadPoolParams {";
@@ -175,12 +176,33 @@ CreateThreadPoolHelper(Env* env, OrtThreadPoolParams options) {
                                       spin_us, /*force_hybrid*/ false, backoff_max);
 }
 
+static constexpr const char* kIntraOpNumThreadsEnvVar = "ORT_INTRA_OP_NUM_THREADS";
+static constexpr const char* kInterOpNumThreadsEnvVar = "ORT_INTER_OP_NUM_THREADS";
+
+// Determines the default thread count from the environment for a pool whose size was not set
+// programmatically (thread_pool_size == 0). Returns 0 when the environment does not specify one,
+// leaving the machine-sized default in place.
+//
+// ORT_INTRA_OP_NUM_THREADS sizes the intra-op pool and ORT_INTER_OP_NUM_THREADS the inter-op pool.
+// Parsing is strict (a negative or non-integer value fails loudly); an explicit value of 0 requests
+// the machine-sized default. This lets CPU-limited containers bound ORT's pools without reaching
+// every InferenceSession call site: physical-core detection cannot see the cgroup CPU reservation,
+// so the machine-sized default otherwise oversubscribes it.
+static int NumThreadsFromEnvironment(ThreadPoolType tpool_type) {
+  const bool is_intra_op = tpool_type == ThreadPoolType::INTRA_OP;
+  const char* ort_env_var = is_intra_op ? kIntraOpNumThreadsEnvVar : kInterOpNumThreadsEnvVar;
+  if (const auto parsed = ParseEnvironmentVariable<int>(ort_env_var); parsed.has_value()) {
+    ORT_ENFORCE(*parsed >= 0, ort_env_var, " must be a non-negative integer, got: ", *parsed);
+    return *parsed;
+  }
+  return 0;
+}
+
 std::unique_ptr<ThreadPool>
 CreateThreadPool(Env* env, OrtThreadPoolParams options, ThreadPoolType tpool_type) {
-  // If openmp is enabled we don't want to create any additional threadpools for sequential execution.
-  // However, parallel execution relies on the existence of a separate threadpool. Hence we allow eigen threadpools
-  // to be created for parallel execution.
-  ORT_UNUSED_PARAMETER(tpool_type);
+  if (options.thread_pool_size == 0) {  // 0 == "use default": consult the environment before the helper sizes to the machine
+    options.thread_pool_size = NumThreadsFromEnvironment(tpool_type);
+  }
   return CreateThreadPoolHelper(env, options);
 }
 

onnxruntime/core/util/thread_utils.h

@@ -8,7 +8,9 @@
 #include <string>
 
 struct OrtThreadPoolParams {
-  // 0: Use default setting. (All the physical cores or half of the logical cores)
+  // 0: Use default setting: the ORT_INTRA_OP_NUM_THREADS / ORT_INTER_OP_NUM_THREADS environment
+  //    variable for the corresponding pool type if set, otherwise all the physical cores or half
+  //    of the logical cores.
   // 1: Don't create thread pool
   // n: Create a thread pool with n threads.
   int thread_pool_size = 0;

onnxruntime/test/platform/threadpool_test.cc

@@ -5,6 +5,7 @@
 #include "core/platform/EigenNonBlockingThreadPool.h"
 #include <mutex>
 #include "core/util/thread_utils.h"
+#include "test/util/include/scoped_env_vars.h"
 #ifdef _WIN32
 #include "test/platform/windows/env.h"
 #include <Windows.h>
@@ -1040,4 +1041,64 @@ TEST(ThreadPoolTest, SpinBackoffWithTimeBoundedSpin) {
   TestSpinBackoffMode(1000, 8U);  // 1ms + backoff cap 8
 }
 
+// Tests for sizing default (thread_pool_size <= 0) pools from the environment.
+// Each test pins both variables so ambient values cannot leak in.
+// DegreeOfParallelism scales by a granularity factor on hybrid CPUs, so tests compare an
+// env-sized pool against a reference pool of the intended explicit size rather than asserting
+// absolute values.
+namespace {
+int PoolDegreeWithEnv(const test::EnvVarMap& env_vars, concurrency::ThreadPoolType tpool_type,
+                      int thread_pool_size = 0) {
+  test::ScopedEnvironmentVariables scoped_env(env_vars);
+  OrtThreadPoolParams tpo;
+  tpo.thread_pool_size = thread_pool_size;
+  auto tp = concurrency::CreateThreadPool(&Env::Default(), tpo, tpool_type);
+  return concurrency::ThreadPool::DegreeOfParallelism(tp.get());
+}
+
+const optional<std::string> kUnset{};
+const test::EnvVarMap kAllUnset{{"ORT_INTRA_OP_NUM_THREADS", kUnset},
+                                {"ORT_INTER_OP_NUM_THREADS", kUnset}};
+
+int PoolDegreeForExplicitSize(int thread_pool_size) {
+  return PoolDegreeWithEnv(kAllUnset, concurrency::ThreadPoolType::INTRA_OP, thread_pool_size);
+}
+}  // namespace
+
+TEST(ThreadPoolTest, DefaultPoolSizeFromOrtEnvVars) {
+  EXPECT_EQ(PoolDegreeWithEnv({{"ORT_INTRA_OP_NUM_THREADS", "3"},
+                               {"ORT_INTER_OP_NUM_THREADS", kUnset}},
+                              concurrency::ThreadPoolType::INTRA_OP),
+            PoolDegreeForExplicitSize(3));
+  EXPECT_EQ(PoolDegreeWithEnv({{"ORT_INTRA_OP_NUM_THREADS", kUnset},
+                               {"ORT_INTER_OP_NUM_THREADS", "3"}},
+                              concurrency::ThreadPoolType::INTER_OP),
+            PoolDegreeForExplicitSize(3));
+}
+
+TEST(ThreadPoolTest, ExplicitPoolSizeWinsOverEnvVars) {
+  EXPECT_EQ(PoolDegreeWithEnv({{"ORT_INTRA_OP_NUM_THREADS", "2"},
+                               {"ORT_INTER_OP_NUM_THREADS", kUnset}},
+                              concurrency::ThreadPoolType::INTRA_OP,
+                              /*thread_pool_size=*/4),
+            PoolDegreeForExplicitSize(4));
+}
+
+TEST(ThreadPoolTest, OrtEnvVarZeroRestoresMachineSizedDefault) {
+  // An explicit 0 opts back into the machine-sized default.
+  EXPECT_EQ(PoolDegreeWithEnv({{"ORT_INTRA_OP_NUM_THREADS", "0"},
+                               {"ORT_INTER_OP_NUM_THREADS", kUnset}},
+                              concurrency::ThreadPoolType::INTRA_OP),
+            PoolDegreeWithEnv(kAllUnset, concurrency::ThreadPoolType::INTRA_OP));
+}
+
+#ifndef ORT_NO_EXCEPTIONS
+TEST(ThreadPoolTest, InvalidOrtEnvVarValueThrows) {
+  EXPECT_THROW(PoolDegreeWithEnv({{"ORT_INTRA_OP_NUM_THREADS", "-1"},
+                                  {"ORT_INTER_OP_NUM_THREADS", kUnset}},
+                                 concurrency::ThreadPoolType::INTRA_OP),
+               OnnxRuntimeException);
+}
+#endif
+
 }  // namespace onnxruntime