File size: 1,880 Bytes
9c0631a eac0082 9c0631a 8c57b9b 9c0631a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import subprocess
import sys
import time
# Function to install necessary dependencies
def install_dependencies():
try:
print("Installing Petals...")
subprocess.check_call([sys.executable, "sudo", "-H", "pip", "install", "--user", "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 = "deca-ai/2-mini"
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)
|