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)