Stable Diffusion: tolerate transient status-poll failures; runpod-lazy v5 in-place model switching (homelab)

4a554f33b396544160b7a277cce2ded01b987627

permissionBRICK <40219477+permissionBRICK@users.noreply.github.com>

Parents: 6276520a70c3b5ac6a13750ec4490ef38a61daf2

homelab/runpod-lazy-proxy.py+111 -34 Viewed

 Models come from a CATALOG pushed by the ST UI: entries of
   {"name": display, "value": model filename, "files": [{"dest": "unet/x.gguf", "url": "..."}]}
 The pod downloads exactly the selected entry's files at boot (boot-models.py
-reads MODEL_MANIFEST). Changing the active model with a running pod recreates
+reads MODEL_MANIFEST). Changing the active model with a RUNNING pod downloads
-it (fresh disk = old model gone, only the new one downloads).
+the new files IN PLACE through the pod's model-manager (port 8189; LRU-evicts
+old models if the disk runs short) - the pod is only recreated when it runs an
+old image without the manager.
 Control API (CORS-enabled; ST warmup UI + runpod-pod.sh CLI):
   GET  /lazy/status    -> {"state": red|orange|green, "model": ..., ...}

 state = {
     'pod_id': None,
-    'model': None,      # catalog `value` the pod was provisioned for
+    'model': None,      # values_key of the models the pod holds
     'phase': 'red',     # red | orange | green
     'gpu': None,
     'since': time.time(),
     'last': time.monotonic(),
+    'ensuring': 0,      # active ensure_pod waiters (status skips probing then)
     'lock': threading.Lock(),
+}

     return '+'.join(sorted(values)) if values else ''
+def key_values(key):
+    return set(k for k in (key or '').split('+') if k)
+def needed_files(values):
+    files = []
+    for value in values or []:
+        entry = catalog_entry(value)
+        if entry:
+            files.extend(entry.get('files') or [])
+    return files
 def api(method, path, body=None):
     req = urllib.request.Request(
         f'https://rest.runpod.io/v1{path}',

         'gpuCount': 1,
         'cloudType': CLOUD_TYPE,
         'containerDiskInGb': 80,
-        'ports': ['8188/http'],
+        'ports': ['8188/http', '8189/http'],
         'env': env,
-        'dockerStartCmd': ['bash', '-c', f'python3 /boot-models.py && cd /comfyui && exec python main.py {COMFY_ARGS}'],
+        # The model-manager subshell tolerates old images without the script.
+        'dockerStartCmd': ['bash', '-c', f'python3 /boot-models.py && (python3 /model-manager.py &) && cd /comfyui && exec python main.py {COMFY_ARGS}'],
+    }
     if DATACENTERS:
         body['dataCenterIds'] = DATACENTERS

     return f'https://{pod_id}-8188.proxy.runpod.net'
+def mm_request(pod_id, method, path, obj=None, timeout=15):
+    """Calls the in-pod model manager (port 8189)."""
+    req = urllib.request.Request(
+        f'https://{pod_id}-8189.proxy.runpod.net{path}',
+        json.dumps(obj).encode() if obj is not None else None,
+        {'User-Agent': UA, 'Content-Type': 'application/json'},
+        method=method)
+    with urllib.request.urlopen(req, timeout=timeout) as resp:
+        return json.loads(resp.read() or b'{}')
 UA = 'Mozilla/5.0 (compatible; runpod-lazy-proxy)'

         log('terminate failed:', err)
-def ensure_pod(values=None, wait=True):
+def _create_pod_with_retries(values):
-    """Ensures a pod provisioned for catalog values (None = active selection)."""
-    values = values or active_values()
-    key = values_key(values)
-    with state['lock']:
-        pod_id, have, gpu = (state['pod_id'], state['model'], state['gpu'])
-        if not pod_id:
-            pod_id, have, gpu = find_pod()
-        # Only recreate when the pod's model set is KNOWN and different. An
-        # unknown set (e.g. env not reported) must not churn pods.
-        if pod_id and key and have and have != key:
-            log(f'pod has models={have}, need {key} - recreating')
-            terminate(pod_id)
-            pod_id, have, gpu = None, None, None
-        if not pod_id:
     state['phase'] = 'orange'
     state['since'] = time.time()
     last_err = None
     for attempt in range(3):
         try:
