From 32e92562ebb2879221c0b45038891af6f02f6f5f Mon Sep 17 00:00:00 2001 From: yoshikisd Date: Sun, 9 Nov 2025 03:53:42 +0000 Subject: [PATCH] Added multi-GPU pytests --- tests/conftest.py | 25 ++- .../multi_gpu_script_plot_and_save.py | 85 +++++++++ tests/multi_gpu/test_multi_gpu.py | 171 ++++++++++++++++++ 3 files changed, 277 insertions(+), 4 deletions(-) create mode 100644 tests/multi_gpu/multi_gpu_script_plot_and_save.py create mode 100644 tests/multi_gpu/test_multi_gpu.py diff --git a/tests/conftest.py b/tests/conftest.py index f0faea5..9de16ca 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -32,21 +32,32 @@ def pytest_addoption(parser): default=False, help="run slow tests, primarily full reconstruction tests." ) + parser.addoption( + "--runmultigpu", + action="store_true", + default=False, + help="Runs tests using 2 NVIDIA CUDA GPUs." + ) def pytest_configure(config): config.addinivalue_line("markers", "slow: mark test as slow to run") + config.addinivalue_line("markers", "multigpu: run the multigpu test using 2 NVIDIA GPUs") def pytest_collection_modifyitems(config, items): - if config.getoption("--runslow"): - # --runslow given in cli: do not skip slow tests - return + # Skip the slow and/or multigpu tests if --runslow and/or --multigpu + # is given in cli. skip_slow = pytest.mark.skip(reason="need --runslow option to run") + skip_multigpu = pytest.mark.skip(reason='need --runmultigpu option to run') + for item in items: - if "slow" in item.keywords: + if "slow" in item.keywords and not config.getoption("--runslow"): item.add_marker(skip_slow) + if "multigpu" in item.keywords and not config.getoption("--runmultigpu"): + item.add_marker(skip_multigpu) + @pytest.fixture def reconstruction_device(request): @@ -415,3 +426,9 @@ def example_nested_dicts(pytestconfig): } return [test_dict_1, test_dict_2, test_dict_3] + + +@pytest.fixture(scope='module') +def multigpu_script(pytestconfig): + return str(pytestconfig.rootpath) + \ + '/tests/multi_gpu/multi_gpu_script_plot_and_save.py' diff --git a/tests/multi_gpu/multi_gpu_script_plot_and_save.py b/tests/multi_gpu/multi_gpu_script_plot_and_save.py new file mode 100644 index 0000000..3f90f61 --- /dev/null +++ b/tests/multi_gpu/multi_gpu_script_plot_and_save.py @@ -0,0 +1,85 @@ +import cdtools +from cdtools.tools import multigpu +import os +from matplotlib import pyplot as plt + +rank = multigpu.get_rank() +world_size = multigpu.get_world_size() +cdtools.tools.multigpu.setup(rank=rank, world_size=world_size) + +filename = os.environ.get('CDTOOLS_TESTING_DATA_PATH') +savedir = os.environ.get('CDTOOLS_TESTING_TMP_PATH') +SHOW_PLOT = bool(int(os.environ.get('CDTOOLS_TESTING_SHOW_PLOT'))) + +print('DONT CLOSE ANY OF THE FIGURES OR THE TEST WILL FAIL!') +dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) + +model = cdtools.models.FancyPtycho.from_dataset( + dataset, + n_modes=3, + oversampling=2, + probe_support_radius=120, + propagation_distance=5e-3, + units='mm', + obj_view_crop=-50, +) + +device = 'cuda' +model.to(device=device) +dataset.get_as(device=device) + +# Test Ptycho2DDataset.inspect +if SHOW_PLOT: + dataset.inspect() + +# Test Ptycho2DDataset.to_cxi +filename_to_cxi = os.path.join(savedir, + f'RANK_{model.rank}_test_to_cxi.h5') +dataset.to_cxi(filename_to_cxi) + +# Test CDIModel.save_to_h5 +filename_save_to_h5 = os.path.join(savedir, + f'RANK_{model.rank}_test_save_to.h5') +model.save_to_h5(filename_save_to_h5, dataset) + +# Test CDIModel.save_on_exit(), CDIModel.inspect() +filename_save_on_exit = os.path.join(savedir, + f'RANK_{model.rank}_test_save_on_exit.h5') + +with model.save_on_exit(filename_save_on_exit, dataset): + for loss in model.Adam_optimize(5, dataset, lr=0.02, batch_size=40): + if rank == 0: + print(model.report()) + if SHOW_PLOT: + model.inspect(dataset) + +if SHOW_PLOT: + # Test CDIModel.compare(dataset) + model.compare(dataset) + + # Test CDIModel.save_figures() + filename_save_figures = os.path.join(savedir, + f'RANK_{model.rank}_test_plot_') + model.save_figures(prefix=filename_save_figures, + extension='.png') + + plt.close('all') + +# Test CDIModel.save_checkpoint +filename_save_checkpoint = \ + os.path.join(savedir, f'RANK_{model.rank}_test_save_checkpoint.pt') +model.save_checkpoint(dataset, checkpoint_file=filename_save_checkpoint) + +# Test CDIModel.save_on_exception() +filename_save_on_except = \ + os.path.join(savedir, f'RANK_{model.rank}_test_save_on_except.h5') + +with model.save_on_exception(filename_save_on_except, dataset): + for loss in model.Adam_optimize(10, dataset, lr=0.02, batch_size=40): + if rank == 0 and model.epoch <= 10: + print(model.report()) + elif model.epoch > 10: + raise Exception('This is a deliberate exception raised to ' + + 'test save on exception') + +cdtools.tools.multigpu.cleanup() diff --git a/tests/multi_gpu/test_multi_gpu.py b/tests/multi_gpu/test_multi_gpu.py new file mode 100644 index 0000000..57be53b --- /dev/null +++ b/tests/multi_gpu/test_multi_gpu.py @@ -0,0 +1,171 @@ +import cdtools +from cdtools.tools import multigpu +import pytest +import os +import subprocess +import torch as t + +""" +This file contains several tests that are relevant to running multi-GPU +operations in CDTools. +""" + + +def reconstruct(rank, world_size, conn): + """ + An example reconstruction script to test the performance of 1 vs 2 GPU + operation. + """ + filename = os.environ.get('CDTOOLS_TESTING_GOLD_BALL_PATH') + cdtools.tools.multigpu.setup(rank=rank, + world_size=world_size) + dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) + + pad = 10 + dataset.pad(pad) + dataset.inspect() + model = cdtools.models.FancyPtycho.from_dataset( + dataset, + n_modes=3, + probe_support_radius=50, + propagation_distance=2e-6, + units='um', + probe_fourier_crop=pad + ) + model.translation_offsets.data += 0.7 * \ + t.randn_like(model.translation_offsets) + model.weights.requires_grad = False + + device = 'cuda' + model.to(device=device) + dataset.get_as(device=device) + + recon = cdtools.reconstructors.AdamReconstructor(model, + dataset, + rank=rank, + world_size=world_size) + + for loss in recon.optimize(10, lr=0.005, batch_size=50): + if rank == 0 and model.epoch == 10: + print(model.report()) + conn.send((model.loss_times, model.loss_history)) + cdtools.tools.multigpu.cleanup() + + +@pytest.mark.multigpu +def test_plotting_saving_torchrun(lab_ptycho_cxi, + multigpu_script, + tmp_path, + show_plot): + """ + Run a multi-GPU test via torchrun on a script that executes several + plotting and file-saving methods from CDIModel and ensure they run + without failure. + + Also, make sure that only 1 GPU is generating the plots. + + If this test fails, one of three things happened: + 1) Either something failed while multigpu_script_2 was called + 2) Somehow, something aside from Rank 0 saved results + 3) multigpu_script_2 was not able to save all the data files + we asked it to save. + """ + # Run the test script, which generates several files that either have + # the prefix + cmd = ['torchrun', + '--standalone', + '--nnodes=1', + '--nproc_per_node=2', + multigpu_script] + + child_env = os.environ.copy() + child_env['CDTOOLS_TESTING_DATA_PATH'] = lab_ptycho_cxi + child_env['CDTOOLS_TESTING_TMP_PATH'] = str(tmp_path) + child_env['CDTOOLS_TESTING_SHOW_PLOT'] = str(int(show_plot)) + + try: + subprocess.run(cmd, check=True, env=child_env) + except subprocess.CalledProcessError: + # The called script is designed to throw an exception. + # TODO: Figure out how to distinguish between the engineered error + # in the script versus any other error. + pass + + # Check if all the generated file names only have the prefix 'RANK_0' + filelist = [f for f in os.listdir(tmp_path) + if os.path.isfile(os.path.join(tmp_path, f))] + + assert all([file.startswith('RANK_0') for file in filelist]) + print('All files have the RANK_0 prefix.') + + # Check if plots have been saved + if show_plot: + print('Plots generated: ' + + f"{sum([file.startswith('RANK_0_test_plot') for file in filelist])}") # noqa + assert any([file.startswith('RANK_0_test_plot') for file in filelist]) + else: + print('--plot not enabled. Checks on plotting and figure saving' + + ' will not be conducted.') + + # Check if we have all five data files saved + file_output_suffix = ('test_save_checkpoint.pt', + 'test_save_on_exit.h5', + 'test_save_on_except.h5', + 'test_save_to.h5', + 'test_to_cxi.h5') + + print(f'{sum([file.endswith(file_output_suffix) for file in filelist])}' + + ' out of 5 data files have been generated.') + assert sum([file.endswith(file_output_suffix) for file in filelist]) \ + == len(file_output_suffix) + + +@pytest.mark.multigpu +def test_reconstruction_quality_spawn(gold_ball_cxi, + show_plot): + """ + Run a multi-GPU speed test based on gold_ball_ptycho_speedtest.py + and make sure the final reconstructed loss using 2 GPUs is similar + to 1 GPU. + + This test requires us to have 2 NVIDIA GPUs and makes use of the + multi-GPU speed test. + + If this test fails, it indicates that the reconstruction quality is + getting noticably worse with increased GPU counts. This may be a symptom + of a synchronization/broadcasting issue between the different GPUs. + """ + # Make the gold_ball_cxi file path visible to the reconstruct function + os.environ['CDTOOLS_TESTING_GOLD_BALL_PATH'] = gold_ball_cxi + + loss_mean_list, loss_std_list, \ + _, _, speed_up_mean_list, speed_up_std_list\ + = multigpu.run_speed_test(fn=reconstruct, + gpu_counts=(1, 2), + runs=3, + show_plot=show_plot) + + # Make sure that the final loss values between the 1 and 2 GPU tests + # are comprable to within 1 std of each other. + single_gpu_loss_mean = loss_mean_list[0][-1] + single_gpu_loss_std = loss_std_list[0][-1] + double_gpu_loss_mean = loss_mean_list[1][-1] + double_gpu_loss_std = loss_std_list[1][-1] + + single_gpu_loss_min = single_gpu_loss_mean - single_gpu_loss_std + single_gpu_loss_max = single_gpu_loss_mean + single_gpu_loss_std + multi_gpu_loss_min = double_gpu_loss_mean - double_gpu_loss_std + multi_gpu_loss_max = double_gpu_loss_mean + double_gpu_loss_std + + has_loss_overlap = \ + min(single_gpu_loss_max, multi_gpu_loss_max)\ + > max(single_gpu_loss_min, multi_gpu_loss_min) + + assert has_loss_overlap + + # Make sure the loss mean falls below 3.2e-4. The values of losses I + # recorded at the time of testing were <3.19 e-4. + assert double_gpu_loss_mean < 3.2e-4 + + # Make sure that we have some speed up... + assert speed_up_mean_list[0] < speed_up_mean_list[1]