Spaces:
Running on Zero
Running on Zero
| import contextlib | |
| import contextvars | |
| import numpy as np | |
| import pickle | |
| import sys | |
| import types | |
| import weakref | |
| import pytest | |
| import torch | |
| import compilation | |
| def test_package_transfer_drops_live_cuda_runner(): | |
| packaged = compilation.PackagedFP16("/runtime/model.pt2") | |
| # A native runner cannot be pickled and must never escape its GPU worker. | |
| packaged.loaded = lambda value: value | |
| restored = pickle.loads(pickle.dumps(packaged)) | |
| assert restored.package_path == "/runtime/model.pt2" | |
| assert restored.loaded is None | |
| def test_package_loads_once_across_zerogpu_requests(monkeypatch): | |
| loads = [] | |
| def load(path): | |
| loads.append(path) | |
| return lambda value: value + 1 | |
| monkeypatch.setattr(torch._inductor, "aoti_load_package", load) | |
| # The real helper reads Linux process maps; it cannot import on Windows. | |
| cleanup = types.SimpleNamespace(_register_aoti_cleanup=contextlib.nullcontext) | |
| monkeypatch.setitem(sys.modules, "spaces.zero.torch.aoti", cleanup) | |
| packaged = compilation.PackagedFP16("/runtime/model.pt2") | |
| # ZeroGPU workers run every request in a fresh copy of the worker context. | |
| results = [contextvars.copy_context().run(packaged, 1) for _ in range(3)] | |
| assert results == [2, 2, 2] | |
| assert loads == ["/runtime/model.pt2"] | |
| def test_disabled_compilation_preserves_working_model(monkeypatch): | |
| monkeypatch.setenv("VRMBG_COMPILE", "0") | |
| model = torch.nn.Linear(2, 1) | |
| expected = model(torch.ones(1, 2)).detach().clone() | |
| assert compilation.enable_compiled_model(model) is False | |
| torch.testing.assert_close(model(torch.ones(1, 2)), expected) | |
| def test_compile_failure_preserves_eager_model(monkeypatch, caplog): | |
| monkeypatch.setenv("VRMBG_COMPILE", "1") | |
| monkeypatch.setenv("SPACES_ZERO_GPU", "true") | |
| monkeypatch.setattr(compilation.spaces, "GPU", lambda **kw: lambda fn: fn) | |
| def fail(model): | |
| raise RuntimeError("Compiler unavailable") | |
| monkeypatch.setattr(compilation, "compile_fp16", fail) | |
| model = torch.nn.Linear(2, 1) | |
| model.register_buffer("integer_indices", torch.arange(4)) | |
| expected = model(torch.ones(1, 2)).detach().clone() | |
| assert compilation.enable_compiled_model(model) is False | |
| torch.testing.assert_close(model(torch.ones(1, 2)), expected) | |
| assert "using eager FP16" in caplog.text | |
| def test_zerogpu_startup_enables_compiled_forward(monkeypatch): | |
| monkeypatch.delenv("VRMBG_COMPILE", raising=False) | |
| monkeypatch.setenv("SPACES_ZERO_GPU", "true") | |
| monkeypatch.setattr(compilation.spaces, "GPU", lambda **kw: lambda fn: fn) | |
| model = torch.nn.Linear(2, 1) | |
| model.register_buffer("integer_indices", torch.arange(4)) | |
| def compiled(value): | |
| return value.sum(dim=1, keepdim=True) | |
| monkeypatch.setattr(compilation, "compile_fp16", lambda model: compiled) | |
| old_parameter = weakref.ref(model.weight) | |
| assert compilation.enable_compiled_model(model) is True | |
| torch.testing.assert_close(model(torch.ones(1, 2)), torch.tensor([[2.0]])) | |
| assert all(tensor.numel() == 0 for tensor in model.state_dict().values()) | |
| assert old_parameter() is None | |
| def test_apply_failure_is_not_reported_as_working_eager_fallback(monkeypatch): | |
| monkeypatch.setenv("VRMBG_COMPILE", "1") | |
| monkeypatch.setattr(compilation.spaces, "GPU", lambda **kw: lambda fn: fn) | |
| monkeypatch.setattr(compilation, "compile_fp16", lambda model: object()) | |
| def fail_apply(candidate, model): | |
| model.weight = None | |
| raise RuntimeError("Apply failed after mutation") | |
| monkeypatch.setattr(compilation, "apply_compiled_model", fail_apply) | |
| with pytest.raises(RuntimeError, match="Apply failed"): | |
| compilation.enable_compiled_model(torch.nn.Linear(2, 1)) | |
| def test_local_compilation_errors_require_a_fresh_process(monkeypatch): | |
| monkeypatch.setenv("VRMBG_COMPILE", "1") | |
| monkeypatch.delenv("SPACES_ZERO_GPU", raising=False) | |
| monkeypatch.setattr(compilation.spaces, "GPU", lambda **kw: lambda fn: fn) | |
| def fail(model): | |
| raise RuntimeError("CUDA failure") | |
| monkeypatch.setattr(compilation, "compile_fp16", fail) | |
| with pytest.raises(RuntimeError, match="CUDA failure"): | |
| compilation.enable_compiled_model(torch.nn.Linear(2, 1)) | |
| def test_quality_check_rejects_invalid_or_changed_mattes(bad): | |
| reference = np.full((3, 8, 8), 0.5, dtype=np.float32) | |
| with pytest.raises(RuntimeError): | |
| compilation.check_mattes(reference, np.full_like(reference, bad)) | |
| def test_quality_check_rejects_shape_changes(): | |
| with pytest.raises(RuntimeError): | |
| compilation.check_mattes(np.zeros((3, 8, 8)), np.zeros((1, 8, 8))) | |
| def test_quality_check_accepts_small_fp16_rounding(): | |
| reference = np.full((3, 8, 8), 0.5, dtype=np.float32) | |
| compilation.check_mattes(reference, reference + 0.0001) | |