-                    pod_id, gpu = create_pod(values)
+            return create_pod(values)
-                    have = key
-                    break
         except urllib.error.HTTPError as err:
             last_err = err.read().decode()[:200]
             log(f'create_pod attempt {attempt + 1} failed:', last_err)
             time.sleep(5)
-            if not pod_id:
     state['phase'] = 'red'
     raise RuntimeError(f'could not create pod: {last_err}')
+def ensure_pod(values=None, wait=True):
+    """Ensures a pod holding the catalog values (None = active selection).
+    An existing pod gets missing models downloaded IN PLACE through the in-pod
+    model manager; recreation only happens for old-image pods without one."""
+    values = values or active_values()
+    key = values_key(values)
+    files = needed_files(values)
+    dests = [f['dest'].lstrip('/') for f in files]
+    created = False
+    with state['lock']:
+        pod_id, have, gpu = (state['pod_id'], state['model'], state['gpu'])
+        if not pod_id:
+            pod_id, have, gpu = find_pod()
+        if not pod_id:
+            pod_id, gpu = _create_pod_with_retries(values)
+            have, created = key, True
         state['pod_id'], state['model'], state['gpu'] = pod_id, have, gpu
         state['last'] = time.monotonic()

         state['phase'] = 'orange'
     if not wait:
         return pod_id
+    # Freshly created pods boot with the right manifest; existing pods need an
+    # /ensure pushed once the manager answers.
+    ensured = created or not files
+    no_manager_strikes = 0
+    state['ensuring'] += 1
+    try:
         deadline = time.monotonic() + START_TIMEOUT
         while time.monotonic() < deadline:
-        # Booting/downloading counts as activity, else the reaper would kill a
+            # Booting/downloading counts as activity, else the reaper would
-        # warming pod mid-download.
+            # kill a warming pod mid-download.
             state['last'] = time.monotonic()
-        if upstream_ready(pod_id):
+            if not ensured:
+                try:
+                    mm_request(pod_id, 'POST', '/ensure', {'files': files})
+                    ensured = True
+                    state['model'] = values_key(key_values(state['model']) | set(values))
+                    log(f'in-place ensure requested: {key}')
+                except Exception:
+                    pass  # manager still booting, or old image without one
+            ready = upstream_ready(pod_id)
+            models_ok = True
+            if not created and files and ensured:
+                try:
+                    mm = mm_request(pod_id, 'GET', '/status', timeout=10)
+                    errs = {d: e for d, e in (mm.get('errors') or {}).items() if d in dests}
+                    if errs:
+                        raise RuntimeError(f'model download failed: {errs}')
+                    models_ok = not (set(dests) - set(mm.get('present') or []))
+                except RuntimeError:
+                    raise
+                except Exception:
+                    models_ok = False
+            if ready and not ensured:
+                no_manager_strikes += 1
+            if (no_manager_strikes >= 3 and key and state['model'] and state['model'] != key):
+                # Comfy is up but there is no manager (old image): last resort.
+                log(f'pod has models={state["model"]}, need {key} - recreating (no model manager)')
+                with state['lock']:
+                    terminate(pod_id)
+                    pod_id, gpu = _create_pod_with_retries(values)
+                    state.update({'pod_id': pod_id, 'model': key, 'gpu': gpu, 'last': time.monotonic()})
+                created, ensured = True, True
+                continue
+            if ready and models_ok:
                 if state['phase'] != 'green':
                     state['phase'] = 'green'
                     state['since'] = time.time()
                     log('pod ready:', pod_id)
                 return pod_id
+            if state['phase'] == 'green':
+                state['phase'] = 'orange'  # in-place download in progress
             if state['pod_id'] != pod_id:
                 raise RuntimeError('pod replaced while waiting')
             time.sleep(5)
         raise RuntimeError('pod not ready within START_TIMEOUT')
+    finally:
+        state['ensuring'] -= 1
 def match_workflow_models(workflow_bytes):

 def status_body():
     pod_id = state['pod_id']
