Allow user to limit the number of threads that OpenMP spawns.

higuoxing · Xuebin Su · xuebinsu · higuoxing · commit 58e2aeca8c84 · 2024-03-08T11:43:14.000+08:00
This patch introduces a new GUC parameter pgml.omp_num_threads to
control control the number of threads that OpenMP can spawn.

Co-authored-by: Xuebin Su &lt;sxuebin@vmware.com&gt;
Co-authored-by: Xuebin Su (苏学斌) &lt;12034000+xuebinsu@users.noreply.github.com&gt;
diff --git a/pgml-extension/src/config.rs b/pgml-extension/src/config.rs
@@ -18,6 +18,8 @@ lazy_static! {
         "pgml.huggingface_trust_remote_code_whitelist",
         GucSetting::<Option<&'static CStr>>::new(None),
     );
+    pub static ref PGML_OMP_NUM_THREADS: (&'static str, GucSetting<i32>) =
+        ("pgml.omp_num_threads", GucSetting::<i32>::new(-1));
 }
 
 pub fn initialize_server_params() {
@@ -53,6 +55,16 @@ pub fn initialize_server_params() {
         GucContext::Userset,
         GucFlags::default(),
     );
+    GucRegistry::define_int_guc(
+        PGML_OMP_NUM_THREADS.0,
+        "Specifies the number of threads used by default of underlying OpenMP library. Only positive integers are valid",
+        "",
+        &PGML_OMP_NUM_THREADS.1,
+        -1,
+        i32::max_value(),
+        GucContext::Backend,
+        GucFlags::default(),
+    );
 }
 
 #[cfg(any(test, feature = "pg_test"))]
@@ -75,4 +87,10 @@ mod tests {
         set_config(name, value).unwrap();
         assert_eq!(PGML_HF_WHITELIST.1.get().unwrap().to_string_lossy(), value);
     }
+
+    #[pg_test]
+    fn omp_num_threads_cannot_be_set_after_startup() {
+        let result = std::panic::catch_unwind(|| set_config("pgml.omp_num_threads", "1"));
+        assert!(result.is_err());
+    }
 }
diff --git a/pgml-extension/src/lib.rs b/pgml-extension/src/lib.rs
@@ -21,10 +21,20 @@ pg_module_magic!();
 
 extension_sql_file!("../sql/schema.sql", name = "schema");
 
+extern "C" {
+    fn omp_set_num_threads(num_threads: i32);
+}
+
 #[cfg(not(feature = "use_as_lib"))]
 #[pg_guard]
 pub extern "C" fn _PG_init() {
     config::initialize_server_params();
+    let omp_num_threads = config::PGML_OMP_NUM_THREADS.1.get();
+    if omp_num_threads > 0 {
+        unsafe {
+            omp_set_num_threads(omp_num_threads);
+        }
+    }
     bindings::python::activate().expect("Error setting python venv");
     orm::project::init();
 }
@@ -54,7 +64,7 @@ pub mod pg_test {
 
     pub fn postgresql_conf_options() -> Vec<&'static str> {
         // return any postgresql.conf settings that are required for your tests
-        let mut options = vec!["shared_preload_libraries = 'pgml'"];
+        let mut options = vec!["shared_preload_libraries = 'pgml'", "pgml.omp_num_threads = '1'"];
         if let Some(venv) = option_env!("PGML_VENV") {
             let option = format!("pgml.venv = '{venv}'");
             options.push(Box::leak(option.into_boxed_str()));
