Merge pull request #1192 from LibrePhotos/feat/improved-service-manag…

…ement

Improved service management

api/management/commands/start_service.py

@@ -0,0 +1,35 @@
+from django.core.management.base import BaseCommand
+from django_q.models import Schedule
+from django_q.tasks import schedule
+
+from api.services import SERVICES, start_service
+
+
+class Command(BaseCommand):
+    help = "Start one of the services."
+
+    # Define all the services that can be started
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "service",
+            type=str,
+            help="The service to start",
+            choices=[
+                SERVICES.keys(),
+                "all",
+            ],
+        )
+
+    def handle(self, *args, **kwargs):
+        service = kwargs["service"]
+        if service == "all":
+            for svc in SERVICES.keys():
+                start_service(svc)
+            if not Schedule.objects.filter(func="api.services.check_services").exists():
+                schedule(
+                    "api.services.check_services",
+                    schedule_type=Schedule.MINUTES,
+                    minutes=1,
+                )
+        else:
+            start_service(service)

api/services.py

@@ -0,0 +1,73 @@
+import subprocess
+
+import requests
+
+from api.util import logger
+
+# Define all the services that can be started, with their respective ports
+SERVICES = {
+    "image_similarity": 8002,
+    "thumbnail": 8003,
+    "face_recognition": 8005,
+    "clip_embeddings": 8006,
+    "llm": 8008,
+    "image_captioning": 8007,
+}
+
+running_processes = {}
+
+
+def check_services():
+    for service in running_processes.keys():
+        if not is_healthy(service):
+            process = running_processes.pop(service)
+            process.terminate()
+            logger.info(f"Starting {service}")
+            start_service(service)
+
+
+def is_healthy(service):
+    port = SERVICES.get(service)
+    try:
+        res = requests.get(f"http://localhost:{port}/health")
+        return res.status_code == 200
+    except BaseException as e:
+        logger.warning(f"Error checking health of {service}: {str(e)}")
+        return False
+
+
+def start_service(service):
+    if service == "image_similarity":
+        process = subprocess.Popen(
+            [
+                "python",
+                "image_similarity/main.py",
+                "2>&1 | tee /logs/image_similarity.log",
+            ]
+        )
+    elif service in SERVICES.keys():
+        process = subprocess.Popen(
+            [
+                "python",
+                f"service/{service}/main.py",
+                "2>&1 | tee /logs/{service}.log",
+            ]
+        )
+    else:
+        logger.warning("Unknown service:", service)
+        return False
+
+    running_processes[service] = process
+    logger.info(f"Service '{service}' started successfully")
+    return True
+
+
+def stop_service(service):
+    if service in running_processes:
+        process = running_processes.pop(service)
+        process.terminate()
+        logger.info(f"Service '{service}' stopped successfully")
+        return True
+    else:
+        logger.warning(f"Service '{service}' is not running")
+        return False

api/views/services.py

@@ -0,0 +1,69 @@
+from rest_framework import status, viewsets
+from rest_framework.decorators import action
+from rest_framework.permissions import IsAdminUser
+from rest_framework.response import Response
+
+from api.services import SERVICES, is_healthy, start_service, stop_service
+
+
+class ServiceViewSet(viewsets.ViewSet):
+    permission_classes = [IsAdminUser]
+
+    def list(self, request):
+        return Response({"services": SERVICES})
+
+    def retrieve(self, request, pk=None):
+        service_name = pk
+
+        if service_name not in SERVICES:
+            return Response(
+                {"error": f"Service {service_name} not found"},
+                status=status.HTTP_404_NOT_FOUND,
+            )
+
+        healthy = is_healthy(service_name)
+        return Response({"service_name": service_name, "healthy": healthy})
+
+    @action(detail=True, methods=["post"])
+    def start(self, request, pk=None):
+        service_name = pk
+
+        if service_name not in SERVICES:
+            return Response(
+                {"error": f"Service {service_name} not found"},
+                status=status.HTTP_404_NOT_FOUND,
+            )
+
+        start_result = start_service(service_name)
+        if start_result:
+            return Response(
+                {"message": f"Service {service_name} started successfully"},
+                status=status.HTTP_200_OK,
+            )
+        else:
+            return Response(
+                {"error": f"Failed to start service {service_name}"},
+                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )
+
+    @action(detail=True, methods=["post"])
+    def stop(self, request, pk=None):
+        service_name = pk
+
+        if service_name not in SERVICES:
+            return Response(
+                {"error": f"Service {service_name} not found"},
+                status=status.HTTP_404_NOT_FOUND,
+            )
+
+        stop_result = stop_service(service_name)
+        if stop_result:
+            return Response(
+                {"message": f"Service {service_name} stopped successfully"},
+                status=status.HTTP_200_OK,
+            )
+        else:
+            return Response(
+                {"error": f"Failed to stop service {service_name}"},
+                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )

