fix: first run error will be logged, log rotation handling revamped, prompt id will be generated first.

custom_routes.py

@@ -94,19 +94,24 @@ async def comfy_deploy_run(request):
         "prompt_id": prompt_id
     }
 
+    prompt_metadata[prompt_id] = {
+        'status_endpoint': data.get('status_endpoint'),
+        'file_upload_endpoint': data.get('file_upload_endpoint'),
+    }
+
     try:
         res = post_prompt(prompt)
     except Exception as e:
         error_type = type(e).__name__
-        stack_trace = traceback.format_exc().strip().split('\n')[-2]
+        stack_trace_short = traceback.format_exc().strip().split('\n')[-2]
+        stack_trace = traceback.format_exc().strip()
         print(f"error: {error_type}, {e}")
-        print(f"stack trace: {stack_trace}")
+        print(f"stack trace: {stack_trace_short}")
-        return web.Response(status=500, reason=f"{error_type}: {e}, {stack_trace}")
+        asyncio.create_task(update_run_with_output(prompt_id, {
-
+            "error_type": error_type,
-    prompt_metadata[res['prompt_id']] = {
+            "stack_trace": stack_trace
-        'status_endpoint': data.get('status_endpoint'),
+        }))
-        'file_upload_endpoint': data.get('file_upload_endpoint'),
+        return web.Response(status=500, reason=f"{error_type}: {e}, {stack_trace_short}")
-    }
 
     status = 200
     if "error" in res:

prestartup_script.py

@@ -7,33 +7,42 @@ import threading
 import logging
 from logging.handlers import RotatingFileHandler
 
-class StreamToLogger(object):
-    def __init__(self, original, logger, log_level):
-        self.original_stdout = original
-        self.logger = logger
-        self.log_level = log_level
-
-    def write(self, buf):
-        self.original_stdout.write(buf)
-        self.original_stdout.flush()
-        for line in buf.rstrip().splitlines():
-            self.logger.log(self.log_level, line.rstrip())
-
-    def flush(self):
-        self.original_stdout.flush()
-
-
-logger = logging.getLogger()
-logger.setLevel(logging.INFO)
-
-# Create a handler that rotates log files every 500KB
 handler = RotatingFileHandler('comfy-deploy.log', maxBytes=500000, backupCount=5)
-logger.addHandler(handler)
 
-# Store original streams
 original_stdout = sys.stdout
 original_stderr = sys.stderr
 
+class StreamToLogger():
+    def __init__(self, log_level):
+        self.log_level = log_level
+
+    def write(self, buf):
+        if (self.log_level == logging.INFO):
+            original_stdout.write(buf)
+            original_stdout.flush()
+        elif (self.log_level == logging.ERROR):
+            original_stderr.write(buf)
+            original_stderr.flush()
+
+        for line in buf.rstrip().splitlines():
+            handler.handle(
+                logging.LogRecord(
+                    name="comfy-deploy",
+                    level=self.log_level,
+                    pathname="prestartup_script.py",
+                    lineno=1,
+                    msg=line.rstrip(),
+                    args=None,
+                    exc_info=None
+                )
+            )
+
+    def flush(self):
+        if (self.log_level == logging.INFO):
+            original_stdout.flush()
+        elif (self.log_level == logging.ERROR):
+            original_stderr.flush()
+
 # Redirect stdout and stderr to the logger
-sys.stdout = StreamToLogger(original_stdout, logger, logging.INFO)
+sys.stdout = StreamToLogger(logging.INFO)
-sys.stderr = StreamToLogger(original_stderr, logger, logging.ERROR)
+sys.stderr = StreamToLogger(logging.ERROR)

web/src/server/createRun.ts

@@ -58,18 +58,38 @@ export const createRun = withServerPromise(
       file_upload_endpoint: `${origin}/api/file-upload`,
     };
 
+    prompt_id = v4();
+
+    // Add to our db
+    const workflow_run = await db
+      .insert(workflowRunsTable)
+      .values({
+        id: prompt_id,
+        workflow_id: workflow_version_data.workflow_id,
+        workflow_version_id: workflow_version_data.id,
+        workflow_inputs: inputs,
+        machine_id,
+        origin: isManualRun ? "manual" : "api",
+      })
+      .returning();
+
+    revalidatePath(`/${workflow_version_data.workflow_id}`);
+
+    try {
       switch (machine.type) {
         case "runpod-serverless":
-        prompt_id = v4();
           const data = {
             input: {
               ...shareData,
               prompt_id: prompt_id,
             },
           };
-        console.log(data);
 
-        if (!machine.auth_token) {
+          if (
+            !machine.auth_token &&
+            !machine.endpoint.includes("localhost") &&
+            !machine.endpoint.includes("127.0.0.1")
+          ) {
             throw new Error("Machine auth token not found");
           }
 
@@ -88,34 +108,42 @@ export const createRun = withServerPromise(
           console.log(data, __result);
           break;
         case "classic":
-        const body = shareData;
+          const body = {
+            ...shareData,
+            prompt_id: prompt_id,
+          };
+          // console.log(body);
           const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`;
           const _result = await fetch(comfyui_endpoint, {
             method: "POST",
             body: JSON.stringify(body),
             cache: "no-store",
           });
-        if (!_result.ok)
+          // console.log(_result);
-          throw new Error(`Error creating run, ${_result.statusText}`);
+
-        const result = await ComfyAPI_Run.parseAsync(await _result.json());
+          if (!_result.ok) {
-        prompt_id = result.prompt_id;
+            let message = `Error creating run, ${_result.statusText}`;
+            try {
+              const result = await ComfyAPI_Run.parseAsync(
+                await _result.json()
+              );
+              message += ` ${result.node_errors}`;
+            } catch (error) {}
+            throw new Error(message);
+          }
+          // prompt_id = result.prompt_id;
           break;
       }
-
+    } catch (e) {
-    // Add to our db
+      console.error(e);
-    const workflow_run = await db
+      await db
-      .insert(workflowRunsTable)
+        .update(workflowRunsTable)
-      .values({
+        .set({
-        id: prompt_id,
+          status: "failed",
-        workflow_id: workflow_version_data.workflow_id,
-        workflow_version_id: workflow_version_data.id,
-        workflow_inputs: inputs,
-        machine_id,
-        origin: isManualRun ? "manual" : "api",
         })
-      .returning();
+        .where(eq(workflowRunsTable.id, workflow_run[0].id));
-
+      throw e;
-    revalidatePath(`/${workflow_version_data.workflow_id}`);
+    }
 
     return {
       workflow_run_id: workflow_run[0].id,