from fastapi import FastAPI
import subprocess

app = FastAPI()

@app.get("/start")
def start_task():
    try:
        # Run each script one by one
        subprocess.run(["python", "model_split.py"], check=True)
        subprocess.run(["python", "ray_node.py"], check=True)
        subprocess.run(["python", "head_node.py"], check=True)
        return {"message": "All scripts executed successfully"}
    except subprocess.CalledProcessError as e:
        return {"error": f"Script failed: {e}"}

@app.get("/stop")
def stop_ray():
    try:
        # This stops all Ray processes
        subprocess.run(["ray", "stop"], check=True)
        return {"message": "Ray stopped successfully"}
    except subprocess.CalledProcessError as e:
        return {"error": f"Failed to stop Ray: {e}"}