"""Tests for diffusion `update_weights_from_disk`.

This module verifies the ability to update model weights in place without restarting
the server, which is critical for RL workflows and iterative fine-tuning scenarios.

Author:

Menyang Liu, https://github.com/dreamyang-liu
Chenyang Zhao, https://github.com/zhaochenyang20

We use two model pairs for testing (base model / instruct model pairs):

- FLUX.2-klein-base-4B / FLUX.2-klein-4B
- Qwen/Qwen-Image / Qwen/Qwen-Image-2512

These model pairs share the same architecture but differ in transformer
weights. The basic testing logic is to refit the instruct model into the
base model and verify the checksum of the transformer weights are the same,
which simulates the real-world RL scenario. However, since these two model
pairs only differ in transformer weights, and we want to verify update a
specific module with update_weights_from_disk API, we need to create a perturbed
instruct model that adds noise to the vae weights. In this sense, the instruct
model differs from the base model in vae and transformer weights, the text
encoder are still the same.

To strictly verify the correctness of the refit API, we compare the checksum in
SHA-256 on the disk and the server.

NOTE and TODO: In the refit a specific module test, we update the transformer
module and keep other modules the same. As described above, the vae's weights
are perturbed. If we add vae as a target module in the future, ideally speaking,
we should assert that the refitted vae's checksum is the same as directly
computed from the perturbed vae weights in the disk. However, since there is
complex weight-name remapping and QKV merge during model loading, it is not easy
to compare the server-disk checksum for vae and text encoder directly. Therefore,
if the target module is vae, we only verify that the refitted vae's checksum is
different from the base model's vae's checksum.

It should be good issue to solve for the community to adds comparison the server-disk
checksum for vae and text encoder in this test.

=============================================================================

Test organization:

5 test cases in 2 classes;
two model pairs are tested locally, one in CI.

=============================================================================

Class 1: TestUpdateWeightsFromDisk                  (4 tests) — API contract, checksum & rollback
Class 2: TestUpdateWeightsFromDiskWithOffload       (1 test) — Offload-aware update + checksum

-----------------------------------------------------------------------------

Class 1: TestUpdateWeightsFromDisk

Validate the update_weights_from_disk API contract, request/response shape,
error handling, checksum verification, and corrupted-weight rollback.

All tests share one class-scoped server (same process, same in-memory weights).
Tests that require "base model then update" should be explicitly reset to
base model first so behavior is order-independent and updates are real
(base -> perturbed), not no-ops (perturbed -> perturbed).

  • test_update_weights_from_disk_default

    base model -> perturbed model with flush_cache=True.
    Verifies after-update transformer checksum == perturbed model's
    transformer disk checksum


  • test_update_weights_specific_modules

    base -> perturbed with flush_cache=False. Updates only transformer as
    target_modules. Verifies that:
    (1) targeted module's in-memory checksum changed;
    (2) non-targeted modules' in-memory checksums are unchanged.

  • test_update_weights_rejects_invalid_requests

    Invalid requests must fail (400, success=False) without mutating live
    weights. Covers nonexistent model, missing model_path, and nonexistent
    target module.

  • test_corrupted_weights_rollback

    All-or-nothing rollback: We first refit the server from base model ->
    perturbed model. We manually truncate the vae weights of the base
    model to get a corrupted model. We then call the refit to update
    the server from the perturbed model -> corrupted model. Verify that:

    1. The update fails due to truncated vae, server should roll back to the
    perturbed model, i.e., server's transformer weights == perturbed model's
    transformer weights != base model's transformer weights.

    2. After the rollback, server's vae weights == perturbed model's vae
    weights != base model's vae weights.

    3. After the rollback, server's text encoder weights == base model's
    text encoder weights == perturbed model's text encoder weights.

-----------------------------------------------------------------------------

Class 2: TestUpdateWeightsFromDiskWithOffload


Ensure weight updates and checksum verification work when layerwise offload is enabled
(--dit-layerwise-offload). With offload, parameters live in CPU buffers and only left
small torch.empty((1,)) as placeholders on GPU; the updater must write into CPU buffers
and update prefetched GPU tensors without shape mismatch.

  • test_update_weights_with_offload_enabled

    Server with --dit-layerwise-offload (base). Load perturbed checkpoint;
    must succeed (200, success=True), no "Shape mismatch". server's transformer checksum
    matches perturbed model's transformer disk checksum.
"""

from __future__ import annotations

import functools
import os
import shutil
import sys
import tempfile
import threading
from collections.abc import Callable

