hao-ai-lab-FastVideo/fastvideo/tests/modal/pr_test.py

"""Dormant manual Modal launcher.

Production PR CI is Slinky/Slurm-only. This module is retained as rollback
code and is not referenced by the Buildkite pipeline or slash commands.
"""

import os
import re
import shutil
import subprocess
import sys
import time

import modal

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
try:
    from modal_image_utils import (  # noqa: E402
        resolve_image_ref, resolve_uv_torch_backend)
except ModuleNotFoundError:
    # Remote Modal containers re-import this module but mount only the
    # entrypoint file; the digest resolution already happened at local
    # launch time, so a passthrough is correct there.
    def resolve_image_ref(image_ref: str) -> str:
        return image_ref

    def resolve_uv_torch_backend(image_tag: str) -> str | None:
        return os.environ.get("UV_TORCH_BACKEND")


app = modal.App()

model_vol = modal.Volume.from_name("hf-model-weights")
image_version = os.getenv("IMAGE_VERSION", "latest")
image_tag = f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{image_version}"
image_ref = resolve_image_ref(image_tag)
print(f"Using image: {image_ref}")

# Mutable tags inherit the registry image's baked backend, keeping a latest-tag
# transition safe. Explicit CUDA tags also work with older images that predate
# the baked setting, and a caller override always wins.
uv_torch_backend_override = resolve_uv_torch_backend(image_tag)

# INVARIANT: this image definition must be byte-identical for every CI job at
# a given base image digest -- one build, shared cache across all concurrent
# lanes. Never put a per-job/per-commit value (BUILDKITE_*, TEST_SCOPE,
# env-derived overrides) into the image via .env()/run_commands: it becomes an
# image layer, so whenever the base digest changes every concurrent job
# rebuilds its own image variant (~15-20 min each), blowing the Buildkite job
# budget. `image_ref` is the only env-derived input allowed here, because it
# *selects* the base digest. Per-job values arrive at runtime via
# `ci_env_secret` below.
image = (modal.Image.from_registry(image_ref, add_python="3.12").run_commands("rm -rf /FastVideo").apt_install(
    "cmake", "pkg-config", "build-essential", "curl", "libssl-dev", "ffmpeg").run_commands(
        "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable").
         run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc").env({
             "PATH": "/root/.cargo/bin:$PATH",
             "HF_REPO_ID": "FastVideo/performance-tracking",
         }))

dreamverse_image = (image.run_commands("curl -fsSL https://deb.nodesource.com/setup_22.x | bash -").apt_install(
    "nodejs").run_commands("node --version && npm --version"))

# Per-job/per-invocation values are injected into the container environment at
# RUNTIME via this secret (attached to every function below), so the image
# stays identical across jobs. Consumers (run_test_command's checkout,
# fastvideo/tests/performance/{compare_baseline,identity}.py, the FA4
# resolver, `uv pip install`) all read os.environ at runtime, so nothing else
# changes.
ci_env_secret = modal.Secret.from_dict({
    "BUILDKITE_REPO":
    os.environ.get("BUILDKITE_REPO", ""),
    "BUILDKITE_COMMIT":
    os.environ.get("BUILDKITE_COMMIT", ""),
    "BUILDKITE_PULL_REQUEST":
    os.environ.get("BUILDKITE_PULL_REQUEST", ""),
    "BUILDKITE_BRANCH":
    os.environ.get("BUILDKITE_BRANCH", ""),
    "BUILDKITE_SOURCE":
    os.environ.get("BUILDKITE_SOURCE", ""),
    "BUILDKITE_BUILD_URL":
    os.environ.get("BUILDKITE_BUILD_URL", ""),
    "BUILDKITE_BUILD_ID":
    os.environ.get("BUILDKITE_BUILD_ID", ""),
    "BUILDKITE_JOB_ID":
    os.environ.get("BUILDKITE_JOB_ID", ""),
    "TEST_SCOPE":
    os.environ.get("TEST_SCOPE", ""),
    "IMAGE_VERSION":
    image_version,
    "FASTVIDEO_CONTAINER_IMAGE_REF":
    image_ref,
    **{
        key: os.environ[key]
        for key in (
            "FASTVIDEO_ATTENTION_BACKEND",
            "FASTVIDEO_PERFORMANCE_PROFILE_VERSION",
        ) if os.environ.get(key)
    },
    **({
        "UV_TORCH_BACKEND": uv_torch_backend_override
    } if uv_torch_backend_override else {}),
    # FA4 is opt-in (FASTVIDEO_FA4). Keep the default enabled for
    # inference/perf parity; model-load and training lanes that do not exercise
    # FA4 explicitly set FASTVIDEO_FA4=0 in their command strings below.
    # Caller override wins.
    "FASTVIDEO_FA4":
    os.environ.get("FASTVIDEO_FA4", "1"),
})

