Skip to content

test: fix tests not working with high nranks and change GA to test that #150

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

test: fix tests not working with high nranks and change GA to test that #150

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
aee9add
test: fix tests not working with high nranks and change GA to test that
mrava87 Jul 1, 2025
12d1708
ci: exclude rank=9 for openmpi as it does not run in CI (works locally)
mrava87 Jul 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion .github/workflows/build.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,12 @@ jobs:
os: [ubuntu-latest, macos-latest]
python-version: ['3.10', '3.11', '3.12', '3.13']
mpi: ['mpich', 'openmpi', 'intelmpi']
rank: ['2', '3', '4']
rank: ['2', '4', '9']
exclude:
- os: macos-latest
mpi: 'intelmpi'
- mpi: 'openmpi'
rank: '9' # works locally, not in CI
runs-on: ${{ matrix.os }}
steps:
- name: Checkout
Expand Down
2 changes: 1 addition & 1 deletion Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
PIP := $(shell command -v pip3 2> /dev/null || command which pip 2> /dev/null)
PYTHON := $(shell command -v python3 2> /dev/null || command which python 2> /dev/null)
NUM_PROCESSES = 3
NUM_PROCESSES = 4

.PHONY: install dev-install dev-install_nccl install_ \
conda install_conda_nccl dev-install_conda dev-install_conda_nccl \
Expand Down
45 changes: 36 additions & 9 deletions tests/test_distributedarray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,12 +199,21 @@ def test_distributed_norm(par):
def test_distributed_masked(par):
"""Test Asarray with masked array"""
# Number of subcommunicators
if MPI.COMM_WORLD.Get_size() % 2 == 0:
size = MPI.COMM_WORLD.Get_size()

