fixing synthax and format

Author: IdirLISN

src/apps/competitions/tasks.py

@@ -3,15 +3,16 @@
 import re
 import traceback
 import zipfile
+import json
 from datetime import timedelta, datetime
-
+from django.conf import settings
 from io import BytesIO
 from tempfile import TemporaryDirectory, NamedTemporaryFile
 
 import oyaml as yaml
 import requests
 from celery._state import app_or_default
-from django.conf import settings
+from django_redis import get_redis_connection
 from django.core.exceptions import ObjectDoesNotExist
 from django.core.files.base import ContentFile
 from django.db.models import Subquery, OuterRef, Count, Case, When, Value, F
@@ -42,9 +43,12 @@
 from asgiref.sync import async_to_sync
 
 import logging
-
 logger = logging.getLogger(__name__)
 
+r = get_redis_connection("default")
+WORKERS_REGISTRY_KEY = "compute_workers_registry"
+WORKER_HEARTBEAT_TTL = 35
+
 COMPETITION_FIELDS = [
     "title",
     "description",
@@ -914,36 +918,61 @@ def refresh_compute_worker_health():
         stats = inspector.stats() or {}
         active = inspector.active() or {}
         reserved = inspector.reserved() or {}
+        active_queues = inspector.active_queues() or {}
     except Exception:
         logger.exception("Unable to inspect Celery workers")
         return
 
     for worker_name in stats.keys():
-        if not worker_name.startswith("compute-worker"):
+        queues = active_queues.get(worker_name, []) or []
+        queue_names = []
+
+        for q in queues:
+            if isinstance(q, dict) and q.get("name"):
+                queue_names.append(q["name"])
+
+        is_compute_worker = (
+            "compute-worker" in queue_names
+            or worker_name.startswith("compute-worker")
+            or worker_name.startswith("CW")
+        )
+
+        if not is_compute_worker:
             continue
 
-        raw_running_jobs = len(active.get(worker_name, [])) + len(
-            reserved.get(worker_name, [])
+        running_jobs = (
+            len(active.get(worker_name, []))
+            + len(reserved.get(worker_name, []))
         )
-        status = "busy" if raw_running_jobs > 0 else "available"
+        status = "busy" if running_jobs > 0 else "available"
 
         payload = {
             "hostname": worker_name,
             "status": status,
-            "running_jobs": raw_running_jobs,
+            "running_jobs": running_jobs,
             "timestamp": now().timestamp(),
         }
 
-        r.set(f"worker:{worker_name}:heartbeat", json.dumps(payload), ex=35)
+        heartbeat_key = f"worker:{worker_name}:heartbeat"
+
+        r.set(
+            heartbeat_key,
+            json.dumps(payload),
+            ex=WORKER_HEARTBEAT_TTL,
+        )
+
         r.hset(
             WORKERS_REGISTRY_KEY,
             worker_name,
-            json.dumps(
-                {
-                    "hostname": worker_name,
-                    "last_seen": payload["timestamp"],
-                }
-            ),
+            json.dumps({
+                "hostname": worker_name,
+                "status": status,
+                "running_jobs": running_jobs,
+                "last_seen": payload["timestamp"],
+            }),
         )
 
         _broadcast_worker_state(payload)
+        logger.info(
+            f"[WORKER-HEALTH] {worker_name} status={status} jobs={running_jobs}"
+        )

src/celery_config.py

@@ -42,10 +42,9 @@ def app_for_vhost(vhost):
         _vhost_apps[vhost] = vhost_app
     return _vhost_apps[vhost]
 
-
 app.conf.beat_schedule = {
     "refresh-compute-worker-health": {
-        "task": "chemin.vers.refresh_compute_worker_health",
-        "schedule": 5.0,
+        "task": "apps.competitions.tasks.refresh_compute_worker_health",
+        "schedule": 3.0,
     },
-}
+}

src/utils/consumers.py

@@ -5,22 +5,33 @@
 from channels.generic.websocket import AsyncJsonWebsocketConsumer
 from celery._state import app_or_default
 
+import logging
+
+logger = logging.getLogger(__name__)
+
 
 class ComputeWorkersConsumer(AsyncJsonWebsocketConsumer):
     async def connect(self):
-        user = self.scope["user"]
+        user = self.scope.get("user")
 
         if user is None or user.is_anonymous:
             await self.close()
             return
+
         await self.accept()
         self._running = True
         self._task = asyncio.create_task(self._push_workers_loop())
 
     async def disconnect(self, close_code):
         self._running = False
-        if hasattr(self, "_task"):
-            self._task.cancel()
+
+        task = getattr(self, "_task", None)
+        if task:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
 
     async def _push_workers_loop(self):
         while self._running:
@@ -46,13 +57,15 @@ def _load_snapshot(self):
             reserved = inspector.reserved() or {}
             active_queues = inspector.active_queues() or {}
         except Exception:
+            logger.exception("Unable to inspect Celery workers")
             return []
 
         workers = []
 
         for worker_name in stats.keys():
             queues = active_queues.get(worker_name, []) or []
             queue_names = []
+
             for q in queues:
                 if isinstance(q, dict) and q.get("name"):
                     queue_names.append(q["name"])
@@ -66,8 +79,9 @@ def _load_snapshot(self):
             if not is_compute_worker:
                 continue
 
-            running_jobs = len(active.get(worker_name, [])) + len(
-                reserved.get(worker_name, [])
+            running_jobs = (
+                len(active.get(worker_name, []))
+                + len(reserved.get(worker_name, []))
             )
             status = "busy" if running_jobs > 0 else "available"
 
@@ -80,4 +94,4 @@ def _load_snapshot(self):
                 }
             )
 
-        return workers
+        return workers