remove redis pubsub from pyworker

lib/backend.py

@@ -30,7 +30,7 @@ from lib.data_types import (
     BenchmarkResult
 )
 
-VERSION = "0.2.0"
+VERSION = "0.1.0"
 
 MSG_HISTORY_LEN = 100
 log = logging.getLogger(__file__)
@@ -66,17 +66,10 @@ 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
@@ -111,19 +104,23 @@ class Backend:
 
     #######################################Private#######################################
     def _fetch_pubkey(self):
-        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")
-       
+        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
 
     async def __handle_request(
         self,
@@ -318,10 +315,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
@@ -396,7 +393,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,7 +286,6 @@ class AutoScalerData:
     """Data that is reported to autoscaler"""
 
     id: int
-    mtoken: str
     version: str
     loadtime: float
     cur_load: float

lib/metrics.py

@@ -28,7 +28,6 @@ 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
@@ -143,16 +142,12 @@ 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,
             }
@@ -214,7 +209,6 @@ 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,
@@ -234,25 +228,17 @@ class Metrics:
 
         async def send_data(report_addr: str) -> bool:
             data = compute_autoscaler_data()
-            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"))
+            full_path = report_addr.rstrip("/") + "/worker_status/"
             log.debug(
                 "\n".join(
                     [
                         "#" * 60,
                         f"sending data to autoscaler",
-                        f"{json.dumps(log_data, indent=2)}",
+                        f"{json.dumps((asdict(data)), indent=2)}",
                         "#" * 60,
                     ]
                 )
             )
-
-            full_path = report_addr.rstrip("/") + "/worker_status/"
             for attempt in range(1, 4):
                 try:
                     session = await self.http()

workers/comfyui-json/client.py

@@ -98,7 +98,6 @@ 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/client.py

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