hf_secret = modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})
wandb_secret = modal.Secret.from_dict({"WANDB_API_KEY": os.environ.get("WANDB_API_KEY", "")})


def _run_git_with_retries(command: list[str], *, cwd: str, cleanup_path: str | None = None) -> None:
    last_returncode = 1
    for attempt in range(1, 4):
        if cleanup_path is not None:
            shutil.rmtree(cleanup_path, ignore_errors=True)

        result = subprocess.run(command, cwd=cwd, check=False)
        if result.returncode == 0:
            return

        last_returncode = result.returncode
        if attempt < 3:
            sleep_seconds = 5 * attempt
            print(f"Git command failed (attempt {attempt}/3, exit {last_returncode}); "
                  f"retrying in {sleep_seconds}s",
                  flush=True)
            time.sleep(sleep_seconds)

    raise RuntimeError(f"Git command failed after 3 attempts with exit code {last_returncode}: " + " ".join(command))


def _checkout_repository(git_repo: str, git_commit: str, pr_number: str | None, repo_root: str = "/FastVideo") -> None:
    if not git_repo or git_repo.startswith("-"):
        raise RuntimeError("BUILDKITE_REPO must be a non-empty repository URL.")

    if pr_number and pr_number != "false":
        try:
            pr_id = int(pr_number)
        except ValueError as error:
            raise RuntimeError(f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}") from error
        if pr_id <= 0:
            raise RuntimeError(f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}")
        target = f"refs/pull/{pr_id}/head"
        print(f"Using PR ref for checkout: {target}")
    else:
        if not git_commit or re.fullmatch(r"[0-9a-fA-F]{7,64}", git_commit) is None:
            raise RuntimeError(f"Invalid BUILDKITE_COMMIT value: {git_commit}")
        target = git_commit
        print(f"Using direct commit checkout: {target}")

    clone_command = [
        "git",
        "-c",
        "http.version=HTTP/1.1",
        "clone",
        "--config",
        "http.version=HTTP/1.1",
        "--depth=1",
        "--filter=blob:none",
        "--no-checkout",
        git_repo,
        repo_root,
    ]
    _run_git_with_retries(clone_command, cwd="/", cleanup_path=repo_root)

    git_prefix = ["git", "-c", "http.version=HTTP/1.1"]
    _run_git_with_retries(git_prefix + [
        "fetch",
        "--prune",
        "--no-tags",
        "--depth=1",
        "--filter=blob:none",
        "origin",
        target,
    ],
                          cwd=repo_root)
    _run_git_with_retries(git_prefix + ["checkout", "--detach", "FETCH_HEAD"], cwd=repo_root)
    _run_git_with_retries(git_prefix + ["submodule", "update", "--init", "--recursive"], cwd=repo_root)


def run_test(pytest_command: str):
    """Helper function to run a test suite with custom pytest command"""
    run_test_command(pytest_command, build_kernel=True)


