feat(builder): update builder for updating time log and migrate to class objects

builder/modal-builder/src/template/app.py

@@ -1,14 +1,14 @@
 from config import config
 import modal
-from modal import Image, Mount, web_endpoint, Stub, asgi_app 
+from modal import Image, Mount, web_endpoint, Stub, asgi_app, method, enter, exit
 import json
 import urllib.request
 import urllib.parse
 from pydantic import BaseModel
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import HTMLResponse
 from volume_setup import volumes
-
+from datetime import datetime
 # deploy_test = False
 
 import os
@@ -158,10 +158,11 @@ image = Image.debian_slim()
 
 target_image = image if deploy_test else dockerfile_image
 
-@stub.function(image=target_image, gpu=config["gpu"]
-   ,volumes=volumes 
-)
-def run(input: Input):
+@stub.cls(image=target_image, gpu=config["gpu"] ,volumes=volumes, timeout=60 * 10, container_idle_timeout=60 * 5)
+class ComfyDeployRunner:
+
+    @enter()
+    def setup(self):
         import subprocess
         import time
         # Make sure that the ComfyUI API is available
@@ -170,7 +171,7 @@ def run(input: Input):
         command = ["python", "main.py",
                 "--disable-auto-launch", "--disable-metadata"]
 
-    server_process = subprocess.Popen(command, cwd="/comfyui")
+        self.server_process = subprocess.Popen(command, cwd="/comfyui")
 
         check_server(
             f"http://{COMFY_HOST}",
@@ -178,15 +179,24 @@ def run(input: Input):
             COMFY_API_AVAILABLE_INTERVAL_MS,
         )
 
+    @exit()
+    def cleanup(self, exc_type, exc_value, traceback):
+        self.server_process.terminate()
+
+    @method()
+    def run(self, input: Input):
+        data = json.dumps({
+            "run_id": input.prompt_id,
+            "status": "started",
+            "time": datetime.now().isoformat()
+        }).encode('utf-8')
+        req = urllib.request.Request(input.status_endpoint, data=data, method='POST')
+        urllib.request.urlopen(req)
+
         job_input = input
 
-    # print(f"comfy-modal - got input {job_input}")
-
-    # Queue the workflow
         try:
-        # job_input is the json input
-        queued_workflow = queue_workflow_comfy_deploy(
-            job_input)  # queue_workflow(workflow)
+            queued_workflow = queue_workflow_comfy_deploy(job_input)  # queue_workflow(workflow)
             prompt_id = queued_workflow["prompt_id"]
             print(f"comfy-modal - queued workflow with ID {prompt_id}")
         except Exception as e:
@@ -223,24 +233,31 @@ def run(input: Input):
             return {"error": f"Error waiting for image generation: {str(e)}"}
 
         print(f"comfy-modal - Finished, turning off")
-    server_process.terminate()
 
-    # Get the generated image and return it as URL in an AWS bucket or as base64
-    # images_result = process_output_images(history[prompt_id].get("outputs"), job["id"])
-    # result = {**images_result, "refresh_worker": REFRESH_WORKER}
         result = {"status": status}
 
         return result
-    print("Running remotely on Modal!")
     
 
 @web_app.post("/run")
-async def bar(request_input: RequestInput):
-    # print(request_input)
+async def post_run(request_input: RequestInput):
     if not deploy_test:
-        return run.remote(request_input.input)
-    # pass
+        # print(request_input.input.prompt_id, request_input.input.status_endpoint)
+        data = json.dumps({
+            "run_id": request_input.input.prompt_id,
+            "status": "queued",
+            "time": datetime.now().isoformat()
+        }).encode('utf-8')
+        req = urllib.request.Request(request_input.input.status_endpoint, data=data, method='POST')
+        urllib.request.urlopen(req)
 
+        model = ComfyDeployRunner()
+        call = model.run.spawn(request_input.input)
+
+        # call = run.spawn()
+        return {"call_id": call.object_id}
+    
+    return {"call_id": None}
 
 @stub.function(image=image
    ,volumes=volumes

builder/modal-builder/src/template/config.py

@@ -1,6 +1,6 @@
 config = {
     "name": "my-app",
-    "deploy_test": "True",
+    "deploy_test": "False",
     "gpu": "T4", 
     "public_model_volume": "model-store",
     "private_model_volume": "private-model-store",

builder/modal-builder/src/template/data/snapshot.json

@@ -1,5 +1,10 @@
 {
   "comfyui": "d0165d819afe76bd4e6bdd710eb5f3e571b6a804",
-    "git_custom_nodes": {},
+  "git_custom_nodes": {
+    "https://github.com/BennyKok/comfyui-deploy.git": {
+      "hash": "a838cb7ad425e5652c3931fbafdc886b53c48a22",
+      "disabled": false
+    }
+  },
   "file_custom_nodes": []
 }