# Exclude not handled cases
shape_axis = par['x'].shape[par['axis']]
print('shape_axis, size', shape_axis, size, shape_axis % size != 0)
if shape_axis % size != 0:
pytest.skip(f"Array dimension to distributed ({shape_axis}) is not "
f"divisible by the number of processes ({size})...")
if size % 2 == 0:
nsub = 2
elif MPI.COMM_WORLD.Get_size() % 3 == 0:
elif size % 3 == 0:
nsub = 3
else:
pass
pytest.skip(f"Number of processes ({size}) is not divisible "
"by 2 or 3...")
subsize = max(1, MPI.COMM_WORLD.Get_size() // nsub)
mask = np.repeat(np.arange(nsub), subsize)

Expand Down Expand Up @@ -236,12 +245,21 @@ def test_distributed_masked(par):
def test_distributed_maskeddot(par1, par2):
"""Test Distributed Dot product with masked array"""
# Number of subcommunicators
if MPI.COMM_WORLD.Get_size() % 2 == 0:
size = MPI.COMM_WORLD.Get_size()

# Exclude not handled cases
shape_axis = par1['x'].shape[par1['axis']]
print('shape_axis, size', shape_axis, size, shape_axis % size != 0)
if shape_axis % size != 0:
pytest.skip(f"Array dimension to distributed ({shape_axis}) is not "
f"divisible by the number of processes ({size})...")
if size % 2 == 0:
nsub = 2
elif MPI.COMM_WORLD.Get_size() % 3 == 0:
elif size % 3 == 0:
nsub = 3
else:
pass
pytest.skip(f"Number of processes ({size}) is not divisible "
"by 2 or 3...")
subsize = max(1, MPI.COMM_WORLD.Get_size() // nsub)
mask = np.repeat(np.arange(nsub), subsize)

Expand Down Expand Up @@ -271,12 +289,21 @@ def test_distributed_maskeddot(par1, par2):
def test_distributed_maskednorm(par):
"""Test Distributed numpy.linalg.norm method with masked array"""
# Number of subcommunicators
if MPI.COMM_WORLD.Get_size() % 2 == 0:
size = MPI.COMM_WORLD.Get_size()

# Exclude not handled cases
shape_axis = par['x'].shape[par['axis']]
print('shape_axis, size', shape_axis, size, shape_axis % size != 0)
if shape_axis % size != 0:
pytest.skip(f"Array dimension to distributed ({shape_axis}) is not "
f"divisible by the number of processes ({size})...")
if size % 2 == 0:
nsub = 2
elif MPI.COMM_WORLD.Get_size() % 3 == 0:
elif size % 3 == 0:
nsub = 3
else:
pass
pytest.skip(f"Number of processes ({size}) is not divisible "
"by 2 or 3...")
subsize = max(1, MPI.COMM_WORLD.Get_size() // nsub)
mask = np.repeat(np.arange(nsub), subsize)
# Replicate x as required in masked arrays
Expand Down
12 changes: 6 additions & 6 deletions tests/test_fredholm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
size = MPI.COMM_WORLD.Get_size()

par1 = {
"nsl": 12,
"nsl": 21,
"ny": 6,
"nx": 4,
"nz": 5,
Expand All @@ -30,7 +30,7 @@
"dtype": "float32",
} # real, saved Gt
par2 = {
"nsl": 12,
"nsl": 21,
"ny": 6,
"nx": 4,
"nz": 5,
Expand All @@ -40,7 +40,7 @@
"dtype": "float32",
} # real, unsaved Gt
par3 = {
"nsl": 12,
"nsl": 21,
"ny": 6,
"nx": 4,
"nz": 5,
Expand All @@ -50,7 +50,7 @@
"dtype": "complex64",
} # complex, saved Gt
par4 = {
"nsl": 12,
"nsl": 21,
"ny": 6,
"nx": 4,
"nz": 5,
Expand All @@ -60,7 +60,7 @@
"dtype": "complex64",
} # complex, unsaved Gt
par5 = {
"nsl": 12,
"nsl": 21,
"ny": 6,
"nx": 4,
"nz": 1,
Expand All @@ -70,7 +70,7 @@
"dtype": "float32",
} # real, saved Gt, nz=1
par6 = {
"nsl": 12,
"nsl": 21,
"ny": 6,
"nx": 4,
"nz": 1,
Expand Down
13 changes: 12 additions & 1 deletion tests/test_solver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
StackedDistributedArray
)

np.random.seed(42)
size = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()

Expand Down Expand Up @@ -94,6 +93,8 @@
)
def test_cg(par):
"""CG with MPIBlockDiag"""
np.random.seed(42)

A = np.ones((par["ny"], par["nx"])) + par[
"imag"] * np.ones((par["ny"], par["nx"]))
Aop = MatrixMult(np.conj(A.T) @ A, dtype=par['dtype'])
Expand Down Expand Up @@ -139,6 +140,8 @@ def test_cg(par):
)
def test_cgls(par):
"""CGLS with MPIBlockDiag"""
np.random.seed(42)

A = np.ones((par["ny"], par["nx"])) + par[
"imag"] * np.ones((par["ny"], par["nx"]))
Aop = MatrixMult(np.conj(A.T) @ A + 1e-5 * np.eye(par["nx"], dtype=par['dtype']),
Expand Down Expand Up @@ -186,6 +189,8 @@ def test_cgls(par):
)
def test_cgls_broadcastdata(par):
"""CGLS with broadcasted data vector"""
np.random.seed(42)

A = (rank + 1) * np.ones((par["ny"], par["nx"])) + (rank + 2) * par[
"imag"
] * np.ones((par["ny"], par["nx"]))
Expand Down Expand Up @@ -232,6 +237,8 @@ def test_cgls_broadcastdata(par):
)
def test_cgls_broadcastmodel(par):
"""CGLS with broadcasted model vector"""
np.random.seed(42)

A = np.ones((par["ny"], par["nx"])) + par[
"imag"] * np.ones((par["ny"], par["nx"]))
Aop = MatrixMult(np.conj(A.T) @ A + 1e-5 * np.eye(par["nx"], dtype=par['dtype']),
Expand Down Expand Up @@ -281,6 +288,8 @@ def test_cgls_broadcastmodel(par):
)
def test_cg_stacked(par):
"""CG with MPIStackedBlockDiag"""
np.random.seed(42)

A = np.ones((par["ny"], par["nx"])) + par[
"imag"] * np.ones((par["ny"], par["nx"]))
Aop = MatrixMult(np.conj(A.T) @ A + 1e-5 * np.eye(par["nx"], dtype=par['dtype']),
Expand Down Expand Up @@ -344,6 +353,8 @@ def test_cg_stacked(par):
)
def test_cgls_stacked(par):
"""CGLS with MPIStackedBlockDiag"""
np.random.seed(42)

A = np.ones((par["ny"], par["nx"])) + par[
"imag"] * np.ones((par["ny"], par["nx"]))
Aop = MatrixMult(np.conj(A.T) @ A + 1e-5 * np.eye(par["nx"], dtype=par['dtype']),
Expand Down