Fix clear_gpu_cache (#898)

levkk · web-flow · commit 160fe576c482 · 2023-08-07T16:54:14.000-07:00
diff --git a/pgml-extension/requirements.txt b/pgml-extension/requirements.txt
@@ -25,3 +25,4 @@ transformers==4.31.0
 xgboost==1.7.6
 langchain==0.0.237
 einops==0.6.1
+pynvml==11.5.0
diff --git a/pgml-extension/src/bindings/transformers/mod.rs b/pgml-extension/src/bindings/transformers/mod.rs
@@ -348,6 +348,8 @@ pub fn load_dataset(
 }
 
 pub fn clear_gpu_cache(memory_usage: Option<f32>) -> Result<bool> {
+    crate::bindings::venv::activate();
+
     Python::with_gil(|py| -> Result<bool> {
         let clear_gpu_cache: Py<PyAny> = PY_MODULE.getattr(py, "clear_gpu_cache")?;
         let success = clear_gpu_cache

Original file line number	Diff line number	Diff line change
`@@ -348,6 +348,8 @@ pub fn load_dataset(`
`348`	`348`	`}`
`349`	`349`
`350`	`350`	`pub fn clear_gpu_cache(memory_usage: Option<f32>) -> Result<bool> {`
	`351`	`+ crate::bindings::venv::activate();`
	`352`	`+`
`351`	`353`	`Python::with_gil(\|py\| -> Result<bool> {`
`352`	`354`	`let clear_gpu_cache: Py<PyAny> = PY_MODULE.getattr(py, "clear_gpu_cache")?;`
`353`	`355`	`let success = clear_gpu_cache`