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

4a554f33b396544160b7a277cce2ded01b987627

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

2 files changed, +128 -40Showing whitespace changes
homelab/runpod-lazy-proxy.py+111 -34
@@ -5,8 +5,10 @@ SillyTavern's fallback chain points at this proxy as a normal ComfyUI server.
55Models come from a CATALOG pushed by the ST UI: entries of
66 {"name": display, "value": model filename, "files": [{"dest": "unet/x.gguf", "url": "..."}]}
77The pod downloads exactly the selected entry's files at boot (boot-models.py
88reads MODEL_MANIFEST). Changing the active model with a runningRUNNING pod recreatesdownloads
9-it (fresh disk = old model gone, only the new one downloads).
9+the new files IN PLACE through the pod's model-manager (port 8189; LRU-evicts
10+old models if the disk runs short) - the pod is only recreated when it runs an
11+old image without the manager.
1012
1113Control API (CORS-enabled; ST warmup UI + runpod-pod.sh CLI):
1214 GET /lazy/status -> {"state": red|orange|green, "model": ..., ...}
@@ -45,11 +47,12 @@ CATALOG_FILE = os.environ.get('CATALOG_FILE', '/app/catalog.json')
4547
4648state = {
4749 'pod_id': None,
4850 'model': None, # catalogvalues_key `value`of the podmodels wasthe provisionedpod forholds
4951 'phase': 'red', # red | orange | green
5052 'gpu': None,
5153 'since': time.time(),
5254 'last': time.monotonic(),
55+ 'ensuring': 0, # active ensure_pod waiters (status skips probing then)
5356 'lock': threading.Lock(),
5457}
5558
@@ -91,6 +94,19 @@ def values_key(values):
9194 return '+'.join(sorted(values)) if values else ''
9295
9396
97+def key_values(key):
98+ return set(k for k in (key or '').split('+') if k)
99+
100+
101+def needed_files(values):
102+ files = []
103+ for value in values or []:
104+ entry = catalog_entry(value)
105+ if entry:
106+ files.extend(entry.get('files') or [])
107+ return files
108+
109+
94110def api(method, path, body=None):
95111 req = urllib.request.Request(
96112 f'https://rest.runpod.io/v1{path}',
@@ -131,9 +147,10 @@ def create_pod(values):
131147 'gpuCount': 1,
132148 'cloudType': CLOUD_TYPE,
133149 'containerDiskInGb': 80,
134150 'ports': ['8188/http', '8189/http'],
135151 'env': env,
136- 'dockerStartCmd': ['bash', '-c', f'python3 /boot-models.py && cd /comfyui && exec python main.py {COMFY_ARGS}'],
152+ # The model-manager subshell tolerates old images without the script.
153+ 'dockerStartCmd': ['bash', '-c', f'python3 /boot-models.py && (python3 /model-manager.py &) && cd /comfyui && exec python main.py {COMFY_ARGS}'],
137154 }
138155 if DATACENTERS:
139156 body['dataCenterIds'] = DATACENTERS
@@ -146,6 +163,17 @@ def pod_url(pod_id):
146163 return f'https://{pod_id}-8188.proxy.runpod.net'
147164
148165
166+def mm_request(pod_id, method, path, obj=None, timeout=15):
167+ """Calls the in-pod model manager (port 8189)."""
168+ req = urllib.request.Request(
169+ f'https://{pod_id}-8189.proxy.runpod.net{path}',
170+ json.dumps(obj).encode() if obj is not None else None,
171+ {'User-Agent': UA, 'Content-Type': 'application/json'},
172+ method=method)
173+ with urllib.request.urlopen(req, timeout=timeout) as resp:
174+ return json.loads(resp.read() or b'{}')
175+
176+
149177UA = 'Mozilla/5.0 (compatible; runpod-lazy-proxy)'
150178
151179
@@ -171,36 +199,38 @@ def terminate(pod_id):
171199 log('terminate failed:', err)
172200
173201
174202def ensure_pod_create_pod_with_retries(values=None, wait=True):
175- """Ensures a pod provisioned for catalog values (None = active selection)."""
176- values = values or active_values()
177- key = values_key(values)
178- with state['lock']:
179- pod_id, have, gpu = (state['pod_id'], state['model'], state['gpu'])
180- if not pod_id:
181- pod_id, have, gpu = find_pod()
182- # Only recreate when the pod's model set is KNOWN and different. An
183- # unknown set (e.g. env not reported) must not churn pods.
184- if pod_id and key and have and have != key:
185- log(f'pod has models={have}, need {key} - recreating')
186- terminate(pod_id)
187- pod_id, have, gpu = None, None, None
188- if not pod_id:
189203 state['phase'] = 'orange'
190204 state['since'] = time.time()
191205 last_err = None
192206 for attempt in range(3):
193207 try:
194208 pod_id, gpu = return create_pod(values)
195- have = key
196- break
197209 except urllib.error.HTTPError as err:
198210 last_err = err.read().decode()[:200]
199211 log(f'create_pod attempt {attempt + 1} failed:', last_err)
200212 time.sleep(5)
201- if not pod_id:
202213 state['phase'] = 'red'
203214 raise RuntimeError(f'could not create pod: {last_err}')
215+
216+
217+def ensure_pod(values=None, wait=True):
218+ """Ensures a pod holding the catalog values (None = active selection).
219+
220+ An existing pod gets missing models downloaded IN PLACE through the in-pod
221+ model manager; recreation only happens for old-image pods without one."""
222+ values = values or active_values()
223+ key = values_key(values)
224+ files = needed_files(values)
225+ dests = [f['dest'].lstrip('/') for f in files]
226+ created = False
227+ with state['lock']:
228+ pod_id, have, gpu = (state['pod_id'], state['model'], state['gpu'])
229+ if not pod_id:
230+ pod_id, have, gpu = find_pod()
231+ if not pod_id:
232+ pod_id, gpu = _create_pod_with_retries(values)
233+ have, created = key, True
204234 state['pod_id'], state['model'], state['gpu'] = pod_id, have, gpu
205235 state['last'] = time.monotonic()
206236
@@ -208,21 +238,63 @@ def ensure_pod(values=None, wait=True):
208238 state['phase'] = 'orange'
209239 if not wait:
210240 return pod_id
241+ # Freshly created pods boot with the right manifest; existing pods need an
242+ # /ensure pushed once the manager answers.
243+ ensured = created or not files
244+ no_manager_strikes = 0
245+ state['ensuring'] += 1
246+ try:
211247 deadline = time.monotonic() + START_TIMEOUT
212248 while time.monotonic() < deadline:
213249 # Booting/downloading counts as activity, else the reaper would kill a
214250 # kill a warming pod mid-download.
215251 state['last'] = time.monotonic()
216- if upstream_ready(pod_id):
252+ if not ensured:
253+ try:
254+ mm_request(pod_id, 'POST', '/ensure', {'files': files})
255+ ensured = True
256+ state['model'] = values_key(key_values(state['model']) | set(values))
257+ log(f'in-place ensure requested: {key}')
258+ except Exception:
259+ pass # manager still booting, or old image without one
260+ ready = upstream_ready(pod_id)
261+ models_ok = True
262+ if not created and files and ensured:
263+ try:
264+ mm = mm_request(pod_id, 'GET', '/status', timeout=10)
265+ errs = {d: e for d, e in (mm.get('errors') or {}).items() if d in dests}
266+ if errs:
267+ raise RuntimeError(f'model download failed: {errs}')
268+ models_ok = not (set(dests) - set(mm.get('present') or []))
269+ except RuntimeError:
270+ raise
271+ except Exception:
272+ models_ok = False
273+ if ready and not ensured:
274+ no_manager_strikes += 1
275+ if (no_manager_strikes >= 3 and key and state['model'] and state['model'] != key):
276+ # Comfy is up but there is no manager (old image): last resort.
277+ log(f'pod has models={state["model"]}, need {key} - recreating (no model manager)')
278+ with state['lock']:
279+ terminate(pod_id)
280+ pod_id, gpu = _create_pod_with_retries(values)
281+ state.update({'pod_id': pod_id, 'model': key, 'gpu': gpu, 'last': time.monotonic()})
282+ created, ensured = True, True
283+ continue
284+ if ready and models_ok:
217285 if state['phase'] != 'green':
218286 state['phase'] = 'green'
219287 state['since'] = time.time()
220288 log('pod ready:', pod_id)
221289 return pod_id
290+ if state['phase'] == 'green':
291+ state['phase'] = 'orange' # in-place download in progress
222292 if state['pod_id'] != pod_id:
223293 raise RuntimeError('pod replaced while waiting')
224294 time.sleep(5)
225295 raise RuntimeError('pod not ready within START_TIMEOUT')
296+ finally:
297+ state['ensuring'] -= 1
226298
227299
228300def match_workflow_models(workflow_bytes):
@@ -254,10 +326,15 @@ def idle_reaper():
254326
255327def status_body():
256328 pod_id = state['pod_id']
257329 if pod_id and state['phaseensuring'] != 'green' and> upstream_ready(pod_id)0:
330+ # An ensure_pod waiter owns the phase; probing here could report green
331+ # mid-download or block the status reply on a dead pod's proxy URL.
332+ pass
333+ elif pod_id:
334+ if state['phase'] != 'green' and upstream_ready(pod_id):
258335 state['phase'] = 'green'
259336 state['since'] = time.time()
260- if not pod_id:
337+ else:
261338 found, have, gpu = find_pod()
262339 if found:
263340 state.update({'pod_id': found, 'model': have, 'gpu': gpu})
@@ -327,10 +404,10 @@ class Proxy(BaseHTTPRequestHandler):
327404 catalog = {'models': payload.get('models', []), 'active': [a for a in (active or []) if a]}
328405 save_catalog(catalog)
329406 log(f"catalog updated: {len(catalog['models'])} models, active={catalog['active']}")
330407 # Re-provisionFetch missing models proactively when the active modelset changed under a live pod
331408 # (known-modelunder podsa onlylive -pod never(in churnplace; arecreate podonly offor unknownold provenanceimages).
332409 if state['pod_id'] and catalog['active'] and state['model'] and statenot (set(catalog['modelactive']) !<= values_keykey_values(catalogstate['activemodel'])):
333410 log('active model set changed - re-provisioningensuring models on pod')
334411 threading.Thread(target=self._safe_ensure, args=(catalog['active'],), daemon=True).start()
335412 return self._reply(200, status_body())
336413 except Exception as err:
@@ -395,7 +472,7 @@ def main():
395472 log('adoption wait failed:', err)
396473 threading.Thread(target=_adopt, daemon=True).start()
397474 server = ThreadingHTTPServer(('0.0.0.0', LISTEN_PORT), Proxy)
398475 log(f'runpod-lazy v4v5 (catalog+lorain-place model switching) on :{LISTEN_PORT} (pod "{POD_NAME}", DCs {DATACENTERS or "any"}, idle {IDLE_SECONDS}s)')
399476 server.serve_forever()
400477
401478
public/scripts/extensions/stable-diffusion/index.js+17 -6
@@ -1580,7 +1580,7 @@ function renderRunpodStatus(status) {
15801580 const left = status.idle_seconds_left ? ` (idle stop in ${Math.round(status.idle_seconds_left / 60)}m)` : '';
15811581 text.textContent = `ready on ${status.gpu ?? 'GPU'}${left}`;
15821582 } else if (phase === 'orange') {
15831583 text.textContent = `starting (models: ${status.modelsmodel || (status.active ?? []).join(' + ') || '?'})…`;
15841584 } else {
15851585 text.textContent = 'off';
15861586 }
@@ -1610,20 +1610,31 @@ function ensureRunpodBarDot() {
16101610 dot.style.display = getRunpodLazyUrl() ? '' : 'none';
16111611}
16121612
1613+let runpodPollFailures = 0;
1614+
16131615async function pollRunpodStatus() {
16141616 const url = getRunpodLazyUrl();
16151617 if (!url) {
16161618 return;
16171619 }
16181620 let phase = 'red'runpodLastPhase;
16191621 try {
16201622 const result = await fetch(`${url}/lazy/status`, { signal: AbortSignal.timeout(50008000) });
1621- phase = renderRunpodStatus(result.ok ? await result.json() : null);
1623+ if (!result.ok) {
1624+ throw new Error(`status ${result.status}`);
1625+ }
1626+ runpodPollFailures = 0;
1627+ phase = renderRunpodStatus(await result.json());
16221628 } catch {
1623- renderRunpodStatus(null);
1629+ // A single slow/failed poll (e.g. while the proxy is provisioning a pod)
1630+ // must not flip the dot to red; only sustained unreachability does.
1631+ runpodPollFailures++;
1632+ if (runpodPollFailures >= 3) {
1633+ phase = renderRunpodStatus(null);
1634+ }
16241635 }
16251636 clearTimeout(runpodStatusTimer);
16261637 runpodStatusTimer = setTimeout(pollRunpodStatus, phase === 'orange' || runpodPollFailures ? RUNPOD_POLL_BUSY_MS : RUNPOD_POLL_IDLE_MS);
16271638}
16281639
16291640async function runpodControl(action) {