Add Flask app for managing AI stacks
This commit is contained in:
@@ -0,0 +1,117 @@
|
||||
import os
|
||||
import subprocess
|
||||
import threading
|
||||
from flask import Flask, render_template, redirect, jsonify
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
IMAGE_STACK = os.environ.get("IMAGE_STACK", "/docker/ai-image-gen")
|
||||
LLM_STACK = os.environ.get("LLM_STACK", "/docker/ai-llm")
|
||||
|
||||
# Track current active stack and swap status
|
||||
state = {
|
||||
"active": None, # "images" | "chat" | None
|
||||
"swapping": False,
|
||||
"message": "No stack active"
|
||||
}
|
||||
|
||||
def run_compose(stack_dir, command):
|
||||
"""Run a docker compose command in a stack directory."""
|
||||
subprocess.run(
|
||||
["docker", "compose", "-f", f"{stack_dir}/docker-compose.yml"] + command,
|
||||
capture_output=True
|
||||
)
|
||||
|
||||
def is_stack_running(stack_dir):
|
||||
"""Check if the GPU container for this stack is running."""
|
||||
container = "ai-image-gen-invokeai-1" if "image" in stack_dir else "ai-llm-ollama-1"
|
||||
result = subprocess.run(
|
||||
["docker", "inspect", "--format", "{{.State.Running}}", container],
|
||||
capture_output=True,
|
||||
text=True
|
||||
)
|
||||
return result.stdout.strip() == "true"
|
||||
|
||||
def do_swap(target):
|
||||
"""Stop the GPU-heavy container of inactive stack, start target."""
|
||||
state["swapping"] = True
|
||||
|
||||
if target == "images":
|
||||
state["message"] = "Stopping Ollama to free VRAM..."
|
||||
subprocess.run(["docker", "stop", "ai-llm-ollama-1"], capture_output=True)
|
||||
state["message"] = "Starting InvokeAI..."
|
||||
subprocess.run(
|
||||
["docker", "compose", "-f", f"{IMAGE_STACK}/docker-compose.yml", "up", "-d"],
|
||||
capture_output=True
|
||||
)
|
||||
else:
|
||||
state["message"] = "Stopping InvokeAI to free VRAM..."
|
||||
subprocess.run(["docker", "stop", "ai-image-gen-invokeai-1"], capture_output=True)
|
||||
state["message"] = "Starting Ollama..."
|
||||
subprocess.run(["docker", "start", "ai-llm-ollama-1"], capture_output=True)
|
||||
|
||||
state["active"] = target
|
||||
state["swapping"] = False
|
||||
state["message"] = f"{target.capitalize()} is ready."
|
||||
|
||||
@app.route("/")
|
||||
@app.route("/portal")
|
||||
def portal():
|
||||
image_running = is_stack_running(IMAGE_STACK)
|
||||
llm_running = is_stack_running(LLM_STACK)
|
||||
return render_template(
|
||||
"portal.html",
|
||||
active=state["active"],
|
||||
swapping=state["swapping"],
|
||||
message=state["message"],
|
||||
image_running=image_running,
|
||||
llm_running=llm_running
|
||||
)
|
||||
|
||||
@app.route("/switch/<target>")
|
||||
def switch(target):
|
||||
if target not in ("images", "chat"):
|
||||
return "Invalid target", 400
|
||||
|
||||
if state["swapping"]:
|
||||
return redirect("/portal")
|
||||
|
||||
# Warn if same stack requested
|
||||
if state["active"] == target:
|
||||
urls = {"chat": "https://chat.mydomain.com", "images": "https://images.mydomain.com"}
|
||||
return redirect(urls[target])
|
||||
|
||||
# Run swap in background so browser gets immediate response
|
||||
thread = threading.Thread(target=do_swap, args=(target,))
|
||||
thread.daemon = True
|
||||
thread.start()
|
||||
|
||||
return render_template(
|
||||
"switching.html",
|
||||
target=target,
|
||||
message=f"Switching to {target}... this takes about 10 seconds."
|
||||
)
|
||||
|
||||
@app.route("/chat")
|
||||
def chat():
|
||||
return redirect(URLS["chat"])
|
||||
|
||||
@app.route("/images")
|
||||
def images():
|
||||
return redirect(URLS["images"])
|
||||
|
||||
@app.route("/status")
|
||||
def status():
|
||||
"""Polled by the switching page to know when swap is done."""
|
||||
return jsonify(state)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Detect what's already running on startup
|
||||
if is_stack_running(IMAGE_STACK):
|
||||
state["active"] = "images"
|
||||
state["message"] = "Images stack is active."
|
||||
elif is_stack_running(LLM_STACK):
|
||||
state["active"] = "chat"
|
||||
state["message"] = "Chat stack is active."
|
||||
|
||||
app.run(host="0.0.0.0", port=8080)
|
||||
Reference in New Issue
Block a user