import pytest
import requests
from safetensors.torch import load_file, save_file

from sglang.multimodal_gen.runtime.loader.utils import (
    _list_safetensors_files,
)
from sglang.multimodal_gen.runtime.loader.weight_utils import (
    compute_weights_checksum,
    safetensors_weights_iterator,
)
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.test.server.test_server_utils import (
    ServerManager,
)
from sglang.multimodal_gen.test.test_utils import (
    DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST,
    DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST,
    DEFAULT_QWEN_IMAGE_2512_MODEL_NAME_FOR_TEST,
    DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
    get_dynamic_server_port,
    is_in_ci,
)

logger = init_logger(__name__)


_TRANSFORMER_MODULE = "transformer"
_VAE_MODULE = "vae"
_TEXT_ENCODER_MODULE_PREFIX = "text_encoder"


# Modules whose weights differ between the base model and the perturbed
# perturbed checkpoint
_DIFFERING_MODULES: list[str] = [_TRANSFORMER_MODULE, _VAE_MODULE]

_ALL_MODEL_PAIRS: list[tuple[str, str]] = [
    (
        DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST,
        DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST,
    ),
    (
        DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
        DEFAULT_QWEN_IMAGE_2512_MODEL_NAME_FOR_TEST,
    ),
]


_CI_MODEL_PAIR_ENV = "SGLANG_MMGEN_UPDATE_WEIGHTS_PAIR"
_CI_MODEL_PAIR_IDS = ("FLUX.2-klein-base-4B",)
_UPDATE_TIMEOUT_SECONDS = 900
_CHECKSUM_TIMEOUT_SECONDS = 600


def _resolve_active_model_pairs() -> list[tuple[str, str]]:
    if not is_in_ci():
        return _ALL_MODEL_PAIRS

    pair_by_id = {pair[0].split("/")[-1]: pair for pair in _ALL_MODEL_PAIRS}
    selected_pair_id = os.environ.get(_CI_MODEL_PAIR_ENV)
    if selected_pair_id is None:
        return [pair_by_id[_CI_MODEL_PAIR_IDS[0]]]

    selected_pair = pair_by_id.get(selected_pair_id)
    if selected_pair is None:
        valid_ids = ", ".join(sorted(pair_by_id))
        raise ValueError(
            f"Invalid {_CI_MODEL_PAIR_ENV}={selected_pair_id!r}. "
            f"Expected one of: {valid_ids}."
        )
    return [selected_pair]


_ACTIVE_MODEL_PAIRS = _resolve_active_model_pairs()
_PAIR_IDS = [p[0].split("/")[-1] for p in _ACTIVE_MODEL_PAIRS]


@functools.lru_cache(maxsize=None)
def _compute_checksum_from_disk(model_path: str, module_name: str) -> str:
    """Compute SHA-256 checksum from safetensors files on disk.

    Uses the same compute_weights_checksum function as the server,
    so the checksums are directly comparable.

    Results are cached (keyed on model_path and module_name) because the
    same disk checksum is requested multiple times across tests.
    """
    local_path = maybe_download_model(model_path)
    weights_dir = os.path.join(local_path, module_name)
    assert os.path.exists(
        weights_dir
    ), f"No weights dir for {module_name} in {local_path}"

    safetensors_files = _list_safetensors_files(weights_dir)
    assert safetensors_files, f"No safetensors files in {weights_dir}"

    return compute_weights_checksum(safetensors_weights_iterator(safetensors_files))


def _clone_model_with_modified_module(
    src_model: str,
    dst_model: str,
    target_module: str,
    transform_safetensor: Callable[[str, str], None],
) -> None:
    # Symlink root-level files (model_index.json, etc.).
    for fname in os.listdir(src_model):
        src_path = os.path.join(src_model, fname)
        dst_path = os.path.join(dst_model, fname)
        if os.path.isfile(src_path) and not os.path.exists(dst_path):
            os.symlink(src_path, dst_path)

    for module_dir in sorted(os.listdir(src_model)):
        src_dir = os.path.join(src_model, module_dir)
        dst_dir = os.path.join(dst_model, module_dir)
        if not os.path.isdir(src_dir):
            continue

        if module_dir != target_module:
            if not os.path.exists(dst_dir):
                os.symlink(src_dir, dst_dir)
            continue

        os.makedirs(dst_dir, exist_ok=True)
        transformed = False
        for fname in sorted(os.listdir(src_dir)):
            src_file = os.path.join(src_dir, fname)
            dst_file = os.path.join(dst_dir, fname)
            if not os.path.isfile(src_file):
                continue

            if not fname.endswith(".safetensors") or transformed:
                if not os.path.exists(dst_file):
                    os.symlink(src_file, dst_file)
                continue

            transform_safetensor(src_file, dst_file)
            transformed = True


