remove test bugs

Use mtoken

lib/backend.py

@@ -30,7 +30,7 @@ from lib.data_types import (
     BenchmarkResult
 )
 
-VERSION = "0.1.0"
+VERSION = "0.2.0"
 
 MSG_HISTORY_LEN = 100
 log = logging.getLogger(__file__)
@@ -66,10 +66,17 @@ class Backend:
     unsecured: bool = dataclasses.field(
         default_factory=lambda: bool(strtobool(os.environ.get("UNSECURED", "false"))),
     )
+    report_addr: str = dataclasses.field(
+        default_factory=lambda: os.environ.get("REPORT_ADDR", "https://run.vast.ai")
+    )
+    mtoken: str = dataclasses.field(
+        default_factory=lambda: os.environ.get("MASTER_TOKEN", "")
+    )
 
     def __post_init__(self):
         self.metrics = Metrics()
         self.metrics._set_version(self.version)
+        self.metrics._set_mtoken(self.mtoken)
         self._total_pubkey_fetch_errors = 0
         self._pubkey = self._fetch_pubkey()
         self.__start_healthcheck: bool = False
@@ -104,23 +111,19 @@ class Backend:
 
     #######################################Private#######################################
     def _fetch_pubkey(self):
-        command = ["curl", "-X", "GET", "https://run.vast.ai/pubkey/"]
-        result = subprocess.check_output(command, universal_newlines=True)
-        log.debug("public key:")
-        log.debug(result)
-        key = None
-        for _ in range(5):
-            try:
-                key = RSA.import_key(result)
-                break
-            except ValueError as e:
-                log.debug(f"Error downloading key: {e}")
-                time.sleep(15)
-        if key is None:
-            self._total_pubkey_fetch_errors += 1
-            if self._total_pubkey_fetch_errors >= MAX_PUBKEY_FETCH_ATTEMPTS:
-                self.backend_errored("Failed to get autoscaler pubkey")
-        return key
+        report_addr = self.report_addr.rstrip("/")
+        command = ["curl", "-X", "GET", f"{report_addr}/pubkey/"]
+        try:
+            result = subprocess.check_output(command, universal_newlines=True)
+            log.debug("public key:")
+            log.debug(result)
+            key = RSA.import_key(result)
+            if key is not None:
+                return key
+        except (ValueError , subprocess.CalledProcessError) as e:
+            log.debug(f"Error downloading key: {e}")
+        self.backend_errored("Failed to get autoscaler pubkey")
+       
 
     async def __handle_request(
         self,
@@ -315,10 +318,10 @@ class Backend:
                 with open(BENCHMARK_INDICATOR_FILE, "r") as f:
                     log.debug("already ran benchmark")
                     # trigger model load
-                    payload = self.benchmark_handler.make_benchmark_payload()
-                    _ = await self.__call_api(
-                        handler=self.benchmark_handler, payload=payload
-                    )
+                    # payload = self.benchmark_handler.make_benchmark_payload()
+                    # _ = await self.__call_api(
+                    #     handler=self.benchmark_handler, payload=payload
+                    # )
                     return float(f.readline())
             except FileNotFoundError:
                 pass
@@ -393,7 +396,7 @@ class Backend:
                         )
                         # some backends need a few seconds after logging successful startup before
                         # they can begin accepting requests
-                        await sleep(5)
+                        # await sleep(5)
                         try:
                             max_throughput = await run_benchmark()
                             self.__start_healthcheck = True

lib/data_types.py

@@ -286,6 +286,7 @@ class AutoScalerData:
     """Data that is reported to autoscaler"""
 
     id: int
+    mtoken: str
     version: str
     loadtime: float
     cur_load: float

lib/metrics.py

@@ -28,6 +28,7 @@ def get_url() -> str:
 @dataclass
 class Metrics:
     version: str = "0"
+    mtoken: str = ""
     last_metric_update: float = 0.0
     last_request_served: float = 0.0
     update_pending: bool = False