image_similarity/main.py

@@ -1,6 +1,5 @@
 import json
 
-import gevent
 from flask import Flask, jsonify, request
 from flask_restful import Api, Resource
 from gevent.pywsgi import WSGIServer
@@ -59,11 +58,21 @@ def post(self):
             return jsonify({"status": False, "result": []}), 500
 
 
+class Health(Resource):
+    def get(self):
+        return jsonify({"status": True})
+
+
 api.add_resource(BuildIndex, "/build/")
 api.add_resource(SearchIndex, "/search/")
+api.add_resource(Health, "/health/")
 
-if __name__ == "__main__":
-    logger.info("starting server")
+
+def start_server():
+    logger.info("Starting server")
     server = WSGIServer(("0.0.0.0", 8002), app)
-    server_thread = gevent.spawn(server.serve_forever)
-    gevent.joinall([server_thread])
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    start_server()

librephotos/urls.py

@@ -33,6 +33,7 @@
     jobs,
     photos,
     search,
+    services,
     sharing,
     timezone,
     upload,
@@ -170,8 +171,9 @@ def post(self, request, *args, **kwargs):
 router.register(r"api/faces", faces.FaceListView, basename="faces")
 
 router.register(r"api/exists", upload.UploadPhotoExists, basename="photo_exists")
-
 router.register(r"api/jobs", jobs.LongRunningJobViewSet, basename="jobs")
+router.register(r"api/services", services.ServiceViewSet, basename="service")
+
 urlpatterns = [
     re_path(r"^", include(router.urls)),
     re_path(r"^api/django-admin/", admin.site.urls),

service/clip_embeddings/main.py

@@ -41,6 +41,11 @@ def calculate_query_embeddings():
     return {"emb": emb, "magnitude": magnitude}, 201
 
 
+@app.route("/health", methods=["GET"])
+def health():
+    return "OK", 200
+
+
 if __name__ == "__main__":
     log("service starting")
     server = WSGIServer(("0.0.0.0", 8006), app)

service/face_recognition/main.py

@@ -48,6 +48,11 @@ def create_face_locations():
     return {"face_locations": face_locations}, 201
 
 
+@app.route("/health", methods=["GET"])
+def health():
+    return "OK", 200
+
+
 if __name__ == "__main__":
     log("service starting")
     server = WSGIServer(("0.0.0.0", 8005), app)

service/image_captioning/main.py

@@ -64,6 +64,11 @@ def export_onnx():
     return "", 200
 
 
+@app.route("/health", methods=["GET"])
+def health():
+    return "OK", 200
+
+
 def check_inactivity():
     global last_request_time
     idle_threshold = 120

service/llm/main.py

@@ -35,6 +35,11 @@ def generate_prompt():
     return {"prompt": output}, 201
 
 
+@app.route("/health", methods=["GET"])
+def health():
+    return "OK", 200
+
+
 if __name__ == "__main__":
     log("service starting")
     server = WSGIServer(("0.0.0.0", 8008), app)

service/thumbnail/main.py

@@ -31,6 +31,11 @@ def create_thumbnail():
     return {"thumbnail": destination}, 201
 
 
+@app.route("/health", methods=["GET"])
+def health():
+    return "OK", 200
+
+
 if __name__ == "__main__":
     log("service starting")
     server = WSGIServer(("0.0.0.0", 8003), app)