Renames services. Move entrypoints to a path

Author: Sergio Soto Núñez

.env.docker

@@ -1,3 +1,4 @@
 ENV=docker
-DB_URL=mongodb://admin:admin@mongodb-service/
-DB_NAME=ai_platform
+DB_URL=mongodb://admin:admin@mongodb-service:27017/
+DB_NAME=ai_platform
+REDIS_URL=redis://redis-service:6379/0

.env.local

@@ -1,3 +1,4 @@
 ENV=local
 DB_URL=mongodb://admin:[email protected]/
-DB_NAME=ai_platform
+DB_NAME=ai_platform
+REDIS_URL=redis://127.0.0.1:6379/0

Dockerfile.api

@@ -11,7 +11,9 @@ COPY --from=requirements-stage /tmp/requirements.txt /code/requirements.txt
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
 COPY ./ai_platform /code/ai_platform
-COPY ./api_entrypoint.sh /code/api_entrypoint.sh
+COPY ./scripts/api_entrypoint.sh /code/api_entrypoint.sh
 RUN chmod +x /code/api_entrypoint.sh
 
+COPY .env.* /code/
+
 CMD ["./api_entrypoint.sh"]

Dockerfile.workers

@@ -20,10 +20,12 @@ RUN poetry install --no-interaction --no-ansi
 COPY ./ai_platform /code/ai_platform
 COPY .env.* /code/
 
-COPY ./worker_initialize.sh /code/worker_initialize.sh
-COPY ./worker_entrypoint.sh /code/worker_entrypoint.sh
+COPY ./scripts/worker_initialize.sh /code/worker_initialize.sh
+COPY ./scripts/worker_entrypoint.sh /code/worker_entrypoint.sh
+
 RUN chmod +x /code/worker_entrypoint.sh
 RUN chmod +x /code/worker_initialize.sh
+
 # Execute worker initialization script
 RUN ./worker_initialize.sh
 

README.md

@@ -0,0 +1,26 @@
+
+# Project overview
+This project is a simple API that allows you to create, update, delete and list machine learning models. It also allows you to create, update, delete and list machine learning model versions. It is a simple API that allows you to manage machine learning models and their versions.
+
+## FastAPI
+
+## Celery
+
+## Redis
+
+## MongoDB
+
+
+# Project local setup
+Clone the project:
+```bash
+git clone https://github.com/ssoto/ai-platform-api.git
+cd ai-platform-api
+```
+You are going to need to have [Docker](https://www.docker.com/) installed in your machine or something similar like [Rancher](https://rancher.com/) to run the project.
+
+Then start compiling proper images and running the project:
+```bash
+docker-compose build  --build-arg ENV=docker workers ai-platform-api 
+docker-compose up -d
+```

ai_platform/api/main.py

@@ -1,23 +1,26 @@
 import logging
 from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
 from motor.motor_asyncio import AsyncIOMotorClient
+
 from ai_platform.api.public.image_tasks.router import router as image_tasks_router
 from ai_platform.api.public.images.router import router as images_router
-from ai_platform.sandbox.images_creation import startup_pipeline
+from ai_platform.task_queue.main import app as celery_app
+from ai_platform.utils import is_redis_ok, is_mongo_ok
 from ai_platform.config import settings
 
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    handlers=[logging.StreamHandler()]
 )
 
 
 def initialize_app(app: FastAPI):
     logging.info("Initializing app...")
-    # do some initialization here
-    # startup_pipeline()
     app.mongodb_client = AsyncIOMotorClient(settings.DB_URL)
     app.mongodb = app.mongodb_client[settings.DB_NAME]
+    app.celery_app = celery_app
     logging.info("App initialized")
     yield
     app.mongodb_client.close()
@@ -26,14 +29,22 @@ def initialize_app(app: FastAPI):
 app = FastAPI(
     title="AI Platform",
     root_path=settings.API_ROOT_PATH,
+    debug=True,
     lifespan=initialize_app
 )
 
 
 @app.get("/health")
 async def health(request: Request):