def run_test_command(test_command: str, build_kernel: bool, install_command: str = 'uv pip install -e ".[test]"'):
    """Helper function to run a test suite with custom test command.

    Most FastVideo CI suites need the custom kernel build. App-level tests like
    DreamVerse's mock-backend UI checks do not, so keep the kernel build
    optional to avoid unrelated CUDA/kernel setup in that CI path.

    The dependency install runs BEFORE the kernel build: pyproject pins the
    PyPI fastvideo-kernel wheel, so an install after the build silently
    replaces the just-built in-tree kernel with the (older) wheel -- every
    lane would then test stale kernels. Pass install_command="" for commands
    that manage their own installs.
    """
    import os
    import subprocess
    import sys

    git_repo = os.environ.get("BUILDKITE_REPO", "")
    git_commit = os.environ.get("BUILDKITE_COMMIT", "")
    pr_number = os.environ.get("BUILDKITE_PULL_REQUEST")

    print(f"Cloning repository: {git_repo}")
    print(f"Target commit: {git_commit}")
    if pr_number:
        print(f"PR number: {pr_number}")
    _checkout_repository(git_repo, git_commit, pr_number)

    setup_steps = [
        "source $HOME/.local/bin/env",
        "source /opt/venv/bin/activate",
        "cd /FastVideo",
    ]
    if install_command:
        setup_steps.append(install_command)
    if build_kernel:
        setup_steps.append("python fastvideo/tests/modal/kernel_build_cache.py install")
    setup_command = " &&\n    ".join(setup_steps)

    setup_result = subprocess.run(["/bin/bash", "-c", setup_command], stdout=sys.stdout, stderr=sys.stderr, check=False)
    if setup_result.returncode != 0:
        raise RuntimeError(f"Setup command failed with exit code {setup_result.returncode}")

    command = " &&\n    ".join([
        "source $HOME/.local/bin/env",
        "source /opt/venv/bin/activate",
        "cd /FastVideo",
        test_command,
    ])
    result = subprocess.run(["/bin/bash", "-c", command], stdout=sys.stdout, stderr=sys.stderr, check=False)

    # Modal containers crash on sys.exit(0); raise on failure, return on success.
    if result.returncode != 0:
        raise RuntimeError(f"Test command failed with exit code {result.returncode}")


