From 088497c7cbcd069cb81c81ca897e3fbdf03f8265 Mon Sep 17 00:00:00 2001 From: Hirotaka Ishihara <38371297+JerryIshihara@users.noreply.github.com> Date: Tue, 30 Jun 2026 23:01:57 +0900 Subject: [PATCH] Add GPU NMF CLI wiring tests --- tests/test_nmf_gpu.py | 134 +++++++++++++++++++++++++++++++++++++++++- tests/test_prepare.py | 50 ++++++++++++++++ 2 files changed, 183 insertions(+), 1 deletion(-) diff --git a/tests/test_nmf_gpu.py b/tests/test_nmf_gpu.py index 0c1b9d4..5aee6b1 100644 --- a/tests/test_nmf_gpu.py +++ b/tests/test_nmf_gpu.py @@ -42,6 +42,7 @@ """ from types import SimpleNamespace +import argparse import builtins import numpy as np @@ -57,7 +58,6 @@ small_nonnegative_matrix, ) - # --------------------------------------------------------------------- # Test harness # --------------------------------------------------------------------- @@ -507,6 +507,138 @@ def test_resolve_gpu_opts_floors_check_every_and_compile_block_to_at_least_one(k assert opts["compile_block"] == 1 +# --------------------------------------------------------------------- +# GPU CLI and cNMF engine wiring +# --------------------------------------------------------------------- +def test_parse_gpu_args_defaults_to_none_until_user_sets_engine(kernel): + """CLI flags should default to None so absent options are distinguishable from explicit values.""" + parser = argparse.ArgumentParser() + kernel.parse_gpu_args(parser) + + args = parser.parse_args([]) + + assert args.engine is None + for name in ("gpu_device", "gpu_dtype", "gpu_allow_tf32", "gpu_compile", + "gpu_eps", "gpu_check_every", "gpu_compile_block"): + assert getattr(args, name) is None, f"{name} should default to None" + + +def test_gpu_kwargs_from_args_rejects_gpu_options_without_gpu_engine(kernel): + """GPU-specific CLI options should raise unless `--engine gpu` was explicitly selected.""" + parser = argparse.ArgumentParser() + kernel.parse_gpu_args(parser) + + for argv in (["--gpu-device", "cuda"], ["--engine", "cpu", "--gpu-device", "cuda"]): + with pytest.raises(ValueError, match="require --engine gpu"): + kernel.gpu_kwargs_from_args(parser.parse_args(argv)) + + +def test_gpu_kwargs_from_args_fills_defaults_when_gpu_engine_selected(kernel): + """`--engine gpu` alone should resolve missing GPU options from DEFAULT_GPU.""" + parser = argparse.ArgumentParser() + kernel.parse_gpu_args(parser) + + args = parser.parse_args(["--engine", "gpu"]) + + assert kernel.gpu_kwargs_from_args(args) == kernel.DEFAULT_GPU + + +def test_gpu_kwargs_from_args_normalizes_cli_overrides(kernel): + """GPU CLI override values should normalize through the same resolver as config dict values.""" + parser = argparse.ArgumentParser() + kernel.parse_gpu_args(parser) + + args = parser.parse_args([ + "--engine", "gpu", + "--gpu-device", "CUDA:0", # device is lower-cased by _resolve_gpu_opts + "--gpu-dtype", "FP32", + "--gpu-allow-tf32", # store_const flags -> True + "--gpu-compile", + "--gpu-eps", "1e-8", + "--gpu-check-every", "5", + "--gpu-compile-block", "100", + ]) + + assert kernel.gpu_kwargs_from_args(args) == { + "device": "cuda:0", + "dtype": "fp32", + "allow_tf32": True, + "compile": True, + "eps": 1e-8, + "check_every": 5, + "compile_block": 100, + } + + +def test_validate_engine_args_for_command_rejects_non_factorize_gpu_options(kernel): + """Engine/GPU options should be accepted only for commands that support GPU factorization.""" + parser = argparse.ArgumentParser() + parser.add_argument("command") + kernel.parse_gpu_args(parser) + supported = ("factorize",) + + # factorize accepts engine/GPU options (no raise) + kernel.validate_engine_args_for_command(parser.parse_args(["factorize", "--engine", "gpu"]), supported) + + # a non-factorize command carrying engine/GPU options is rejected + for argv in (["prepare", "--engine", "gpu"], ["consensus", "--gpu-device", "cuda"]): + with pytest.raises(ValueError, match="only valid with"): + kernel.validate_engine_args_for_command(parser.parse_args(argv), supported) + + # a non-factorize command without engine/GPU options is fine + kernel.validate_engine_args_for_command(parser.parse_args(["prepare"]), supported) + + +def test_configure_nmf_engine_cpu_leaves_cnmf_instance_unchanged(kernel): + """The default CPU engine should be a no-op so existing sklearn behavior is preserved.""" + class DummyCNMF: + def _nmf(self, X, nmf_kwargs): + return "sklearn-path" + + obj = DummyCNMF() + result = kernel.configure_nmf_engine(obj, engine="cpu", gpu_kwargs={"device": "cuda"}) + + assert result is obj + assert "_nmf" not in vars(obj) # no instance override added + assert obj._nmf("X", {}) == "sklearn-path" # original class method intact + + +def test_configure_nmf_engine_rejects_unknown_engine(kernel): + """Unknown engine names should fail loudly instead of silently using CPU.""" + with pytest.raises(ValueError, match="engine must be 'cpu' or 'gpu'"): + kernel.configure_nmf_engine(object(), engine="tpu") + + +def test_configure_nmf_engine_gpu_patches_instance_nmf_with_adapter(kernel, monkeypatch): + """The GPU engine should replace the instance `_nmf` callable with the GPU adapter path.""" + captured = {} + + def fake_nmf_gpu(self, X, nmf_kwargs): + captured["self"], captured["X"], captured["nmf_kwargs"] = self, X, dict(nmf_kwargs) + return ("spectra", "usages") + + monkeypatch.setattr(kernel, "_nmf_gpu", fake_nmf_gpu) + + class DummyCNMF: + def _nmf(self, X, nmf_kwargs): + return "sklearn-path" + + obj = DummyCNMF() + gpu_kwargs = {"device": "cuda", "dtype": "fp32"} + result = kernel.configure_nmf_engine(obj, engine="gpu", gpu_kwargs=gpu_kwargs) + + assert result is obj + assert "_nmf" in vars(obj) # instance _nmf is now overridden + + out = obj._nmf("Xdata", {"n_components": 5}) + + assert out == ("spectra", "usages") # dispatched through the GPU adapter + assert captured["self"] is obj and captured["X"] == "Xdata" + assert captured["nmf_kwargs"]["engine"] == "gpu" # engine + gpu kwargs embedded first + assert captured["nmf_kwargs"]["gpu"] == gpu_kwargs + assert captured["nmf_kwargs"]["n_components"] == 5 # original kwargs preserved + + # --------------------------------------------------------------------- # Device selection # --------------------------------------------------------------------- diff --git a/tests/test_prepare.py b/tests/test_prepare.py index dcca8e8..2cbb020 100644 --- a/tests/test_prepare.py +++ b/tests/test_prepare.py @@ -58,6 +58,56 @@ def generate_counts_file(tmp_path, file_format, dtype=np.int64, zero_count=False return str(counts_fn) + +# --------------------------------------------------------------------- +# GPU factorize wiring integration +# --------------------------------------------------------------------- +def test_get_nmf_iter_params_default_cpu_engine_does_not_change_sklearn_kwargs(mock_cnmf): + """Default CPU runs should not add engine/gpu keys to sklearn NMF kwargs.""" + _replicate_params, run_params = mock_cnmf.get_nmf_iter_params(ks=[5, 7], n_iter=3, random_state_seed=14) + + # GPU wiring must not leak into the sklearn factorization kwargs. + assert "engine" not in run_params + assert "gpu" not in run_params + # Only the existing cNMF/sklearn factorization keys are present. + assert set(run_params) == {"alpha_W", "alpha_H", "l1_ratio", "beta_loss", "solver", "tol", "max_iter", "init"} + + +def test_factorize_gpu_engine_passes_seed_components_run_params_and_gpu_kwargs(mock_cnmf, monkeypatch, tmp_path): + """cNMF factorize should pass n_components, random_state, run params, and GPU kwargs to `_nmf_gpu`.""" + import cnmf.nmf_gpu as gpu_mod + + counts_fn = generate_counts_file(tmp_path, "txt", np.int64) + mock_cnmf.prepare(counts_fn, components=[5], n_iter=2, densify=True, seed=14) + + captured = [] + + def fake_nmf_gpu(self, X, nmf_kwargs): + captured.append(dict(nmf_kwargs)) + k = nmf_kwargs["n_components"] + return np.zeros((k, X.shape[1])), np.zeros((X.shape[0], k)) # (spectra, usages) + + monkeypatch.setattr(gpu_mod, "_nmf_gpu", fake_nmf_gpu) + gpu_kwargs = {"device": "cpu", "dtype": "fp64"} + gpu_mod.configure_nmf_engine(mock_cnmf, engine="gpu", gpu_kwargs=gpu_kwargs) + + mock_cnmf.factorize(worker_i=0, total_workers=1) + + assert len(captured) == 2 # n_iter=2 replicates for k=5 + replicate_params = load_df_from_npz(mock_cnmf.paths["nmf_replicate_parameters"]) + expected_seeds = set(replicate_params["nmf_seed"]) + observed_seeds = set() + for kw in captured: + assert kw["n_components"] == 5 # set per replicate by factorize + assert kw["engine"] == "gpu" # adapter embedded the engine + assert kw["gpu"] == gpu_kwargs # ...and the resolved GPU kwargs + assert "beta_loss" in kw and "init" in kw # original run params forwarded + observed_seeds.add(kw["random_state"]) + assert observed_seeds == expected_seeds # exact seeds from prepared replicate params + for iter_i in replicate_params["iter"]: + assert os.path.exists(mock_cnmf.paths["iter_spectra"] % (5, iter_i)) + + @pytest.mark.parametrize("file_format", ["txt", "npz", "h5ad"]) @pytest.mark.parametrize("dtype", [np.int64, np.float32, np.float64]) @pytest.mark.parametrize("densify", [True, False])