-    result = await request.app.mongodb_client.server_info()
-    return result
+
+    result = {
+        "mongodb": await is_mongo_ok(request.app.mongodb_client),
+        "redis": is_redis_ok(),
+    }
+    return JSONResponse(
+        content=result,
+        status_code=200 if all(result.values()) else 500
+    )
 
 app.include_router(image_tasks_router)
 app.include_router(images_router)

ai_platform/api/public/image_tasks/router.py

@@ -4,9 +4,10 @@
 from fastapi.encoders import jsonable_encoder
 
 from ai_platform.domain.image_tasks.models import ImageTask
-from ai_platform.sandbox.images_creation import generate_image
-from ai_platform.domain.image_tasks.use_cases import create_image_task, find_image_task_by_id
-from ai_platform.domain.image_repository.use_cases import  get_image_url
+from ai_platform.domain.image_tasks.use_cases import create_image_task, afind_image_task_by_id
+from ai_platform.domain.image_repository.use_cases import get_image_url
+from ai_platform.domain.task_queues.use_cases import send_generation_message
+
 
 router = APIRouter(
     prefix="/imageTasks",
@@ -18,7 +19,7 @@
 @router.get("/")
 async def retrieve(request: Request, id_task: str):
 
-    result = await find_image_task_by_id(id_task, request.app.mongodb["image_tasks"])
+    result = await afind_image_task_by_id(id_task, request.app.mongodb["image_tasks"])
     if not result:
         return JSONResponse(
             status_code=status.HTTP_404_NOT_FOUND,
@@ -40,17 +41,17 @@ async def generate(
 
     image_task = ImageTask(
         prompt=prompt,
-        status="processing"
     )
     # FIXME: this image service is a local endpoint, it should be a service
     image_task.url = get_image_url(image_task.id)
     await create_image_task(
         image_task,
         request.app.mongodb["image_tasks"]
     )
-
-    background_tasks.add_task(generate_image, image_task)
-    logging.info(f"Task {image_task.id} added to the queue")
+    background_tasks.add_task(
+        send_generation_message, image_task, request.app.celery_app,
+    )
+    logging.info(f"Task {image_task.id} send to the queue")
     return JSONResponse(
         status_code=status.HTTP_201_CREATED,
         content=image_task.dict()

ai_platform/api/public/images/router.py

@@ -1,5 +1,8 @@
-from fastapi import APIRouter, Request, Response
-from ai_platform.domain.image_tasks.use_cases import find_image_task_by_id
+import os.path
+
+from fastapi import APIRouter, Request
+from fastapi.responses import Response
+from ai_platform.domain.image_tasks.use_cases import afind_image_task_by_id
 
 router = APIRouter(
     prefix="/images",
@@ -19,10 +22,16 @@
     response_class=Response,
 )
 async def retrieve(request: Request, image_id: str):
-    image_task = await find_image_task_by_id(
+    image_task = await afind_image_task_by_id(
         image_id,
         request.app.mongodb["image_tasks"]
     )
+    if not os.path.exists(image_task.image_path):
+        return Response(
+            status_code=404,
+            content=f"Image {image_id}.png not found"
+        )
+
     with open(image_task.image_path, "rb") as image_file:
         image_bytes = image_file.read()
     return Response(

ai_platform/config/__init__.py

@@ -33,7 +33,11 @@ class DatabaseSettings(BaseSettings):
     DB_NAME: str
 
 
-class Settings(CommonSettings, ServerSettings, DatabaseSettings):
+class RedisSettings(BaseSettings):
+    REDIS_URL: str
+
+
+class Settings(CommonSettings, ServerSettings, DatabaseSettings, RedisSettings):
     pass
 
 

ai_platform/domain/image_tasks/use_cases.py

@@ -10,13 +10,20 @@ async def create_image_task(image_task: ImageTask, db_collection):
     return await db_collection.insert_one(image_task.dict(by_alias=True))
 
 
-async def find_image_task_by_id(id_task: str, db_collection):
+async def afind_image_task_by_id(id_task: str, db_collection):
     result = await db_collection.find_one({"_id": id_task})
     if result:
         return ImageTask(**result)
     return None
 
 
+def find_image_task_by_id(id_task: str, db_collection):
+    result = db_collection.find_one({"_id": id_task})
+    if result:
+        return ImageTask(**result)
+    return None
+
+
 def update_image_task(image_task: ImageTask):
     client = None
     try:

ai_platform/sandbox/__init__.py renamed to ai_platform/domain/task_queues/__init__.py

@@ -0,0 +1,30 @@
+from pydantic import BaseModel, Field
+from ai_platform.domain.image_tasks.models import ImageTask
+
+
+class ImageTaskMesage(BaseModel):
+
+    image: ImageTask = Field(..., alias="imageTask")
+
+    class Config:
+        allow_population_by_field_name = True
+        fields = {
+            "imageTask": "image"
+        }
+        json_encoders = {
+            ImageTask: lambda v: v.dict(by_alias=True)
+        }
+
+
+class TaskQueueMessage(BaseModel):
+
+    task: ImageTaskMesage = Field(..., alias="taskQueue")
+
+    class Config:
+        allow_population_by_field_name = True
+        fields = {
+            "taskQueue": "task"
+        }
+        json_encoders = {
+            ImageTaskMesage: lambda v: v.dict(by_alias=True)
+        }

ai_platform/domain/task_queues/models.py

@@ -0,0 +1,13 @@
+import logging
+from celery import Celery
+from ai_platform.domain.image_tasks.use_cases import ImageTask
+
+
+def send_generation_message(image_task: ImageTask, celery_app: Celery):
+    result = celery_app.send_task(
+        "generate_image",
+        kwargs=image_task.dict(),
+        queue="default"
+    )
+    logging.info(f"Task {result.id} sent to the queue")
+

ai_platform/domain/task_queues/use_cases.py

@@ -5,15 +5,10 @@
 import torch
 
 from ai_platform.domain.image_tasks.models import ImageTask
-from ai_platform.domain.image_tasks.use_cases import update_image_task
-from ai_platform.domain.image_repository.use_cases import update_image_task
-
-_pipe = None
 
 
 def startup_pipeline(only_download=False):
-    global _pipe
-    _pipe = DiffusionPipeline.from_pretrained(
+    pipe = DiffusionPipeline.from_pretrained(
         "runwayml/stable-diffusion-v1-5",
         torch_dtype=torch.float16,
         variant="fp16",
@@ -22,35 +17,29 @@ def startup_pipeline(only_download=False):
         return
 
     if 'macOS' in platform.platform():
-        _pipe.to("mps")
+        pipe.to("mps")
 
-    _pipe.safety_checker = None
-    _pipe.requires_safety_checker = False
+    pipe.safety_checker = None
+    pipe.requires_safety_checker = False
 
     prompt = "fake prompt to wwarmup the pipeline"
     logging.info("Warming up the pipeline")
     # First-time "warmup" pass if PyTorch version is 1.13
-    _ = _pipe(prompt, num_inference_steps=1)
+    _ = pipe(prompt, num_inference_steps=1)
     logging.info("Pipeline warmed up")
 
+    return pipe
+
 
-def generate_image(image: ImageTask):
+def create_image(pipe: DiffusionPipeline, image: ImageTask):
     # Results match those from the CPU device after the warmup pass.
-    global _pipe
-    try:
-        result = _pipe(
-            image.prompt,
-            num_inference_steps=image.generation_steps
-        )
-        image_file = result.images[0]
-        logging.info(f"Image generated: {image_file.shape}")
-        image_file.save(image.image_path)
-        image.set_completed()
-    except Exception as e:
-        image.set_failed(reason=repr(e))
-    finally:
-        update_image_task(image)
-        logging.info(f"Process finished for image {image.id}")
+    result = pipe(
+        image.prompt,
+        num_inference_steps=image.generation_steps
+    )
+    image_file = result.images[0]
+    logging.info(f"Image generated: {image_file}")
+    return image_file
 
 
 def parse_args():
@@ -79,14 +68,6 @@ def main():
     if args.download_model:
         return
 
-    # Testing the image generation
-    image_task = ImageTask(
-        prompt="A cat in the snow",
-        generation_steps=10,
-    )
-    image_task.set_processing()
-    generate_image(image_task)
-
 
 if __name__ == "__main__":
     main()