Create setup-petals-swarm.py

setup-petals-swarm.py

@@ -0,0 +1,59 @@
+import subprocess
+import sys
+import time
+
+# Function to install necessary dependencies
+def install_dependencies():
+    try:
+        print("Installing Petals...")
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/bigscience-workshop/petals"])
+        print("Petals installation successful.")
+    except subprocess.CalledProcessError:
+        print("Error during Petals installation.")
+        sys.exit(1)
+
+# Function to start the Petals server
+def start_petal_server(model_name, port=31337):
+    try:
+        command = [
+            sys.executable,
+            "-m",
+            "petals.cli.run_server",
+            model_name,
+            "--port",
+            str(port)
+        ]
+        server_process = subprocess.Popen(command)
+        print(f"Petals server started for model {model_name} on port {port}.")
+        return server_process
+    except Exception as e:
+        print(f"Error starting Petals server: {e}")
+        sys.exit(1)
+
+# Function to keep the server alive
+def keep_server_alive(server_process):
+    try:
+        while True:
+            time.sleep(60)  # Check every minute
+            if server_process.poll() is not None:
+                print("Petals server has stopped unexpectedly.")
+                break
+    except KeyboardInterrupt:
+        print("Keep-alive interrupted. Terminating server...")
+        server_process.terminate()
+        server_process.wait()
+        print("Petals server terminated.")
+
+if __name__ == "__main__":
+    # Define the model name and port
+    model_name = "meta-llama/Meta-Llama-3.1-405B-Instruct"
+    port = 31337
+
+    # Install dependencies (optional, but should be included for Docker consistency)
+    install_dependencies()
+
+    # Start the Petals server
+    server_process = start_petal_server(model_name, port)
+
+    # Keep the server alive
+    keep_server_alive(server_process)