def _truncate_safetensor(src_file: str, dst_file: str) -> None:
    shutil.copy2(src_file, dst_file)
    size = os.path.getsize(dst_file)
    with open(dst_file, "r+b") as f:
        f.truncate(size - 2)
    logger.info(
        "Created corrupted safetensors: %s (%d -> %d bytes)",
        dst_file,
        size,
        size - 2,
    )


def _perturb_safetensor(src_file: str, dst_file: str) -> None:
    tensors = load_file(src_file)
    perturbed = {
        k: (t + 0.01 if t.is_floating_point() else t) for k, t in tensors.items()
    }
    save_file(perturbed, dst_file)
    logger.info("Created perturbed safetensors: %s", dst_file)


class _UpdateWeightsApiMixin:
    def _update_weights(
        self,
        base_url: str,
        model_path: str,
        flush_cache: bool = True,
        target_modules: list[str] | None = None,
        timeout: int = _UPDATE_TIMEOUT_SECONDS,
    ) -> tuple[dict, int]:
        payload = {"model_path": model_path, "flush_cache": flush_cache}
        if target_modules is not None:
            payload["target_modules"] = target_modules
        response = requests.post(
            f"{base_url}/update_weights_from_disk",
            json=payload,
            timeout=timeout,
        )
        return response.json(), response.status_code

    def _get_weights_checksum(
        self,
        base_url: str,
        module_names: list[str] | None = None,
        timeout: int = _CHECKSUM_TIMEOUT_SECONDS,
    ) -> dict:
        payload = {}
        if module_names is not None:
            payload["module_names"] = module_names
        response = requests.post(
            f"{base_url}/get_weights_checksum",
            json=payload,
            timeout=timeout,
        )
        assert (
            response.status_code == 200
        ), f"get_weights_checksum failed: {response.status_code} {response.text}"
        return response.json()

    def _assert_server_matches_model(
        self,
        base_url: str,
        expected_model: str,
    ) -> None:
        server_checksums = self._get_weights_checksum(
            base_url, module_names=[_TRANSFORMER_MODULE]
        )
        expected_cs = _compute_checksum_from_disk(expected_model, _TRANSFORMER_MODULE)
        server_cs = server_checksums.get(_TRANSFORMER_MODULE)
        assert server_cs == expected_cs, (
            f"Checksum mismatch on '{_TRANSFORMER_MODULE}'\n"
            f"  expected({expected_model}): {expected_cs}\n"
            f"  server: {server_cs}"
        )