@app.function(gpu="H100:1",
              image=image,
              timeout=1200,
              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_encoder_tests():
    run_test(
        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && bash .buildkite/scripts/lanes/encoder.sh"
    )


@app.function(gpu="L40S:1",
              image=image,
              timeout=1200,
              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_vae_tests():
    run_test(
        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && bash .buildkite/scripts/lanes/vae.sh")


@app.function(gpu="L40S:1",
              image=image,
              timeout=900,
              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_golden_gate_tests():
    # Single-layer bitwise DiT fingerprints (~40s/model on GPU): a green gate
    # means the compute path is bit-identical to the golden, so the expensive
    # SSIM generation for that model cannot have regressed. Downloads only the
    # shards holding the gated layer, never full checkpoints.
    run_test(
        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && bash .buildkite/scripts/lanes/golden_gate.sh"
    )


@app.function(gpu="L40S:1",
              image=image,
              timeout=900,
              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_transformer_tests():
    run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
             "FASTVIDEO_FA4=0 bash .buildkite/scripts/lanes/transformer.sh")


@app.function(gpu="L40S:4",
              cpu=8.0,
              memory=32768,
              image=image,
              timeout=900,
              secrets=[wandb_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_training_tests():
    run_test("export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && "
             "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/Vanilla -srP")


@app.function(gpu="L40S:2",
              cpu=8.0,
              memory=32768,
              image=image,
              timeout=900,
              secrets=[wandb_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_training_lora_tests():
    run_test("export HF_HOME='/root/data/.cache' && wandb login $WANDB_API_KEY && "
             "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP")


@app.function(gpu="H100!:2", image=image, timeout=900, secrets=[wandb_secret, ci_env_secret])
def run_training_tests_VSA():
    run_test("wandb login $WANDB_API_KEY && FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/VSA -srP")


@app.function(gpu="H100:1", image=image, timeout=900, secrets=[ci_env_secret])
def run_kernel_tests():
    run_test("bash .buildkite/scripts/lanes/kernel_tests.sh")


# @app.function(gpu="H100:1", image=image, timeout=900, secrets=[ci_env_secret])
# def run_precision_tests_VSA():
#     # VSA correctness is covered by the same file now
#     run_test("pytest fastvideo-kernel/tests/test_correctness.py")

# @app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
# def run_precision_tests_vmoba():
#     run_test("pytest fastvideo-kernel/tests/test_vmoba_correctness.py")


@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
def run_inference_tests_vmoba():
    run_test("bash .buildkite/scripts/lanes/inference_vmoba.sh")


@app.function(gpu="L40S:1", image=image, timeout=1200, secrets=[ci_env_secret])
def run_inference_lora_tests():
    run_test("bash .buildkite/scripts/lanes/inference_lora.sh")


@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[ci_env_secret])
def run_distill_dmd_tests():
    run_test("FASTVIDEO_FA4=0 bash .buildkite/scripts/lanes/distillation_dmd.sh")


@app.function(gpu="L40S:2", image=image, timeout=900, secrets=[wandb_secret, ci_env_secret])
def run_self_forcing_tests():
    run_test("wandb login $WANDB_API_KEY && "
             "FASTVIDEO_FA4=0 pytest ./fastvideo/tests/training/self-forcing/test_self_forcing.py -vs")


@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
def run_unit_test():
    run_test("bash .buildkite/scripts/unit_test.sh")


# TODO: David: GPU only used to resolve import time requirement (not needed for this test). Maybe make those imports lazy?
@app.function(gpu="L40S:1", image=dreamverse_image, timeout=1800, secrets=[ci_env_secret])
def run_dreamverse_app_tests():
    run_test_command(install_command="",
                     build_kernel=False,
                     test_command="""
        uv pip install -e ".[test,dreamverse]" &&
        export PYTHONPATH=/FastVideo/apps/dreamverse:$PYTHONPATH &&
        pytest apps/dreamverse/dreamverse/tests -q &&
        cd apps/dreamverse/web &&
        npm ci &&
        npm run typecheck &&
        npm test &&
        npx playwright install --with-deps chromium webkit firefox &&
        bash -c '
            set -e
            BACKEND_PORT="${BACKEND_PORT:-8009}"
            python -m uvicorn dreamverse.mock_server:app --host 127.0.0.1 --port "$BACKEND_PORT" &
            MOCK_SERVER_PID=$!
            trap "kill $MOCK_SERVER_PID 2>/dev/null || true" EXIT
            for i in {1..30}; do
                curl -fsS "http://127.0.0.1:$BACKEND_PORT/healthz" && break
                sleep 1
            done
            curl -fsS "http://127.0.0.1:$BACKEND_PORT/healthz"
            BACKEND_HOST=127.0.0.1 BACKEND_PORT="$BACKEND_PORT" CI=1 \
                npm run e2e -- \
                    --project=chromium \
                    --project=webkit \
                    --project=firefox \
                    --project=mobile-safari \
                    --project=mobile-chromium
        '
        """)


@app.function(gpu="L40S:1",
              cpu=8.0,
              memory=32768,
              image=image,
              timeout=1800,
              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_train_framework_tests():
    run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
             "FASTVIDEO_FA4=0 bash .buildkite/scripts/lanes/train_framework.sh")


@app.function(gpu="L40S:1",
              image=image,
              timeout=1800,
              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def seed_grad_norm_references():
    """Record the legacy L40S per-method grad-norm reference.

    Phase 2 / 5a-ii one-off seeding entrypoint. Pinned to ``gpu="L40S:1"`` (the
    dormant Modal rollback runtime), so this function only seeds the ``L40S`` key in
    ``fastvideo/tests/train/methods/grad_norm_refs.json``.

    ``FASTVIDEO_GRADNORM_UPDATE=1`` makes ``check_grad_norm_regression`` record
    the measured norm instead of asserting; ``-rs`` surfaces the recorded value
    in the log so it can be copied into the JSON.

    To seed any other device (e.g. our local Blackwell dev box → ``GB200``
    key), run the same env-var + pytest invocation directly on that
    workstation — see the module docstring of ``grad_norm_regression.py`` for
    the local command and the ``_DEVICE_MAPPINGS`` table.
    """
    run_test("export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && "
             "FASTVIDEO_FA4=0 FASTVIDEO_GRADNORM_UPDATE=1 pytest ./fastvideo/tests/train/methods -vs -rs")


@app.function(gpu="L40S:1",
              image=image,
              timeout=3600,
              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_eval_tests():
    # Eval metric regression: drives the high-level fastvideo.eval API on a
    # fixed asset and asserts each score matches the upstream reference number
    # checked into fastvideo/tests/eval/reference_scores/. Pulls several scorer
    # checkpoints (VideoScore2 VLM, VBench nets, audio models) on first run;
    # they cache on the hf-model-weights volume thereafter.
    #
    # Installs [eval-full] (eval + vbench + audio extras) on top of [test]:
    # the dev image only ships [dev], and without the extras skip_missing_deps
    # in conftest would silently drop nearly every metric and the lane would
    # pass vacuously. detectron2-backed vbench metrics remain skipped by
    # design (not pip-installable; see fastvideo/eval/README.md).
    run_test_command(
        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && bash .buildkite/scripts/lanes/eval.sh",
        build_kernel=True,
        install_command='uv pip install -e ".[test,eval-full]"')


@app.function(gpu="L40S:1", image=image, timeout=3600, secrets=[hf_secret, ci_env_secret])
def run_lora_extraction_tests():
    run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/lora_extraction/test_lora_extraction.py")


@app.function(gpu="L40S:2",
              cpu=8.0,
              memory=32768,
              image=image,
              timeout=1800,
              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_performance_tests():
    # PR/direct records are uploaded only on pass; scheduled main uploads pass
    # and fail so the dashboard records every canonical baseline attempt.
    run_test(
        "export HF_HOME='/root/data/.cache' && "
        "export PERFORMANCE_TRACKING_ROOT='/tmp/perf-tracking' && "
        "hf auth login --token $HF_API_KEY && "
        "if [ -n \"${BUILDKITE_PULL_REQUEST:-}\" ] && [ \"${BUILDKITE_PULL_REQUEST:-false}\" != 'false' ]; then "
        "export PERF_RUN_SOURCE='pr'; "
        "export PERF_UPLOAD_POLICY='pass'; "
        "elif [ \"${BUILDKITE_BRANCH:-}\" = 'main' ] && ( [ \"${BUILDKITE_SOURCE:-}\" = 'schedule' ] || [ \"${TEST_SCOPE:-}\" = 'full' ] ); then "
        "export PERF_RUN_SOURCE='scheduled_main'; "
        "export PERF_UPLOAD_POLICY='always'; "
        "elif [ \"${TEST_SCOPE:-}\" = 'direct' ]; then "
        "export PERF_RUN_SOURCE='unknown'; "
        "export PERF_UPLOAD_POLICY='pass'; "
        "else "
        "export PERF_RUN_SOURCE='unknown'; "
        "export PERF_UPLOAD_POLICY='never'; "
        "fi; "
        "(nvidia-smi --query-gpu=index,timestamp,clocks.sm,clocks.max.sm,power.draw,power.limit,temperature.gpu "
        "--format=csv -l 10 > /tmp/gpu_telemetry.csv 2>/dev/null &); "
        "pytest ./fastvideo/tests/performance -vs; "
        "PYTEST_RC=$?; "
        "PERF_RC=0; "
        "if [ $PYTEST_RC -eq 0 ] || [ \"$PERF_UPLOAD_POLICY\" = 'always' ]; then "
        "PERF_PYTEST_RC=$PYTEST_RC python ./fastvideo/tests/performance/compare_baseline.py; "
        "PERF_RC=$?; "
        "fi; "
        "python ./fastvideo/tests/performance/dashboard.py || true; "
        "echo '--- GPU telemetry (clocks.sm vs clocks.max.sm reveals capped hosts) ---'; "
        "cat /tmp/gpu_telemetry.csv || true; "
        "FINAL_RC=$PYTEST_RC; "
        "if [ $FINAL_RC -eq 0 ]; then FINAL_RC=$PERF_RC; fi; "
        "exit $FINAL_RC")


@app.function(gpu="L40S:1",
              image=image,
              timeout=1800,
              secrets=[hf_secret, ci_env_secret],
              volumes={"/root/data": model_vol})
def run_api_server_tests():
    run_test(
        "export HF_HOME='/root/data/.cache' && hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/entrypoints/test_openai_api_integration.py -vs"
    )