debug: print num_threads per worker

open-atmos · Mar 17, 2024 · 87e696b · 87e696b
1 parent e32a7e1
commit 87e696b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 13 deletions.
diff --git a/.github/workflows/tests+pypi.yml b/.github/workflows/tests+pypi.yml
@@ -141,7 +141,7 @@ jobs:
             export COV_ARGS="--cov=./ --cov-report=xml"
             pip install pytest-cov
           fi
-          NUMBA_NUM_THREADS=3 mpiexec $_mpiexec_args -n ${{ matrix.mpi-np }} pytest $COV_ARGS --timeout=600 --timeout_method=thread -s -vv -We tests/local;
+          mpiexec $_mpiexec_args -n ${{ matrix.mpi-np }} pytest $COV_ARGS --timeout=600 --timeout_method=thread -s -vv -We tests/local;
       - uses: actions/upload-artifact@v2
         with:
           name: plots

diff --git a/tests/local/contract_tests/test_single_vs_multi_node.py b/tests/local/contract_tests/test_single_vs_multi_node.py
@@ -61,6 +61,7 @@ def test_single_vs_multi_node(  # pylint: disable=too-many-arguments,too-many-br
     (which is simulation performed on single node environment)
 
     """
+    print("numba.NUMBA_NUM_THREADS", numba.config.NUMBA_NUM_THREADS)
     # pylint: disable=too-many-locals
     if scenario_class is SphericalScenario and options_kwargs["n_iters"] > 1:
         pytest.skip("TODO #56")
@@ -149,18 +150,15 @@ def test_single_vs_multi_node(  # pylint: disable=too-many-arguments,too-many-br
 
                 # plot
                 if plot:
-                    if MPI_DIM != INNER:
-                        print("Plotting supports only INNER dimension")
-                    else:
-                        tmp = np.empty_like(dataset[:, :, -1])
-                        for i, _ in enumerate(output_steps):
-                            tmp[:] = np.nan
-                            tmp[:, mpi_range] = dataset[:, mpi_range, i]
-                            simulation.quick_look(tmp)
-                            filename = f"step={i:04d}.svg"
-                            pyplot.savefig(plot_path / filename)
-                            print("Saving figure")
-                            pyplot.close()
+                    tmp = np.empty_like(dataset[:, :, -1])
+                    for i, _ in enumerate(output_steps):
+                        tmp[:] = np.nan
+                        tmp[:, mpi_range] = dataset[:, mpi_range, i]
+                        simulation.quick_look(tmp)
+                        filename = f"step={i:04d}.svg"
+                        pyplot.savefig(plot_path / filename)
+                        print("Saving figure")
+                        pyplot.close()
 
     # assert
     with barrier_enclosed():