@@ -142,12 +143,16 @@ class Metrics:
     def _set_version(self, version: str) -> None:
         self.version = version
 
+    def _set_mtoken(self, mtoken: str) -> None:
+        self.mtoken = mtoken
+
     #######################################Private#######################################
 
     async def __send_delete_requests_and_reset(self):
         async def post(report_addr: str, idxs: list[int], success_flag: bool) -> bool:
             data = {
                 "worker_id": self.id,
+                "mtoken": self.mtoken,
                 "request_idxs": idxs,
                 "success": success_flag,
             }
@@ -209,6 +214,7 @@ class Metrics:
         def compute_autoscaler_data() -> AutoScalerData:
             return AutoScalerData(
                 id=self.id,
+                mtoken=self.mtoken,
                 version=self.version,
                 loadtime=(loadtime_snapshot or 0.0), 
                 new_load=self.model_metrics.workload_processing,
@@ -228,17 +234,25 @@ class Metrics:
 
         async def send_data(report_addr: str) -> bool:
             data = compute_autoscaler_data()
-            full_path = report_addr.rstrip("/") + "/worker_status/"
+            log_data = asdict(data)
+            def obfuscate(secret: str) -> str:
+                if secret is None:
+                    return ""
+                return secret[:7] + "..." if len(secret) > 7 else ("*" * len(secret))
+            
+            log_data["mtoken"] = obfuscate(log_data.get("mtoken"))
             log.debug(
                 "\n".join(
                     [
                         "#" * 60,
                         f"sending data to autoscaler",
-                        f"{json.dumps((asdict(data)), indent=2)}",
+                        f"{json.dumps(log_data, indent=2)}",
                         "#" * 60,
                     ]
                 )
             )
+
+            full_path = report_addr.rstrip("/") + "/worker_status/"
             for attempt in range(1, 4):
                 try:
                     session = await self.http()

lib/server.py

@@ -3,38 +3,58 @@ import logging
 from typing import List
 import ssl
 from asyncio import run, gather
-
+import asyncio
 
 from lib.backend import Backend
+from lib.metrics import Metrics
 from aiohttp import web
 
 log = logging.getLogger(__file__)
 
 
 def start_server(backend: Backend, routes: List[web.RouteDef], **kwargs):
-    log.debug("getting certificate...")
-    use_ssl = os.environ.get("USE_SSL", "false") == "true"
-    if use_ssl is True:
-        ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
-        ssl_context.load_cert_chain(
-            certfile="/etc/instance.crt",
-            keyfile="/etc/instance.key",
-        )
-    else:
-        ssl_context = None
+    try:
+        log.debug("getting certificate...")
+        use_ssl = os.environ.get("USE_SSL", "false") == "true"
+        if use_ssl is True:
+            ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+            ssl_context.load_cert_chain(
+                certfile="/etc/instance.crt",
+                keyfile="/etc/instance.key",
+            )
+        else:
+            ssl_context = None
 
-    async def main():
-        log.debug("starting server...")
-        app = web.Application()
-        app.add_routes(routes)
-        runner = web.AppRunner(app)
-        await runner.setup()
-        site = web.TCPSite(
-            runner,
-            ssl_context=ssl_context,
-            port=int(os.environ["WORKER_PORT"]),
-            **kwargs
-        )
-        await gather(site.start(), backend._start_tracking())
+        async def main():
+            log.debug("starting server...")
+            app = web.Application()
+            app.add_routes(routes)
+            runner = web.AppRunner(app)
+            await runner.setup()
+            site = web.TCPSite(
+                runner,
+                ssl_context=ssl_context,
+                port=int(os.environ["WORKER_PORT"]),
+                **kwargs
+            )
+            await gather(site.start(), backend._start_tracking())
 
-    run(main())
+        run(main())
+
+    except Exception as e:
+        err_msg = f"PyWorker failed to launch: {e}"
+        log.error(err_msg)
+
+        async def beacon():
+            metrics = Metrics()
+            metrics._set_version(getattr(backend, "version", "0"))
+            metrics._set_mtoken(getattr(backend, "mtoken", ""))
+            try:
+                while True:
+                    metrics._model_errored(err_msg)
+                    await metrics._Metrics__send_metrics_and_reset()
+                    await asyncio.sleep(10)
+            finally:
+                await metrics.aclose()
+
+        run(beacon())

start_server.sh

@@ -12,7 +12,6 @@ PYWORKER_LOG="$WORKSPACE_DIR/pyworker.log"
 REPORT_ADDR="${REPORT_ADDR:-https://run.vast.ai}"
 USE_SSL="${USE_SSL:-true}"
 WORKER_PORT="${WORKER_PORT:-3000}"
-MODEL_TYPE="${MODEL_TYPE:-image}"
 mkdir -p "$WORKSPACE_DIR"
 cd "$WORKSPACE_DIR"
 
@@ -129,5 +128,44 @@ echo "launching PyWorker server"
 # from the run prior to reboot. past logs are saved in $MODEL_LOG.old for debugging only
 [ -e "$MODEL_LOG" ] && cat "$MODEL_LOG" >> "$MODEL_LOG.old" && : > "$MODEL_LOG"
 
-(python3 -m "workers.$BACKEND.server" |& tee -a "$PYWORKER_LOG") &
-echo "launching PyWorker server done"
+
+set +e
+python3 -m "workers.$BACKEND.server" |& tee -a "$PYWORKER_LOG"
+PY_STATUS=${PIPESTATUS[0]}
+set -e
+
+if [ "${PY_STATUS}" -ne 0 ]; then
+  echo "PyWorker exited with status ${PY_STATUS}; notifying autoscaler..."
+  ERROR_MSG="PyWorker exited: code ${PY_STATUS}"
+  MTOKEN="${MASTER_TOKEN:-}"
+  VERSION="${PYWORKER_VERSION:-0}"
+
+  IFS=',' read -r -a REPORT_ADDRS <<< "${REPORT_ADDR}"
+  for addr in "${REPORT_ADDRS[@]}"; do
+    curl -sS -X POST -H 'Content-Type: application/json' \
+      -d "$(cat <<JSON
+{
+  "id": ${CONTAINER_ID:-0},
+  "mtoken": "${MTOKEN}",
+  "version": "${VERSION}",
+  "loadtime": 0,
+  "new_load": 0,
+  "cur_load": 0,
+  "rej_load": 0,
+  "max_perf": 0,
+  "cur_perf": 0,
+  "error_msg": "${ERROR_MSG}",
+  "num_requests_working": 0,
+  "num_requests_recieved": 0,
+  "additional_disk_usage": 0,
+  "working_request_idxs": [],
+  "cur_capacity": 0,
+  "max_capacity": 0,
+  "url": "${URL}"
+}
+JSON
+)" "${addr%/}/worker_status/" || true
+  done
+fi
+
+echo "launching PyWorker server done"

