From e5ecc9a3a0cfc46488f6c59c23fce18f8427b100 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 03:18:38 +0100 Subject: [PATCH 01/32] Add RTX PRO 6000 RunPod deployment --- .dockerignore | 22 ++++++++ Dockerfile | 80 +++++++++++++++++++++++++++++ runpod/download_models.py | 58 +++++++++++++++++++++ runpod/health_server.py | 58 +++++++++++++++++++++ runpod/start.py | 105 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 323 insertions(+) create mode 100644 .dockerignore create mode 100644 Dockerfile create mode 100644 runpod/download_models.py create mode 100644 runpod/health_server.py create mode 100644 runpod/start.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..f65c83b --- /dev/null +++ b/.dockerignore @@ -0,0 +1,22 @@ +.git +.github +.vscode +.idea + +__pycache__ +*.pyc +*.pyo +*.log + +.venv +venv +env + +deploy/deps +deploy/tmp +deploy/recordings + +*.pth +*.pt +*.safetensors +*.onnx \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..14620fe --- /dev/null +++ b/Dockerfile @@ -0,0 +1,80 @@ +# syntax=docker/dockerfile:1.7 + +FROM nvidia/cuda:12.8.1-devel-ubuntu22.04 + +ARG DEBIAN_FRONTEND=noninteractive +ARG MAX_JOBS=8 + +ENV PYTHONUNBUFFERED=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PIP_NO_CACHE_DIR=1 \ + CUDA_HOME=/usr/local/cuda \ + TORCH_CUDA_ARCH_LIST=12.0 \ + JOYOMNI_OPS_CUDA_ARCHS=120a \ + MAX_JOBS=${MAX_JOBS} + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + ffmpeg \ + git \ + libglib2.0-0 \ + libgl1 \ + ninja-build \ + python3 \ + python3-dev \ + python3-pip \ + python3-venv \ + python-is-python3 \ + && rm -rf /var/lib/apt/lists/* \ + && python3 -m pip install --upgrade pip setuptools wheel + +WORKDIR /opt/joyai + +COPY deploy/requirements.txt /tmp/requirements.txt + +RUN python3 -m pip install -r /tmp/requirements.txt + +# Install patched SageAttention for RTX PRO 6000 Blackwell. +COPY deploy/sageattention-cudagraph-stream.patch /tmp/sageattention.patch + +RUN git clone https://github.com/thu-ml/SageAttention.git /tmp/SageAttention \ + && git -C /tmp/SageAttention checkout d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5 \ + && git -C /tmp/SageAttention apply /tmp/sageattention.patch \ + && cd /tmp/SageAttention \ + && EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" python3 setup.py install \ + && rm -rf /tmp/SageAttention + +# Build the repository's original FP8 CUDA operations. +COPY deploy/joyomni_ops /opt/joyai/deploy/joyomni_ops + +RUN git clone https://github.com/NVIDIA/cutlass.git /tmp/cutlass \ + && git -C /tmp/cutlass checkout dcf215af \ + && JOYOMNI_OPS_CUTLASS_DIR=/tmp/cutlass \ + python3 -m pip install --no-build-isolation /opt/joyai/deploy/joyomni_ops \ + && rm -rf /tmp/cutlass /root/.cache/pip + +COPY . /opt/joyai + +RUN chmod +x /opt/joyai/deploy/run_server.sh \ + && mkdir -p /runpod-volume/joyai \ + && mkdir -p /tmp/joyomni-recordings + +ENV JOYOMNI_DEVICE=cuda:0 \ + JOYOMNI_HOST=0.0.0.0 \ + JOYOMNI_PORT=8080 \ + JOYOMNI_CKPT_ROOT=/runpod-volume/joyai/checkpoints \ + JOYOMNI_WIDTH=840 \ + JOYOMNI_HEIGHT=480 \ + JOYOMNI_FPS=24 \ + JOYOMNI_FP8_IMG=1 \ + JOYOMNI_FP8_TXT=1 \ + JOYOMNI_CUDA_GRAPH=1 \ + JOYOMNI_SAGE_ATTN=1 \ + JOYOMNI_TXT_PARALLEL=1 \ + JOYOMNI_RECORD_DIR=/tmp/joyomni-recordings + +EXPOSE 8080 8081 + +CMD ["python3", "/opt/joyai/runpod/start.py"] \ No newline at end of file diff --git a/runpod/download_models.py b/runpod/download_models.py new file mode 100644 index 0000000..d0acdb7 --- /dev/null +++ b/runpod/download_models.py @@ -0,0 +1,58 @@ +import os +import urllib.request +from pathlib import Path + +from huggingface_hub import snapshot_download + + +checkpoint_root = Path( + os.getenv( + "JOYOMNI_CKPT_ROOT", + "/runpod-volume/joyai/checkpoints", + ) +) + +checkpoint_root.mkdir(parents=True, exist_ok=True) + +hf_token = os.getenv("HF_TOKEN") or None + +print("Downloading JoyAI 0811 model and VAE...") + +snapshot_download( + repo_id="jdopensource/JoyAI-Video-Edit", + local_dir=checkpoint_root / "JoyAI-Video-Edit", + allow_patterns=[ + "dit/joyai_video_edit_dit_0811.pth", + "vae/*", + ], + token=hf_token, + max_workers=8, +) + +print("Downloading MiMo-VL text and vision encoder...") + +snapshot_download( + repo_id="XiaomiMiMo/MiMo-VL-7B-RL-2508", + local_dir=checkpoint_root / "MiMo-VL-7B-RL-2508", + token=hf_token, + max_workers=8, +) + +face_model = checkpoint_root / "face_detection_yunet_2023mar.onnx" + +if not face_model.exists(): + print("Downloading YuNet face detector...") + + temporary_file = face_model.with_suffix(".download") + + urllib.request.urlretrieve( + "https://media.githubusercontent.com/media/opencv/" + "opencv_zoo/main/models/face_detection_yunet/" + "face_detection_yunet_2023mar.onnx", + temporary_file, + ) + + temporary_file.replace(face_model) + +print("All required model files have been downloaded.") +print(f"Checkpoint location: {checkpoint_root}") \ No newline at end of file diff --git a/runpod/health_server.py b/runpod/health_server.py new file mode 100644 index 0000000..24c3711 --- /dev/null +++ b/runpod/health_server.py @@ -0,0 +1,58 @@ +import json +import os +import urllib.error +import urllib.request + +import uvicorn +from fastapi import FastAPI +from fastapi.responses import JSONResponse + + +app = FastAPI() + + +@app.get("/ping") +def ping(): + model_port = os.getenv( + "JOYOMNI_PORT", + os.getenv("PORT", "8080"), + ) + + try: + with urllib.request.urlopen( + f"http://127.0.0.1:{model_port}/health", + timeout=2, + ) as response: + health = json.loads(response.read()) + + if health.get("runtime_loaded") is True: + return { + "status": "healthy", + "model": "ready", + } + + except ( + urllib.error.URLError, + TimeoutError, + ValueError, + ): + pass + + return JSONResponse( + status_code=503, + content={ + "status": "initializing", + "model": "loading", + }, + ) + + +if __name__ == "__main__": + health_port = int(os.getenv("PORT_HEALTH", "8081")) + + uvicorn.run( + app, + host="0.0.0.0", + port=health_port, + log_level="info", + ) \ No newline at end of file diff --git a/runpod/start.py b/runpod/start.py new file mode 100644 index 0000000..f453b6a --- /dev/null +++ b/runpod/start.py @@ -0,0 +1,105 @@ +import os +import subprocess +import sys +from pathlib import Path + + +repository_root = Path(__file__).resolve().parents[1] + +checkpoint_root = Path( + os.getenv( + "JOYOMNI_CKPT_ROOT", + "/runpod-volume/joyai/checkpoints", + ) +) + +required_items = [ + checkpoint_root + / "JoyAI-Video-Edit" + / "dit" + / "joyai_video_edit_dit_0811.pth", + + checkpoint_root + / "JoyAI-Video-Edit" + / "vae" + / "config.json", + + checkpoint_root + / "JoyAI-Video-Edit" + / "vae" + / "diffusion_pytorch_model.safetensors", + + checkpoint_root + / "MiMo-VL-7B-RL-2508", +] + +missing_items = [ + str(item) + for item in required_items + if not item.exists() +] + +if missing_items: + print("Required model files are missing:", flush=True) + + for item in missing_items: + print(f" - {item}", flush=True) + + print( + "\nRun this command once on the attached RunPod volume:", + flush=True, + ) + print( + "python3 /opt/joyai/runpod/download_models.py", + flush=True, + ) + + raise SystemExit(1) + +environment = os.environ.copy() + +# RunPod supplies PORT for public traffic. +model_port = environment.get( + "PORT", + environment.get("JOYOMNI_PORT", "8080"), +) + +environment["JOYOMNI_PORT"] = model_port + +print( + f"Starting JoyAI on port {model_port}...", + flush=True, +) + +health_process = subprocess.Popen( + [ + sys.executable, + str(repository_root / "runpod" / "health_server.py"), + ], + cwd=repository_root, + env=environment, +) + +exit_code = 1 + +try: + model_process = subprocess.Popen( + [ + "bash", + str(repository_root / "deploy" / "run_server.sh"), + ], + cwd=repository_root, + env=environment, + ) + + exit_code = model_process.wait() + +finally: + health_process.terminate() + + try: + health_process.wait(timeout=10) + except subprocess.TimeoutExpired: + health_process.kill() + +raise SystemExit(exit_code) \ No newline at end of file From 0a2b48189b4048891ef1a0abcaaa67551b03a1e1 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 08:53:59 +0100 Subject: [PATCH 02/32] Add manual RunPod container build workflow --- .github/workflows/build-runpod-image.yml | 65 ++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 .github/workflows/build-runpod-image.yml diff --git a/.github/workflows/build-runpod-image.yml b/.github/workflows/build-runpod-image.yml new file mode 100644 index 0000000..f2ab5b1 --- /dev/null +++ b/.github/workflows/build-runpod-image.yml @@ -0,0 +1,65 @@ +name: Build RunPod container image + +on: + workflow_dispatch: + +permissions: + contents: read + packages: write + +concurrency: + group: runpod-image-${{ github.ref }} + cancel-in-progress: false + +env: + REGISTRY: ghcr.io + IMAGE_NAME: samuellucky2424-afk/joyai-video-edit + +jobs: + build-and-push: + name: Build and publish RTX PRO 6000 image + runs-on: ubuntu-24.04 + timeout-minutes: 330 + + steps: + - name: Free disk space for the CUDA image + uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be + with: + tool-cache: false + android: true + dotnet: true + haskell: true + large-packages: true + docker-images: true + swap-storage: false + + - name: Check out the RunPod branch + uses: actions/checkout@v6 + + - name: Log in to GitHub Container Registry + uses: docker/login-action@65b78e6e13532edd9afa3aa52ac7964289d1a9c1 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build and publish the container image + id: build + uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4 + with: + context: . + file: ./Dockerfile + platforms: linux/amd64 + push: true + build-args: | + MAX_JOBS=2 + tags: | + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:runpod-rtx-pro-6000 + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:sha-${{ github.sha }} + labels: | + org.opencontainers.image.source=${{ github.server_url }}/${{ github.repository }} + org.opencontainers.image.revision=${{ github.sha }} + org.opencontainers.image.description=JoyAI Video Edit for RunPod RTX PRO 6000 + + - name: Show immutable deployment reference + run: echo "${REGISTRY}/${IMAGE_NAME}@${{ steps.build.outputs.digest }}" From f32133fbcd5466fab2b0c696970076799d2352d0 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 12:05:57 +0100 Subject: [PATCH 03/32] Add H200 RunPod container build --- Dockerfile.h200 | 81 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 Dockerfile.h200 diff --git a/Dockerfile.h200 b/Dockerfile.h200 new file mode 100644 index 0000000..4865cab --- /dev/null +++ b/Dockerfile.h200 @@ -0,0 +1,81 @@ +# syntax=docker/dockerfile:1.7 + +FROM nvidia/cuda:12.8.1-devel-ubuntu22.04 + +ARG DEBIAN_FRONTEND=noninteractive +ARG MAX_JOBS=8 + +ENV PYTHONUNBUFFERED=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PIP_NO_CACHE_DIR=1 \ + CUDA_HOME=/usr/local/cuda \ + TORCH_CUDA_ARCH_LIST=9.0 \ + JOYOMNI_OPS_CUDA_ARCHS=90 \ + MAX_JOBS=${MAX_JOBS} + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + ffmpeg \ + git \ + libglib2.0-0 \ + libgl1 \ + ninja-build \ + python3 \ + python3-dev \ + python3-pip \ + python3-venv \ + python-is-python3 \ + && rm -rf /var/lib/apt/lists/* \ + && python3 -m pip install --upgrade pip setuptools wheel + +WORKDIR /opt/joyai + +COPY deploy/requirements.txt /tmp/requirements.txt + +RUN python3 -m pip install -r /tmp/requirements.txt + +# SageAttention supports Hopper GPUs. Compile its CUDA extension for H200 (sm_90). +COPY deploy/sageattention-cudagraph-stream.patch /tmp/sageattention.patch + +RUN git clone https://github.com/thu-ml/SageAttention.git /tmp/SageAttention \ + && git -C /tmp/SageAttention checkout d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5 \ + && git -C /tmp/SageAttention apply /tmp/sageattention.patch \ + && cd /tmp/SageAttention \ + && EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" python3 setup.py install \ + && rm -rf /tmp/SageAttention + +# Build JoyAI's FP8 CUDA operations for Hopper (sm_90). +COPY deploy/joyomni_ops /opt/joyai/deploy/joyomni_ops + +RUN git clone https://github.com/NVIDIA/cutlass.git /tmp/cutlass \ + && git -C /tmp/cutlass checkout dcf215af \ + && JOYOMNI_OPS_CUTLASS_DIR=/tmp/cutlass \ + python3 -m pip install --no-build-isolation /opt/joyai/deploy/joyomni_ops \ + && rm -rf /tmp/cutlass /root/.cache/pip + +COPY . /opt/joyai + +RUN chmod +x /opt/joyai/deploy/run_server.sh \ + && mkdir -p /runpod-volume/joyai \ + && mkdir -p /tmp/joyomni-recordings + +ENV JOYOMNI_DEVICE=cuda:0 \ + JOYOMNI_HOST=0.0.0.0 \ + JOYOMNI_PORT=8080 \ + JOYOMNI_CKPT_ROOT=/runpod-volume/joyai/checkpoints \ + JOYOMNI_WIDTH=840 \ + JOYOMNI_HEIGHT=480 \ + JOYOMNI_FPS=24 \ + JOYOMNI_FP8_IMG=1 \ + JOYOMNI_FP8_TXT=1 \ + JOYOMNI_CUDA_GRAPH=1 \ + JOYOMNI_SAGE_ATTN=1 \ + JOYOMNI_TXT_PARALLEL=1 \ + JOYOMNI_RECORD_ENABLED=0 \ + JOYOMNI_RECORD_DIR=/tmp/joyomni-recordings + +EXPOSE 8080 8081 + +CMD ["python3", "/opt/joyai/runpod/start.py"] From becaa73b5af4c97141144a0001933303358125bc Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 12:06:02 +0100 Subject: [PATCH 04/32] Add H200 container build workflow --- .github/workflows/build-runpod-h200.yml | 65 +++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 .github/workflows/build-runpod-h200.yml diff --git a/.github/workflows/build-runpod-h200.yml b/.github/workflows/build-runpod-h200.yml new file mode 100644 index 0000000..37d30d1 --- /dev/null +++ b/.github/workflows/build-runpod-h200.yml @@ -0,0 +1,65 @@ +name: Build RunPod H200 container image + +on: + workflow_dispatch: + +permissions: + contents: read + packages: write + +concurrency: + group: runpod-h200-image-${{ github.ref }} + cancel-in-progress: false + +env: + REGISTRY: ghcr.io + IMAGE_NAME: samuellucky2424-afk/joyai-video-edit + +jobs: + build-and-push: + name: Build and publish H200 image + runs-on: ubuntu-24.04 + timeout-minutes: 330 + + steps: + - name: Free disk space for the CUDA image + uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be + with: + tool-cache: false + android: true + dotnet: true + haskell: true + large-packages: true + docker-images: true + swap-storage: false + + - name: Check out the repository + uses: actions/checkout@v6 + + - name: Log in to GitHub Container Registry + uses: docker/login-action@65b78e6e13532edd9afa3aa52ac7964289d1a9c1 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build and publish the H200 container image + id: build + uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4 + with: + context: . + file: ./Dockerfile.h200 + platforms: linux/amd64 + push: true + build-args: | + MAX_JOBS=2 + tags: | + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:runpod-h200 + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:h200-sha-${{ github.sha }} + labels: | + org.opencontainers.image.source=${{ github.server_url }}/${{ github.repository }} + org.opencontainers.image.revision=${{ github.sha }} + org.opencontainers.image.description=JoyAI Video Edit for RunPod H200 + + - name: Show immutable deployment reference + run: echo "${REGISTRY}/${IMAGE_NAME}@${{ steps.build.outputs.digest }}" From c171b1aa46d98f9b03b23402062988686f740b4d Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 16:30:43 +0100 Subject: [PATCH 05/32] Reduce H200 runtime image size --- .github/workflows/build-runpod-h200.yml | 1 + Dockerfile.h200 | 68 +++++++++++++++++-------- deploy/requirements-h200-runtime.txt | 16 ++++++ 3 files changed, 63 insertions(+), 22 deletions(-) create mode 100644 deploy/requirements-h200-runtime.txt diff --git a/.github/workflows/build-runpod-h200.yml b/.github/workflows/build-runpod-h200.yml index 37d30d1..84f9c5e 100644 --- a/.github/workflows/build-runpod-h200.yml +++ b/.github/workflows/build-runpod-h200.yml @@ -51,6 +51,7 @@ jobs: file: ./Dockerfile.h200 platforms: linux/amd64 push: true + provenance: false build-args: | MAX_JOBS=2 tags: | diff --git a/Dockerfile.h200 b/Dockerfile.h200 index 4865cab..5868565 100644 --- a/Dockerfile.h200 +++ b/Dockerfile.h200 @@ -1,12 +1,15 @@ # syntax=docker/dockerfile:1.7 -FROM nvidia/cuda:12.8.1-devel-ubuntu22.04 +ARG PYTORCH_IMAGE_VERSION=2.9.1-cuda12.8-cudnn9 + +# Compile the Hopper-specific CUDA extensions in a development image. Nothing +# from this stage is shipped except the two finished Python wheels. +FROM pytorch/pytorch:${PYTORCH_IMAGE_VERSION}-devel AS extensions ARG DEBIAN_FRONTEND=noninteractive ARG MAX_JOBS=8 -ENV PYTHONUNBUFFERED=1 \ - PIP_DISABLE_PIP_VERSION_CHECK=1 \ +ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \ PIP_NO_CACHE_DIR=1 \ CUDA_HOME=/usr/local/cuda \ TORCH_CUDA_ARCH_LIST=9.0 \ @@ -16,45 +19,66 @@ ENV PYTHONUNBUFFERED=1 \ RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential \ ca-certificates \ - curl \ - ffmpeg \ git \ - libglib2.0-0 \ - libgl1 \ ninja-build \ - python3 \ - python3-dev \ - python3-pip \ - python3-venv \ - python-is-python3 \ && rm -rf /var/lib/apt/lists/* \ - && python3 -m pip install --upgrade pip setuptools wheel + && python -m pip install --upgrade pip setuptools wheel WORKDIR /opt/joyai +RUN mkdir -p /wheels -COPY deploy/requirements.txt /tmp/requirements.txt - -RUN python3 -m pip install -r /tmp/requirements.txt - -# SageAttention supports Hopper GPUs. Compile its CUDA extension for H200 (sm_90). +# SageAttention supports Hopper GPUs. Compile its CUDA extension for H200 +# (sm_90) and keep only the resulting wheel. COPY deploy/sageattention-cudagraph-stream.patch /tmp/sageattention.patch RUN git clone https://github.com/thu-ml/SageAttention.git /tmp/SageAttention \ && git -C /tmp/SageAttention checkout d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5 \ && git -C /tmp/SageAttention apply /tmp/sageattention.patch \ && cd /tmp/SageAttention \ - && EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" python3 setup.py install \ + && EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" \ + python setup.py bdist_wheel --dist-dir /wheels \ && rm -rf /tmp/SageAttention -# Build JoyAI's FP8 CUDA operations for Hopper (sm_90). +# Build JoyAI's FP8 CUDA operations for Hopper (sm_90) as a wheel. COPY deploy/joyomni_ops /opt/joyai/deploy/joyomni_ops RUN git clone https://github.com/NVIDIA/cutlass.git /tmp/cutlass \ && git -C /tmp/cutlass checkout dcf215af \ && JOYOMNI_OPS_CUTLASS_DIR=/tmp/cutlass \ - python3 -m pip install --no-build-isolation /opt/joyai/deploy/joyomni_ops \ + python -m pip wheel --no-build-isolation --no-deps \ + --wheel-dir /wheels /opt/joyai/deploy/joyomni_ops \ && rm -rf /tmp/cutlass /root/.cache/pip + +# The final image already contains CUDA 12.8, cuDNN 9, PyTorch 2.9.1, and +# torchvision. It intentionally excludes compilers, source trees, and the +# CUDA development toolchain from the published RunPod image. +FROM pytorch/pytorch:${PYTORCH_IMAGE_VERSION}-runtime + +ARG DEBIAN_FRONTEND=noninteractive + +ENV PYTHONUNBUFFERED=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PIP_NO_CACHE_DIR=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates \ + ffmpeg \ + libglib2.0-0 \ + libgl1 \ + libgomp1 \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /opt/joyai + +COPY deploy/requirements-h200-runtime.txt /tmp/requirements.txt +RUN python -m pip install -r /tmp/requirements.txt \ + && rm -f /tmp/requirements.txt + +COPY --from=extensions /wheels /tmp/wheels +RUN python -m pip install --no-deps /tmp/wheels/*.whl \ + && rm -rf /tmp/wheels + COPY . /opt/joyai RUN chmod +x /opt/joyai/deploy/run_server.sh \ @@ -78,4 +102,4 @@ ENV JOYOMNI_DEVICE=cuda:0 \ EXPOSE 8080 8081 -CMD ["python3", "/opt/joyai/runpod/start.py"] +CMD ["python", "/opt/joyai/runpod/start.py"] diff --git a/deploy/requirements-h200-runtime.txt b/deploy/requirements-h200-runtime.txt new file mode 100644 index 0000000..50955f6 --- /dev/null +++ b/deploy/requirements-h200-runtime.txt @@ -0,0 +1,16 @@ +transformers>=4.57.1,<4.58 +accelerate==1.10.1 +diffusers==0.36.0 +einops==0.8.2 +numpy==2.2.6 +pillow==12.2.0 +opencv-python-headless==4.13.0.92 +av==13.1.0 +imageio-ffmpeg==0.6.0 +fastapi==0.117.1 +uvicorn==0.37.0 +uvloop==0.22.1 +ninja==1.13.0 +websockets==16.0 +openai==2.41.0 +loguru==0.7.3 From 4caeccda18dc27f08d02e25bec502be98e442dd9 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 17:58:50 +0100 Subject: [PATCH 06/32] Add runtime compiler for Torch Inductor --- Dockerfile.h200 | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/Dockerfile.h200 b/Dockerfile.h200 index 5868565..4e5f8c5 100644 --- a/Dockerfile.h200 +++ b/Dockerfile.h200 @@ -51,17 +51,21 @@ RUN git clone https://github.com/NVIDIA/cutlass.git /tmp/cutlass \ # The final image already contains CUDA 12.8, cuDNN 9, PyTorch 2.9.1, and -# torchvision. It intentionally excludes compilers, source trees, and the -# CUDA development toolchain from the published RunPod image. +# torchvision. Keep the small host compiler toolchain because Torch Inductor +# compiles CPU launcher code while warming dynamic VAE shapes at runtime. +# CUDA extensions and the full CUDA development toolchain stay in the builder. FROM pytorch/pytorch:${PYTORCH_IMAGE_VERSION}-runtime ARG DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 \ PIP_DISABLE_PIP_VERSION_CHECK=1 \ - PIP_NO_CACHE_DIR=1 + PIP_NO_CACHE_DIR=1 \ + CC=gcc \ + CXX=g++ RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ ca-certificates \ ffmpeg \ libglib2.0-0 \ From 2e690ac48ce1b3b58eff746359623d9a7c9b8bd3 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 19:29:18 +0100 Subject: [PATCH 07/32] Fix RunPod initializing health status --- runpod/health_server.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/runpod/health_server.py b/runpod/health_server.py index 24c3711..2d40c84 100644 --- a/runpod/health_server.py +++ b/runpod/health_server.py @@ -4,8 +4,7 @@ import urllib.request import uvicorn -from fastapi import FastAPI -from fastapi.responses import JSONResponse +from fastapi import FastAPI, Response app = FastAPI() @@ -38,13 +37,9 @@ def ping(): ): pass - return JSONResponse( - status_code=503, - content={ - "status": "initializing", - "model": "loading", - }, - ) + # RunPod load-balancer health contract: + # 200 = ready, 204 = still initializing, anything else = unhealthy. + return Response(status_code=204) if __name__ == "__main__": @@ -55,4 +50,4 @@ def ping(): host="0.0.0.0", port=health_port, log_level="info", - ) \ No newline at end of file + ) From 561db994a57ab8bce42cb79572a5936762980002 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 20:08:36 +0100 Subject: [PATCH 08/32] Skip VAE autotune when serverless compile is disabled --- deploy/xvideo/models/vae/vae_compile.py | 40 ++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/deploy/xvideo/models/vae/vae_compile.py b/deploy/xvideo/models/vae/vae_compile.py index ff0401a..e7f43da 100644 --- a/deploy/xvideo/models/vae/vae_compile.py +++ b/deploy/xvideo/models/vae/vae_compile.py @@ -1,21 +1,47 @@ from __future__ import annotations +import os + import torch import torch.nn as nn from xvideo.inductor_autotune_fix import install as _install_autotune_fix +def compile_enabled() -> bool: + """Return whether Torch Inductor VAE compilation is enabled. + + Compilation remains enabled by default for non-serverless deployments. + RunPod H200 images disable it explicitly so cold workers become healthy + without spending minutes autotuning dozens of VAE shapes. + """ + value = os.getenv("JOYOMNI_VAE_COMPILE", "1").strip().lower() + return value not in {"0", "false", "no", "off"} + + # Must run before any compiled function executes, so warm restarts reuse the # on-disk autotune results instead of re-running coordinate descent. -_install_autotune_fix() +if compile_enabled(): + _install_autotune_fix() _configured: set[int] = set() _configured_encode: set[int] = set() _configured_encode_dynamic: set[int] = set() +_skip_notices: set[str] = set() + + +def _skip_compile(stage: str) -> bool: + if compile_enabled(): + return False + if stage not in _skip_notices: + print(f"[vae_compile] disabled by JOYOMNI_VAE_COMPILE=0; skipping {stage}") + _skip_notices.add(stage) + return True def maybe_setup_decode(vae) -> None: + if _skip_compile("decode compilation"): + return if id(vae) in _configured: return n_conv = 0 @@ -36,10 +62,14 @@ def maybe_setup_decode(vae) -> None: def prep_input(z: torch.Tensor) -> torch.Tensor: + if not compile_enabled(): + return z return z.to(memory_format=torch.channels_last_3d) def maybe_setup_encode(vae) -> None: + if _skip_compile("encode compilation"): + return if id(vae) in _configured_encode: return n_conv = 0 @@ -63,6 +93,8 @@ def warmup_encode(vae, in_channels: int, h_px: int, w_px: int, device: torch.device, dtype: torch.dtype, temporal_lens: tuple[int, ...] = (1, 9), autocast: bool = False) -> None: + if _skip_compile("encode warmup"): + return maybe_setup_encode(vae) from contextlib import nullcontext dev_type = torch.device(device).type @@ -85,6 +117,8 @@ def warmup_encode(vae, in_channels: int, h_px: int, w_px: int, def maybe_setup_encode_dynamic(vae) -> None: + if _skip_compile("dynamic encode compilation"): + return if id(vae) in _configured_encode_dynamic: return if hasattr(vae, "_encode"): @@ -113,6 +147,8 @@ def encode_via_dynamic(vae, x: torch.Tensor): def warmup_encode_dynamic(vae, in_channels: int, hw_list, device: torch.device, dtype: torch.dtype, temporal_lens: tuple[int, ...] = (1,), autocast: bool = False) -> None: + if _skip_compile("dynamic encode warmup"): + return fn = getattr(vae, "_encode_dynamic", None) if fn is None: print("[vae_compile] warmup_encode_dynamic skipped: _encode_dynamic not set up") @@ -144,6 +180,8 @@ def warmup_decode(vae, latent_channels: int, h_lat: int, w_lat: int, device: torch.device, dtype: torch.dtype, temporal_lens: tuple[int, ...] = (1, 2), autocast: bool = True) -> None: + if _skip_compile("decode warmup"): + return maybe_setup_decode(vae) from contextlib import nullcontext dev_type = torch.device(device).type From 9058d6b7fbf45ddba11192a2ba35599cfae8f0d2 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 20:08:57 +0100 Subject: [PATCH 09/32] Disable VAE compilation in H200 serverless image --- Dockerfile.h200 | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/Dockerfile.h200 b/Dockerfile.h200 index 4e5f8c5..ba5ae50 100644 --- a/Dockerfile.h200 +++ b/Dockerfile.h200 @@ -51,8 +51,8 @@ RUN git clone https://github.com/NVIDIA/cutlass.git /tmp/cutlass \ # The final image already contains CUDA 12.8, cuDNN 9, PyTorch 2.9.1, and -# torchvision. Keep the small host compiler toolchain because Torch Inductor -# compiles CPU launcher code while warming dynamic VAE shapes at runtime. +# torchvision. Keep the small host compiler toolchain as an escape hatch for +# deployments that explicitly re-enable Torch Inductor VAE compilation. # CUDA extensions and the full CUDA development toolchain stay in the builder. FROM pytorch/pytorch:${PYTORCH_IMAGE_VERSION}-runtime @@ -101,6 +101,7 @@ ENV JOYOMNI_DEVICE=cuda:0 \ JOYOMNI_CUDA_GRAPH=1 \ JOYOMNI_SAGE_ATTN=1 \ JOYOMNI_TXT_PARALLEL=1 \ + JOYOMNI_VAE_COMPILE=0 \ JOYOMNI_RECORD_ENABLED=0 \ JOYOMNI_RECORD_DIR=/tmp/joyomni-recordings From b2b3aeac64e79044a93aea0ccb7a7d82364522b4 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 20:12:13 +0100 Subject: [PATCH 10/32] Build H200 repair image on branch updates --- .github/workflows/build-runpod-h200.yml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/.github/workflows/build-runpod-h200.yml b/.github/workflows/build-runpod-h200.yml index 84f9c5e..63b77c1 100644 --- a/.github/workflows/build-runpod-h200.yml +++ b/.github/workflows/build-runpod-h200.yml @@ -2,6 +2,14 @@ name: Build RunPod H200 container image on: workflow_dispatch: + push: + branches: + - agent/reduce-h200-image + paths: + - .github/workflows/build-runpod-h200.yml + - Dockerfile.h200 + - deploy/requirements-h200-runtime.txt + - deploy/xvideo/models/vae/vae_compile.py permissions: contents: read From 3c10e73e07c6934c60ba1f742f1f099bf593a95f Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 20:48:59 +0100 Subject: [PATCH 11/32] Fix H200 FP8 CUTLASS target to sm_90a --- Dockerfile.h200 | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/Dockerfile.h200 b/Dockerfile.h200 index ba5ae50..fc4ebe5 100644 --- a/Dockerfile.h200 +++ b/Dockerfile.h200 @@ -13,7 +13,7 @@ ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \ PIP_NO_CACHE_DIR=1 \ CUDA_HOME=/usr/local/cuda \ TORCH_CUDA_ARCH_LIST=9.0 \ - JOYOMNI_OPS_CUDA_ARCHS=90 \ + JOYOMNI_OPS_CUDA_ARCHS=90a \ MAX_JOBS=${MAX_JOBS} RUN apt-get update && apt-get install -y --no-install-recommends \ @@ -39,7 +39,9 @@ RUN git clone https://github.com/thu-ml/SageAttention.git /tmp/SageAttention \ python setup.py bdist_wheel --dist-dir /wheels \ && rm -rf /tmp/SageAttention -# Build JoyAI's FP8 CUDA operations for Hopper (sm_90) as a wheel. +# Build JoyAI's FP8 CUDA operations for Hopper (sm_90a) as a wheel. +# CUTLASS WGMMA instructions are architecture-accelerated and abort at runtime +# when the extension is compiled only for the generic sm_90 target. COPY deploy/joyomni_ops /opt/joyai/deploy/joyomni_ops RUN git clone https://github.com/NVIDIA/cutlass.git /tmp/cutlass \ From 64973d103707e265cba719d9ab1eafec184fd92f Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 21:25:39 +0100 Subject: [PATCH 12/32] Fix VAE attention fallback on H200 --- deploy/xvideo/models/vae/vae.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/deploy/xvideo/models/vae/vae.py b/deploy/xvideo/models/vae/vae.py index 1f8d751..6a75dd2 100644 --- a/deploy/xvideo/models/vae/vae.py +++ b/deploy/xvideo/models/vae/vae.py @@ -6,6 +6,7 @@ import torch from torch import nn, Tensor import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel from diffusers.utils.torch_utils import randn_tensor from diffusers.configuration_utils import ConfigMixin, register_to_config @@ -127,11 +128,15 @@ def attention(self, x: Tensor) -> Tensor: k = self.k(x) v = self.v(x) - q = rearrange(q, "b c t h w -> (b t) 1 (h w) c") - k = rearrange(k, "b c t h w -> (b t) 1 (h w) c") - v = rearrange(v, "b c t h w -> (b t) 1 (h w) c") + q = rearrange(q, "b c t h w -> (b t) 1 (h w) c").contiguous() + k = rearrange(k, "b c t h w -> (b t) 1 (h w) c").contiguous() + v = rearrange(v, "b c t h w -> (b t) 1 (h w) c").contiguous() - x = F.scaled_dot_product_attention(q, k, v) + # This VAE uses a 1024-wide attention head. CUDA flash attention only + # supports head dimensions up to 256, so fused-only SDPA aborts on H200. + # Keep this fallback local to the VAE instead of changing DiT attention. + with sdpa_kernel(SDPBackend.MATH): + x = F.scaled_dot_product_attention(q, k, v) x = rearrange(x, "(b t) 1 (h w) c -> b c t h w", b=b, t=t, h=h, w=w) x = self.proj_out(x) From f199c31f9c636ea9b3f35fd5dc605db152c97801 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 21:26:07 +0100 Subject: [PATCH 13/32] Build H200 image for VAE attention changes --- .github/workflows/build-runpod-h200.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/build-runpod-h200.yml b/.github/workflows/build-runpod-h200.yml index 63b77c1..fcdbd99 100644 --- a/.github/workflows/build-runpod-h200.yml +++ b/.github/workflows/build-runpod-h200.yml @@ -10,6 +10,7 @@ on: - Dockerfile.h200 - deploy/requirements-h200-runtime.txt - deploy/xvideo/models/vae/vae_compile.py + - deploy/xvideo/models/vae/vae.py permissions: contents: read From 9ca25437b0ae750230c5817abb426d467374805a Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Sun, 16 Aug 2026 22:32:16 +0100 Subject: [PATCH 14/32] Fix RunPod connection lifecycle and realtime proxy --- Dockerfile.h200 | 1 + runpod/README.md | 65 +++++++++ runpod/Start-JoyAI-Realtime-Test.ps1 | 70 +++++++++ runpod/health_server.py | 7 +- runpod/local_proxy.py | 183 ++++++++++++++++++++++++ runpod/start.py | 206 ++++++++++++++++----------- tests/test_runpod_connection.py | 55 +++++++ 7 files changed, 502 insertions(+), 85 deletions(-) create mode 100644 runpod/README.md create mode 100644 runpod/Start-JoyAI-Realtime-Test.ps1 create mode 100644 runpod/local_proxy.py create mode 100644 tests/test_runpod_connection.py diff --git a/Dockerfile.h200 b/Dockerfile.h200 index fc4ebe5..e85f046 100644 --- a/Dockerfile.h200 +++ b/Dockerfile.h200 @@ -95,6 +95,7 @@ ENV JOYOMNI_DEVICE=cuda:0 \ JOYOMNI_HOST=0.0.0.0 \ JOYOMNI_PORT=8080 \ JOYOMNI_CKPT_ROOT=/runpod-volume/joyai/checkpoints \ + JOYOMNI_PRELOAD=1 \ JOYOMNI_WIDTH=840 \ JOYOMNI_HEIGHT=480 \ JOYOMNI_FPS=24 \ diff --git a/runpod/README.md b/runpod/README.md new file mode 100644 index 0000000..78456a0 --- /dev/null +++ b/runpod/README.md @@ -0,0 +1,65 @@ +# RunPod connection contract + +JoyAI uses two HTTP servers inside every RunPod load-balancer worker. The +public application server listens on `PORT` (normally `8080`). A small internal +readiness server listens on `PORT_HEALTH` (normally `8081`). + +| Purpose | Correct route | Caller | +| --- | --- | --- | +| Web interface | `GET /` | Browser or client | +| JoyAI server health | `GET /health` | Browser, client, or diagnostics | +| Load/warm the model | `POST /load` | Startup or authenticated client | +| Real-time video stream | `WS /ws` | Web interface | +| RunPod worker readiness | `GET /ping` on `PORT_HEALTH` | RunPod only | + +Do **not** call `/ping` through the public `*.api.runpod.ai` address. That +address routes to the JoyAI server on `PORT`, where the health route is +`/health`. RunPod calls the separate `/ping` route internally on +`PORT_HEALTH`. + +## Startup sequence + +1. `runpod/start.py` starts the internal readiness server on `PORT_HEALTH`. +2. It starts JoyAI with `--preload`, so the model is loaded immediately from + `JOYOMNI_CKPT_ROOT`. +3. While the model loads, internal `GET /ping` returns `204` (initializing). +4. When public `GET /health` reports `runtime_loaded: true`, internal + `GET /ping` returns `200` and RunPod begins routing public traffic. + +This avoids the deadlock where the worker waited for a public `/load` request +while RunPod waited for the worker to become ready before routing that request. + +## RunPod endpoint settings + +Use these values for a load-balancer endpoint: + +| Setting | Value | +| --- | --- | +| `PORT` | `8080` | +| `PORT_HEALTH` | `8081` | +| `JOYOMNI_PRELOAD` | `1` | +| `JOYOMNI_CKPT_ROOT` | The checkpoint path on the mounted network volume | +| Exposed HTTP ports | `8080,8081` | +| Health-check path | `/ping` | +| Active/min workers | `0` while testing to avoid idle GPU charges | + +The H200 image default checkpoint path is +`/runpod-volume/joyai/checkpoints`. Mount the model network volume so that the +downloaded checkpoint tree is available at that path. + +## Windows real-time test + +The browser cannot add a RunPod bearer token to normal page navigation or the +WebSocket constructor. Use the included local authenticated proxy: + +```powershell +cd path\to\JoyAI-Video-Edit\runpod +powershell.exe -NoProfile -ExecutionPolicy Bypass -File .\Start-JoyAI-Realtime-Test.ps1 +``` + +The script securely asks for the API key, warms through `POST /load`, verifies +`GET /health`, starts the local HTTP/WebSocket proxy, and opens the interface. +The key stays in the local process and is removed when the script exits. + +Stop the test with `Ctrl+C`. With zero active workers and a short idle timeout, +RunPod can scale the worker back to zero after the connection closes. diff --git a/runpod/Start-JoyAI-Realtime-Test.ps1 b/runpod/Start-JoyAI-Realtime-Test.ps1 new file mode 100644 index 0000000..4fe372f --- /dev/null +++ b/runpod/Start-JoyAI-Realtime-Test.ps1 @@ -0,0 +1,70 @@ +param( + [string]$EndpointId = "ex9647vtulowka", + [int]$LocalPort = 9000, + [int]$WarmTimeoutSeconds = 900 +) + +$ErrorActionPreference = "Stop" +$proxyProcess = $null +$secret = $null +$apiKey = $null + +try { + $secret = Read-Host "Paste your RunPod API key" -AsSecureString + $apiKey = [System.Net.NetworkCredential]::new("", $secret).Password + $headers = @{ Authorization = "Bearer $apiKey" } + $baseUrl = "https://$EndpointId.api.runpod.ai" + + Write-Host "Warming the JoyAI model through POST /load..." + $loadResponse = Invoke-RestMethod ` + -Uri "$baseUrl/load" ` + -Headers $headers ` + -Method Post ` + -ContentType "application/json" ` + -Body "{}" ` + -TimeoutSec $WarmTimeoutSeconds + + Write-Host "Checking the public JoyAI health route..." + $health = Invoke-RestMethod ` + -Uri "$baseUrl/health" ` + -Headers $headers ` + -Method Get ` + -TimeoutSec 30 + + if (-not $health.ok -or -not $health.runtime_loaded) { + throw "JoyAI reported unhealthy or the model is not loaded." + } + + python.exe -m pip install --user aiohttp + + $env:RUNPOD_API_KEY = $apiKey + $proxyScript = Join-Path $PSScriptRoot "local_proxy.py" + $proxyProcess = Start-Process ` + -FilePath "python.exe" ` + -ArgumentList @( + ('"{0}"' -f $proxyScript), + "--endpoint-id", $EndpointId, + "--port", $LocalPort + ) ` + -NoNewWindow ` + -PassThru + + Start-Sleep -Seconds 2 + $localUrl = "http://127.0.0.1:$LocalPort/" + Write-Host "JoyAI is ready: $localUrl" + Write-Host "Keep this window open. Press Ctrl+C when the test is finished." + Start-Process $localUrl + Wait-Process -Id $proxyProcess.Id +} +catch { + Write-Host "TEST FAILED: $($_.Exception.Message)" -ForegroundColor Red + exit 1 +} +finally { + if ($null -ne $proxyProcess -and -not $proxyProcess.HasExited) { + Stop-Process -Id $proxyProcess.Id -Force -ErrorAction SilentlyContinue + } + Remove-Item Env:RUNPOD_API_KEY -ErrorAction SilentlyContinue + $apiKey = $null + $secret = $null +} diff --git a/runpod/health_server.py b/runpod/health_server.py index 2d40c84..2226194 100644 --- a/runpod/health_server.py +++ b/runpod/health_server.py @@ -12,6 +12,11 @@ @app.get("/ping") def ping(): + """RunPod-only readiness probe served on PORT_HEALTH. + + The public JoyAI server owns /health. A worker remains in RunPod's + initializing state (204) until that endpoint confirms the runtime loaded. + """ model_port = os.getenv( "JOYOMNI_PORT", os.getenv("PORT", "8080"), @@ -24,7 +29,7 @@ def ping(): ) as response: health = json.loads(response.read()) - if health.get("runtime_loaded") is True: + if health.get("ok") is True and health.get("runtime_loaded") is True: return { "status": "healthy", "model": "ready", diff --git a/runpod/local_proxy.py b/runpod/local_proxy.py new file mode 100644 index 0000000..5d2dc73 --- /dev/null +++ b/runpod/local_proxy.py @@ -0,0 +1,183 @@ +"""Local authenticated proxy for the JoyAI RunPod load-balancer endpoint. + +Browsers cannot attach a secret RunPod bearer token to normal navigation or a +WebSocket constructor. This proxy keeps the token in the local process and +adds it to upstream HTTP and WebSocket requests. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os + +from aiohttp import ClientError, ClientSession, ClientTimeout, WSMsgType, web + + +HOP_BY_HOP_HEADERS = { + "connection", + "content-length", + "host", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", +} + + +def filtered_headers(headers) -> dict[str, str]: + return { + name: value + for name, value in headers.items() + if name.lower() not in HOP_BY_HOP_HEADERS + and name.lower() != "authorization" + } + + +def upstream_url(request: web.Request) -> str: + return f"{request.app['upstream']}{request.rel_url}" + + +async def proxy_websocket(request: web.Request) -> web.WebSocketResponse: + downstream = web.WebSocketResponse(heartbeat=30) + await downstream.prepare(request) + + session: ClientSession = request.app["session"] + try: + async with session.ws_connect( + upstream_url(request), + headers=request.app["auth_headers"], + heartbeat=30, + max_msg_size=0, + ) as upstream: + + async def browser_to_runpod() -> None: + async for message in downstream: + if message.type == WSMsgType.TEXT: + await upstream.send_str(message.data) + elif message.type == WSMsgType.BINARY: + await upstream.send_bytes(message.data) + elif message.type == WSMsgType.ERROR: + break + + async def runpod_to_browser() -> None: + async for message in upstream: + if message.type == WSMsgType.TEXT: + await downstream.send_str(message.data) + elif message.type == WSMsgType.BINARY: + await downstream.send_bytes(message.data) + elif message.type == WSMsgType.ERROR: + break + + tasks = [ + asyncio.create_task(browser_to_runpod()), + asyncio.create_task(runpod_to_browser()), + ] + done, pending = await asyncio.wait( + tasks, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + await asyncio.gather(*done, *pending, return_exceptions=True) + + except (ClientError, asyncio.TimeoutError, OSError) as error: + if not downstream.closed: + await downstream.send_json( + { + "type": "proxy_error", + "error": f"RunPod connection failed: {error}", + } + ) + finally: + if not downstream.closed: + await downstream.close() + + return downstream + + +async def proxy_http(request: web.Request) -> web.Response: + session: ClientSession = request.app["session"] + headers = filtered_headers(request.headers) + headers.update(request.app["auth_headers"]) + + try: + async with session.request( + request.method, + upstream_url(request), + headers=headers, + data=await request.read(), + allow_redirects=False, + ) as response: + body = await response.read() + return web.Response( + body=body, + status=response.status, + reason=response.reason, + headers=filtered_headers(response.headers), + ) + except (ClientError, asyncio.TimeoutError, OSError) as error: + return web.Response( + status=503, + content_type="application/json", + text=json.dumps( + { + "status": 503, + "title": "RunPod connection unavailable", + "detail": str(error), + } + ), + ) + + +async def route_request(request: web.Request): + if request.path == "/ws": + return await proxy_websocket(request) + return await proxy_http(request) + + +async def create_session(app: web.Application) -> None: + app["session"] = ClientSession( + timeout=ClientTimeout(total=None, connect=60, sock_connect=60), + ) + + +async def close_session(app: web.Application) -> None: + await app["session"].close() + + +def create_app(endpoint_id: str, api_key: str) -> web.Application: + app = web.Application(client_max_size=1024**3) + app["upstream"] = f"https://{endpoint_id}.api.runpod.ai" + app["auth_headers"] = {"Authorization": f"Bearer {api_key}"} + app.on_startup.append(create_session) + app.on_cleanup.append(close_session) + app.router.add_route("*", "/{tail:.*}", route_request) + return app + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--endpoint-id", required=True) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", default=9000, type=int) + args = parser.parse_args() + + api_key = os.getenv("RUNPOD_API_KEY") + if not api_key: + raise SystemExit("RUNPOD_API_KEY is required") + + web.run_app( + create_app(args.endpoint_id, api_key), + host=args.host, + port=args.port, + print=lambda message: print(message, flush=True), + ) + + +if __name__ == "__main__": + main() diff --git a/runpod/start.py b/runpod/start.py index f453b6a..f351fd2 100644 --- a/runpod/start.py +++ b/runpod/start.py @@ -4,102 +4,140 @@ from pathlib import Path -repository_root = Path(__file__).resolve().parents[1] - -checkpoint_root = Path( - os.getenv( - "JOYOMNI_CKPT_ROOT", - "/runpod-volume/joyai/checkpoints", +TRUE_VALUES = {"1", "true", "yes", "on"} +FALSE_VALUES = {"0", "false", "no", "off"} + + +def env_enabled(name: str, *, default: bool) -> bool: + """Read a strict boolean environment flag. + + A typo must fail during startup instead of silently disabling model preload + and leaving RunPod's load balancer waiting forever for a ready worker. + """ + value = os.getenv(name) + if value is None: + return default + + normalized = value.strip().lower() + if normalized in TRUE_VALUES: + return True + if normalized in FALSE_VALUES: + return False + + choices = ", ".join(sorted(TRUE_VALUES | FALSE_VALUES)) + raise ValueError(f"{name} must be one of: {choices}") + + +def required_checkpoint_items(checkpoint_root: Path) -> list[Path]: + return [ + checkpoint_root + / "JoyAI-Video-Edit" + / "dit" + / "joyai_video_edit_dit_0811.pth", + checkpoint_root + / "JoyAI-Video-Edit" + / "vae" + / "config.json", + checkpoint_root + / "JoyAI-Video-Edit" + / "vae" + / "diffusion_pytorch_model.safetensors", + checkpoint_root / "MiMo-VL-7B-RL-2508", + ] + + +def build_model_command(repository_root: Path, *, preload: bool) -> list[str]: + return [ + "bash", + str(repository_root / "deploy" / "run_server.sh"), + "--preload" if preload else "--no-preload", + ] + + +def main() -> int: + repository_root = Path(__file__).resolve().parents[1] + checkpoint_root = Path( + os.getenv( + "JOYOMNI_CKPT_ROOT", + "/runpod-volume/joyai/checkpoints", + ) ) -) - -required_items = [ - checkpoint_root - / "JoyAI-Video-Edit" - / "dit" - / "joyai_video_edit_dit_0811.pth", - - checkpoint_root - / "JoyAI-Video-Edit" - / "vae" - / "config.json", - - checkpoint_root - / "JoyAI-Video-Edit" - / "vae" - / "diffusion_pytorch_model.safetensors", - - checkpoint_root - / "MiMo-VL-7B-RL-2508", -] - -missing_items = [ - str(item) - for item in required_items - if not item.exists() -] -if missing_items: - print("Required model files are missing:", flush=True) + missing_items = [ + str(item) + for item in required_checkpoint_items(checkpoint_root) + if not item.exists() + ] + + if missing_items: + print("Required model files are missing:", flush=True) + for item in missing_items: + print(f" - {item}", flush=True) + + print( + "\nRun this command once on the attached RunPod volume:", + flush=True, + ) + print( + "python3 /opt/joyai/runpod/download_models.py", + flush=True, + ) + return 1 + + environment = os.environ.copy() + + # RunPod supplies PORT for public traffic. PORT_HEALTH is reserved for the + # internal /ping server and must not be used by clients. + model_port = environment.get( + "PORT", + environment.get("JOYOMNI_PORT", "8080"), + ) + environment["JOYOMNI_PORT"] = model_port - for item in missing_items: - print(f" - {item}", flush=True) + try: + preload = env_enabled("JOYOMNI_PRELOAD", default=True) + except ValueError as error: + print(f"Invalid startup configuration: {error}", flush=True) + return 1 + print(f"Starting JoyAI on public port {model_port}...", flush=True) print( - "\nRun this command once on the attached RunPod volume:", - flush=True, - ) - print( - "python3 /opt/joyai/runpod/download_models.py", + "Routes: GET /, GET /health, POST /load, WS /ws; " + "RunPod health: GET /ping on PORT_HEALTH.", flush=True, ) + if preload: + print( + "Preloading the JoyAI runtime before RunPod marks this worker ready...", + flush=True, + ) - raise SystemExit(1) - -environment = os.environ.copy() - -# RunPod supplies PORT for public traffic. -model_port = environment.get( - "PORT", - environment.get("JOYOMNI_PORT", "8080"), -) - -environment["JOYOMNI_PORT"] = model_port - -print( - f"Starting JoyAI on port {model_port}...", - flush=True, -) - -health_process = subprocess.Popen( - [ - sys.executable, - str(repository_root / "runpod" / "health_server.py"), - ], - cwd=repository_root, - env=environment, -) - -exit_code = 1 - -try: - model_process = subprocess.Popen( + health_process = subprocess.Popen( [ - "bash", - str(repository_root / "deploy" / "run_server.sh"), + sys.executable, + str(repository_root / "runpod" / "health_server.py"), ], cwd=repository_root, env=environment, ) - exit_code = model_process.wait() - -finally: - health_process.terminate() - + exit_code = 1 try: - health_process.wait(timeout=10) - except subprocess.TimeoutExpired: - health_process.kill() - -raise SystemExit(exit_code) \ No newline at end of file + model_process = subprocess.Popen( + build_model_command(repository_root, preload=preload), + cwd=repository_root, + env=environment, + ) + exit_code = model_process.wait() + finally: + health_process.terminate() + try: + health_process.wait(timeout=10) + except subprocess.TimeoutExpired: + health_process.kill() + + return exit_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py new file mode 100644 index 0000000..a42d1b8 --- /dev/null +++ b/tests/test_runpod_connection.py @@ -0,0 +1,55 @@ +import importlib.util +import os +import unittest +from pathlib import Path +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +START_PATH = ROOT / "runpod" / "start.py" + +SPEC = importlib.util.spec_from_file_location("runpod_start", START_PATH) +START = importlib.util.module_from_spec(SPEC) +assert SPEC.loader is not None +SPEC.loader.exec_module(START) + + +class RunPodConnectionContractTests(unittest.TestCase): + def test_preload_defaults_to_enabled(self): + with patch.dict(os.environ, {}, clear=True): + self.assertTrue(START.env_enabled("JOYOMNI_PRELOAD", default=True)) + + def test_preload_can_be_disabled_explicitly(self): + with patch.dict(os.environ, {"JOYOMNI_PRELOAD": "0"}, clear=True): + self.assertFalse(START.env_enabled("JOYOMNI_PRELOAD", default=True)) + + def test_invalid_preload_value_fails_fast(self): + with patch.dict(os.environ, {"JOYOMNI_PRELOAD": "maybe"}, clear=True): + with self.assertRaises(ValueError): + START.env_enabled("JOYOMNI_PRELOAD", default=True) + + def test_model_command_makes_preload_explicit(self): + command = START.build_model_command(ROOT, preload=True) + self.assertEqual(command[-1], "--preload") + + def test_runpod_ping_is_internal_and_checks_public_health(self): + text = (ROOT / "runpod" / "health_server.py").read_text() + self.assertIn('@app.get("/ping")', text) + self.assertIn('/health', text) + self.assertIn('status_code=204', text) + + def test_windows_script_uses_public_routes(self): + text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + self.assertIn('$baseUrl/load', text) + self.assertIn('$baseUrl/health', text) + self.assertNotIn('$baseUrl/ping', text) + self.assertIn('(\'"{0}"\' -f $proxyScript)', text) + + def test_proxy_handles_websocket_and_does_not_target_public_ping(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn('request.path == "/ws"', text) + self.assertNotIn('api.runpod.ai/ping', text) + + +if __name__ == "__main__": + unittest.main() From fbf5343b35cb841945fe0c90786b4cdb1e4e1e3e Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 09:05:00 +0100 Subject: [PATCH 15/32] Fix bounded RunPod connection warm-up --- runpod/Start-JoyAI-Realtime-Test.ps1 | 144 +++++++++++++++++++++++---- 1 file changed, 123 insertions(+), 21 deletions(-) diff --git a/runpod/Start-JoyAI-Realtime-Test.ps1 b/runpod/Start-JoyAI-Realtime-Test.ps1 index 4fe372f..1b8bda1 100644 --- a/runpod/Start-JoyAI-Realtime-Test.ps1 +++ b/runpod/Start-JoyAI-Realtime-Test.ps1 @@ -1,7 +1,8 @@ param( [string]$EndpointId = "ex9647vtulowka", [int]$LocalPort = 9000, - [int]$WarmTimeoutSeconds = 900 + [int]$WarmTimeoutSeconds = 300, + [int]$RetryDelaySeconds = 10 ) $ErrorActionPreference = "Stop" @@ -9,33 +10,110 @@ $proxyProcess = $null $secret = $null $apiKey = $null +function Get-HttpStatusCode { + param([System.Management.Automation.ErrorRecord]$ErrorRecord) + + try { + if ($null -ne $ErrorRecord.Exception.Response.StatusCode) { + return [int]$ErrorRecord.Exception.Response.StatusCode + } + } + catch { + return $null + } + return $null +} + +function Get-ErrorDetail { + param([System.Management.Automation.ErrorRecord]$ErrorRecord) + + if ($null -ne $ErrorRecord.ErrorDetails -and $ErrorRecord.ErrorDetails.Message) { + return $ErrorRecord.ErrorDetails.Message + } + return $ErrorRecord.Exception.Message +} + try { + Write-Host "Checking the local proxy dependency..." + python.exe -c "import aiohttp" 2>$null + if ($LASTEXITCODE -ne 0) { + python.exe -m pip install --user aiohttp + if ($LASTEXITCODE -ne 0) { + throw "Could not install aiohttp for the local proxy." + } + } + $secret = Read-Host "Paste your RunPod API key" -AsSecureString $apiKey = [System.Net.NetworkCredential]::new("", $secret).Password $headers = @{ Authorization = "Bearer $apiKey" } $baseUrl = "https://$EndpointId.api.runpod.ai" - Write-Host "Warming the JoyAI model through POST /load..." - $loadResponse = Invoke-RestMethod ` - -Uri "$baseUrl/load" ` - -Headers $headers ` - -Method Post ` - -ContentType "application/json" ` - -Body "{}" ` - -TimeoutSec $WarmTimeoutSeconds - - Write-Host "Checking the public JoyAI health route..." - $health = Invoke-RestMethod ` - -Uri "$baseUrl/health" ` - -Headers $headers ` - -Method Get ` - -TimeoutSec 30 - - if (-not $health.ok -or -not $health.runtime_loaded) { - throw "JoyAI reported unhealthy or the model is not loaded." + Write-Host "Starting the JoyAI worker through GET /health..." + Write-Host "The H200 must load the 31 GB DiT checkpoint before RunPod routes traffic." + Write-Host "A RunPod 400 'timed out waiting for worker' response is retried until the $WarmTimeoutSeconds-second safety limit." + + $warmTimer = [System.Diagnostics.Stopwatch]::StartNew() + $health = $null + $workerReady = $false + $attempt = 0 + + while (-not $workerReady -and $warmTimer.Elapsed.TotalSeconds -lt $WarmTimeoutSeconds) { + $attempt += 1 + $remainingSeconds = $WarmTimeoutSeconds - [int]$warmTimer.Elapsed.TotalSeconds + if ($remainingSeconds -le 0) { + break + } + + # RunPod's load balancer can spend about two minutes waiting for a cold + # worker before returning HTTP 400. Keep each request slightly above + # that window, then retry without creating another worker. + $requestTimeout = [Math]::Min(135, $remainingSeconds) + Write-Host "Readiness attempt $attempt (elapsed $([int]$warmTimer.Elapsed.TotalSeconds)s)..." + + try { + $health = Invoke-RestMethod ` + -Uri "$baseUrl/health" ` + -Headers $headers ` + -Method Get ` + -TimeoutSec $requestTimeout + + if ($health.ok -and $health.runtime_loaded) { + $workerReady = $true + break + } + + Write-Host "Worker answered, but the runtime is still initializing." + } + catch { + $statusCode = Get-HttpStatusCode $_ + $detail = Get-ErrorDetail $_ + $retryableStatusCodes = @(400, 408, 425, 429, 500, 502, 503, 504) + $isColdStartResponse = ( + $retryableStatusCodes -contains $statusCode -or + $detail -match "timed out waiting for worker|cold start|no worker|temporarily unavailable|semaphore timeout|forcibly closed" + ) + + if (-not $isColdStartResponse) { + throw + } + + Write-Host "Worker is still cold-starting: $detail" + } + + $remainingSeconds = $WarmTimeoutSeconds - [int]$warmTimer.Elapsed.TotalSeconds + if ($remainingSeconds -gt 0) { + $sleepSeconds = [Math]::Min($RetryDelaySeconds, $remainingSeconds) + Write-Host "Retrying in $sleepSeconds seconds..." + Start-Sleep -Seconds $sleepSeconds + } } - python.exe -m pip install --user aiohttp + $warmTimer.Stop() + if (-not $workerReady) { + throw "JoyAI did not become ready within $WarmTimeoutSeconds seconds. Stop here and inspect the worker log; do not rebuild the image or start another Pod." + } + + Write-Host "JoyAI runtime is ready after $([int]$warmTimer.Elapsed.TotalSeconds) seconds." $env:RUNPOD_API_KEY = $apiKey $proxyScript = Join-Path $PSScriptRoot "local_proxy.py" @@ -49,8 +127,32 @@ try { -NoNewWindow ` -PassThru - Start-Sleep -Seconds 2 $localUrl = "http://127.0.0.1:$LocalPort/" + $localHealthUrl = "${localUrl}health" + $localReady = $false + for ($attempt = 1; $attempt -le 10; $attempt++) { + if ($proxyProcess.HasExited) { + throw "The local JoyAI proxy stopped unexpectedly." + } + try { + $localHealth = Invoke-RestMethod ` + -Uri $localHealthUrl ` + -Method Get ` + -TimeoutSec 30 + if ($localHealth.ok -and $localHealth.runtime_loaded) { + $localReady = $true + break + } + } + catch { + Start-Sleep -Seconds 1 + } + } + + if (-not $localReady) { + throw "The local proxy could not reach the ready JoyAI worker." + } + Write-Host "JoyAI is ready: $localUrl" Write-Host "Keep this window open. Press Ctrl+C when the test is finished." Start-Process $localUrl From eca57f9c079a3fc8d026d43a41eb57c734183508 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 09:05:06 +0100 Subject: [PATCH 16/32] Handle proxy compression and Windows resets --- runpod/local_proxy.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/runpod/local_proxy.py b/runpod/local_proxy.py index 5d2dc73..e81d67a 100644 --- a/runpod/local_proxy.py +++ b/runpod/local_proxy.py @@ -29,6 +29,20 @@ } +def proxy_exception_handler(loop: asyncio.AbstractEventLoop, context: dict) -> None: + """Ignore the harmless Windows Proactor reset logged after a peer closes.""" + error = context.get("exception") + message = context.get("message", "") + if ( + os.name == "nt" + and isinstance(error, ConnectionResetError) + and getattr(error, "winerror", None) == 10054 + and "_call_connection_lost" in message + ): + return + loop.default_exception_handler(context) + + def filtered_headers(headers) -> dict[str, str]: return { name: value @@ -141,8 +155,14 @@ async def route_request(request: web.Request): async def create_session(app: web.Application) -> None: + if os.name == "nt": + asyncio.get_running_loop().set_exception_handler(proxy_exception_handler) app["session"] = ClientSession( timeout=ClientTimeout(total=None, connect=60, sock_connect=60), + # Relay compressed bytes and their Content-Encoding header unchanged. + # The browser can decode Brotli itself; the local Python proxy should + # not require the optional Brotli package just to forward RunPod HTML. + auto_decompress=False, ) From 398ad2b31debee73ae5974aaa68d2658a21fba26 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 09:05:11 +0100 Subject: [PATCH 17/32] Document RunPod connection safeguards --- runpod/README.md | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/runpod/README.md b/runpod/README.md index 78456a0..ab9150c 100644 --- a/runpod/README.md +++ b/runpod/README.md @@ -57,9 +57,17 @@ cd path\to\JoyAI-Video-Edit\runpod powershell.exe -NoProfile -ExecutionPolicy Bypass -File .\Start-JoyAI-Realtime-Test.ps1 ``` -The script securely asks for the API key, warms through `POST /load`, verifies -`GET /health`, starts the local HTTP/WebSocket proxy, and opens the interface. -The key stays in the local process and is removed when the script exits. +The script securely asks for the API key and polls `GET /health` until the +preloaded runtime is ready. RunPod's load balancer may return HTTP `400` with +`timed out waiting for worker` after its own two-minute wait while the worker +is still initializing; the script treats that response as retryable up to its +five-minute safety limit. It then starts the local HTTP/WebSocket proxy and +opens the interface. The key stays in the local process and is removed when +the script exits. + +`POST /load` is not used by this test flow because the container already starts +with `JOYOMNI_PRELOAD=1`, and RunPod cannot route that request until the worker +has passed readiness anyway. Stop the test with `Ctrl+C`. With zero active workers and a short idle timeout, RunPod can scale the worker back to zero after the connection closes. From 68215a5ce68feee04198bf18ec37426bfa7afea6 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 09:05:17 +0100 Subject: [PATCH 18/32] Test RunPod connection safeguards --- tests/test_runpod_connection.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index a42d1b8..b033543 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -40,16 +40,46 @@ def test_runpod_ping_is_internal_and_checks_public_health(self): def test_windows_script_uses_public_routes(self): text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() - self.assertIn('$baseUrl/load', text) self.assertIn('$baseUrl/health', text) + self.assertNotIn('$baseUrl/load', text) self.assertNotIn('$baseUrl/ping', text) self.assertIn('(\'"{0}"\' -f $proxyScript)', text) + def test_windows_script_prepares_proxy_dependency_before_warming_worker(self): + text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + dependency_check = text.index('python.exe -c "import aiohttp"') + warm_request = text.index('$baseUrl/health') + self.assertLess(dependency_check, warm_request) + self.assertIn("$localReady", text) + + def test_windows_script_retries_runpod_cold_start_timeout(self): + text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + self.assertIn("while (-not $workerReady", text) + self.assertIn("timed out waiting for worker", text) + self.assertIn("$retryableStatusCodes", text) + self.assertIn("$WarmTimeoutSeconds", text) + + def test_windows_script_opens_browser_only_after_local_health_is_ready(self): + text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + local_health_check = text.index('$localHealthUrl') + browser_open = text.index('Start-Process $localUrl') + self.assertLess(local_health_check, browser_open) + def test_proxy_handles_websocket_and_does_not_target_public_ping(self): text = (ROOT / "runpod" / "local_proxy.py").read_text() self.assertIn('request.path == "/ws"', text) self.assertNotIn('api.runpod.ai/ping', text) + def test_proxy_relays_brotli_without_a_python_decoder(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn("auto_decompress=False", text) + + def test_proxy_suppresses_only_windows_cleanup_resets(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn("proxy_exception_handler", text) + self.assertIn('getattr(error, "winerror", None) == 10054', text) + self.assertIn('"_call_connection_lost" in message', text) + if __name__ == "__main__": unittest.main() From 7b143dc7ef2e2da1e4008d4ed1018c72b3ef1269 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 16:53:27 +0100 Subject: [PATCH 19/32] Trigger verified H200 image build From 6c08352ff94da760c01e88891e5d93596cf069cd Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 16:56:07 +0100 Subject: [PATCH 20/32] Report H200 container build status --- .github/workflows/build-runpod-h200.yml | 47 ++++++++++++++++++++++++- 1 file changed, 46 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build-runpod-h200.yml b/.github/workflows/build-runpod-h200.yml index 0e6edb5..c54eb12 100644 --- a/.github/workflows/build-runpod-h200.yml +++ b/.github/workflows/build-runpod-h200.yml @@ -16,10 +16,11 @@ on: permissions: contents: read packages: write + statuses: write concurrency: group: runpod-h200-image-${{ github.ref }} - cancel-in-progress: false + cancel-in-progress: true env: REGISTRY: ghcr.io @@ -53,6 +54,24 @@ jobs: username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} + - name: Report H200 image build pending + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + payload=$(jq -nc \ + --arg state pending \ + --arg context runpod-h200-image \ + --arg description "Building the immutable RunPod H200 image" \ + --arg target_url "${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID}" \ + '{state: $state, context: $context, description: $description, target_url: $target_url}') + curl --fail-with-body --retry 3 \ + -X POST \ + -H "Accept: application/vnd.github+json" \ + -H "Authorization: Bearer ${GH_TOKEN}" \ + -H "X-GitHub-Api-Version: 2022-11-28" \ + "${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/statuses/${GITHUB_SHA}" \ + -d "${payload}" + - name: Build and publish the H200 container image id: build uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4 @@ -74,3 +93,29 @@ jobs: - name: Show immutable deployment reference run: echo "${REGISTRY}/${IMAGE_NAME}@${{ steps.build.outputs.digest }}" + + - name: Report H200 image build result + if: always() + env: + BUILD_OUTCOME: ${{ steps.build.outcome }} + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + state=failure + description="RunPod H200 image build failed" + if [ "${BUILD_OUTCOME}" = "success" ]; then + state=success + description="Immutable RunPod H200 image published" + fi + payload=$(jq -nc \ + --arg state "${state}" \ + --arg context runpod-h200-image \ + --arg description "${description}" \ + --arg target_url "${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID}" \ + '{state: $state, context: $context, description: $description, target_url: $target_url}') + curl --fail-with-body --retry 3 \ + -X POST \ + -H "Accept: application/vnd.github+json" \ + -H "Authorization: Bearer ${GH_TOKEN}" \ + -H "X-GitHub-Api-Version: 2022-11-28" \ + "${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/statuses/${GITHUB_SHA}" \ + -d "${payload}" From fd21f4d36884dc4c220c6bf2e3877fbc7ff79d14 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 17:07:58 +0100 Subject: [PATCH 21/32] Remove fragile GitHub status reporting --- .github/workflows/build-runpod-h200.yml | 45 ------------------------- 1 file changed, 45 deletions(-) diff --git a/.github/workflows/build-runpod-h200.yml b/.github/workflows/build-runpod-h200.yml index c54eb12..faf94fd 100644 --- a/.github/workflows/build-runpod-h200.yml +++ b/.github/workflows/build-runpod-h200.yml @@ -16,7 +16,6 @@ on: permissions: contents: read packages: write - statuses: write concurrency: group: runpod-h200-image-${{ github.ref }} @@ -54,24 +53,6 @@ jobs: username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} - - name: Report H200 image build pending - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - payload=$(jq -nc \ - --arg state pending \ - --arg context runpod-h200-image \ - --arg description "Building the immutable RunPod H200 image" \ - --arg target_url "${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID}" \ - '{state: $state, context: $context, description: $description, target_url: $target_url}') - curl --fail-with-body --retry 3 \ - -X POST \ - -H "Accept: application/vnd.github+json" \ - -H "Authorization: Bearer ${GH_TOKEN}" \ - -H "X-GitHub-Api-Version: 2022-11-28" \ - "${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/statuses/${GITHUB_SHA}" \ - -d "${payload}" - - name: Build and publish the H200 container image id: build uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4 @@ -93,29 +74,3 @@ jobs: - name: Show immutable deployment reference run: echo "${REGISTRY}/${IMAGE_NAME}@${{ steps.build.outputs.digest }}" - - - name: Report H200 image build result - if: always() - env: - BUILD_OUTCOME: ${{ steps.build.outcome }} - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - state=failure - description="RunPod H200 image build failed" - if [ "${BUILD_OUTCOME}" = "success" ]; then - state=success - description="Immutable RunPod H200 image published" - fi - payload=$(jq -nc \ - --arg state "${state}" \ - --arg context runpod-h200-image \ - --arg description "${description}" \ - --arg target_url "${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}/actions/runs/${GITHUB_RUN_ID}" \ - '{state: $state, context: $context, description: $description, target_url: $target_url}') - curl --fail-with-body --retry 3 \ - -X POST \ - -H "Accept: application/vnd.github+json" \ - -H "Authorization: Bearer ${GH_TOKEN}" \ - -H "X-GitHub-Api-Version: 2022-11-28" \ - "${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/statuses/${GITHUB_SHA}" \ - -d "${payload}" From 9c5e87f6e92aed3ef065979cef825503d80e3a56 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 17:43:42 +0100 Subject: [PATCH 22/32] Prevent browser codec probes from locking startup --- deploy/static/index.html | 27 +++++++++++++++++++++++---- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/deploy/static/index.html b/deploy/static/index.html index 0da32c0..3587c18 100644 --- a/deploy/static/index.html +++ b/deploy/static/index.html @@ -214,7 +214,6 @@ - - - -
@@ -1139,8 +1138,17 @@ let h264DecodeOk = false; let h264DecodeReady = Promise.resolve(false); +const CODEC_PROBE_TIMEOUT_MS = 1200; +function codecProbeWithTimeout(probe, fallback = false) { + return Promise.race([ + Promise.resolve(probe).catch(() => fallback), + new Promise((resolve) => setTimeout(() => resolve(fallback), CODEC_PROBE_TIMEOUT_MS)), + ]); +} if (typeof VideoDecoder !== "undefined") { - h264DecodeReady = VideoDecoder.isConfigSupported({ codec: "avc1.640028", optimizeForLatency: true }) + h264DecodeReady = codecProbeWithTimeout( + VideoDecoder.isConfigSupported({ codec: "avc1.640028", optimizeForLatency: true }) + ) .then((s) => { h264DecodeOk = !!(s && s.supported); return h264DecodeOk; }) .catch(() => false); } @@ -1630,7 +1638,10 @@ upCodecH264 = false; if (!usingVideoFile && typeof VideoEncoder !== "undefined") { try { - const s = await VideoEncoder.isConfigSupported(upEncoderConfig(width, height)); + const s = await codecProbeWithTimeout( + VideoEncoder.isConfigSupported(upEncoderConfig(width, height)), + null + ); upCodecH264 = !!(s && s.supported); } catch (e) {} } @@ -1698,7 +1709,12 @@ sessionGranted = true; setSendBusy(true, t("busy_starting")); showOutputStart(true, t("busy_starting")); - beginSession(); + try { + const started = await beginSession(); + if (!started) throw new Error("session start was not sent"); + } catch (err) { + stopActiveRun(); + } return; } if (msg.type === "started") { @@ -1976,6 +1992,9 @@ currentWs.onclose = null; try { currentWs.close(); } catch (err) {} } + startingRun = false; + sessionGranted = false; + setSendBusy(false, ""); } async function send() { From 71e68e98e40b8aec6d9e0b9d7b432d0c16e5192d Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 17:43:43 +0100 Subject: [PATCH 23/32] Add disabled-send-button regression coverage --- tests/test_runpod_connection.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index 19ecb33..2373a73 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -100,10 +100,21 @@ def test_h200_live_mode_disables_recording_and_uses_low_bandwidth_defaults(self) self.assertNotIn(' --record-dir "$RECORD_DIR" \\\n', launcher) self.assertNotIn('id="downloadBubble"', html) self.assertNotIn('id="outputStartOverlay"', html) + self.assertNotIn('id="outputIdleOverlay"', html) + self.assertNotIn('id="outputWaitOverlay"', html) + self.assertNotIn('id="outputToast"', html) + self.assertNotIn('id="sendHint"', html) self.assertIn('data-upq="0.2" class="on"', html) self.assertIn('data-fps="16" class="on"', html) self.assertIn("let autoQTier = 0;", html) + def test_browser_codec_probes_cannot_lock_the_send_button(self): + html = (ROOT / "deploy" / "static" / "index.html").read_text() + self.assertIn("const CODEC_PROBE_TIMEOUT_MS = 1200;", html) + self.assertGreaterEqual(html.count("codecProbeWithTimeout("), 3) + self.assertIn('if (!started) throw new Error("session start was not sent")', html) + self.assertIn('setSendBusy(false, "");', html) + if __name__ == "__main__": unittest.main() From 89d094c4e9bbd3ea96e567418fdd3bf64e50d5d3 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 20:05:51 +0100 Subject: [PATCH 24/32] Tune live streaming for 20 FPS low latency --- Dockerfile.h200 | 2 +- deploy/static/index.html | 24 ++++++++++++------------ runpod/README.md | 8 ++++---- runpod/local_proxy.py | 31 ++++++++++++++++++------------- tests/test_runpod_connection.py | 12 ++++++++++-- 5 files changed, 45 insertions(+), 32 deletions(-) diff --git a/Dockerfile.h200 b/Dockerfile.h200 index 0361454..8484a9b 100644 --- a/Dockerfile.h200 +++ b/Dockerfile.h200 @@ -97,7 +97,7 @@ ENV JOYOMNI_DEVICE=cuda:0 \ JOYOMNI_PRELOAD=1 \ JOYOMNI_WIDTH=840 \ JOYOMNI_HEIGHT=480 \ - JOYOMNI_FPS=16 \ + JOYOMNI_FPS=20 \ JOYOMNI_FP8_IMG=1 \ JOYOMNI_FP8_TXT=1 \ JOYOMNI_CUDA_GRAPH=1 \ diff --git a/deploy/static/index.html b/deploy/static/index.html index 3587c18..f81324a 100644 --- a/deploy/static/index.html +++ b/deploy/static/index.html @@ -227,7 +227,7 @@ - +
-
-
+
+
@@ -274,8 +274,8 @@
-
-
+
+
@@ -617,8 +617,8 @@ let pendingPeCacheKey = null; let outputPlaybackTimer = null; let outputPlaybackIntervalMs = null; -const DEFAULT_TARGET_OUTPUT_QUEUE_DELAY_MS = 200; -const DEFAULT_MAX_OUTPUT_QUEUE_DELAY_MS = 300; +const DEFAULT_TARGET_OUTPUT_QUEUE_DELAY_MS = 100; +const DEFAULT_MAX_OUTPUT_QUEUE_DELAY_MS = 200; function targetOutputQueueDelayMs() { const el = document.getElementById("targetQueueDelayMs"); const raw = el && el.value !== "" ? el.value : null; @@ -632,7 +632,7 @@ const v = raw === null ? NaN : Number(raw); return Number.isFinite(v) && v >= 0 ? v : DEFAULT_MAX_OUTPUT_QUEUE_DELAY_MS; } -const MAX_BACKEND_PENDING_FRAMES = 32; +const MAX_BACKEND_PENDING_FRAMES = 16; let sendRealtimeSkips = 0; let sentFrames = 0; let backendAckedFrames = 0; @@ -1159,7 +1159,7 @@ let upCodecH264 = false; let upEncoder = null; -const UPLINK_KEYFRAME_INTERVAL = 8; +const UPLINK_KEYFRAME_INTERVAL = 20; let upSeq = 0; const upPendingByTs = new Map(); @@ -2593,7 +2593,7 @@ b.addEventListener("click", () => { const d = b.getAttribute("data-q"); if (d === "auto") { - autoQuality = true; autoQTier = 2; autoQLowStreak = 0; autoQHoldTicks = 0; + autoQuality = true; autoQTier = 0; autoQLowStreak = 0; autoQHoldTicks = 0; highlightQTier("auto"); setDownQuality(Q_TIERS[autoQTier], { live: true }); } else { diff --git a/runpod/README.md b/runpod/README.md index fb07914..2dcff9e 100644 --- a/runpod/README.md +++ b/runpod/README.md @@ -77,10 +77,10 @@ rebuild the image during that inspection. The H200 image defaults to live-only operation. Recording and download finalization are disabled, the presence gate is off, and the browser starts at -16 FPS with low upload/downlink quality. The adaptive downlink can increase -quality after the connection proves stable. Refreshing the browser replaces -the previous WebSocket session instead of waiting behind its stale session -ticket. +20 FPS with low upload/downlink quality. Its low-latency playback buffer targets +100 ms, and the adaptive downlink can increase quality after the connection +proves stable. Refreshing the browser replaces the previous WebSocket session +instead of waiting behind its stale session ticket. Stop the test with `Ctrl+C`. With zero active workers and a short idle timeout, RunPod can scale the worker back to zero after the connection closes. diff --git a/runpod/local_proxy.py b/runpod/local_proxy.py index bcffae9..39a6b5a 100644 --- a/runpod/local_proxy.py +++ b/runpod/local_proxy.py @@ -61,6 +61,7 @@ async def proxy_websocket(request: web.Request) -> web.WebSocketResponse: await downstream.prepare(request) session: ClientSession = request.app["session"] + proxy_state = request.app["proxy_state"] upstream = None try: upstream = await session.ws_connect( @@ -73,11 +74,11 @@ async def proxy_websocket(request: web.Request) -> web.WebSocketResponse: # The local proxy is intentionally single-viewer. Replacing a stale # browser socket explicitly tells the server to release its session # ticket, avoiding a long queue after a page refresh. - async with request.app["websocket_lock"]: - previous_upstream = request.app.get("active_upstream") - previous_downstream = request.app.get("active_downstream") - request.app["active_upstream"] = upstream - request.app["active_downstream"] = downstream + async with proxy_state["websocket_lock"]: + previous_upstream = proxy_state["active_upstream"] + previous_downstream = proxy_state["active_downstream"] + proxy_state["active_upstream"] = upstream + proxy_state["active_downstream"] = downstream # Close the replaced pair outside the lock. The old handler also uses # this lock during cleanup, so awaiting its close while holding the @@ -142,11 +143,11 @@ async def runpod_to_browser() -> None: except (ClientError, asyncio.TimeoutError, OSError): pass await upstream.close() - async with request.app["websocket_lock"]: - if request.app.get("active_upstream") is upstream: - request.app["active_upstream"] = None - if request.app.get("active_downstream") is downstream: - request.app["active_downstream"] = None + async with proxy_state["websocket_lock"]: + if proxy_state["active_upstream"] is upstream: + proxy_state["active_upstream"] = None + if proxy_state["active_downstream"] is downstream: + proxy_state["active_downstream"] = None if not downstream.closed: await downstream.close() @@ -207,9 +208,13 @@ async def create_session(app: web.Application) -> None: # not require the optional Brotli package just to forward RunPod HTML. auto_decompress=False, ) - app["websocket_lock"] = asyncio.Lock() - app["active_upstream"] = None - app["active_downstream"] = None + # aiohttp warns when top-level application state is mutated after startup. + # Keep live socket state inside one mutable object registered at startup. + app["proxy_state"] = { + "websocket_lock": asyncio.Lock(), + "active_upstream": None, + "active_downstream": None, + } async def close_session(app: web.Application) -> None: diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index 2373a73..83435b3 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -77,6 +77,8 @@ def test_proxy_handles_websocket_and_does_not_target_public_ping(self): self.assertNotIn('api.runpod.ai/ping', text) self.assertIn('await previous_upstream.send_json({"type": "stop"})', text) self.assertIn('message=b"replaced by a refreshed browser"', text) + self.assertIn('proxy_state = request.app["proxy_state"]', text) + self.assertNotIn('request.app["active_upstream"] =', text) def test_proxy_relays_brotli_without_a_python_decoder(self): text = (ROOT / "runpod" / "local_proxy.py").read_text() @@ -95,7 +97,7 @@ def test_h200_live_mode_disables_recording_and_uses_low_bandwidth_defaults(self) html = (ROOT / "deploy" / "static" / "index.html").read_text() self.assertIn("JOYOMNI_RECORD_ENABLED=0", dockerfile) self.assertIn("JOYOMNI_ONLINE_GATE_ENABLED=0", dockerfile) - self.assertIn("JOYOMNI_FPS=16", dockerfile) + self.assertIn("JOYOMNI_FPS=20", dockerfile) self.assertIn('EXTRA_ARGS+=(--record-dir "$RECORD_DIR")', launcher) self.assertNotIn(' --record-dir "$RECORD_DIR" \\\n', launcher) self.assertNotIn('id="downloadBubble"', html) @@ -104,9 +106,15 @@ def test_h200_live_mode_disables_recording_and_uses_low_bandwidth_defaults(self) self.assertNotIn('id="outputWaitOverlay"', html) self.assertNotIn('id="outputToast"', html) self.assertNotIn('id="sendHint"', html) + self.assertIn('class="kvreset-field keep-min" style="display:none"', html) self.assertIn('data-upq="0.2" class="on"', html) - self.assertIn('data-fps="16" class="on"', html) + self.assertIn('data-fps="20" class="on"', html) self.assertIn("let autoQTier = 0;", html) + self.assertIn("const DEFAULT_TARGET_OUTPUT_QUEUE_DELAY_MS = 100;", html) + self.assertIn("const DEFAULT_MAX_OUTPUT_QUEUE_DELAY_MS = 200;", html) + self.assertIn("const MAX_BACKEND_PENDING_FRAMES = 16;", html) + self.assertIn("const UPLINK_KEYFRAME_INTERVAL = 20;", html) + self.assertIn('autoQuality = true; autoQTier = 0;', html) def test_browser_codec_probes_cannot_lock_the_send_button(self): html = (ROOT / "deploy" / "static" / "index.html").read_text() From 89001e8b26fef150ad828376ee0d20506947a3ca Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 22:05:16 +0100 Subject: [PATCH 25/32] Enable persistent compiled VAE on RunPod H200 --- Dockerfile.h200 | 8 +- deploy/run_server.sh | 12 ++- deploy/xvideo/models/vae/vae_compile.py | 59 +++++++++++++-- deploy/xvideo/serving/joyomni_streaming.py | 73 +++++++++++++++---- .../xvideo/serving/serve_joyomni_streaming.py | 24 ++++++ runpod/README.md | 13 +++- runpod/Start-JoyAI-Realtime-Test.ps1 | 11 ++- tests/test_runpod_connection.py | 37 ++++++++++ 8 files changed, 210 insertions(+), 27 deletions(-) diff --git a/Dockerfile.h200 b/Dockerfile.h200 index 8484a9b..19f23d7 100644 --- a/Dockerfile.h200 +++ b/Dockerfile.h200 @@ -94,6 +94,8 @@ ENV JOYOMNI_DEVICE=cuda:0 \ JOYOMNI_HOST=0.0.0.0 \ JOYOMNI_PORT=8080 \ JOYOMNI_CKPT_ROOT=/runpod-volume/joyai/checkpoints \ + JOYOMNI_CACHE_ROOT=/runpod-volume/joyai/cache/h200-torch291-cu128 \ + JOYOMNI_CACHE_READY_MARKER=/runpod-volume/joyai/cache/h200-torch291-cu128/ready.json \ JOYOMNI_PRELOAD=1 \ JOYOMNI_WIDTH=840 \ JOYOMNI_HEIGHT=480 \ @@ -103,7 +105,11 @@ ENV JOYOMNI_DEVICE=cuda:0 \ JOYOMNI_CUDA_GRAPH=1 \ JOYOMNI_SAGE_ATTN=1 \ JOYOMNI_TXT_PARALLEL=1 \ - JOYOMNI_VAE_COMPILE=0 \ + JOYOMNI_VAE_COMPILE=1 \ + JOYOMNI_VAE_COMPILE_STRICT=1 \ + JOYOMNI_LOAD_WARMUP_STRICT=1 \ + JOYOMNI_WARMUP_BOTH_ORIENTATIONS=0 \ + JOYOMNI_WARMUP_REFERENCE_BUCKETS=0 \ JOYOMNI_RECORD_ENABLED=0 \ JOYOMNI_ONLINE_GATE_ENABLED=0 diff --git a/deploy/run_server.sh b/deploy/run_server.sh index c54ead9..9c26c96 100755 --- a/deploy/run_server.sh +++ b/deploy/run_server.sh @@ -26,13 +26,17 @@ fi cd "$HERE" -export TORCHINDUCTOR_CACHE_DIR="$HERE/deps/cache/torchinductor" -export TRITON_CACHE_DIR="$HERE/deps/cache/triton" -export CUDA_CACHE_PATH="$HERE/deps/cache/nv_compute" -export TORCHINDUCTOR_FX_GRAPH_CACHE=1 +CACHE_ROOT="${JOYOMNI_CACHE_ROOT:-$HERE/deps/cache}" +export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-$CACHE_ROOT/torchinductor}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-$CACHE_ROOT/triton}" +export CUDA_CACHE_PATH="${CUDA_CACHE_PATH:-$CACHE_ROOT/nv_compute}" +export TORCHINDUCTOR_FX_GRAPH_CACHE="${TORCHINDUCTOR_FX_GRAPH_CACHE:-1}" export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" mkdir -p "$TORCHINDUCTOR_CACHE_DIR" "$TRITON_CACHE_DIR" "$CUDA_CACHE_PATH" +echo "JoyAI compile cache root: $CACHE_ROOT" +echo "JoyAI VAE compile: ${JOYOMNI_VAE_COMPILE:-1} (strict=${JOYOMNI_VAE_COMPILE_STRICT:-0})" + export PYTHONUNBUFFERED=1 export PYTHONPATH="$HERE" diff --git a/deploy/xvideo/models/vae/vae_compile.py b/deploy/xvideo/models/vae/vae_compile.py index e7f43da..ace4eba 100644 --- a/deploy/xvideo/models/vae/vae_compile.py +++ b/deploy/xvideo/models/vae/vae_compile.py @@ -10,14 +10,20 @@ def compile_enabled() -> bool: """Return whether Torch Inductor VAE compilation is enabled. - Compilation remains enabled by default for non-serverless deployments. - RunPod H200 images disable it explicitly so cold workers become healthy - without spending minutes autotuning dozens of VAE shapes. + Compilation remains enabled by default. Serverless deployments should put + the Inductor/Triton caches on persistent storage so later workers reuse the + first worker's autotuning artifacts. """ value = os.getenv("JOYOMNI_VAE_COMPILE", "1").strip().lower() return value not in {"0", "false", "no", "off"} +def strict_enabled() -> bool: + """Fail startup when an explicitly required compiled path cannot warm.""" + value = os.getenv("JOYOMNI_VAE_COMPILE_STRICT", "0").strip().lower() + return value in {"1", "true", "yes", "on"} + + # Must run before any compiled function executes, so warm restarts reuse the # on-disk autotune results instead of re-running coordinate descent. if compile_enabled(): @@ -28,6 +34,47 @@ def compile_enabled() -> bool: _configured_encode: set[int] = set() _configured_encode_dynamic: set[int] = set() _skip_notices: set[str] = set() +_compile_failures: list[str] = [] + + +def _warmup_failed(stage: str, exc: BaseException) -> None: + message = f"{stage} failed: {exc!r}" + _compile_failures.append(message) + print(f"[vae_compile] {message}", flush=True) + if strict_enabled(): + raise RuntimeError(message) from exc + + +def runtime_status() -> dict[str, object]: + """Return a JSON-safe snapshot used by health checks and launch logs.""" + ready = ( + compile_enabled() + and len(_configured_encode) >= 2 + and len(_configured) >= 1 + and not _compile_failures + ) + return { + "enabled": compile_enabled(), + "strict": strict_enabled(), + "ready": ready, + "encode_instances": len(_configured_encode), + "decode_instances": len(_configured), + "dynamic_encode_instances": len(_configured_encode_dynamic), + "failures": list(_compile_failures), + "cache": { + "torchinductor": os.getenv("TORCHINDUCTOR_CACHE_DIR"), + "triton": os.getenv("TRITON_CACHE_DIR"), + "cuda": os.getenv("CUDA_CACHE_PATH"), + }, + } + + +def assert_runtime_ready() -> None: + """Reject the slow eager path when the deployment requires compilation.""" + status = runtime_status() + if status["ready"]: + return + raise RuntimeError(f"compiled VAE did not become ready: {status}") def _skip_compile(stage: str) -> bool: @@ -113,7 +160,7 @@ def warmup_encode(vae, in_channels: int, h_px: int, w_px: int, torch.cuda.synchronize(device) print(f"[vae_compile] warmup compiled encode shape (1,{in_channels},{t},{h_px},{w_px}) autocast={autocast}") except Exception as exc: # noqa: BLE001 - print(f"[vae_compile] encode warmup failed for t={t}: {exc!r}") + _warmup_failed(f"encode warmup t={t}", exc) def maybe_setup_encode_dynamic(vae) -> None: @@ -172,7 +219,7 @@ def warmup_encode_dynamic(vae, in_channels: int, hw_list, device: torch.device, torch.cuda.synchronize(device) n_ok += 1 except Exception as exc: # noqa: BLE001 - print(f"[vae_compile] dynamic encode warmup failed for ({h_px},{w_px},t={t}): {exc!r}") + _warmup_failed(f"dynamic encode warmup ({h_px},{w_px},t={t})", exc) print(f"[vae_compile] dynamic encode warmup done: {n_ok}/{len(hw_list) * len(temporal_lens)} shapes autocast={autocast}") @@ -200,4 +247,4 @@ def warmup_decode(vae, latent_channels: int, h_lat: int, w_lat: int, torch.cuda.synchronize(device) print(f"[vae_compile] warmup compiled decode shape (1,{latent_channels},{t},{h_lat},{w_lat}) autocast={autocast}") except Exception as exc: # noqa: BLE001 - print(f"[vae_compile] warmup failed for t={t}: {exc!r}") + _warmup_failed(f"decode warmup t={t}", exc) diff --git a/deploy/xvideo/serving/joyomni_streaming.py b/deploy/xvideo/serving/joyomni_streaming.py index d514521..845bdc1 100644 --- a/deploy/xvideo/serving/joyomni_streaming.py +++ b/deploy/xvideo/serving/joyomni_streaming.py @@ -41,6 +41,13 @@ def _autocast_ctx(device_type: str, dtype: torch.dtype, enabled: bool): _GRAPH_CACHE_CAP = 1 _GRAPH_CAPTURE_MAX_FAILS = 2 + +def _env_on(name: str, default: bool = False) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + def _vae_compile_module(): from xvideo.models.vae import vae_compile as module return module @@ -179,6 +186,20 @@ def __init__( self.graph_runners: dict[tuple, StreamingGraphRunner] = {} self.graph_capture_failures: dict[tuple, int] = {} + def optimization_status(self) -> dict[str, Any]: + runners = list(self.graph_runners.values()) + return { + "vae_compile": _vae_compile_module().runtime_status(), + "cuda_graph": { + "enabled": graph_env_enabled(), + "ready": any(getattr(runner, "ready", False) for runner in runners), + "ready_runners": sum( + bool(getattr(runner, "ready", False)) for runner in runners + ), + "capture_failures": sum(self.graph_capture_failures.values()), + }, + } + @classmethod def load( cls, @@ -251,7 +272,10 @@ def load( pipeline.transformer.eval() _orientations = [(warmup_height, warmup_width)] - if (warmup_width, warmup_height) != (warmup_height, warmup_width): + if ( + _env_on("JOYOMNI_WARMUP_BOTH_ORIENTATIONS", True) + and (warmup_width, warmup_height) != (warmup_height, warmup_width) + ): _orientations.append((warmup_width, warmup_height)) _stem_mod = pipeline.vae.stem @@ -272,6 +296,8 @@ def load( ) except Exception as _vc_exc: print(f"#####[STREAM] VAE compile warmup skipped: {_vc_exc!r}") + if _env_on("JOYOMNI_VAE_COMPILE_STRICT"): + raise RuntimeError("required VAE decode compilation failed") from _vc_exc try: _vc = _vae_compile_module() @@ -293,20 +319,28 @@ def load( autocast=_vae_ac, ) - _ref_basesize = getattr(cfg, "ref_image_basesize", DEFAULT_REFERENCE_IMG_IV2V_BASESIZE) - _ref_cfgs = generate_video_image_bucket( - img_basesize=_ref_basesize, bs_img=1, bs_vid=0, bs_mimg=0, bs_mvid=0, - ) - _ref_hw = sorted({(c[3], c[4]) for c in _ref_cfgs}) - _vc.maybe_setup_encode_dynamic(_src_vae) - _vc.warmup_encode_dynamic( - _src_vae, 3, _ref_hw, - device=_module_device(_src_vae), dtype=_vae_dt, - temporal_lens=(1,), - autocast=False, - ) + if _env_on("JOYOMNI_WARMUP_REFERENCE_BUCKETS", True): + _ref_basesize = getattr(cfg, "ref_image_basesize", DEFAULT_REFERENCE_IMG_IV2V_BASESIZE) + _ref_cfgs = generate_video_image_bucket( + img_basesize=_ref_basesize, bs_img=1, bs_vid=0, bs_mimg=0, bs_mvid=0, + ) + _ref_hw = sorted({(c[3], c[4]) for c in _ref_cfgs}) + _vc.maybe_setup_encode_dynamic(_src_vae) + _vc.warmup_encode_dynamic( + _src_vae, 3, _ref_hw, + device=_module_device(_src_vae), dtype=_vae_dt, + temporal_lens=(1,), + autocast=False, + ) + else: + print("#####[STREAM] reference-bucket VAE warmup disabled; reference shapes compile lazily") except Exception as _ve_exc: print(f"#####[STREAM] VAE encode compile warmup skipped: {_ve_exc!r}") + if _env_on("JOYOMNI_VAE_COMPILE_STRICT"): + raise RuntimeError("required VAE encode compilation failed") from _ve_exc + + if _env_on("JOYOMNI_VAE_COMPILE_STRICT"): + _vae_compile_module().assert_runtime_ready() runtime = cls( cfg=cfg, @@ -326,6 +360,13 @@ def load( runtime.warmup_full_pipeline(height=_wh, width=_ww) except Exception as _wexc: print(f"#####[STREAM] full-pipeline warmup error (non-fatal): {_wexc!r}") + if _env_on("JOYOMNI_LOAD_WARMUP_STRICT"): + raise RuntimeError("required full-pipeline warmup failed") from _wexc + + if _env_on("JOYOMNI_LOAD_WARMUP_STRICT") and graph_env_enabled(): + ready_graphs = [runner for runner in runtime.graph_runners.values() if getattr(runner, "ready", False)] + if not ready_graphs: + raise RuntimeError("required CUDA graph was not captured during full-pipeline warmup") return runtime @@ -395,9 +436,15 @@ def warmup_full_pipeline( r = session.wait_async_result(timeout=0.5) if r is not None: completed += 1 + if completed < num_chunks and _env_on("JOYOMNI_LOAD_WARMUP_STRICT"): + raise RuntimeError( + f"full-pipeline warmup returned only {completed}/{num_chunks} chunks" + ) print(f"#####[STREAM] full-pipeline warmup done: {completed} chunks in {time.time() - t0:.1f}s") except Exception as exc: print(f"#####[STREAM] full-pipeline warmup skipped/failed: {exc!r}") + if _env_on("JOYOMNI_LOAD_WARMUP_STRICT"): + raise finally: if session is not None: try: diff --git a/deploy/xvideo/serving/serve_joyomni_streaming.py b/deploy/xvideo/serving/serve_joyomni_streaming.py index 976a547..00c51ca 100644 --- a/deploy/xvideo/serving/serve_joyomni_streaming.py +++ b/deploy/xvideo/serving/serve_joyomni_streaming.py @@ -603,6 +603,25 @@ def get_runtime() -> JoyOmniRuntime: warmup_height=args.height, warmup_width=args.width, ) + optimization_status = app.state.runtime.optimization_status() + print( + f"#####[OPTIMIZATIONS] {json.dumps(optimization_status, sort_keys=True)}", + flush=True, + ) + marker_path = os.getenv("JOYOMNI_CACHE_READY_MARKER") + if marker_path: + marker = Path(marker_path) + marker.parent.mkdir(parents=True, exist_ok=True) + marker.write_text( + json.dumps( + { + "ready_at": time.time(), + "optimizations": optimization_status, + }, + indent=2, + sort_keys=True, + ) + ) return app.state.runtime @asynccontextmanager @@ -683,6 +702,11 @@ def health() -> JSONResponse: "max_temporal_ids": args.max_temporal_ids, "freeze_kv_on_static": args.freeze_kv_on_static, "static_diff_thresh": args.static_diff_thresh, + "optimizations": ( + app.state.runtime.optimization_status() + if app.state.runtime is not None + else None + ), } ) diff --git a/runpod/README.md b/runpod/README.md index 2dcff9e..f154489 100644 --- a/runpod/README.md +++ b/runpod/README.md @@ -61,7 +61,7 @@ The script securely asks for the API key and polls `GET /health` until the preloaded runtime is ready. RunPod's load balancer may return HTTP `400` with `timed out waiting for worker` after its own two-minute wait while the worker is still initializing; the script treats that response as retryable up to its -five-minute safety limit. It then starts the local HTTP/WebSocket proxy and +15-minute safety limit. It then starts the local HTTP/WebSocket proxy and opens the interface. The key stays in the local process and is removed when the script exits. @@ -69,7 +69,7 @@ the script exits. with `JOYOMNI_PRELOAD=1`, and RunPod cannot route that request until the worker has passed readiness anyway. -The five-minute limit bounds how long the local test waits; it is not a hard +The 15-minute limit bounds how long the local test waits; it is not a hard RunPod billing cutoff. RunPod bills from worker start until the worker fully stops. If the limit expires, stop the current worker (or temporarily set max workers to zero) before investigating the log. Do not start a second worker or @@ -82,5 +82,14 @@ finalization are disabled, the presence gate is off, and the browser starts at proves stable. Refreshing the browser replaces the previous WebSocket session instead of waiting behind its stale session ticket. +The H200 image enables the upstream compiled/autotuned VAE path. TorchInductor, +Triton, and CUDA caches are stored under +`/runpod-volume/joyai/cache/h200-torch291-cu128`, so the expensive first compile +is reused by later workers with the same pinned image stack. The health payload +reports `optimizations.vae_compile` and `optimizations.cuda_graph`; the strict +H200 defaults fail startup rather than silently serving the slow eager path. +For the live-only profile, startup precompiles the 840×480 landscape stream and +does not spend GPU time warming unused portrait or reference-image shapes. + Stop the test with `Ctrl+C`. With zero active workers and a short idle timeout, RunPod can scale the worker back to zero after the connection closes. diff --git a/runpod/Start-JoyAI-Realtime-Test.ps1 b/runpod/Start-JoyAI-Realtime-Test.ps1 index 500e2f7..84a671c 100644 --- a/runpod/Start-JoyAI-Realtime-Test.ps1 +++ b/runpod/Start-JoyAI-Realtime-Test.ps1 @@ -1,7 +1,7 @@ param( [string]$EndpointId = "ex9647vtulowka", [int]$LocalPort = 9000, - [int]$WarmTimeoutSeconds = 300, + [int]$WarmTimeoutSeconds = 900, [int]$RetryDelaySeconds = 10 ) @@ -50,6 +50,7 @@ try { Write-Host "Starting the JoyAI worker through GET /health..." Write-Host "The H200 must load the 32.5 GB DiT checkpoint before RunPod routes traffic." + Write-Host "The first compiled-VAE start can take longer while it creates the persistent optimization cache." Write-Host "A RunPod 400 'timed out waiting for worker' response is retried until the $WarmTimeoutSeconds-second safety limit." $warmTimer = [System.Diagnostics.Stopwatch]::StartNew() @@ -114,6 +115,14 @@ try { } Write-Host "JoyAI runtime is ready after $([int]$warmTimer.Elapsed.TotalSeconds) seconds." + if ( + $null -eq $health.optimizations -or + -not $health.optimizations.vae_compile.ready -or + -not $health.optimizations.cuda_graph.ready + ) { + throw "JoyAI became reachable without the required compiled VAE and CUDA graph. Stop the worker and inspect its optimization log before testing." + } + Write-Host "Verified: compiled VAE and CUDA graph are active." $env:RUNPOD_API_KEY = $apiKey $proxyScript = Join-Path $PSScriptRoot "local_proxy.py" diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index 83435b3..c223a08 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -54,6 +54,11 @@ def test_windows_script_prepares_proxy_dependency_before_warming_worker(self): def test_windows_script_retries_runpod_cold_start_timeout(self): text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + self.assertIn("[int]$WarmTimeoutSeconds = 900", text) + self.assertIn("first compiled-VAE start", text) + self.assertIn("compiled VAE and CUDA graph are active", text) + self.assertIn("$health.optimizations.vae_compile.ready", text) + self.assertIn("$health.optimizations.cuda_graph.ready", text) self.assertIn("while (-not $workerReady", text) self.assertIn("timed out waiting for worker", text) self.assertIn("operation has timed out", text) @@ -98,6 +103,19 @@ def test_h200_live_mode_disables_recording_and_uses_low_bandwidth_defaults(self) self.assertIn("JOYOMNI_RECORD_ENABLED=0", dockerfile) self.assertIn("JOYOMNI_ONLINE_GATE_ENABLED=0", dockerfile) self.assertIn("JOYOMNI_FPS=20", dockerfile) + self.assertIn("JOYOMNI_VAE_COMPILE=1", dockerfile) + self.assertIn("JOYOMNI_VAE_COMPILE_STRICT=1", dockerfile) + self.assertIn("JOYOMNI_LOAD_WARMUP_STRICT=1", dockerfile) + self.assertIn("JOYOMNI_WARMUP_BOTH_ORIENTATIONS=0", dockerfile) + self.assertIn("JOYOMNI_WARMUP_REFERENCE_BUCKETS=0", dockerfile) + self.assertIn( + "JOYOMNI_CACHE_ROOT=/runpod-volume/joyai/cache/h200-torch291-cu128", + dockerfile, + ) + self.assertIn('CACHE_ROOT="${JOYOMNI_CACHE_ROOT:-$HERE/deps/cache}"', launcher) + self.assertIn('$CACHE_ROOT/torchinductor', launcher) + self.assertIn('$CACHE_ROOT/triton', launcher) + self.assertIn('$CACHE_ROOT/nv_compute', launcher) self.assertIn('EXTRA_ARGS+=(--record-dir "$RECORD_DIR")', launcher) self.assertNotIn(' --record-dir "$RECORD_DIR" \\\n', launcher) self.assertNotIn('id="downloadBubble"', html) @@ -116,6 +134,25 @@ def test_h200_live_mode_disables_recording_and_uses_low_bandwidth_defaults(self) self.assertIn("const UPLINK_KEYFRAME_INTERVAL = 20;", html) self.assertIn('autoQuality = true; autoQTier = 0;', html) + def test_health_reports_required_runtime_optimizations(self): + server = ( + ROOT / "deploy" / "xvideo" / "serving" / "serve_joyomni_streaming.py" + ).read_text() + runtime = ( + ROOT / "deploy" / "xvideo" / "serving" / "joyomni_streaming.py" + ).read_text() + vae_compile = ( + ROOT / "deploy" / "xvideo" / "models" / "vae" / "vae_compile.py" + ).read_text() + + self.assertIn('"optimizations": (', server) + self.assertIn("app.state.runtime.optimization_status()", server) + self.assertIn("JOYOMNI_CACHE_READY_MARKER", server) + self.assertIn('"vae_compile": _vae_compile_module().runtime_status()', runtime) + self.assertIn('"cuda_graph": {', runtime) + self.assertIn("assert_runtime_ready", vae_compile) + self.assertIn("JOYOMNI_VAE_COMPILE_STRICT", vae_compile) + def test_browser_codec_probes_cannot_lock_the_send_button(self): html = (ROOT / "deploy" / "static" / "index.html").read_text() self.assertIn("const CODEC_PROBE_TIMEOUT_MS = 1200;", html) From 168b47ba28bc71bd5708e2cd915d1c0f286d9cc0 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 23:13:59 +0100 Subject: [PATCH 26/32] Allow compiled pipeline warmup to finish --- Dockerfile.h200 | 1 + deploy/xvideo/serving/joyomni_streaming.py | 24 ++++++++++++++++++---- runpod/README.md | 7 +++++-- runpod/Start-JoyAI-Realtime-Test.ps1 | 4 ++-- tests/test_runpod_connection.py | 6 +++++- 5 files changed, 33 insertions(+), 9 deletions(-) diff --git a/Dockerfile.h200 b/Dockerfile.h200 index 19f23d7..81b09f3 100644 --- a/Dockerfile.h200 +++ b/Dockerfile.h200 @@ -108,6 +108,7 @@ ENV JOYOMNI_DEVICE=cuda:0 \ JOYOMNI_VAE_COMPILE=1 \ JOYOMNI_VAE_COMPILE_STRICT=1 \ JOYOMNI_LOAD_WARMUP_STRICT=1 \ + JOYOMNI_FULL_WARMUP_TIMEOUT_SECONDS=300 \ JOYOMNI_WARMUP_BOTH_ORIENTATIONS=0 \ JOYOMNI_WARMUP_REFERENCE_BUCKETS=0 \ JOYOMNI_RECORD_ENABLED=0 \ diff --git a/deploy/xvideo/serving/joyomni_streaming.py b/deploy/xvideo/serving/joyomni_streaming.py index 845bdc1..b8d4aba 100644 --- a/deploy/xvideo/serving/joyomni_streaming.py +++ b/deploy/xvideo/serving/joyomni_streaming.py @@ -38,6 +38,7 @@ def _autocast_ctx(device_type: str, dtype: torch.dtype, enabled: bool): _FULL_WARMUP_CHUNKS = 4 +_FULL_WARMUP_TIMEOUT_SECONDS = 300.0 _GRAPH_CACHE_CAP = 1 _GRAPH_CAPTURE_MAX_FAILS = 2 @@ -359,8 +360,10 @@ def load( for (_wh, _ww) in _orientations: runtime.warmup_full_pipeline(height=_wh, width=_ww) except Exception as _wexc: - print(f"#####[STREAM] full-pipeline warmup error (non-fatal): {_wexc!r}") - if _env_on("JOYOMNI_LOAD_WARMUP_STRICT"): + _warmup_strict = _env_on("JOYOMNI_LOAD_WARMUP_STRICT") + _severity = "fatal" if _warmup_strict else "non-fatal" + print(f"#####[STREAM] full-pipeline warmup error ({_severity}): {_wexc!r}") + if _warmup_strict: raise RuntimeError("required full-pipeline warmup failed") from _wexc if _env_on("JOYOMNI_LOAD_WARMUP_STRICT") and graph_env_enabled(): @@ -431,8 +434,21 @@ def warmup_full_pipeline( if completed >= num_chunks: break - deadline = time.time() + 120.0 - while completed < num_chunks and time.time() < deadline: + timeout_seconds = max( + 30.0, + float( + os.environ.get( + "JOYOMNI_FULL_WARMUP_TIMEOUT_SECONDS", + str(_FULL_WARMUP_TIMEOUT_SECONDS), + ) + ), + ) + print( + f"#####[STREAM] waiting up to {timeout_seconds:.0f}s for " + f"{num_chunks - completed} warmup chunk(s)" + ) + deadline = time.monotonic() + timeout_seconds + while completed < num_chunks and time.monotonic() < deadline: r = session.wait_async_result(timeout=0.5) if r is not None: completed += 1 diff --git a/runpod/README.md b/runpod/README.md index f154489..88a3b86 100644 --- a/runpod/README.md +++ b/runpod/README.md @@ -61,7 +61,7 @@ The script securely asks for the API key and polls `GET /health` until the preloaded runtime is ready. RunPod's load balancer may return HTTP `400` with `timed out waiting for worker` after its own two-minute wait while the worker is still initializing; the script treats that response as retryable up to its -15-minute safety limit. It then starts the local HTTP/WebSocket proxy and +30-minute safety limit. It then starts the local HTTP/WebSocket proxy and opens the interface. The key stays in the local process and is removed when the script exits. @@ -69,7 +69,7 @@ the script exits. with `JOYOMNI_PRELOAD=1`, and RunPod cannot route that request until the worker has passed readiness anyway. -The 15-minute limit bounds how long the local test waits; it is not a hard +The 30-minute limit bounds how long the local test waits; it is not a hard RunPod billing cutoff. RunPod bills from worker start until the worker fully stops. If the limit expires, stop the current worker (or temporarily set max workers to zero) before investigating the log. Do not start a second worker or @@ -90,6 +90,9 @@ reports `optimizations.vae_compile` and `optimizations.cuda_graph`; the strict H200 defaults fail startup rather than silently serving the slow eager path. For the live-only profile, startup precompiles the 840×480 landscape stream and does not spend GPU time warming unused portrait or reference-image shapes. +After compilation, the server allows up to 300 seconds for its four asynchronous +full-pipeline warm-up chunks. This keeps strict readiness without rejecting a +healthy first compile when later chunks finish shortly after two minutes. Stop the test with `Ctrl+C`. With zero active workers and a short idle timeout, RunPod can scale the worker back to zero after the connection closes. diff --git a/runpod/Start-JoyAI-Realtime-Test.ps1 b/runpod/Start-JoyAI-Realtime-Test.ps1 index 84a671c..b1aee13 100644 --- a/runpod/Start-JoyAI-Realtime-Test.ps1 +++ b/runpod/Start-JoyAI-Realtime-Test.ps1 @@ -1,7 +1,7 @@ param( [string]$EndpointId = "ex9647vtulowka", [int]$LocalPort = 9000, - [int]$WarmTimeoutSeconds = 900, + [int]$WarmTimeoutSeconds = 1800, [int]$RetryDelaySeconds = 10 ) @@ -50,7 +50,7 @@ try { Write-Host "Starting the JoyAI worker through GET /health..." Write-Host "The H200 must load the 32.5 GB DiT checkpoint before RunPod routes traffic." - Write-Host "The first compiled-VAE start can take longer while it creates the persistent optimization cache." + Write-Host "The first compiled-VAE start can take up to 30 minutes while it creates the persistent optimization cache." Write-Host "A RunPod 400 'timed out waiting for worker' response is retried until the $WarmTimeoutSeconds-second safety limit." $warmTimer = [System.Diagnostics.Stopwatch]::StartNew() diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index c223a08..03527d8 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -54,7 +54,8 @@ def test_windows_script_prepares_proxy_dependency_before_warming_worker(self): def test_windows_script_retries_runpod_cold_start_timeout(self): text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() - self.assertIn("[int]$WarmTimeoutSeconds = 900", text) + self.assertIn("[int]$WarmTimeoutSeconds = 1800", text) + self.assertIn("up to 30 minutes", text) self.assertIn("first compiled-VAE start", text) self.assertIn("compiled VAE and CUDA graph are active", text) self.assertIn("$health.optimizations.vae_compile.ready", text) @@ -106,6 +107,7 @@ def test_h200_live_mode_disables_recording_and_uses_low_bandwidth_defaults(self) self.assertIn("JOYOMNI_VAE_COMPILE=1", dockerfile) self.assertIn("JOYOMNI_VAE_COMPILE_STRICT=1", dockerfile) self.assertIn("JOYOMNI_LOAD_WARMUP_STRICT=1", dockerfile) + self.assertIn("JOYOMNI_FULL_WARMUP_TIMEOUT_SECONDS=300", dockerfile) self.assertIn("JOYOMNI_WARMUP_BOTH_ORIENTATIONS=0", dockerfile) self.assertIn("JOYOMNI_WARMUP_REFERENCE_BUCKETS=0", dockerfile) self.assertIn( @@ -150,6 +152,8 @@ def test_health_reports_required_runtime_optimizations(self): self.assertIn("JOYOMNI_CACHE_READY_MARKER", server) self.assertIn('"vae_compile": _vae_compile_module().runtime_status()', runtime) self.assertIn('"cuda_graph": {', runtime) + self.assertIn('"JOYOMNI_FULL_WARMUP_TIMEOUT_SECONDS"', runtime) + self.assertIn("time.monotonic() + timeout_seconds", runtime) self.assertIn("assert_runtime_ready", vae_compile) self.assertIn("JOYOMNI_VAE_COMPILE_STRICT", vae_compile) From 2d0d27bd2d91b144fa3ef1b7672e7b4fabcf8683 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Mon, 17 Aug 2026 23:57:22 +0100 Subject: [PATCH 27/32] Serialize compiled VAE calls across streaming threads --- deploy/xvideo/models/pipeline.py | 3 +- deploy/xvideo/models/vae/vae_compile.py | 115 +++++++++++++-------- deploy/xvideo/serving/joyomni_streaming.py | 9 +- tests/test_runpod_connection.py | 6 ++ 4 files changed, 84 insertions(+), 49 deletions(-) diff --git a/deploy/xvideo/models/pipeline.py b/deploy/xvideo/models/pipeline.py index d9da262..3b3759f 100644 --- a/deploy/xvideo/models/pipeline.py +++ b/deploy/xvideo/models/pipeline.py @@ -305,7 +305,8 @@ def _encode_vae_single( from xvideo.models.vae import vae_compile as _vc inputs = _vc.prep_input(inputs) - latents = self.vae.encode(inputs).latent_dist.sample() + with _vc.call_guard(): + latents = self.vae.encode(inputs).latent_dist.sample() if enable_denormalization: latents = self.normalize_latents(latents) return latents.to(original_device) diff --git a/deploy/xvideo/models/vae/vae_compile.py b/deploy/xvideo/models/vae/vae_compile.py index ace4eba..8d89d83 100644 --- a/deploy/xvideo/models/vae/vae_compile.py +++ b/deploy/xvideo/models/vae/vae_compile.py @@ -1,6 +1,8 @@ from __future__ import annotations import os +import threading +from contextlib import contextmanager import torch import torch.nn as nn @@ -35,6 +37,25 @@ def strict_enabled() -> bool: _configured_encode_dynamic: set[int] = set() _skip_notices: set[str] = set() _compile_failures: list[str] = [] +_call_lock = threading.RLock() + + +@contextmanager +def call_guard(): + """Serialize entry into Dynamo/Inductor from streaming worker threads. + + PyTorch's compiler tracing state is not reliably isolated across Python + threads. The streaming pipeline deliberately runs encode, decode, and + pseudo-encode in parallel threads, so unguarded first calls can race while + a cached graph is being restored or specialized. The lock only covers + Python-side dispatch/kernel submission; CUDA streams may continue running + asynchronously after the guarded call returns. + """ + if not compile_enabled(): + yield + return + with _call_lock: + yield def _warmup_failed(stage: str, exc: BaseException) -> None: @@ -60,6 +81,7 @@ def runtime_status() -> dict[str, object]: "encode_instances": len(_configured_encode), "decode_instances": len(_configured), "dynamic_encode_instances": len(_configured_encode_dynamic), + "thread_call_guard": compile_enabled(), "failures": list(_compile_failures), "cache": { "torchinductor": os.getenv("TORCHINDUCTOR_CACHE_DIR"), @@ -89,22 +111,23 @@ def _skip_compile(stage: str) -> bool: def maybe_setup_decode(vae) -> None: if _skip_compile("decode compilation"): return - if id(vae) in _configured: - return - n_conv = 0 - for m in vae.modules(): - if isinstance(m, nn.Conv3d): - m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) - n_conv += 1 - if hasattr(vae, "_decode"): - vae._decode = torch.compile(vae._decode, mode="max-autotune-no-cudagraphs", dynamic=False) - target = "_decode" - elif hasattr(vae, "decode"): - vae.decode = torch.compile(vae.decode, mode="max-autotune-no-cudagraphs", dynamic=False) - target = "decode" - else: - raise RuntimeError("VAE has neither _decode nor decode; cannot compile") - _configured.add(id(vae)) + with call_guard(): + if id(vae) in _configured: + return + n_conv = 0 + for m in vae.modules(): + if isinstance(m, nn.Conv3d): + m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) + n_conv += 1 + if hasattr(vae, "_decode"): + vae._decode = torch.compile(vae._decode, mode="max-autotune-no-cudagraphs", dynamic=False) + target = "_decode" + elif hasattr(vae, "decode"): + vae.decode = torch.compile(vae.decode, mode="max-autotune-no-cudagraphs", dynamic=False) + target = "decode" + else: + raise RuntimeError("VAE has neither _decode nor decode; cannot compile") + _configured.add(id(vae)) print(f"[vae_compile] converted {n_conv} Conv3d weights to channels_last_3d + compiled vae.{target}") @@ -117,22 +140,23 @@ def prep_input(z: torch.Tensor) -> torch.Tensor: def maybe_setup_encode(vae) -> None: if _skip_compile("encode compilation"): return - if id(vae) in _configured_encode: - return - n_conv = 0 - for m in vae.modules(): - if isinstance(m, nn.Conv3d): - m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) - n_conv += 1 - if hasattr(vae, "_encode"): - vae._encode = torch.compile(vae._encode, mode="max-autotune-no-cudagraphs", dynamic=False) - target = "_encode" - elif hasattr(vae, "encode"): - vae.encode = torch.compile(vae.encode, mode="max-autotune-no-cudagraphs", dynamic=False) - target = "encode" - else: - raise RuntimeError("VAE has neither _encode nor encode; cannot compile") - _configured_encode.add(id(vae)) + with call_guard(): + if id(vae) in _configured_encode: + return + n_conv = 0 + for m in vae.modules(): + if isinstance(m, nn.Conv3d): + m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) + n_conv += 1 + if hasattr(vae, "_encode"): + vae._encode = torch.compile(vae._encode, mode="max-autotune-no-cudagraphs", dynamic=False) + target = "_encode" + elif hasattr(vae, "encode"): + vae.encode = torch.compile(vae.encode, mode="max-autotune-no-cudagraphs", dynamic=False) + target = "encode" + else: + raise RuntimeError("VAE has neither _encode nor encode; cannot compile") + _configured_encode.add(id(vae)) print(f"[vae_compile] converted {n_conv} Conv3d weights to channels_last_3d + compiled vae.{target} (encode)") @@ -154,7 +178,7 @@ def warmup_encode(vae, in_channels: int, h_px: int, w_px: int, if use_ac else nullcontext() ) try: - with torch.no_grad(), ctx: + with torch.no_grad(), ctx, call_guard(): _ = vae.encode(x) if torch.cuda.is_available(): torch.cuda.synchronize(device) @@ -166,16 +190,17 @@ def warmup_encode(vae, in_channels: int, h_px: int, w_px: int, def maybe_setup_encode_dynamic(vae) -> None: if _skip_compile("dynamic encode compilation"): return - if id(vae) in _configured_encode_dynamic: - return - if hasattr(vae, "_encode"): - core = getattr(vae, "_encode") - elif hasattr(vae, "encode"): - core = getattr(vae, "encode") - else: - raise RuntimeError("VAE has neither _encode nor encode; cannot compile") - vae._encode_dynamic = torch.compile(core, mode="max-autotune-no-cudagraphs", dynamic=True) - _configured_encode_dynamic.add(id(vae)) + with call_guard(): + if id(vae) in _configured_encode_dynamic: + return + if hasattr(vae, "_encode"): + core = getattr(vae, "_encode") + elif hasattr(vae, "encode"): + core = getattr(vae, "encode") + else: + raise RuntimeError("VAE has neither _encode nor encode; cannot compile") + vae._encode_dynamic = torch.compile(core, mode="max-autotune-no-cudagraphs", dynamic=True) + _configured_encode_dynamic.add(id(vae)) print("[vae_compile] compiled vae._encode_dynamic (dynamic=True, reference-image path)") @@ -213,7 +238,7 @@ def warmup_encode_dynamic(vae, in_channels: int, hw_list, device: torch.device, if use_ac else nullcontext() ) try: - with torch.no_grad(), ctx: + with torch.no_grad(), ctx, call_guard(): _ = fn(x) if torch.cuda.is_available(): torch.cuda.synchronize(device) @@ -241,7 +266,7 @@ def warmup_decode(vae, latent_channels: int, h_lat: int, w_lat: int, if use_ac else nullcontext() ) try: - with torch.no_grad(), ctx: + with torch.no_grad(), ctx, call_guard(): _ = vae.decode(z, return_dict=False)[0] if torch.cuda.is_available(): torch.cuda.synchronize(device) diff --git a/deploy/xvideo/serving/joyomni_streaming.py b/deploy/xvideo/serving/joyomni_streaming.py index b8d4aba..1870a3b 100644 --- a/deploy/xvideo/serving/joyomni_streaming.py +++ b/deploy/xvideo/serving/joyomni_streaming.py @@ -735,7 +735,8 @@ def _encode_ref_image_latent(self) -> torch.Tensor | None: ref_img_encoded = ref_img_tensor.to(device=encode_device, dtype=self.vae_dtype) _vc = _vae_compile_module() - encoded = _vc.encode_via_dynamic(self.pipeline.vae, ref_img_encoded) + with _vc.call_guard(): + encoded = _vc.encode_via_dynamic(self.pipeline.vae, ref_img_encoded) if not hasattr(encoded, "latent_dist"): raise TypeError(f"Unsupported VAE encode output type for ref image: {type(encoded)}") ref_img_latent = encoded.latent_dist.sample() @@ -1733,7 +1734,8 @@ def _decode_chunk_pixels( vae_ctx = _autocast_ctx(vae_device_type, self.vae_dtype, self.vae_autocast_enabled) with vae_ctx: started = self._timer_start(vae_device) - chunk_decoded = decode_vae.decode(decode_input, return_dict=False)[0] + with _vc.call_guard(): + chunk_decoded = decode_vae.decode(decode_input, return_dict=False)[0] if profile is not None: self._timer_record(profile, "vae_decode_s", started) @@ -1760,7 +1762,8 @@ def _encode_next_decode_pseudo_latent( prev_pixels = _vc.prep_input(prev_pixels) vae_ctx = _autocast_ctx(pseudo_device_type, self.vae_dtype, self.vae_autocast_enabled) with vae_ctx: - pseudo_enc = pseudo_vae.encode(prev_pixels) + with _vc.call_guard(): + pseudo_enc = pseudo_vae.encode(prev_pixels) if hasattr(pseudo_enc, "latent_dist"): pseudo_latent = pseudo_enc.latent_dist.sample() else: diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index 03527d8..7bfbade 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -156,6 +156,12 @@ def test_health_reports_required_runtime_optimizations(self): self.assertIn("time.monotonic() + timeout_seconds", runtime) self.assertIn("assert_runtime_ready", vae_compile) self.assertIn("JOYOMNI_VAE_COMPILE_STRICT", vae_compile) + self.assertIn("def call_guard():", vae_compile) + self.assertIn('"thread_call_guard": compile_enabled()', vae_compile) + self.assertIn("with _vc.call_guard():", runtime) + self.assertIn("with _vc.call_guard():", ( + ROOT / "deploy" / "xvideo" / "models" / "pipeline.py" + ).read_text()) def test_browser_codec_probes_cannot_lock_the_send_button(self): html = (ROOT / "deploy" / "static" / "index.html").read_text() From 317664d0b087784e0a018e22b0ab3b4adac6cf8e Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Tue, 18 Aug 2026 00:45:50 +0100 Subject: [PATCH 28/32] Recover live streams from websocket stalls --- deploy/static/index.html | 29 +++++++++++++++++ .../xvideo/serving/serve_joyomni_streaming.py | 16 +++++++++- runpod/local_proxy.py | 8 +++++ tests/test_runpod_connection.py | 31 +++++++++++++++++++ 4 files changed, 83 insertions(+), 1 deletion(-) diff --git a/deploy/static/index.html b/deploy/static/index.html index f81324a..dd5b18c 100644 --- a/deploy/static/index.html +++ b/deploy/static/index.html @@ -589,6 +589,11 @@ try { new ResizeObserver(() => updateCropMask()).observe(camera); } catch (e) {} } let ws = null; +let streamingWanted = false; +let reconnectTimer = null; +let reconnectAttempt = 0; +const RECONNECT_BASE_DELAY_MS = 750; +const RECONNECT_MAX_DELAY_MS = 5000; let stream = null; let videoFileUrl = null; let usingVideoFile = false; @@ -788,6 +793,25 @@ }, 1000); } +function cancelReconnect() { + if (reconnectTimer) clearTimeout(reconnectTimer); + reconnectTimer = null; +} + +function scheduleReconnect() { + if (!streamingWanted || reconnectTimer) return; + const delay = Math.min( + RECONNECT_MAX_DELAY_MS, + RECONNECT_BASE_DELAY_MS * Math.pow(2, Math.min(reconnectAttempt, 3)) + ); + reconnectAttempt += 1; + reconnectTimer = setTimeout(async () => { + reconnectTimer = null; + if (!streamingWanted || ws) return; + await start(); + }, delay); +} + function readRefImageDataUrl() { const el = document.getElementById("refImage"); const file = el && el.files && el.files[0]; @@ -1718,6 +1742,7 @@ return; } if (msg.type === "started") { + reconnectAttempt = 0; if (Number(msg.width) > 0 && Number(msg.height) > 0) { document.getElementById("width").value = msg.width; document.getElementById("height").value = msg.height; @@ -1959,6 +1984,7 @@ setSendBusy(false, t("busy_disconnected")); showOutputIdle(true); updateMetrics(); + scheduleReconnect(); }; socket.onerror = () => { if (ws === socket) { @@ -1969,6 +1995,8 @@ } function stopActiveRun() { + streamingWanted = false; + cancelReconnect(); if (timer) clearInterval(timer); timer = null; stopVideoPull(); @@ -2015,6 +2043,7 @@ sessionGranted = false; } stopActiveRun(); + streamingWanted = true; await start(); } diff --git a/deploy/xvideo/serving/serve_joyomni_streaming.py b/deploy/xvideo/serving/serve_joyomni_streaming.py index 00c51ca..ce326a4 100644 --- a/deploy/xvideo/serving/serve_joyomni_streaming.py +++ b/deploy/xvideo/serving/serve_joyomni_streaming.py @@ -73,6 +73,10 @@ def release(self, ticket: int) -> None: WS_SEND_TIMEOUT_S = 10.0 +HOLDER_IDLE_TIMEOUT_S = max( + 10.0, + float(os.environ.get("JOYOMNI_SESSION_IDLE_TIMEOUT_SECONDS", "60")), +) REF_IMAGE_DIR = REPO_ROOT / "rv2v_reference" REF_IMAGE_FILES = { @@ -1319,7 +1323,6 @@ async def _reset_session(reason: str) -> None: app.state.ws_debug = ws_debug - HOLDER_IDLE_TIMEOUT_S = 10.0 last_activity = time.monotonic() last_frames_out = frames_out while True: @@ -1327,6 +1330,11 @@ async def _reset_session(reason: str) -> None: last_frames_out = frames_out last_activity = time.monotonic() if time.monotonic() - last_activity >= HOLDER_IDLE_TIMEOUT_S: + print( + f"#####[WS-GUARD] releasing session after " + f"{HOLDER_IDLE_TIMEOUT_S:.0f}s without browser or output activity", + flush=True, + ) try: await _send_json( { @@ -1344,6 +1352,12 @@ async def _reset_session(reason: str) -> None: if "text" in message and message["text"] is not None: payload = json.loads(message["text"]) msg_type = payload.get("type") + # ACK and ping messages prove that the live browser and its + # camera loop are still reachable. The old 10-second guard + # ignored them, so a short frame/backpressure pause could + # release a healthy session and leave the last frame frozen. + last_activity = time.monotonic() + ws_debug["last_client_activity_at"] = time.time() if msg_type == "start": print(f"#####[RESTART] 'start' received (session {'live' if session is not None else 'none'})", flush=True) last_activity = time.monotonic() diff --git a/runpod/local_proxy.py b/runpod/local_proxy.py index 39a6b5a..77cbfef 100644 --- a/runpod/local_proxy.py +++ b/runpod/local_proxy.py @@ -70,6 +70,7 @@ async def proxy_websocket(request: web.Request) -> web.WebSocketResponse: heartbeat=30, max_msg_size=0, ) + print("JoyAI WebSocket connected through RunPod.", flush=True) # The local proxy is intentionally single-viewer. Replacing a stale # browser socket explicitly tells the server to release its session @@ -137,6 +138,13 @@ async def runpod_to_browser() -> None: } ) finally: + upstream_code = getattr(upstream, "close_code", None) + downstream_code = getattr(downstream, "close_code", None) + print( + "JoyAI WebSocket closed " + f"(browser={downstream_code}, RunPod={upstream_code}).", + flush=True, + ) if upstream is not None and not upstream.closed: try: await upstream.send_json({"type": "stop"}) diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index 7bfbade..004b446 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -170,6 +170,37 @@ def test_browser_codec_probes_cannot_lock_the_send_button(self): self.assertIn('if (!started) throw new Error("session start was not sent")', html) self.assertIn('setSendBusy(false, "");', html) + def test_live_session_pings_prevent_false_idle_disconnects(self): + server = ( + ROOT / "deploy" / "xvideo" / "serving" / "serve_joyomni_streaming.py" + ).read_text() + self.assertIn('JOYOMNI_SESSION_IDLE_TIMEOUT_SECONDS", "60"', server) + self.assertIn('ws_debug["last_client_activity_at"] = time.time()', server) + ping_handler = server.index('elif msg_type == "ping":') + client_activity = server.rindex( + 'last_activity = time.monotonic()', 0, ping_handler + ) + payload_decode = server.rindex( + 'payload = json.loads(message["text"])', 0, ping_handler + ) + self.assertGreater(client_activity, payload_decode) + self.assertIn("releasing session after", server) + + def test_browser_recovers_an_unexpected_websocket_close(self): + html = (ROOT / "deploy" / "static" / "index.html").read_text() + self.assertIn("let streamingWanted = false;", html) + self.assertIn("function scheduleReconnect()", html) + self.assertIn("RECONNECT_MAX_DELAY_MS", html) + self.assertIn("scheduleReconnect();", html) + self.assertIn("streamingWanted = false;\n cancelReconnect();", html) + self.assertIn("streamingWanted = true;\n await start();", html) + + def test_proxy_reports_websocket_close_codes(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn("JoyAI WebSocket connected through RunPod.", text) + self.assertIn("JoyAI WebSocket closed", text) + self.assertIn('f"(browser={downstream_code}, RunPod={upstream_code})."', text) + if __name__ == "__main__": unittest.main() From 8d868b455e3bb371145ab8ff56719447f1325835 Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Tue, 18 Aug 2026 06:55:42 +0100 Subject: [PATCH 29/32] Keep live streams pinned to one RunPod worker (#11) Use RunPod strict-resume worker affinity and authenticated HTTP keepalives so WebSocket-only sessions are not scaled down or rerouted to a second cold H200. --- runpod/README.md | 15 +++++ runpod/local_proxy.py | 107 +++++++++++++++++++++++++++++--- tests/test_runpod_connection.py | 21 +++++++ 3 files changed, 135 insertions(+), 8 deletions(-) diff --git a/runpod/README.md b/runpod/README.md index 88a3b86..7ec704f 100644 --- a/runpod/README.md +++ b/runpod/README.md @@ -42,6 +42,16 @@ Use these values for a load-balancer endpoint: | Exposed HTTP ports | `8080,8081` | | Health-check path | `/ping` | | Active/min workers | `0` while testing to avoid idle GPU charges | +| Max workers | `1` for one-viewer testing and cost safety | +| Idle timeout | `60` seconds | + +RunPod does not reliably treat a quiet WebSocket as autoscaling activity. The +local proxy therefore sends an authenticated `GET /health` every three seconds +while the stream is open and pins every follow-up request with +`X-Runpod-Worker-Id: strict-resume `. This keeps the selected worker +alive during a session without paying for an always-on active worker. Limiting +the endpoint to one worker also prevents a reconnect or page refresh from +starting a second H200 while the first one is still usable. The H200 image default checkpoint path is `/runpod-volume/joyai/checkpoints`. Mount the model network volume so that the @@ -82,6 +92,11 @@ finalization are disabled, the presence gate is off, and the browser starts at proves stable. Refreshing the browser replaces the previous WebSocket session instead of waiting behind its stale session ticket. +The proxy records the `X-Runpod-Worker-Id` returned by the initial local health +check and uses RunPod strict-resume affinity for the page, assets, and WebSocket. +This is required because the video session and loaded model state live inside +one worker; an unpinned reconnect can otherwise land on a different cold worker. + The H200 image enables the upstream compiled/autotuned VAE path. TorchInductor, Triton, and CUDA caches are stored under `/runpod-volume/joyai/cache/h200-torch291-cu128`, so the expensive first compile diff --git a/runpod/local_proxy.py b/runpod/local_proxy.py index 77cbfef..90ec06b 100644 --- a/runpod/local_proxy.py +++ b/runpod/local_proxy.py @@ -12,7 +12,15 @@ import json import os -from aiohttp import ClientError, ClientSession, ClientTimeout, WSCloseCode, WSMsgType, web +from aiohttp import ( + ClientError, + ClientSession, + ClientTimeout, + WSCloseCode, + WSMsgType, + WSServerHandshakeError, + web, +) HOP_BY_HOP_HEADERS = { @@ -28,6 +36,12 @@ "upgrade", } +RUNPOD_WORKER_HEADER = "X-Runpod-Worker-Id" +# RunPod's load balancer may not count an otherwise active WebSocket as worker +# activity. Keep this below the platform's 5-second default idle timeout so a +# live stream cannot be scaled down underneath the browser. +RUNPOD_KEEPALIVE_SECONDS = 3.0 + def proxy_exception_handler(loop: asyncio.AbstractEventLoop, context: dict) -> None: """Ignore the harmless Windows Proactor reset logged after a peer closes.""" @@ -56,6 +70,57 @@ def upstream_url(request: web.Request) -> str: return f"{request.app['upstream']}{request.rel_url}" +def upstream_headers(app: web.Application) -> dict[str, str]: + """Authenticate and pin follow-up traffic to the selected RunPod worker.""" + headers = dict(app["auth_headers"]) + worker_id = app["proxy_state"].get("worker_id") + if worker_id: + headers[RUNPOD_WORKER_HEADER] = f"strict-resume {worker_id}" + return headers + + +def remember_worker(app: web.Application, headers) -> str | None: + """Remember the worker chosen by RunPod for later HTTP and WS requests.""" + worker_id = str(headers.get(RUNPOD_WORKER_HEADER, "")).strip() + if not worker_id: + return None + state = app["proxy_state"] + if state.get("worker_id") != worker_id: + state["worker_id"] = worker_id + print(f"JoyAI pinned to RunPod worker {worker_id}.", flush=True) + return worker_id + + +async def keep_worker_active(app: web.Application) -> None: + """Reset RunPod's idle timer while the browser has a live WebSocket.""" + session: ClientSession = app["session"] + failures = 0 + while True: + await asyncio.sleep(RUNPOD_KEEPALIVE_SECONDS) + try: + async with session.get( + f"{app['upstream']}/health", + headers=upstream_headers(app), + timeout=ClientTimeout(total=10), + ) as response: + await response.read() + remember_worker(app, response.headers) + if response.status != 200: + raise ClientError( + f"RunPod keepalive returned HTTP {response.status}" + ) + failures = 0 + except asyncio.CancelledError: + raise + except (ClientError, asyncio.TimeoutError, OSError) as error: + failures += 1 + if failures == 1 or failures % 5 == 0: + print( + f"JoyAI RunPod keepalive failed ({failures}): {error}", + flush=True, + ) + + async def proxy_websocket(request: web.Request) -> web.WebSocketResponse: downstream = web.WebSocketResponse(heartbeat=30) await downstream.prepare(request) @@ -63,14 +128,35 @@ async def proxy_websocket(request: web.Request) -> web.WebSocketResponse: session: ClientSession = request.app["session"] proxy_state = request.app["proxy_state"] upstream = None + keepalive_task = None try: - upstream = await session.ws_connect( - upstream_url(request), - headers=request.app["auth_headers"], - heartbeat=30, - max_msg_size=0, - ) + # The local readiness check captures X-Runpod-Worker-Id before the + # browser opens this socket. Strict-resume prevents a reconnect from + # silently landing on a different, cold H200 worker. + try: + upstream = await session.ws_connect( + upstream_url(request), + headers=upstream_headers(request.app), + heartbeat=10, + max_msg_size=0, + ) + except WSServerHandshakeError as error: + if error.status != 404 or not proxy_state.get("worker_id"): + raise + # A redeploy gives the worker a new ID. Retry normal routing once; + # the successful handshake response becomes the new affinity. + proxy_state["worker_id"] = None + upstream = await session.ws_connect( + upstream_url(request), + headers=upstream_headers(request.app), + heartbeat=10, + max_msg_size=0, + ) + response = getattr(upstream, "_response", None) + if response is not None: + remember_worker(request.app, response.headers) print("JoyAI WebSocket connected through RunPod.", flush=True) + keepalive_task = asyncio.create_task(keep_worker_active(request.app)) # The local proxy is intentionally single-viewer. Replacing a stale # browser socket explicitly tells the server to release its session @@ -138,6 +224,9 @@ async def runpod_to_browser() -> None: } ) finally: + if keepalive_task is not None: + keepalive_task.cancel() + await asyncio.gather(keepalive_task, return_exceptions=True) upstream_code = getattr(upstream, "close_code", None) downstream_code = getattr(downstream, "close_code", None) print( @@ -169,7 +258,7 @@ async def proxy_http(request: web.Request) -> web.Response: # upstream load balancer for an identity response so both PowerShell's # readiness probe and the browser receive bytes they can consume directly. headers["Accept-Encoding"] = "identity" - headers.update(request.app["auth_headers"]) + headers.update(upstream_headers(request.app)) try: async with session.request( @@ -180,6 +269,7 @@ async def proxy_http(request: web.Request) -> web.Response: allow_redirects=False, ) as response: body = await response.read() + remember_worker(request.app, response.headers) return web.Response( body=body, status=response.status, @@ -222,6 +312,7 @@ async def create_session(app: web.Application) -> None: "websocket_lock": asyncio.Lock(), "active_upstream": None, "active_downstream": None, + "worker_id": None, } diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index 004b446..593b07b 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -201,6 +201,27 @@ def test_proxy_reports_websocket_close_codes(self): self.assertIn("JoyAI WebSocket closed", text) self.assertIn('f"(browser={downstream_code}, RunPod={upstream_code})."', text) + def test_proxy_pins_reconnects_to_one_runpod_worker(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn('RUNPOD_WORKER_HEADER = "X-Runpod-Worker-Id"', text) + self.assertIn('f"strict-resume {worker_id}"', text) + self.assertIn("remember_worker(request.app, response.headers)", text) + self.assertIn('"worker_id": None', text) + + def test_proxy_keeps_worker_active_during_websocket_stream(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn("RUNPOD_KEEPALIVE_SECONDS = 3.0", text) + self.assertIn("async def keep_worker_active", text) + self.assertIn('f"{app[\'upstream\']}/health"', text) + self.assertIn("asyncio.create_task(keep_worker_active(request.app))", text) + self.assertIn("keepalive_task.cancel()", text) + + def test_runpod_settings_limit_one_viewer_tests_to_one_worker(self): + text = (ROOT / "runpod" / "README.md").read_text() + self.assertIn("| Max workers | `1`", text) + self.assertIn("| Idle timeout | `60` seconds |", text) + self.assertIn("X-Runpod-Worker-Id: strict-resume", text) + if __name__ == "__main__": unittest.main() From 9beab8c8e2251f99d8eec8176932cce05b56c14a Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Tue, 18 Aug 2026 07:14:12 +0100 Subject: [PATCH 30/32] Unblock the RunPod WebSocket handshake (#12) Use normal routing for the WebSocket upgrade and stream keepalive while Max workers=1 prevents duplicate H200 workers. --- runpod/README.md | 16 +++++----- runpod/local_proxy.py | 53 ++++++++++++++++----------------- tests/test_runpod_connection.py | 13 ++++++-- 3 files changed, 44 insertions(+), 38 deletions(-) diff --git a/runpod/README.md b/runpod/README.md index 7ec704f..2040d50 100644 --- a/runpod/README.md +++ b/runpod/README.md @@ -47,11 +47,11 @@ Use these values for a load-balancer endpoint: RunPod does not reliably treat a quiet WebSocket as autoscaling activity. The local proxy therefore sends an authenticated `GET /health` every three seconds -while the stream is open and pins every follow-up request with -`X-Runpod-Worker-Id: strict-resume `. This keeps the selected worker -alive during a session without paying for an always-on active worker. Limiting -the endpoint to one worker also prevents a reconnect or page refresh from -starting a second H200 while the first one is still usable. +while the stream is open. Ordinary HTTP follow-up requests use soft +`X-Runpod-Worker-Id` affinity, while the WebSocket upgrade uses normal routing: +strict affinity can be held outside the worker when RunPod considers a +long-lived connection at capacity. Limiting the endpoint to one worker is the +hard guarantee that normal routing cannot start a second H200. The H200 image default checkpoint path is `/runpod-volume/joyai/checkpoints`. Mount the model network volume so that the @@ -93,9 +93,9 @@ proves stable. Refreshing the browser replaces the previous WebSocket session instead of waiting behind its stale session ticket. The proxy records the `X-Runpod-Worker-Id` returned by the initial local health -check and uses RunPod strict-resume affinity for the page, assets, and WebSocket. -This is required because the video session and loaded model state live inside -one worker; an unpinned reconnect can otherwise land on a different cold worker. +check and uses it as a soft preference for page and asset requests. Keep +`Max workers` at `1` during one-viewer testing so the unblocked WebSocket route +still reaches that same worker and cannot launch a different cold worker. The H200 image enables the upstream compiled/autotuned VAE path. TorchInductor, Triton, and CUDA caches are stored under diff --git a/runpod/local_proxy.py b/runpod/local_proxy.py index 90ec06b..c5c94fb 100644 --- a/runpod/local_proxy.py +++ b/runpod/local_proxy.py @@ -18,7 +18,6 @@ ClientTimeout, WSCloseCode, WSMsgType, - WSServerHandshakeError, web, ) @@ -70,12 +69,19 @@ def upstream_url(request: web.Request) -> str: return f"{request.app['upstream']}{request.rel_url}" -def upstream_headers(app: web.Application) -> dict[str, str]: - """Authenticate and pin follow-up traffic to the selected RunPod worker.""" +def upstream_headers( + app: web.Application, + *, + affinity: bool = True, +) -> dict[str, str]: + """Authenticate requests and optionally prefer the selected worker.""" headers = dict(app["auth_headers"]) worker_id = app["proxy_state"].get("worker_id") - if worker_id: - headers[RUNPOD_WORKER_HEADER] = f"strict-resume {worker_id}" + if affinity and worker_id: + # Soft affinity never waits behind a worker RunPod considers at + # capacity. Max workers=1 remains the hard guarantee that a fallback + # cannot create a second H200 during one-viewer testing. + headers[RUNPOD_WORKER_HEADER] = worker_id return headers @@ -100,7 +106,10 @@ async def keep_worker_active(app: web.Application) -> None: try: async with session.get( f"{app['upstream']}/health", - headers=upstream_headers(app), + # RunPod can hold an affinity request while a long-lived WS is + # occupying the selected worker. With max workers=1, normal + # routing still reaches the only worker without that wait. + headers=upstream_headers(app, affinity=False), timeout=ClientTimeout(total=10), ) as response: await response.read() @@ -130,28 +139,16 @@ async def proxy_websocket(request: web.Request) -> web.WebSocketResponse: upstream = None keepalive_task = None try: - # The local readiness check captures X-Runpod-Worker-Id before the - # browser opens this socket. Strict-resume prevents a reconnect from - # silently landing on a different, cold H200 worker. - try: - upstream = await session.ws_connect( - upstream_url(request), - headers=upstream_headers(request.app), - heartbeat=10, - max_msg_size=0, - ) - except WSServerHandshakeError as error: - if error.status != 404 or not proxy_state.get("worker_id"): - raise - # A redeploy gives the worker a new ID. Retry normal routing once; - # the successful handshake response becomes the new affinity. - proxy_state["worker_id"] = None - upstream = await session.ws_connect( - upstream_url(request), - headers=upstream_headers(request.app), - heartbeat=10, - max_msg_size=0, - ) + # Do not put strict/soft affinity on the WebSocket upgrade. RunPod can + # hold an affinity upgrade outside the worker for several minutes even + # when ordinary HTTP requests reach it immediately. Max workers=1 pins + # normal routing to the only worker without blocking the handshake. + upstream = await session.ws_connect( + upstream_url(request), + headers=upstream_headers(request.app, affinity=False), + heartbeat=10, + max_msg_size=0, + ) response = getattr(upstream, "_response", None) if response is not None: remember_worker(request.app, response.headers) diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index 593b07b..b6e2ae2 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -204,10 +204,19 @@ def test_proxy_reports_websocket_close_codes(self): def test_proxy_pins_reconnects_to_one_runpod_worker(self): text = (ROOT / "runpod" / "local_proxy.py").read_text() self.assertIn('RUNPOD_WORKER_HEADER = "X-Runpod-Worker-Id"', text) - self.assertIn('f"strict-resume {worker_id}"', text) + self.assertIn("headers[RUNPOD_WORKER_HEADER] = worker_id", text) self.assertIn("remember_worker(request.app, response.headers)", text) self.assertIn('"worker_id": None', text) + def test_websocket_and_keepalive_do_not_wait_on_worker_affinity(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertGreaterEqual( + text.count("upstream_headers(request.app, affinity=False)"), + 1, + ) + self.assertIn("upstream_headers(app, affinity=False)", text) + self.assertNotIn('f"strict-resume {worker_id}"', text) + def test_proxy_keeps_worker_active_during_websocket_stream(self): text = (ROOT / "runpod" / "local_proxy.py").read_text() self.assertIn("RUNPOD_KEEPALIVE_SECONDS = 3.0", text) @@ -220,7 +229,7 @@ def test_runpod_settings_limit_one_viewer_tests_to_one_worker(self): text = (ROOT / "runpod" / "README.md").read_text() self.assertIn("| Max workers | `1`", text) self.assertIn("| Idle timeout | `60` seconds |", text) - self.assertIn("X-Runpod-Worker-Id: strict-resume", text) + self.assertIn("soft", text) if __name__ == "__main__": From eb179bb58a2b16fbbd0ce844472e5682806c312b Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Tue, 18 Aug 2026 07:46:42 +0100 Subject: [PATCH 31/32] Stabilize reference-person switching over WebSocket (#13) Use application-level liveness without an aggressive protocol heartbeat and add an explicit reference-person replacement preset. --- deploy/static/index.html | 3 ++- runpod/local_proxy.py | 6 +++++- tests/test_runpod_connection.py | 17 +++++++++++++++++ 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/deploy/static/index.html b/deploy/static/index.html index dd5b18c..75f3391 100644 --- a/deploy/static/index.html +++ b/deploy/static/index.html @@ -2063,8 +2063,9 @@ }, { key: "ref_tryon", - label: "参考图换装", label_en: "Try-On from Reference", + label: "参考图编辑", label_en: "Reference Image", cases: [ + { title: "我的参考人物", title_en: "My Reference Person", desc: "用上传图片中的人物替换主角", desc_en: "Replace the subject with your uploaded person", text: "将视频中的主角替换为上传的参考图中的人物。保持参考人物的身份、面部、发型和服装，同时保留原视频的姿势、动作、镜头和背景。", text_en: "Replace the main subject in the video with the person shown in the uploaded reference image. Preserve the reference person's identity, face, hairstyle, and clothing while keeping the source pose, motion, camera, and background unchanged." }, { title: "NY 棒球帽", title_en: "NY Cap", desc: "戴上参考图里的黑色 NY 棒球帽", desc_en: "Black NY cap from the reference image", ref: "hat", text: "将参考图中的黑色 NY 棒球帽戴到主角头上，帽子正面朝前。", text_en: "Put the black NY baseball cap from the reference image on the subject's head, facing forward." }, { title: "红色围巾", title_en: "Red Scarf", desc: "围上参考图里的红色流苏围巾", desc_en: "Red tasseled scarf from the reference image", ref: "scarf", text: "将参考图中的红色流苏围巾围到主角的颈部。", text_en: "Wrap the red tasseled scarf from the reference image around the subject's neck." }, { title: "粉色T恤", title_en: "Pink Tee", desc: "换上参考图里的粉色短袖 T 恤", desc_en: "Pink short-sleeve tee from the reference image", ref: "pink_tee", text: "将主角的衣服换成参考图里的粉色短袖 T 恤。", text_en: "Replace the subject's clothing with the pink short-sleeve T-shirt from the reference image." }, diff --git a/runpod/local_proxy.py b/runpod/local_proxy.py index c5c94fb..2e71b45 100644 --- a/runpod/local_proxy.py +++ b/runpod/local_proxy.py @@ -146,7 +146,11 @@ async def proxy_websocket(request: web.Request) -> web.WebSocketResponse: upstream = await session.ws_connect( upstream_url(request), headers=upstream_headers(request.app, affinity=False), - heartbeat=10, + # The browser already sends an application-level JSON ping every + # second. A protocol heartbeat adds a second liveness mechanism + # that can falsely close with code 1006 when RunPod's load + # balancer does not relay a Pong during graph/session rebuilds. + heartbeat=None, max_msg_size=0, ) response = getattr(upstream, "_response", None) diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index b6e2ae2..c49ff37 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -201,6 +201,23 @@ def test_proxy_reports_websocket_close_codes(self): self.assertIn("JoyAI WebSocket closed", text) self.assertIn('f"(browser={downstream_code}, RunPod={upstream_code})."', text) + def test_proxy_does_not_duplicate_browser_application_heartbeat(self): + proxy = (ROOT / "runpod" / "local_proxy.py").read_text() + html = (ROOT / "deploy" / "static" / "index.html").read_text() + self.assertIn("heartbeat=None", proxy) + self.assertNotIn("heartbeat=10", proxy) + self.assertIn('{ type: "ping", t: Date.now()', html) + + def test_reference_person_replacement_has_an_explicit_prompt(self): + html = (ROOT / "deploy" / "static" / "index.html").read_text() + self.assertIn('label_en: "Reference Image"', html) + self.assertIn('title_en: "My Reference Person"', html) + self.assertIn( + "Replace the main subject in the video with the person shown in the " + "uploaded reference image.", + html, + ) + def test_proxy_pins_reconnects_to_one_runpod_worker(self): text = (ROOT / "runpod" / "local_proxy.py").read_text() self.assertIn('RUNPOD_WORKER_HEADER = "X-Runpod-Worker-Id"', text) From 3a3e31d3b0eabd496549d8490a8b1c7b59a08dfb Mon Sep 17 00:00:00 2001 From: samuellucky2424-afk Date: Tue, 18 Aug 2026 09:10:24 +0100 Subject: [PATCH 32/32] Verify and pin the upgraded RV2V checkpoint (#14) Pin downloads to the official upgraded 0811 RV2V release, verify checkpoint metadata before model load, expose verification in health, and add an optional full-hash fallback. --- deploy/xvideo/checkpoint_status.py | 93 +++++++++++++++++++ .../xvideo/serving/serve_joyomni_streaming.py | 2 + runpod/Start-JoyAI-Realtime-Test.ps1 | 12 +++ runpod/download_models.py | 6 +- runpod/start.py | 38 +++++++- runpod/verify_checkpoint.py | 42 +++++++++ tests/test_runpod_connection.py | 70 ++++++++++++++ 7 files changed, 261 insertions(+), 2 deletions(-) create mode 100644 deploy/xvideo/checkpoint_status.py create mode 100644 runpod/verify_checkpoint.py diff --git a/deploy/xvideo/checkpoint_status.py b/deploy/xvideo/checkpoint_status.py new file mode 100644 index 0000000..a40d052 --- /dev/null +++ b/deploy/xvideo/checkpoint_status.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import Any + + +JOYAI_DIT_FILENAME = "dit/joyai_video_edit_dit_0811.pth" +JOYAI_DIT_RELEASE_COMMIT = "eda14f342ef99c52485bbb8dc271c29b42298089" +JOYAI_DIT_XET_HASH = "86a577acfe936e9b56ae7c89f04d2db61d5b69f97f6fa6496da8f7cfcc47305f" +JOYAI_DIT_SHA256 = "b3904b6fda53d13b230918bb616f322d12cfb2337b0e8d9dc203cdabc36605ba" + + +def _metadata_path(checkpoint_path: Path) -> Path: + """Return huggingface_hub's local_dir metadata path for the DiT file.""" + local_dir = checkpoint_path.parent.parent + return local_dir / ".cache" / "huggingface" / "download" / "dit" / ( + checkpoint_path.name + ".metadata" + ) + + +def _read_metadata(checkpoint_path: Path) -> tuple[str | None, str | None, str | None]: + metadata_path = _metadata_path(checkpoint_path) + if not metadata_path.is_file(): + return None, None, None + + try: + with metadata_path.open(encoding="utf-8") as metadata: + revision = metadata.readline().strip() or None + etag = metadata.readline().strip().strip('"') or None + metadata.readline() # timestamp; validate only by successful reading + except (OSError, UnicodeError) as error: + return None, None, f"{type(error).__name__}: {error}" + + return revision, etag, None + + +def checkpoint_status( + checkpoint: str | Path, + *, + full_hash: bool = False, +) -> dict[str, Any]: + """Inspect the RV2V DiT without reading 32.5 GB unless explicitly requested.""" + checkpoint_path = Path(checkpoint) + report: dict[str, Any] = { + "status": "missing", + "path": str(checkpoint_path), + "expected_release_commit": JOYAI_DIT_RELEASE_COMMIT, + "expected_xet_hash": JOYAI_DIT_XET_HASH, + "expected_sha256": JOYAI_DIT_SHA256, + "metadata_revision": None, + "metadata_etag": None, + "sha256": None, + "size_bytes": None, + "verification": "none", + } + + if not checkpoint_path.is_file(): + return report + + try: + report["size_bytes"] = checkpoint_path.stat().st_size + except OSError: + pass + + revision, etag, metadata_error = _read_metadata(checkpoint_path) + report["metadata_revision"] = revision + report["metadata_etag"] = etag + if metadata_error is not None: + report["metadata_error"] = metadata_error + + if etag in {JOYAI_DIT_XET_HASH, JOYAI_DIT_SHA256}: + report["status"] = "current" + report["verification"] = "huggingface_metadata" + elif etag: + report["status"] = "stale" + report["verification"] = "huggingface_metadata" + else: + report["status"] = "unknown" + + if full_hash: + digest = hashlib.sha256() + with checkpoint_path.open("rb") as checkpoint_file: + for block in iter(lambda: checkpoint_file.read(16 * 1024 * 1024), b""): + digest.update(block) + actual_sha256 = digest.hexdigest() + report["sha256"] = actual_sha256 + report["verification"] = "sha256" + report["status"] = ( + "current" if actual_sha256 == JOYAI_DIT_SHA256 else "stale" + ) + + return report diff --git a/deploy/xvideo/serving/serve_joyomni_streaming.py b/deploy/xvideo/serving/serve_joyomni_streaming.py index ce326a4..3de2325 100644 --- a/deploy/xvideo/serving/serve_joyomni_streaming.py +++ b/deploy/xvideo/serving/serve_joyomni_streaming.py @@ -26,6 +26,7 @@ import uvicorn from xvideo.serving.pe import DEFAULT_MODEL as DEFAULT_PE_MODEL +from xvideo.checkpoint_status import checkpoint_status from xvideo.serving.joyomni_streaming import ( JoyOmniRuntime, StreamingSettings, @@ -678,6 +679,7 @@ def health() -> JSONResponse: "ok": True, "runtime_loaded": app.state.runtime is not None, "dit_ckpt": args.dit_ckpt, + "checkpoint": checkpoint_status(args.dit_ckpt), "device": str(app.state.runtime.device) if app.state.runtime is not None else args.device, "vae_device": str(_module_device(app.state.runtime.pipeline.vae)) if app.state.runtime is not None else args.vae_device, "vae_encode_device": ( diff --git a/runpod/Start-JoyAI-Realtime-Test.ps1 b/runpod/Start-JoyAI-Realtime-Test.ps1 index b1aee13..c052b1d 100644 --- a/runpod/Start-JoyAI-Realtime-Test.ps1 +++ b/runpod/Start-JoyAI-Realtime-Test.ps1 @@ -115,6 +115,18 @@ try { } Write-Host "JoyAI runtime is ready after $([int]$warmTimer.Elapsed.TotalSeconds) seconds." + if ( + $null -eq $health.checkpoint -or + $health.checkpoint.status -ne "current" + ) { + $checkpointStatus = if ($null -eq $health.checkpoint) { + "missing from the health response" + } else { + [string]$health.checkpoint.status + } + throw "The upgraded RV2V checkpoint was not verified ($checkpointStatus). Stop this worker. Verify the mounted volume with: python3 /opt/joyai/runpod/verify_checkpoint.py --full-hash" + } + Write-Host "Verified: the upgraded RV2V checkpoint is active." if ( $null -eq $health.optimizations -or -not $health.optimizations.vae_compile.ready -or diff --git a/runpod/download_models.py b/runpod/download_models.py index d0acdb7..ff2fe7a 100644 --- a/runpod/download_models.py +++ b/runpod/download_models.py @@ -5,6 +5,9 @@ from huggingface_hub import snapshot_download +JOYAI_RV2V_REVISION = "eda14f342ef99c52485bbb8dc271c29b42298089" + + checkpoint_root = Path( os.getenv( "JOYOMNI_CKPT_ROOT", @@ -20,6 +23,7 @@ snapshot_download( repo_id="jdopensource/JoyAI-Video-Edit", + revision=JOYAI_RV2V_REVISION, local_dir=checkpoint_root / "JoyAI-Video-Edit", allow_patterns=[ "dit/joyai_video_edit_dit_0811.pth", @@ -55,4 +59,4 @@ temporary_file.replace(face_model) print("All required model files have been downloaded.") -print(f"Checkpoint location: {checkpoint_root}") \ No newline at end of file +print(f"Checkpoint location: {checkpoint_root}") diff --git a/runpod/start.py b/runpod/start.py index f351fd2..51c1b0b 100644 --- a/runpod/start.py +++ b/runpod/start.py @@ -4,6 +4,14 @@ from pathlib import Path +REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +DEPLOY_ROOT = REPOSITORY_ROOT / "deploy" +if str(DEPLOY_ROOT) not in sys.path: + sys.path.insert(0, str(DEPLOY_ROOT)) + +from xvideo.checkpoint_status import checkpoint_status + + TRUE_VALUES = {"1", "true", "yes", "on"} FALSE_VALUES = {"0", "false", "no", "off"} @@ -55,7 +63,7 @@ def build_model_command(repository_root: Path, *, preload: bool) -> list[str]: def main() -> int: - repository_root = Path(__file__).resolve().parents[1] + repository_root = REPOSITORY_ROOT checkpoint_root = Path( os.getenv( "JOYOMNI_CKPT_ROOT", @@ -84,6 +92,34 @@ def main() -> int: ) return 1 + dit_checkpoint = required_checkpoint_items(checkpoint_root)[0] + dit_status = checkpoint_status(dit_checkpoint) + status = dit_status["status"] + print( + "JoyAI RV2V checkpoint status: " + f"{status} ({dit_status['verification']}).", + flush=True, + ) + if status == "stale": + print( + "The mounted volume contains a different 0811 DiT than the " + "upgraded RV2V release. Refusing to load stale weights.", + flush=True, + ) + print( + "Update the volume with: " + "python3 /opt/joyai/runpod/download_models.py", + flush=True, + ) + return 1 + if status == "unknown": + print( + "Hugging Face metadata is unavailable, so this checkpoint cannot " + "be identified without a one-time full hash. Run: " + "python3 /opt/joyai/runpod/verify_checkpoint.py --full-hash", + flush=True, + ) + environment = os.environ.copy() # RunPod supplies PORT for public traffic. PORT_HEALTH is reserved for the diff --git a/runpod/verify_checkpoint.py b/runpod/verify_checkpoint.py new file mode 100644 index 0000000..c868541 --- /dev/null +++ b/runpod/verify_checkpoint.py @@ -0,0 +1,42 @@ +import argparse +import json +import os +import sys +from pathlib import Path + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +DEPLOY_ROOT = REPOSITORY_ROOT / "deploy" +if str(DEPLOY_ROOT) not in sys.path: + sys.path.insert(0, str(DEPLOY_ROOT)) + +from xvideo.checkpoint_status import checkpoint_status + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Verify the mounted JoyAI upgraded RV2V DiT checkpoint." + ) + parser.add_argument( + "--full-hash", + action="store_true", + help="Read and SHA-256 hash the complete 32.5 GB file.", + ) + args = parser.parse_args() + + checkpoint_root = Path( + os.getenv("JOYOMNI_CKPT_ROOT", "/runpod-volume/joyai/checkpoints") + ) + checkpoint = ( + checkpoint_root + / "JoyAI-Video-Edit" + / "dit" + / "joyai_video_edit_dit_0811.pth" + ) + report = checkpoint_status(checkpoint, full_hash=args.full_hash) + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if report["status"] == "current" else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py index c49ff37..1c7f93f 100644 --- a/tests/test_runpod_connection.py +++ b/tests/test_runpod_connection.py @@ -2,6 +2,7 @@ import os import unittest from pathlib import Path +from tempfile import TemporaryDirectory from unittest.mock import patch @@ -13,8 +14,77 @@ assert SPEC.loader is not None SPEC.loader.exec_module(START) +CHECKPOINT_STATUS_PATH = ROOT / "deploy" / "xvideo" / "checkpoint_status.py" +CHECKPOINT_SPEC = importlib.util.spec_from_file_location( + "checkpoint_status", CHECKPOINT_STATUS_PATH +) +CHECKPOINT_STATUS = importlib.util.module_from_spec(CHECKPOINT_SPEC) +assert CHECKPOINT_SPEC.loader is not None +CHECKPOINT_SPEC.loader.exec_module(CHECKPOINT_STATUS) + class RunPodConnectionContractTests(unittest.TestCase): + def test_current_rv2v_checkpoint_is_pinned_and_reported_by_health(self): + downloader = (ROOT / "runpod" / "download_models.py").read_text() + launcher = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + server = ( + ROOT / "deploy" / "xvideo" / "serving" / "serve_joyomni_streaming.py" + ).read_text() + self.assertIn(CHECKPOINT_STATUS.JOYAI_DIT_RELEASE_COMMIT, downloader) + self.assertIn('"checkpoint": checkpoint_status(args.dit_ckpt)', server) + self.assertIn('$health.checkpoint.status -ne "current"', launcher) + self.assertIn("the upgraded RV2V checkpoint is active", launcher) + + def test_checkpoint_metadata_identifies_current_and_stale_weights(self): + with TemporaryDirectory() as temporary_directory: + local_dir = Path(temporary_directory) / "JoyAI-Video-Edit" + checkpoint = local_dir / "dit" / "joyai_video_edit_dit_0811.pth" + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"test checkpoint") + metadata = ( + local_dir + / ".cache" + / "huggingface" + / "download" + / "dit" + / "joyai_video_edit_dit_0811.pth.metadata" + ) + metadata.parent.mkdir(parents=True) + metadata.write_text( + CHECKPOINT_STATUS.JOYAI_DIT_RELEASE_COMMIT + + "\n" + + CHECKPOINT_STATUS.JOYAI_DIT_XET_HASH + + "\n0\n" + ) + + report = CHECKPOINT_STATUS.checkpoint_status(checkpoint) + self.assertEqual(report["status"], "current") + self.assertEqual(report["verification"], "huggingface_metadata") + + metadata.write_text("old-revision\nold-etag\n0\n") + report = CHECKPOINT_STATUS.checkpoint_status(checkpoint) + self.assertEqual(report["status"], "stale") + + def test_checkpoint_without_metadata_is_unknown_until_full_hash(self): + with TemporaryDirectory() as temporary_directory: + checkpoint = ( + Path(temporary_directory) + / "JoyAI-Video-Edit" + / "dit" + / "joyai_video_edit_dit_0811.pth" + ) + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"not the official checkpoint") + + report = CHECKPOINT_STATUS.checkpoint_status(checkpoint) + self.assertEqual(report["status"], "unknown") + + report = CHECKPOINT_STATUS.checkpoint_status( + checkpoint, full_hash=True + ) + self.assertEqual(report["status"], "stale") + self.assertEqual(report["verification"], "sha256") + def test_preload_defaults_to_enabled(self): with patch.dict(os.environ, {}, clear=True): self.assertTrue(START.env_enabled("JOYOMNI_PRELOAD", default=True))