enhancement on swarm processes regex

src-tauri/src/swarm.rs

@@ -1,7 +1,7 @@
 use reqwest::get;
 use serde::Deserialize;
 use tauri::api::process::Command;
-use std::{env, collections::HashMap, str};
+use std::{env, collections::HashMap, str, path::PathBuf};
 
 #[derive(Deserialize)]
 struct PetalsModelInfo {
@@ -125,23 +125,6 @@ pub fn run_swarm(num_blocks: i32, model: String, public_name: String) {
         .expect("🙈 Failed to run swarm");
 }
 
-// #[tauri::command(async)]
-// pub fn run_swarm(handle: tauri::AppHandle, num_blocks: i32, model: String, public_name: String) {
-//     let petals_path = get_petals_path(handle);
-//     let env = get_environment();
-//     println!("🚀 Starting the Swarm...");
-//     let _ = Command::new("sh")
-//         .args([
-//             format!("{petals_path}/run_swarm.sh").as_str(), 
-//             &num_blocks.to_string(),
-//             &public_name,
-//             &model
-//         ])
-//         .envs(env.clone())
-//         .spawn()
-//         .expect("🙈 Failed to run swarm");
-// }
-
 
 fn get_environment() -> HashMap<String, String> {
 
@@ -190,9 +173,14 @@ pub fn get_swarm_processes() -> String {
         return "".to_string();
     }
 
+    let env = get_environment();
+    let python_path = PathBuf::from(env.get(&String::from("PREM_PYTHON")).unwrap().to_string());
+    let env_path = python_path.parent().unwrap().parent().unwrap().to_str().unwrap();
+    let regex = format!("https://github.com/bigscience-workshop/petals|{env_path}.*(petals.cli.run_server|multiprocessing.resource_tracker|from multiprocessing.spawn)");
+
     // If create_env.sh is not running, get the processes from petals
     let output = Command::new("/usr/bin/pgrep")
-        .args(&["-f", "https://github.com/bigscience-workshop/petals|petals.cli.run_server|multiprocessing.resource_tracker|from multiprocessing.spawn"])
+        .args(&["-f", &regex])
         .output()
         .map_err(|e| {
             println!("🙈 Failed to execute command: {}", e);