workers/comfyui-json/client.py

@@ -98,6 +98,7 @@ def call_text2image_workflow(
         endpoint=route_response["endpoint"],
         reqnum=route_response["reqnum"],
         url=route_response["url"],
+        request_idx=route_response["request_idx"],
     )
     
     # Build the payload for the worker request

workers/comfyui-json/data_types.py

@@ -1,5 +1,3 @@
-import dataclasses
-from enum import Enum
 import os
 import sys
 import random
@@ -15,12 +13,6 @@ from lib.data_types import ApiPayload, JsonDataException
 
 log = logging.getLogger(__file__)
 
-class ModelType(Enum):
-    image = "image"
-    audio = "audio"
-    video = "video"
-
-
 def count_workload() -> float:
     # Always 100.0 where there is a single instance of ComfyUI handling requests
     # Results will indicate % or a job completed per second.  Avoids sub 0.1 sec performance indication
@@ -29,11 +21,6 @@ def count_workload() -> float:
 @dataclasses.dataclass
 class ComfyWorkflowData(ApiPayload):
     input: dict
-    model_type: ModelType = dataclasses.field(
-        default_factory=lambda: ModelType(
-            os.environ.get("MODEL_TYPE", "image").lower()
-        )
-    )
 
     @classmethod
     def for_test(cls):