class TestUpdateWeightsFromDisk(_UpdateWeightsApiMixin):
    @pytest.fixture(
        scope="class",
        params=_ACTIVE_MODEL_PAIRS,
        ids=_PAIR_IDS,
    )
    def diffusion_server_no_offload(self, request):
        default_model, source_model = request.param
        port = get_dynamic_server_port()
        wait_deadline = float(os.environ.get("SGLANG_TEST_WAIT_SECS", "600"))

        manager = ServerManager(
            model=default_model,
            port=port,
            wait_deadline=wait_deadline,
            extra_args="--num-gpus 1",
        )

        # Ensure models are local before spawning threads that need the paths.
        local_default = maybe_download_model(default_model)
        local_source = maybe_download_model(source_model)

        perturbed_vae_model_dir = tempfile.mkdtemp(prefix="sglang_perturbed_vae_")
        corrupted_vae_model_dir = tempfile.mkdtemp(prefix="sglang_corrupted_")

        # Run all disk I/O in background while the server boots.
        bg_threads = [
            threading.Thread(
                target=_compute_checksum_from_disk, args=(default_model, module)
            )
            for module in _DIFFERING_MODULES
        ] + [
            threading.Thread(
                target=_clone_model_with_modified_module,
                args=(
                    local_source,
                    perturbed_vae_model_dir,
                    _VAE_MODULE,
                    _perturb_safetensor,
                ),
            ),
            threading.Thread(
                target=_clone_model_with_modified_module,
                args=(
                    local_default,
                    corrupted_vae_model_dir,
                    _VAE_MODULE,
                    _truncate_safetensor,
                ),
            ),
        ]
        for t in bg_threads:
            t.start()

        ctx = manager.start()
        for t in bg_threads:
            t.join()

        # Sanity: all _DIFFERING_MODULES should differ between base and perturbed.
        for module in _DIFFERING_MODULES:
            assert _compute_checksum_from_disk(
                default_model, module
            ) != _compute_checksum_from_disk(perturbed_vae_model_dir, module), (
                f"Assumption violated: {module} should differ between "
                f"{default_model} and {perturbed_vae_model_dir}"
            )

        try:
            yield ctx, default_model, perturbed_vae_model_dir, corrupted_vae_model_dir
        finally:
            ctx.cleanup()
            shutil.rmtree(perturbed_vae_model_dir, ignore_errors=True)
            shutil.rmtree(corrupted_vae_model_dir, ignore_errors=True)

    def test_update_weights_from_disk_default(self, diffusion_server_no_offload):
        """Default update (target_modules=None, flush_cache=True): all changed modules updated."""
        ctx, default_model, perturbed_model_dir, _ = diffusion_server_no_offload
        base_url = f"http://localhost:{ctx.port}"

        self._update_weights(base_url, default_model, flush_cache=True)

        result, status_code = self._update_weights(
            base_url, perturbed_model_dir, flush_cache=True
        )
        assert status_code == 200
        assert result.get("success", False), f"Update failed: {result.get('message')}"

        self._assert_server_matches_model(base_url, perturbed_model_dir)

    def test_update_weights_specific_modules(self, diffusion_server_no_offload):
        ctx, default_model, perturbed_model_dir, _ = diffusion_server_no_offload
        base_url = f"http://localhost:{ctx.port}"

        # Reset server to default_model.
        self._update_weights(base_url, default_model)
        before_checksums = self._get_weights_checksum(
            base_url, module_names=_DIFFERING_MODULES
        )

        target_modules = [_TRANSFORMER_MODULE]
        result, status_code = self._update_weights(
            base_url,
            perturbed_model_dir,
            target_modules=target_modules,
            flush_cache=False,
        )
        assert status_code == 200, f"Update failed: {result}"
        assert result.get("success", False), f"Update failed: {result.get('message')}"

        after_checksums = self._get_weights_checksum(
            base_url, module_names=_DIFFERING_MODULES
        )

        # Targeted module should have changed.
        for name in target_modules:
            assert after_checksums.get(name) != before_checksums.get(name), (
                f"Targeted module '{name}' checksum should change after update\n"
                f"  before: {before_checksums.get(name)}\n"
                f"  after:  {after_checksums.get(name)}"
            )

        # Non-targeted modules should be unchanged.
        for name, cs in after_checksums.items():
            if name in target_modules or cs == "not_found":
                continue
            assert cs == before_checksums.get(name), (
                f"Non-targeted module '{name}' should be unchanged\n"
                f"  before: {before_checksums.get(name)}\n"
                f"  after:  {cs}"
            )

    def test_update_weights_rejects_invalid_requests(self, diffusion_server_no_offload):
        ctx, _, perturbed_model_dir, _ = diffusion_server_no_offload
        base_url = f"http://localhost:{ctx.port}"

        before_checksums = self._get_weights_checksum(
            base_url, module_names=_DIFFERING_MODULES
        )

        result, status_code = self._update_weights(
            base_url,
            "/nonexistent/path/to/model",
            timeout=60,
        )
        logger.info(f"Update result for nonexistent model: {result}")

        assert status_code == 400, f"Expected 400, got {status_code}"
        assert not result.get("success", True), "Should fail for nonexistent model"

        response = requests.post(
            f"{base_url}/update_weights_from_disk",
            json={},
            timeout=30,
        )

        assert response.status_code == 400, f"Expected 400, got {response.status_code}"
        result = response.json()
        assert not result.get("success", True), "Should fail when model_path is missing"

        result, status_code = self._update_weights(
            base_url,
            perturbed_model_dir,
            target_modules=["nonexistent_module"],
            timeout=60,
        )
        logger.info(f"Update nonexistent module result: {result}")

        assert status_code == 400, f"Expected 400, got {status_code}"
        assert not result.get("success", True), "Should fail for nonexistent module"
        assert "not found in pipeline" in result.get("message", "")

        after_checksums = self._get_weights_checksum(
            base_url, module_names=_DIFFERING_MODULES
        )
        assert after_checksums == before_checksums

    def test_corrupted_weights_rollback(self, diffusion_server_no_offload):
        ctx, default_model, perturbed_model_dir, corrupted_vae_model_dir = (
            diffusion_server_no_offload
        )
        base_url = f"http://localhost:{ctx.port}"

        # base → perturbed
        self._update_weights(base_url, default_model)
        base_checksums = self._get_weights_checksum(base_url)

        result, status_code = self._update_weights(base_url, perturbed_model_dir)
        assert status_code == 200 and result.get("success")
        perturbed_checksums = self._get_weights_checksum(base_url)

        text_encoder_modules = sorted(
            name
            for name in perturbed_checksums
            if _TEXT_ENCODER_MODULE_PREFIX in name
            and perturbed_checksums.get(name) != "not_found"
            and base_checksums.get(name) != "not_found"
        )
        assert (
            text_encoder_modules
        ), "Expected at least one text encoder module checksum"

        # perturbed → corrupted (should fail and rollback)
        rollback_targets = [_TRANSFORMER_MODULE, _VAE_MODULE]
        result, status_code = self._update_weights(
            base_url,
            corrupted_vae_model_dir,
            target_modules=rollback_targets,
        )
        assert (
            status_code == 400
        ), f"Expected 400 on corrupted weights, got {status_code}"
        assert not result.get("success", True)
        message = result.get("message", "")
        assert "rolled back" in message.lower()
        # The updater reports the first failing module in the error message.
        # With ordered target_modules=[transformer, vae], this makes the
        # failure point explicit: transformer is processed first, then vae fails.
        assert (
            "Failed to update module 'vae'" in message
        ), f"Expected vae to be the explicit failure point, got: {message}"
        rolled_back_checksums = self._get_weights_checksum(base_url)

        # 1) transformer: server == perturbed != base
        transformer_base = base_checksums.get(_TRANSFORMER_MODULE)
        transformer_perturbed = perturbed_checksums.get(_TRANSFORMER_MODULE)
        transformer_rolled_back = rolled_back_checksums.get(_TRANSFORMER_MODULE)
        assert transformer_rolled_back == transformer_perturbed
        assert transformer_rolled_back != transformer_base

        # 2) vae: server == perturbed != base
        vae_base = base_checksums.get(_VAE_MODULE)
        vae_perturbed = perturbed_checksums.get(_VAE_MODULE)
        vae_rolled_back = rolled_back_checksums.get(_VAE_MODULE)
        assert vae_rolled_back == vae_perturbed
        assert vae_rolled_back != vae_base

        # 3) text encoder(s): server == base == perturbed
        for name in text_encoder_modules:
            assert rolled_back_checksums.get(name) == perturbed_checksums.get(
                name
            ), f"Text encoder module '{name}' should stay equal to perturbed"
            assert rolled_back_checksums.get(name) == base_checksums.get(
                name
            ), f"Text encoder module '{name}' should stay equal to base"


