Added support for Swarm Mode

README.md

@@ -95,7 +95,7 @@ Let us know which Model / Feature you wanna see in Prem [here](https://feedback.
 - [Docker](https://docs.docker.com/engine/install/)
 - [Tauri](https://tauri.app/v1/guides/getting-started/prerequisites#installing)
 
-### Run the app with Tauri
+### Run the app with Tauri (Desktop App)
 
 > Make sure to remove any value in the `.env` file.
 
@@ -104,7 +104,7 @@ npm i
 npm run tauri dev
 ```
 
-### Run the app with React
+### Run the app with React (Browser App)
 
 ```bash
 # run the daemon

src-tauri/src/main.rs

@@ -2,10 +2,11 @@
 #![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
 
 use reqwest::blocking::get;
+use reqwest::get as reqwest_get;
 use serde::Deserialize;
-use std::{env, thread};
+use std::{env, thread, str};
 use tauri::{
-    AboutMetadata, api::process::Command, CustomMenuItem, Manager, Menu, MenuItem, RunEvent,
+    api::process::Command, AboutMetadata, CustomMenuItem, Manager, Menu, MenuItem, RunEvent,
     Submenu, SystemTray, SystemTrayEvent, SystemTrayMenu, SystemTrayMenuItem, WindowEvent,
 };
 
@@ -26,6 +27,12 @@ struct Config {
     prem: Prem,
 }
 
+#[derive(Deserialize)]
+struct ModelInfo {
+    name: String,
+    state: String,
+}
+
 #[tauri::command]
 fn run_container() {
     // check if docker is running
@@ -99,6 +106,156 @@ fn is_docker_running() -> bool {
     return false;
 }
 
+fn is_python_installed() -> bool {
+    let output = Command::new("/usr/bin/python3")
+        .args(["--version"])
+        .output()
+        .map_err(|e| {
+            println!("Failed to execute python --version: {}", e);
+            e
+        });
+
+    if !output.unwrap().stdout.is_empty() {
+        println!("🐍 Python is installed");
+        return true;
+    }
+    return false;
+}
+
+#[tauri::command]
+fn is_swarm_supported() -> bool {
+    match env::consts::OS {
+        "macos" => true,
+        _ => false
+    }
+}
+
+#[tauri::command]
+async fn get_petals_models() -> Result<Vec<String>, String> {
+    let url = "https://health.petals.dev/api/v1/state";
+    let response = reqwest_get(url).await.map_err(|err| err.to_string())?;
+
+    if response.status().is_success() {
+        let json_data: serde_json::Value = response.json().await.map_err(|err| err.to_string())?;
+
+        let models: Vec<String> = json_data["model_reports"]
+            .as_array()
+            .unwrap_or(&vec![])
+            .iter()
+            .filter_map(|model_report| {
+                let model_info: Result<ModelInfo, _> =
+                    serde_json::from_value(model_report.clone());
+                match model_info {
+                    Ok(model_info) if model_info.state == "healthy" => Some(model_info.name),
+                    _ => None,
+                }
+            })
+            .collect();
+
+        Ok(models)
+    } else {
+        Err("Request failed".to_string())
+    }
+}
+
+
+
+#[tauri::command]
+fn is_swarm_mode_running() -> bool {
+    let output_value = get_swarm_processes();
+
+    if !output_value.is_empty() {
+        println!(
+            "🏃‍♀️ Processeses running: {}",
+            output_value.replace("\n", " ")
+        );
+        return true;
+    }
+    return false;
+}
+
+fn get_username() -> String {
+    let output = Command::new("whoami").output();
+
+    match output {
+        Ok(output) => {
+            output.stdout.trim().to_string()
+        },
+        Err(_) => "prem-app".to_string(),
+    }
+}
+
+#[tauri::command]
+fn run_swarm_mode(num_blocks: i32, model: String) {
+    if is_python_installed() {
+        thread::spawn(move || {
+            println!("🚀 Starting the Swarm...");
+
+            let _ = Command::new("/usr/bin/python3")
+                .args(&["-m", "pip", "install", "petals==2.2.0"])
+                .output()
+                .expect("🙈 Failed to execute command");
+
+            // Print stdout and stderr
+            // println!("🛠️ Installing the dependencies >>> {}", String::from_utf8_lossy(&output.stdout));
+            // eprintln!("{}", String::from_utf8_lossy(&output.stderr));
+
+            let username = get_username();
+            let _ = Command::new("/usr/bin/python3")
+                .args(&[
+                    "-m",
+                    "petals.cli.run_server",
+                    "--num_blocks",
+                    &num_blocks.to_string(),
+                    "--public_name",
+                    &username,
+                    &model,
+                ])
+                .output()
+                .expect("🙈 Failed to execute command");
+
+            // Print stdout and stderr
+            // println!("🚀 Running the Swarm >>> {}", String::from_utf8_lossy(&output.stdout));
+            // eprintln!("{}", String::from_utf8_lossy(&output.stderr));
+        });
+    } else {
+        println!("🙈 Python is not installed");
+    }
+}
+
+fn get_swarm_processes() -> String {
+    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"])
+        .output()
+        .map_err(|e| {
+            println!("🙈 Failed to execute command: {}", e);
+            e
+        });
+
+    let output_value = output.unwrap().stdout;
+    return output_value;
+}
+
+#[tauri::command]
+fn stop_swarm_mode() {
+    println!("🛑 Stopping the Swarm...");
+    let processes = get_swarm_processes().replace("\n", " ");
+    println!("🛑 Stopping Processes: {}", processes);
+    let processes = processes.split(" ").collect::<Vec<&str>>();
+
+    for process in processes {
+        println!("🛑 Stopping Process: {}", process);
+        let _ = Command::new("kill")
+            .args(&[process.to_string()])
+            .output()
+            .map_err(|e| {
+                println!("🙈 Failed to execute command: {}", e);
+                e
+            });
+    }
+    println!("🛑 Stopped all the Swarm Processes.");
+}
+
 #[tauri::command]
 fn is_container_running() -> Result<bool, String> {
     let output = Command::new("/usr/local/bin/docker")
@@ -131,10 +288,9 @@ fn main() {
                 ))
                 .add_native_item(MenuItem::Minimize)
                 .add_native_item(MenuItem::Hide)
-                .add_item(CustomMenuItem::new(
-                    "quit",
-                    "Quit Prem App",
-                ).accelerator("CommandOrControl+Q")),
+                .add_item(
+                    CustomMenuItem::new("quit", "Quit Prem App").accelerator("CommandOrControl+Q"),
+                ),
         ))
         .add_submenu(Submenu::new(
             "Edit",
@@ -164,11 +320,16 @@ fn main() {
     let system_tray = SystemTray::new().with_menu(tray_menu);
 
     #[allow(unused_mut)]
-        let mut app = tauri::Builder::default()
+    let mut app = tauri::Builder::default()
         .invoke_handler(tauri::generate_handler![
             run_container,
             is_docker_running,
             is_container_running,
+            is_swarm_supported,
+            get_petals_models,
+            run_swarm_mode,
+            stop_swarm_mode,
+            is_swarm_mode_running
         ])
         .menu(menu)
         .on_menu_event(|event| match event.menu_item_id() {

src/modules/service/components/Petals.tsx

@@ -0,0 +1,7 @@
+import clsx from "clsx";
+
+const Petals = ({ className }: React.ButtonHTMLAttributes<HTMLButtonElement>) => {
+  return <sup className={clsx("text-[9px] text-white ml-2", className)}>PETALS</sup>;
+};
+
+export default Petals;

src/modules/service/components/ServiceCard.tsx

@@ -10,6 +10,7 @@ import { SERVICES_KEY } from "shared/hooks/useServices";
 import type { ServiceCardProps } from "../types";
 
 import Beta from "./Beta";
+import Petals from "./Petals";
 import ServiceActions from "./ServiceActions";
 
 const ServiceCard = ({ className, icon, service }: ServiceCardProps) => {
@@ -48,6 +49,7 @@ const ServiceCard = ({ className, icon, service }: ServiceCardProps) => {
       <h3>
         {title}
         {service.beta && <Beta className="md:-top-2 top-[-5px]" />}
+        {service.petals && <Petals className="md:-top-2 top-[-5px]" />}
       </h3>
     </Link>
   );

src/modules/service/types.ts

@@ -45,6 +45,7 @@ export type Service = {
   volumeName: string | null;
   volumePath: string | null;
   beta: boolean;
+  petals: boolean;
   comingSoon: boolean;
   promptTemplate: string;
   invokeMethod: {

src/modules/settings/components/Settings.tsx

@@ -8,6 +8,7 @@ import { useLockedBody, useMediaQuery } from "usehooks-ts";
 
 import AdvancedSettings from "./AdvancedSettings";
 import GPUResources from "./GPUResources";
+import SwarmMode from "./SwarmMode";
 import SystemResources from "./SystemResources";
 
 const Settings = () => {
@@ -49,6 +50,9 @@ const Settings = () => {
         <div className="lg:max-w-[80%] mx-auto settings">
           <AdvancedSettings />
         </div>
+        <div className="lg:max-w-[80%] mx-auto settings">
+          <SwarmMode />
+        </div>
       </div>
     </AppContainer>
   );