-    if pod_id and state['phase'] != 'green' and upstream_ready(pod_id):
+    if state['ensuring'] > 0:
+        # An ensure_pod waiter owns the phase; probing here could report green
+        # mid-download or block the status reply on a dead pod's proxy URL.
+        pass
+    elif pod_id:
+        if state['phase'] != 'green' and upstream_ready(pod_id):
             state['phase'] = 'green'
             state['since'] = time.time()
-    if not pod_id:
+    else:
         found, have, gpu = find_pod()
         if found:
             state.update({'pod_id': found, 'model': have, 'gpu': gpu})

                 catalog = {'models': payload.get('models', []), 'active': [a for a in (active or []) if a]}
                 save_catalog(catalog)
                 log(f"catalog updated: {len(catalog['models'])} models, active={catalog['active']}")
-                # Re-provision proactively when the active model changed under a live pod
+                # Fetch missing models proactively when the active set changed
-                # (known-model pods only - never churn a pod of unknown provenance).
+                # under a live pod (in place; recreate only for old images).
-                if state['pod_id'] and catalog['active'] and state['model'] and state['model'] != values_key(catalog['active']):
+                if state['pod_id'] and catalog['active'] and state['model'] and not (set(catalog['active']) <= key_values(state['model'])):
-                    log('active model set changed - re-provisioning pod')
+                    log('active model set changed - ensuring models on pod')
                     threading.Thread(target=self._safe_ensure, args=(catalog['active'],), daemon=True).start()
                 return self._reply(200, status_body())
             except Exception as err:

                 log('adoption wait failed:', err)
         threading.Thread(target=_adopt, daemon=True).start()
     server = ThreadingHTTPServer(('0.0.0.0', LISTEN_PORT), Proxy)
-    log(f'runpod-lazy v4 (catalog+lora) on :{LISTEN_PORT} (pod "{POD_NAME}", DCs {DATACENTERS or "any"}, idle {IDLE_SECONDS}s)')
+    log(f'runpod-lazy v5 (in-place model switching) on :{LISTEN_PORT} (pod "{POD_NAME}", DCs {DATACENTERS or "any"}, idle {IDLE_SECONDS}s)')
     server.serve_forever()

public/scripts/extensions/stable-diffusion/index.js+17 -6 Viewed

         const left = status.idle_seconds_left ? ` (idle stop in ${Math.round(status.idle_seconds_left / 60)}m)` : '';
         text.textContent = `ready on ${status.gpu ?? 'GPU'}${left}`;
     } else if (phase === 'orange') {
-        text.textContent = `starting (models: ${status.models ?? '?'})…`;
+        text.textContent = `starting (models: ${status.model || (status.active ?? []).join(' + ') || '?'})…`;
     } else {
         text.textContent = 'off';
+    }

     dot.style.display = getRunpodLazyUrl() ? '' : 'none';
+}
+let runpodPollFailures = 0;
 async function pollRunpodStatus() {
     const url = getRunpodLazyUrl();
     if (!url) {
         return;
+    }
-    let phase = 'red';
+    let phase = runpodLastPhase;
     try {
-        const result = await fetch(`${url}/lazy/status`, { signal: AbortSignal.timeout(5000) });
+        const result = await fetch(`${url}/lazy/status`, { signal: AbortSignal.timeout(8000) });
-        phase = renderRunpodStatus(result.ok ? await result.json() : null);
+        if (!result.ok) {
+            throw new Error(`status ${result.status}`);
+        }
+        runpodPollFailures = 0;
+        phase = renderRunpodStatus(await result.json());
     } catch {
-        renderRunpodStatus(null);
+        // A single slow/failed poll (e.g. while the proxy is provisioning a pod)
+        // must not flip the dot to red; only sustained unreachability does.
+        runpodPollFailures++;
+        if (runpodPollFailures >= 3) {
+            phase = renderRunpodStatus(null);
+        }
+    }
     clearTimeout(runpodStatusTimer);
-    runpodStatusTimer = setTimeout(pollRunpodStatus, phase === 'orange' ? RUNPOD_POLL_BUSY_MS : RUNPOD_POLL_IDLE_MS);
+    runpodStatusTimer = setTimeout(pollRunpodStatus, phase === 'orange' || runpodPollFailures ? RUNPOD_POLL_BUSY_MS : RUNPOD_POLL_IDLE_MS);
+}
 async function runpodControl(action) {