Activate venv at server start and remove the check from functions (#1276)

SilasMarvin · web-flow · commit 5fdbd5dc1523 · 2024-01-05T14:42:59.000-08:00
diff --git a/pgml-extension/src/bindings/langchain/mod.rs b/pgml-extension/src/bindings/langchain/mod.rs
@@ -8,8 +8,6 @@ use crate::create_pymodule;
 create_pymodule!("/src/bindings/langchain/langchain.py");
 
 pub fn chunk(splitter: &str, text: &str, kwargs: &serde_json::Value) -> Result<Vec<String>> {
-    crate::bindings::python::activate()?;
-
     let kwargs = serde_json::to_string(kwargs).unwrap();
 
     Python::with_gil(|py| -> Result<Vec<String>> {
diff --git a/pgml-extension/src/bindings/python/mod.rs b/pgml-extension/src/bindings/python/mod.rs
@@ -30,7 +30,6 @@ pub fn activate() -> Result<bool> {
 }
 
 pub fn pip_freeze() -> Result<TableIterator<'static, (name!(package, String),)>> {
-    activate()?;
     let packages = Python::with_gil(|py| -> Result<Vec<String>> {
         let freeze = get_module!(PY_MODULE).getattr(py, "freeze")?;
         let result = freeze.call0(py)?;
@@ -42,7 +41,6 @@ pub fn pip_freeze() -> Result<TableIterator<'static, (name!(package, String),)>>
 }
 
 pub fn validate_dependencies() -> Result<bool> {
-    activate()?;
     Python::with_gil(|py| {
         let sys = PyModule::import(py, "sys").unwrap();
         let version: String = sys.getattr("version").unwrap().extract().unwrap();
@@ -68,7 +66,6 @@ pub fn validate_dependencies() -> Result<bool> {
 }
 
 pub fn version() -> Result<String> {
-    activate()?;
     Python::with_gil(|py| {
         let sys = PyModule::import(py, "sys").unwrap();
         let version: String = sys.getattr("version").unwrap().extract().unwrap();
@@ -77,7 +74,6 @@ pub fn version() -> Result<String> {
 }
 
 pub fn package_version(name: &str) -> Result<String> {
-    activate()?;
     Python::with_gil(|py| {
         let package = py.import(name)?;
         Ok(package.getattr("__version__")?.extract()?)
diff --git a/pgml-extension/src/bindings/transformers/mod.rs b/pgml-extension/src/bindings/transformers/mod.rs
@@ -34,8 +34,6 @@ pub fn get_model_from(task: &Value) -> Result<String> {
 }
 
 pub fn embed(transformer: &str, inputs: Vec<&str>, kwargs: &serde_json::Value) -> Result<Vec<Vec<f32>>> {
-    crate::bindings::python::activate()?;
-
     let kwargs = serde_json::to_string(kwargs)?;
     Python::with_gil(|py| -> Result<Vec<Vec<f32>>> {
         let embed: Py<PyAny> = get_module!(PY_MODULE).getattr(py, "embed").format_traceback(py)?;
@@ -58,8 +56,6 @@ pub fn embed(transformer: &str, inputs: Vec<&str>, kwargs: &serde_json::Value) -
 }
 
 pub fn tune(task: &Task, dataset: TextDataset, hyperparams: &JsonB, path: &Path) -> Result<HashMap<String, f64>> {
-    crate::bindings::python::activate()?;
-
     let task = task.to_string();
     let hyperparams = serde_json::to_string(&hyperparams.0)?;
 
@@ -86,8 +82,6 @@ pub fn tune(task: &Task, dataset: TextDataset, hyperparams: &JsonB, path: &Path)
 }
 
 pub fn generate(model_id: i64, inputs: Vec<&str>, config: JsonB) -> Result<Vec<String>> {
-    crate::bindings::python::activate()?;
-
     Python::with_gil(|py| -> Result<Vec<String>> {
         let generate = get_module!(PY_MODULE).getattr(py, "generate").format_traceback(py)?;
         let config = serde_json::to_string(&config.0)?;
@@ -162,8 +156,6 @@ pub fn load_dataset(
     limit: Option<usize>,
     kwargs: &serde_json::Value,
 ) -> Result<usize> {
-    crate::bindings::python::activate()?;
-
     let kwargs = serde_json::to_string(kwargs)?;
 
     let dataset = Python::with_gil(|py| -> Result<String> {
@@ -313,8 +305,6 @@ pub fn load_dataset(
 }
 
 pub fn clear_gpu_cache(memory_usage: Option<f32>) -> Result<bool> {
-    crate::bindings::python::activate().unwrap();
-
     Python::with_gil(|py| -> Result<bool> {
         let clear_gpu_cache: Py<PyAny> = get_module!(PY_MODULE)
             .getattr(py, "clear_gpu_cache")
diff --git a/pgml-extension/src/bindings/transformers/transform.rs b/pgml-extension/src/bindings/transformers/transform.rs
@@ -46,7 +46,6 @@ pub fn transform<T: serde::Serialize>(
     args: &serde_json::Value,
     inputs: T,
 ) -> Result<serde_json::Value> {
-    crate::bindings::python::activate()?;
     whitelist::verify_task(task)?;
 
     let task = serde_json::to_string(task)?;
@@ -74,7 +73,6 @@ pub fn transform_stream<T: serde::Serialize>(
     args: &serde_json::Value,
     input: T,
 ) -> Result<Py<PyAny>> {
-    crate::bindings::python::activate()?;
     whitelist::verify_task(task)?;
 
     let task = serde_json::to_string(task)?;
diff --git a/pgml-extension/src/lib.rs b/pgml-extension/src/lib.rs
@@ -24,6 +24,7 @@ extension_sql_file!("../sql/schema.sql", name = "schema");
 #[cfg(not(feature = "use_as_lib"))]
 #[pg_guard]
 pub extern "C" fn _PG_init() {
+    bindings::python::activate().expect("Error setting python venv");
     orm::project::init();
 }
 
diff --git a/pgml-extension/src/orm/model.rs b/pgml-extension/src/orm/model.rs
@@ -88,10 +88,6 @@ impl Model {
             },
         };
 
-        if runtime == Runtime::python {
-            let _ = crate::bindings::python::activate();
-        }
-
         let dataset = snapshot.tabular_dataset();
         let status = Status::in_progress;
         // Create the model record.
@@ -334,10 +330,7 @@ impl Model {
                     }
 
                     #[cfg(feature = "python")]
-                    Runtime::python => {
-                        let _ = crate::bindings::python::activate();
-                        crate::bindings::sklearn::Estimator::from_bytes(&data)?
-                    }
+                    Runtime::python => crate::bindings::sklearn::Estimator::from_bytes(&data)?,
 
                     #[cfg(not(feature = "python"))]
                     Runtime::python => {

Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,7 @@ extension_sql_file!("../sql/schema.sql", name = "schema");`
`24`	`24`	`#[cfg(not(feature = "use_as_lib"))]`
`25`	`25`	`#[pg_guard]`
`26`	`26`	`pub extern "C" fn _PG_init() {`
	`27`	`+ bindings::python::activate().expect("Error setting python venv");`
`27`	`28`	`orm::project::init();`
`28`	`29`	`}`
`29`	`30`