|
| 1 | +"""ECS task scale-in protection via the container agent endpoint. |
| 2 | +
|
| 3 | +When at least one GUI session is active the task is marked as protected so |
| 4 | +that ECS auto-scaling and deployment scale-in events cannot terminate it. |
| 5 | +Protection is renewed periodically (before the expiry window closes) and |
| 6 | +cleared when the last GUI session ends. |
| 7 | +
|
| 8 | +The module is a no-op when ``ECS_AGENT_URI`` is not set (local development). |
| 9 | +""" |
| 10 | + |
| 11 | +import json |
| 12 | +import os |
| 13 | +import threading |
| 14 | + |
| 15 | +import requests |
| 16 | + |
| 17 | +from aind_ephys_portal.session_logging import list_gui_sessions |
| 18 | + |
| 19 | +# --- Configuration ----------------------------------------------------------- |
| 20 | + |
| 21 | +PROTECTION_DURATION_MINUTES = 60 |
| 22 | +"""Each protection window lasts this long. Renewed before it expires.""" |
| 23 | + |
| 24 | +_REFRESH_INTERVAL_SECONDS = 50 * 60 # 50 min — well before 60-min expiry |
| 25 | + |
| 26 | +# --- Module state (guarded by _lock) ----------------------------------------- |
| 27 | + |
| 28 | +_lock = threading.Lock() |
| 29 | +_is_protected = False |
| 30 | +_refresh_timer: threading.Timer | None = None |
| 31 | + |
| 32 | + |
| 33 | +# --- Low-level ECS agent call ------------------------------------------------ |
| 34 | + |
| 35 | +def _agent_uri() -> str | None: |
| 36 | + return os.environ.get("ECS_AGENT_URI") |
| 37 | + |
| 38 | + |
| 39 | +def _set_task_protection(enabled: bool, expires_minutes: int = PROTECTION_DURATION_MINUTES) -> bool: |
| 40 | + """PUT to the ECS agent endpoint. Returns True on success.""" |
| 41 | + uri = _agent_uri() |
| 42 | + if not uri: |
| 43 | + action = "protect" if enabled else "unprotect" |
| 44 | + print(f"[ecs_protection] No ECS_AGENT_URI — skipping {action} (local dev)") |
| 45 | + return True # treat as success in local dev |
| 46 | + |
| 47 | + url = f"{uri}/task-protection/v1/state" |
| 48 | + body: dict = {"ProtectionEnabled": enabled} |
| 49 | + if enabled: |
| 50 | + body["ExpiresInMinutes"] = expires_minutes |
| 51 | + |
| 52 | + try: |
| 53 | + resp = requests.put(url, json=body, timeout=5) |
| 54 | + data = resp.json() |
| 55 | + if "error" in data: |
| 56 | + print(f"[ecs_protection] Agent error: {json.dumps(data['error'])}") |
| 57 | + return False |
| 58 | + print(f"[ecs_protection] Task protection set: enabled={enabled}, response={json.dumps(data)}") |
| 59 | + return True |
| 60 | + except Exception as exc: |
| 61 | + print(f"[ecs_protection] Failed to set protection (enabled={enabled}): {exc}") |
| 62 | + return False |
| 63 | + |
| 64 | + |
| 65 | +# --- High-level API ---------------------------------------------------------- |
| 66 | + |
| 67 | +def _cancel_refresh_timer(): |
| 68 | + global _refresh_timer |
| 69 | + if _refresh_timer is not None: |
| 70 | + _refresh_timer.cancel() |
| 71 | + _refresh_timer = None |
| 72 | + |
| 73 | + |
| 74 | +def _schedule_refresh(): |
| 75 | + global _refresh_timer |
| 76 | + _cancel_refresh_timer() |
| 77 | + _refresh_timer = threading.Timer(_REFRESH_INTERVAL_SECONDS, _refresh_protection) |
| 78 | + _refresh_timer.daemon = True |
| 79 | + _refresh_timer.start() |
| 80 | + |
| 81 | + |
| 82 | +def _refresh_protection(): |
| 83 | + """Called by the timer — renew protection if sessions still exist, else clear.""" |
| 84 | + with _lock: |
| 85 | + if len(list_gui_sessions()) > 0: |
| 86 | + print("[ecs_protection] Renewing task protection (sessions still active)") |
| 87 | + _set_task_protection(True) |
| 88 | + _schedule_refresh() |
| 89 | + else: |
| 90 | + print("[ecs_protection] No sessions at refresh time — clearing protection") |
| 91 | + _unprotect_task_locked() |
| 92 | + |
| 93 | + |
| 94 | +def _unprotect_task_locked(): |
| 95 | + """Must be called while holding _lock.""" |
| 96 | + global _is_protected |
| 97 | + _cancel_refresh_timer() |
| 98 | + _set_task_protection(False) |
| 99 | + _is_protected = False |
| 100 | + |
| 101 | + |
| 102 | +def protect_task(): |
| 103 | + """Enable scale-in protection for this task (idempotent while protected).""" |
| 104 | + global _is_protected |
| 105 | + with _lock: |
| 106 | + if _is_protected: |
| 107 | + return |
| 108 | + if _set_task_protection(True): |
| 109 | + _is_protected = True |
| 110 | + _schedule_refresh() |
| 111 | + |
| 112 | + |
| 113 | +def unprotect_task(): |
| 114 | + """Disable scale-in protection for this task.""" |
| 115 | + with _lock: |
| 116 | + if not _is_protected: |
| 117 | + return |
| 118 | + _unprotect_task_locked() |
| 119 | + |
| 120 | + |
| 121 | +def is_protected() -> bool: |
| 122 | + return _is_protected |
| 123 | + |
| 124 | + |
| 125 | +def get_protection_status() -> dict | None: |
| 126 | + """GET the current protection status from the ECS agent. Returns None in local dev.""" |
| 127 | + uri = _agent_uri() |
| 128 | + if not uri: |
| 129 | + return {"local_dev": True, "is_protected": _is_protected} |
| 130 | + try: |
| 131 | + resp = requests.get(f"{uri}/task-protection/v1/state", timeout=5) |
| 132 | + return resp.json() |
| 133 | + except Exception as exc: |
| 134 | + return {"error": str(exc)} |
0 commit comments