@@ -43,17 +30,15 @@ class ComfyWorkflowData(ApiPayload):
         Example: SD1.5, simple image gen 10000 steps, 512px x 512px will run for approximately 9 minutes @ ~18 it/s (RTX 4090)
         """
         # Try to load benchmark.json
-        #Note:  We should cross check with Rob if the audio sample benchmark file is correct
-        model_type = ModelType(os.environ.get("MODEL_TYPE", "image").lower())
-        benchmark_file = Path(f"workers/comfyui-json/misc/benchmark_{model_type.value}.json")
+        benchmark_file = Path("workers/comfyui-json/misc/benchmark.json")
+        
         if benchmark_file.exists():
             try:
                 with open(benchmark_file, "r") as f:
                     benchmark_workflow = json.load(f)
-                log.info(f"using benchmark json file for {model_type.value}")
                 return cls(
                     input={
-                        "request_id": f"{model_type.value}-{random.randint(1000, 99999)}",
+                        "request_id": f"test-{random.randint(1000, 99999)}",
                         "workflow_json": benchmark_workflow
                     }
                 )

workers/comfyui-json/misc/benchmark_audio.json

@@ -1,118 +0,0 @@
-{
-        "3": {
-            "inputs": {
-            "seed": 98942092715729,
-            "steps": 50,
-            "cfg": 4.98,
-            "sampler_name": "dpmpp_3m_sde_gpu",
-            "scheduler": "exponential",
-            "denoise": 1,
-            "model": [
-                "4",
-                0
-            ],
-            "positive": [
-                "6",
-                0
-            ],
-            "negative": [
-                "7",
-                0
-            ],
-            "latent_image": [
-                "11",
-                0
-            ]
-            },
-            "class_type": "KSampler",
-            "_meta": {
-            "title": "KSampler"
-            }
-        },
-        "4": {
-            "inputs": {
-            "ckpt_name": "stable-audio-open-1.0.safetensors"
-            },
-            "class_type": "CheckpointLoaderSimple",
-            "_meta": {
-            "title": "Load Checkpoint"
-            }
-        },
-        "6": {
-            "inputs": {
-            "text": "heaven church electronic dance music",
-            "clip": [
-                "10",
-                0
-            ]
-            },
-            "class_type": "CLIPTextEncode",
-            "_meta": {
-            "title": "CLIP Text Encode (Prompt)"
-            }
-        },
-        "7": {
-            "inputs": {
-            "text": "",
-            "clip": [
-                "10",
-                0
-            ]
-            },
-            "class_type": "CLIPTextEncode",
-            "_meta": {
-            "title": "CLIP Text Encode (Prompt)"
-            }
-        },
-        "10": {
-            "inputs": {
-            "clip_name": "t5-base.safetensors",
-            "type": "stable_audio",
-            "device": "default"
-            },
-            "class_type": "CLIPLoader",
-            "_meta": {
-            "title": "Load CLIP"
-            }
-        },
-        "11": {
-            "inputs": {
-            "seconds": 47.6,
-            "batch_size": 1
-            },
-            "class_type": "EmptyLatentAudio",
-            "_meta": {
-            "title": "EmptyLatentAudio"
-            }
-        },
-        "12": {
-            "inputs": {
-            "samples": [
-                "3",
-                0
-            ],
-            "vae": [
-                "4",
-                2
-            ]
-            },
-            "class_type": "VAEDecodeAudio",
-            "_meta": {
-            "title": "VAEDecodeAudio"
-            }
-        },
-        "13": {
-            "inputs": {
-            "filename_prefix": "audio/ComfyUI",
-            "audioUI": "",
-            "audio": [
-                "12",
-                0
-            ]
-            },
-            "class_type": "SaveAudio",
-            "_meta": {
-            "title": "SaveAudio"
-            }
-        }
-    }

workers/comfyui-json/misc/benchmark_image.json

@@ -1,107 +0,0 @@
-{
-      "3": {
-        "inputs": {
-          "seed": 588445435278533,
-          "steps": 20,
-          "cfg": 8,
-          "sampler_name": "euler",
-          "scheduler": "normal",
-          "denoise": 1,
-          "model": [
-            "4",
-            0
-          ],
-          "positive": [
-            "6",
-            0
-          ],
-          "negative": [
-            "7",
-            0
-          ],
-          "latent_image": [
-            "5",
-            0
-          ]
-        },
-        "class_type": "KSampler",
-        "_meta": {
-          "title": "KSampler"
-        }
-      },
-      "4": {
-        "inputs": {
-          "ckpt_name": "v1-5-pruned-emaonly-fp16.safetensors"
-        },
-        "class_type": "CheckpointLoaderSimple",
-        "_meta": {
-          "title": "Load Checkpoint"
-        }
-      },
-      "5": {
-        "inputs": {
-          "width": 512,
-          "height": 512,
-          "batch_size": 1
-        },
-        "class_type": "EmptyLatentImage",
-        "_meta": {
-          "title": "Empty Latent Image"
-        }
-      },
-      "6": {
-        "inputs": {
-          "text": "beautiful scenery nature glass bottle landscape, , purple galaxy bottle,",
-          "clip": [
-            "4",
-            1
-          ]
-        },
-        "class_type": "CLIPTextEncode",
-        "_meta": {
-          "title": "CLIP Text Encode (Prompt)"
-        }
-      },
-      "7": {
-        "inputs": {
-          "text": "text, watermark",
-          "clip": [
-            "4",
-            1
-          ]
-        },
-        "class_type": "CLIPTextEncode",
-        "_meta": {
-          "title": "CLIP Text Encode (Prompt)"
-        }
-      },
-      "8": {
-        "inputs": {
-          "samples": [
-            "3",
-            0
-          ],
-          "vae": [
-            "4",
-            2
-          ]
-        },
-        "class_type": "VAEDecode",
-        "_meta": {
-          "title": "VAE Decode"
-        }
-      },
-      "9": {
-        "inputs": {
-          "filename_prefix": "ComfyUI",
-          "images": [
-            "8",
-            0
-          ]
-        },
-        "class_type": "SaveImage",
-        "_meta": {
-          "title": "Save Image"
-        }
-      }
-    }

workers/comfyui-json/misc/benchmark_video.json

@@ -1,216 +0,0 @@
-{
-    "90": {
-        "inputs": {
-            "clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors",
-            "type": "wan",
-            "device": "default"
-        },
-        "class_type": "CLIPLoader",
-        "_meta": {
-            "title": "Load CLIP"
-        }
-    },
-    "91": {
-        "inputs": {
-            "text": "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走，裸露，NSFW",
-            "clip": [
-                "90",
-                0
-            ]
-        },
-        "class_type": "CLIPTextEncode",
-        "_meta": {
-            "title": "CLIP Text Encode (Negative Prompt)"
-        }
-    },
-    "92": {
-        "inputs": {
-            "vae_name": "wan_2.1_vae.safetensors"
-        },
-        "class_type": "VAELoader",
-        "_meta": {
-            "title": "Load VAE"
-        }
-    },
-    "93": {
-        "inputs": {
-            "shift": 8.000000000000002,
-            "model": [
-                "101",
-                0
-            ]
-        },
-        "class_type": "ModelSamplingSD3",
-        "_meta": {
-            "title": "ModelSamplingSD3"
-        }
-    },
-    "94": {
-        "inputs": {
-            "shift": 8,
-            "model": [
-                "102",
-                0
-            ]
-        },
-        "class_type": "ModelSamplingSD3",
-        "_meta": {
-            "title": "ModelSamplingSD3"
-        }
-    },
-    "95": {
-        "inputs": {
-            "add_noise": "disable",
-            "noise_seed": 0,
-            "steps": 20,
-            "cfg": 3.5,
-            "sampler_name": "euler",
-            "scheduler": "simple",
-            "start_at_step": 10,
-            "end_at_step": 10000,
-            "return_with_leftover_noise": "disable",
-            "model": [
-                "94",
-                0
-            ],
-            "positive": [
-                "99",
-                0
-            ],
-            "negative": [
-                "91",
-                0
-            ],
-            "latent_image": [
-                "96",
-                0
-            ]
-        },
-        "class_type": "KSamplerAdvanced",
-        "_meta": {
-            "title": "KSampler (Advanced)"
-        }
-    },
-    "96": {
-        "inputs": {
-            "add_noise": "enable",
-            "noise_seed": "__RANDOM_INT__",
-            "steps": 20,
-            "cfg": 3.5,
-            "sampler_name": "euler",
-            "scheduler": "simple",
-            "start_at_step": 0,
-            "end_at_step": 10,
-            "return_with_leftover_noise": "enable",
-            "model": [
-                "93",
-                0
-            ],
-            "positive": [
-                "99",
-                0
-            ],
-            "negative": [
-                "91",
-                0
-            ],
-            "latent_image": [
-                "104",
-                0
-            ]
-        },
-        "class_type": "KSamplerAdvanced",
-        "_meta": {
-            "title": "KSampler (Advanced)"
-        }
-    },
-    "97": {
-        "inputs": {
-            "samples": [
-                "95",
-                0
-            ],
-            "vae": [
-                "92",
-                0
-            ]
-        },
-        "class_type": "VAEDecode",
-        "_meta": {
-            "title": "VAE Decode"
-        }
-    },
-    "98": {
-        "inputs": {
-            "filename_prefix": "video/ComfyUI",
-            "format": "auto",
-            "codec": "auto",
-            "video": [
-                "100",
-                0
-            ]
-        },
-        "class_type": "SaveVideo",
-        "_meta": {
-            "title": "Save Video"
-        }
-    },
-    "99": {
-        "inputs": {
-            "text": "Beautiful young European woman with honey blonde hair gracefully turning her head back over shoulder, gentle smile, bright eyes looking at camera. Hair flowing in slow motion as she turns. Soft natural lighting, clean background, cinematic portrait.",
-            "clip": [
-                "90",
-                0
-            ]
-        },
-        "class_type": "CLIPTextEncode",
-        "_meta": {
-            "title": "CLIP Text Encode (Positive Prompt)"
-        }
-    },
-    "100": {
-        "inputs": {
-            "fps": 16,
-            "images": [
-                "97",
-                0
-            ]
-        },
-        "class_type": "CreateVideo",
-        "_meta": {
-            "title": "Create Video"
-        }
-    },
-    "101": {
-        "inputs": {
-            "unet_name": "wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors",
-            "weight_dtype": "default"
-        },
-        "class_type": "UNETLoader",
-        "_meta": {
-            "title": "Load Diffusion Model"
-        }
-    },
-    "102": {
-        "inputs": {
-            "unet_name": "wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors",
-            "weight_dtype": "default"
-        },
-        "class_type": "UNETLoader",
-        "_meta": {
-            "title": "Load Diffusion Model"
-        }
-    },
-    "104": {
-        "inputs": {
-            "width": 640,
-            "height": 640,
-            "length": 81,
-            "batch_size": 1
-        },
-        "class_type": "EmptyHunyuanLatentVideo",
-        "_meta": {
-            "title": "EmptyHunyuanLatentVideo"
-        }
-    }
-}

workers/comfyui/client.py

@@ -82,6 +82,7 @@ def call_custom_workflow_for_sd3(
         endpoint=message["endpoint"],
         reqnum=message["reqnum"],
         url=message["url"],
+        request_idx=message["request_idx"],
     )
     workflow = {
         "3": {