Added support for Swarm Mode

src-tauri/src/main.rs

@@ -5,7 +5,7 @@ use reqwest::blocking::get;
 use serde::Deserialize;
 use std::{env, thread};
 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,
 };
 
@@ -99,6 +99,106 @@ 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_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;
+}
+
+#[tauri::command]
+fn run_swarm_mode() {
+    if is_python_installed() {
+        thread::spawn(|| {
+            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 _ = Command::new("/usr/bin/python3")
+                .args(&[
+                    "-m",
+                    "petals.cli.run_server",
+                    "--num_blocks",
+                    "10",
+                    "--public_name",
+                    "prem-app",
+                    "petals-team/StableBeluga2",
+                ])
+                .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 +231,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 +263,14 @@ 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,
+            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>
   );

src/modules/settings/components/SwarmMode.tsx

@@ -0,0 +1,61 @@
+import React, { useState, useEffect } from "react";
+import { runSwarmMode, checkSwarmModeRunning, stopSwarmMode } from "shared/helpers/utils";
+
+import PrimaryButton from "../../../shared/components/PrimaryButton";
+
+const SwarmMode = () => {
+  const [swarmMode, setSwarmMode] = useState(false);
+
+  useEffect(() => {
+    const intervalId = setInterval(async () => {
+      const isRunning = await checkSwarmModeRunning();
+      setSwarmMode(isRunning);
+    }, 1000);
+
+    console.log("intervalId", intervalId);
+
+    return () => {
+      clearInterval(intervalId);
+    };
+  }, []);
+
+  const onStart = async () => {
+    try {
+      await runSwarmMode();
+      setSwarmMode(true);
+    } catch (error) {
+      console.error(error);
+    }
+  };
+
+  const onStop = async () => {
+    try {
+      await stopSwarmMode();
+      setSwarmMode(false);
+    } catch (error) {
+      console.error(error);
+    }
+  };
+
+  return (
+    <div className="flex items-end justify-between mr-[45px]">
+      {swarmMode ? (
+        <>
+          <label className="text-grey-300 mr-2 backend-url md:text-lg text-[11px]">
+            Swarm Mode Enabled
+          </label>
+          <PrimaryButton onClick={onStop}>Stop</PrimaryButton>
+        </>
+      ) : (
+        <>
+          <label className="text-grey-300 mr-2 backend-url md:text-lg text-[11px]">
+            Swarm Mode Not Enabled
+          </label>
+          <PrimaryButton onClick={onStart}>Start</PrimaryButton>
+        </>
+      )}
+    </div>
+  );
+};
+
+export default SwarmMode;

src/shared/helpers/utils.ts

@@ -29,10 +29,29 @@ export const checkIsServerRunning = async () => {
 
 export const runDockerContainer = async () => {
   const containerRunning = await checkIsContainerRunning();
-  if (containerRunning) return;
+  if (containerRunning) {
+    return;
+  }
   await invoke("run_container");
 };
 
+export const checkSwarmModeRunning = async () => {
+  const check = await invoke("is_swarm_mode_running");
+  return Boolean(check);
+};
+
+export const stopSwarmMode = async () => {
+  await invoke("stop_swarm_mode");
+};
+
+export const runSwarmMode = async () => {
+  const isSwarmMode = await checkSwarmModeRunning();
+  if (isSwarmMode) {
+    return;
+  }
+  await invoke("run_swarm_mode");
+};
+
 export const isIP = (host: string): boolean => {
   if (host.includes("localhost")) return true;
   /**