Simplify global variables.

higuoxing · higuoxing · commit 3768179e282e · 2024-04-07T00:08:24.000+08:00
diff --git a/pgml-extension/src/bindings/python/mod.rs b/pgml-extension/src/bindings/python/mod.rs
@@ -21,7 +21,7 @@ pub fn activate_venv(venv: &str) -> Result<bool> {
 }
 
 pub fn activate() -> Result<bool> {
-    match PGML_VENV.1.get() {
+    match PGML_VENV.get() {
         Some(venv) => activate_venv(&venv.to_string_lossy()),
         None => Ok(false),
     }
diff --git a/pgml-extension/src/bindings/transformers/whitelist.rs b/pgml-extension/src/bindings/transformers/whitelist.rs
@@ -5,29 +5,29 @@ use pgrx::{pg_schema, pg_test};
 use serde_json::Value;
 use std::ffi::CStr;
 
-use crate::config::{PGML_HF_TRUST_REMOTE_CODE, PGML_HF_TRUST_WHITELIST, PGML_HF_WHITELIST};
+use crate::config::{PGML_HF_TRUST_REMOTE_CODE, PGML_HF_TRUST_REMOTE_CODE_WHITELIST, PGML_HF_WHITELIST};
 
 /// Verify that the model in the task JSON is allowed based on the huggingface whitelists.
 pub fn verify_task(task: &Value) -> Result<(), Error> {
     let task_model = match get_model_name(task) {
         Some(model) => model.to_string(),
         None => return Ok(()),
     };
-    let whitelisted_models = config_csv_list(&PGML_HF_WHITELIST.1);
+    let whitelisted_models = config_csv_list(&PGML_HF_WHITELIST);
 
     let model_is_allowed = whitelisted_models.is_empty() || whitelisted_models.contains(&task_model);
     if !model_is_allowed {
         bail!(
             "model {} is not whitelisted. Consider adding to {} in postgresql.conf",
             task_model,
-            PGML_HF_WHITELIST.0
+            "pgml.huggingface_whitelist"
         );
     }
 
     let task_trust = get_trust_remote_code(task);
-    let trust_remote_code = PGML_HF_TRUST_REMOTE_CODE.1.get();
+    let trust_remote_code = PGML_HF_TRUST_REMOTE_CODE.get();
 
-    let trusted_models = config_csv_list(&PGML_HF_TRUST_WHITELIST.1);
+    let trusted_models = config_csv_list(&PGML_HF_TRUST_REMOTE_CODE_WHITELIST);
 
     let model_is_trusted = trusted_models.is_empty() || trusted_models.contains(&task_model);
 
@@ -36,9 +36,9 @@ pub fn verify_task(task: &Value) -> Result<(), Error> {
         bail!(
             "model {} is not trusted to run remote code. Consider setting {} = 'true' or adding {} to {}",
             task_model,
-            PGML_HF_TRUST_REMOTE_CODE.0,
+            "pgml.huggingface_trust_remote_code",
             task_model,
-            PGML_HF_TRUST_WHITELIST.0
+            "pgml.huggingface_trust_remote_code_whitelist",
         );
     }
 
@@ -129,7 +129,7 @@ mod tests {
     #[pg_test]
     fn test_empty_whitelist() {
         let model = "Salesforce/xgen-7b-8k-inst";
-        set_config(PGML_HF_WHITELIST.0, "").unwrap();
+        set_config("pgml.huggingface_whitelist", "").unwrap();
         let task_json = format!(json_template!(), model, false);
         let task: Value = serde_json::from_str(&task_json).unwrap();
         assert!(verify_task(&task).is_ok());
@@ -138,12 +138,12 @@ mod tests {
     #[pg_test]
     fn test_nonempty_whitelist() {
         let model = "Salesforce/xgen-7b-8k-inst";
-        set_config(PGML_HF_WHITELIST.0, model).unwrap();
+        set_config("pgml.huggingface_whitelist", model).unwrap();
         let task_json = format!(json_template!(), model, false);
         let task: Value = serde_json::from_str(&task_json).unwrap();
         assert!(verify_task(&task).is_ok());
 
-        set_config(PGML_HF_WHITELIST.0, "other_model").unwrap();
+        set_config("pgml.huggingface_whitelist", "other_model").unwrap();
         let task_json = format!(json_template!(), model, false);
         let task: Value = serde_json::from_str(&task_json).unwrap();
         assert!(verify_task(&task).is_err());
@@ -152,8 +152,8 @@ mod tests {
     #[pg_test]
     fn test_trusted_model() {
         let model = "Salesforce/xgen-7b-8k-inst";
-        set_config(PGML_HF_WHITELIST.0, model).unwrap();
-        set_config(PGML_HF_TRUST_WHITELIST.0, model).unwrap();
+        set_config("pgml.huggingface_whitelist", model).unwrap();
+        set_config("pgml.huggingface_trust_remote_code_whitelist", model).unwrap();
 
         let task_json = format!(json_template!(), model, false);
         let task: Value = serde_json::from_str(&task_json).unwrap();
@@ -163,7 +163,7 @@ mod tests {
         let task: Value = serde_json::from_str(&task_json).unwrap();
         assert!(verify_task(&task).is_err());
 
-        set_config(PGML_HF_TRUST_REMOTE_CODE.0, "true").unwrap();
+        set_config("pgml.huggingface_trust_remote_code", "true").unwrap();
         let task_json = format!(json_template!(), model, false);
         let task: Value = serde_json::from_str(&task_json).unwrap();
         assert!(verify_task(&task).is_ok());
@@ -176,8 +176,8 @@ mod tests {
     #[pg_test]
     fn test_untrusted_model() {
         let model = "Salesforce/xgen-7b-8k-inst";
-        set_config(PGML_HF_WHITELIST.0, model).unwrap();
-        set_config(PGML_HF_TRUST_WHITELIST.0, "other_model").unwrap();
+        set_config("pgml.huggingface_whitelist", model).unwrap();
+        set_config("pgml.huggingface_trust_remote_code_whitelist", "other_model").unwrap();
 
         let task_json = format!(json_template!(), model, false);
         let task: Value = serde_json::from_str(&task_json).unwrap();
@@ -187,7 +187,7 @@ mod tests {
         let task: Value = serde_json::from_str(&task_json).unwrap();
         assert!(verify_task(&task).is_err());
 
-        set_config(PGML_HF_TRUST_REMOTE_CODE.0, "true").unwrap();
+        set_config("pgml.huggingface_trust_remote_code", "true").unwrap();
         let task_json = format!(json_template!(), model, false);
         let task: Value = serde_json::from_str(&task_json).unwrap();
         assert!(verify_task(&task).is_ok());
diff --git a/pgml-extension/src/config.rs b/pgml-extension/src/config.rs
@@ -1,67 +1,58 @@
-use once_cell::sync::Lazy;
 use pgrx::{GucContext, GucFlags, GucRegistry, GucSetting};
 use std::ffi::CStr;
 
 #[cfg(any(test, feature = "pg_test"))]
 use pgrx::{pg_schema, pg_test};
 
-pub static PGML_VENV: Lazy<(&'static str, GucSetting<Option<&'static CStr>>)> =
-    Lazy::new(|| ("pgml.venv", GucSetting::<Option<&'static CStr>>::new(None)));
-pub static PGML_HF_WHITELIST: Lazy<(&'static str, GucSetting<Option<&'static CStr>>)> = Lazy::new(|| {
-    (
-        "pgml.huggingface_whitelist",
-        GucSetting::<Option<&'static CStr>>::new(None),
-    )
-});
-pub static PGML_HF_TRUST_REMOTE_CODE: Lazy<(&'static str, GucSetting<bool>)> =
-    Lazy::new(|| ("pgml.huggingface_trust_remote_code", GucSetting::<bool>::new(false)));
-pub static PGML_HF_TRUST_WHITELIST: Lazy<(&'static str, GucSetting<Option<&'static CStr>>)> = Lazy::new(|| {
-    (
-        "pgml.huggingface_trust_remote_code_whitelist",
-        GucSetting::<Option<&'static CStr>>::new(None),
-    )
-});
-pub static PGML_OMP_NUM_THREADS: Lazy<(&'static str, GucSetting<i32>)> =
-    Lazy::new(|| ("pgml.omp_num_threads", GucSetting::<i32>::new(0)));
+pub static PGML_VENV: GucSetting<Option<&'static CStr>> = GucSetting::<Option<&'static CStr>>::new(None);
+pub static PGML_HF_WHITELIST: GucSetting<Option<&'static CStr>> = GucSetting::<Option<&'static CStr>>::new(None);
+pub static PGML_HF_TRUST_REMOTE_CODE: GucSetting<bool> = GucSetting::<bool>::new(false);
+pub static PGML_HF_TRUST_REMOTE_CODE_WHITELIST: GucSetting<Option<&'static CStr>> =
+    GucSetting::<Option<&'static CStr>>::new(None);
+pub static PGML_OMP_NUM_THREADS: GucSetting<i32> = GucSetting::<i32>::new(0);
 
 pub fn initialize_server_params() {
     GucRegistry::define_string_guc(
-        PGML_VENV.0,
+        "pgml.venv",
         "Python's virtual environment path",
         "",
-        &PGML_VENV.1,
+        &PGML_VENV,
         GucContext::Userset,
         GucFlags::default(),
     );
+
     GucRegistry::define_string_guc(
-        PGML_HF_WHITELIST.0,
+        "pgml.huggingface_whitelist",
         "Models allowed to be downloaded from huggingface",
         "",
-        &PGML_HF_WHITELIST.1,
+        &PGML_HF_WHITELIST,
         GucContext::Userset,
         GucFlags::default(),
     );
+
     GucRegistry::define_bool_guc(
-        PGML_HF_TRUST_REMOTE_CODE.0,
+        "pgml.huggingface_trust_remote_code",
         "Whether model can execute remote codes",
         "",
-        &PGML_HF_TRUST_REMOTE_CODE.1,
+        &PGML_HF_TRUST_REMOTE_CODE,
         GucContext::Userset,
         GucFlags::default(),
     );
+
     GucRegistry::define_string_guc(
-        PGML_HF_TRUST_WHITELIST.0,
+        "pgml.huggingface_trust_remote_code_whitelist",
         "Models allowed to execute remote codes when pgml.hugging_face_trust_remote_code = 'on'",
         "",
-        &PGML_HF_TRUST_WHITELIST.1,
+        &PGML_HF_TRUST_REMOTE_CODE_WHITELIST,
         GucContext::Userset,
         GucFlags::default(),
     );
+
     GucRegistry::define_int_guc(
-        PGML_OMP_NUM_THREADS.0,
+        "pgml.omp_num_threads",
         "Specifies the number of threads used by default of underlying OpenMP library. Only positive integers are valid",
         "",
-        &PGML_OMP_NUM_THREADS.1,
+        &PGML_OMP_NUM_THREADS,
         0,
         i32::max_value(),
         GucContext::Backend,
@@ -87,7 +78,8 @@ mod tests {
         let name = "pgml.huggingface_whitelist";
         let value = "meta-llama/Llama-2-7b";
         set_config(name, value).unwrap();
-        assert_eq!(PGML_HF_WHITELIST.1.get().unwrap().to_string_lossy(), value);
+        assert_eq!(PGML_HF_WHITELIST.get().unwrap().to_str().unwrap(), value);
+        //assert_eq!((&*PGML_HF_WHITELIST).get().unwrap().to_str().unwrap(), value);
     }
 
     #[pg_test]
diff --git a/pgml-extension/src/lib.rs b/pgml-extension/src/lib.rs
@@ -29,7 +29,7 @@ extern "C" {
 #[pg_guard]
 pub extern "C" fn _PG_init() {
     config::initialize_server_params();
-    let omp_num_threads = config::PGML_OMP_NUM_THREADS.1.get();
+    let omp_num_threads = config::PGML_OMP_NUM_THREADS.get();
     if omp_num_threads > 0 {
         unsafe {
             omp_set_num_threads(omp_num_threads);

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ pub fn activate_venv(venv: &str) -> Result<bool> {`
`21`	`21`	`}`
`22`	`22`
`23`	`23`	`pub fn activate() -> Result<bool> {`
`24`		`- match PGML_VENV.1.get() {`
	`24`	`+ match PGML_VENV.get() {`
`25`	`25`	`Some(venv) => activate_venv(&venv.to_string_lossy()),`
`26`	`26`	`None => Ok(false),`
`27`	`27`	`}`