class TestUpdateWeightsFromDiskWithOffload(_UpdateWeightsApiMixin):
    """Test update_weights_from_disk with layerwise offload enabled."""

    @pytest.fixture(scope="class", params=_ACTIVE_MODEL_PAIRS, ids=_PAIR_IDS)
    def diffusion_server_with_offload(self, request):
        default_model, source_model = request.param
        port = get_dynamic_server_port()
        wait_deadline = float(os.environ.get("SGLANG_TEST_WAIT_SECS", "600"))

        local_source = maybe_download_model(source_model)
        perturbed_vae_model_dir = tempfile.mkdtemp(prefix="sglang_perturbed_vae_")

        clone_thread = threading.Thread(
            target=_clone_model_with_modified_module,
            args=(
                local_source,
                perturbed_vae_model_dir,
                _VAE_MODULE,
                _perturb_safetensor,
            ),
        )
        clone_thread.start()

        manager = ServerManager(
            model=default_model,
            port=port,
            wait_deadline=wait_deadline,
            extra_args="--num-gpus 1 --dit-layerwise-offload true",
        )

        ctx = manager.start()
        clone_thread.join()

        try:
            yield ctx, default_model, perturbed_vae_model_dir
        finally:
            ctx.cleanup()
            shutil.rmtree(perturbed_vae_model_dir, ignore_errors=True)

    def test_update_weights_with_offload_enabled(self, diffusion_server_with_offload):
        ctx, _, perturbed_model_dir = diffusion_server_with_offload
        base_url = f"http://localhost:{ctx.port}"

        result, status_code = self._update_weights(base_url, perturbed_model_dir)
        assert status_code == 200, f"Expected 200, got {status_code}"
        assert result.get("success", False), f"Update failed: {result.get('message')}"

        message = result.get("message", "")
        assert "Shape mismatch" not in message, f"Shape mismatch detected: {message}"

        self._assert_server_matches_model(base_url, perturbed_model_dir)


if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v", "-s"]))
