Repository: RadioAstronomySoftwareGroup/pyuvsim

Issue:
Python exceptions are not raised in MPI
When run within MPI on Oscar, if a Python exception is raised (by any rank) the stack trace is not printed. The only information given is that MPI_Abort() was called. I've resorted to inserting "print" statements to track down bugs...

This may just be an Oscar/SLURM issue.

Discussion:
coveralls: ## Pull Request Test Coverage Report for [Build 1508](https://coveralls.io/builds/26755043)

* **-2** of **43**   **(95.35%)**  changed or added relevant lines in **4** files are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage decreased (**-0.06%**) to **95.575%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [pyuvsim/mpi.py](https://coveralls.io/builds/26755043/source?filename=pyuvsim%2Fmpi.py#L215) | 8 | 10 | 80.0%
<!-- | **Total:** | **41** | **43** | **95.35%** | -->


|  Totals | [![Coverage Status](https://coveralls.io/builds/26755043/badge)](https://coveralls.io/builds/26755043) |
| :-- | --: |
| Change from base [Build 1490](https://coveralls.io/builds/26477601): |  -0.06% |
| Covered Lines: | 1771 |
| Relevant Lines: | 1853 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 207edc0..274d505 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+
+## Added
+- A function for checking the memory usage on each Node
+
 ### Changed
 - Replaced init_uvdata_out function with complete_uvdata
 - init_uvdata_out function is more modular.
@@ -10,6 +14,8 @@
 - Added functions to read healpix maps and added support for frequency axis
 
 ### Fixed
+- skymodel splitting was not working correctly to avoid running out of memory.
+- Flush stderr in excepthook, so stack trace is printed when an exception is raised in MPI.
 - No longer calling UVBeam.interp with freq_interp_kind, for updated pyuvdata.
 - Circular-installation bug
 - Bug in how integration_time array is set in setup
diff --git a/environment.yml b/environment.yml
index 1909c52..295e3f0 100644
--- a/environment.yml
+++ b/environment.yml
@@ -9,7 +9,7 @@ dependencies:
   - scipy>=1.0.1
   - PyYAML>=5.1
   - mpi4py>=3.0.0
-  - astropy>=2.0
+  - astropy>=2.0.16
   - pip
   - psutil
   - line_profiler
diff --git a/pyuvsim/mpi.py b/pyuvsim/mpi.py
index 4a7d70a..7187ee4 100644
--- a/pyuvsim/mpi.py
+++ b/pyuvsim/mpi.py
@@ -7,6 +7,7 @@ from __future__ import absolute_import, division, print_function
 import sys
 from array import array
 from threading import Thread
+import resource
 
 import mpi4py
 import numpy as np
@@ -16,7 +17,6 @@ from mpi4py import MPI
 
 rank = 0  # COMM_WORLD rank
 Npus = 1
-Npus_node = 1
 world_comm = None
 node_comm = None
 rank_comm = None
@@ -27,15 +27,23 @@ def set_mpi_excepthook(mpi_comm):
 
     def mpi_excepthook(exctype, value, traceback):  # pragma: no cover
         sys.__excepthook__(exctype, value, traceback)
+        sys.stderr.flush()
         mpi_comm.Abort(1)
 
     sys.excepthook = mpi_excepthook
 
 
-def start_mpi():
+def start_mpi(block_nonroot_stdout=True):
     """
     Check if MPI has already been initialized. If so, just set the communicators,
     Npus, and rank variables.
+
+    Parameters
+    ----------
+
+    block_nonroot_stdout : bool (True)
+        Redirect stdout on nonzero ranks to /dev/null, for cleaner output.
+
     """
     global world_comm, node_comm, rank_comm, rank, Npus
     if not MPI.Is_initialized():
@@ -51,10 +59,9 @@ def start_mpi():
 
     world_comm.Barrier()
 
-    if not rank == 0:  # pragma: no cover
+    if (not rank == 0) and block_nonroot_stdout:  # pragma: no cover
         # For non-root ranks, do not print to stdout.
         # (Uncovered until we have multi-rank tests)
-        global stdout
         sys.stdout = open('/dev/null', 'w')
 
 
@@ -97,7 +104,7 @@ def shared_mem_bcast(arr, root=0):
 
     if node_comm.rank == root:
         # Now fill the window on each node with the data.
-        sh_arr[:] = arr
+        sh_arr[:] = arr[()]
 
     sh_arr.flags['WRITEABLE'] = False  # Do not want ranks overwriting the data.
 
@@ -180,6 +187,40 @@ class Counter(object):
         return val
 
 
+def get_max_node_rss(return_per_node=False):
+    """
+    Find the maximum memory usage on any node in the job in bytes.
+
+    Parameters
+    ----------
+
+    return_per_node : bool (optional)
+        Return the total memory on the node to each rank on
+        that node. (Default is False)
+
+    Returns
+    -------
+
+    max_mem : float
+        Maximum memory usage in GiB across the job.
+        Only returns to the zero-th rank on the world_comm.
+    """
+
+    # On linux, getrusage returns in kiB
+    # On Mac systems, getrusage returns in B
+    scale = 1.0
+    if 'linux' in sys.platform:
+        scale = 2**10
+
+    memory_usage_GiB = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * scale / 2**30
+    node_mem_tot = node_comm.allreduce(memory_usage_GiB, op=MPI.SUM)
+    if return_per_node:
+        return node_mem_tot
+
+    max_mem = world_comm.reduce(node_mem_tot, op=MPI.MAX, root=0)
+    return max_mem
+
+
 def get_rank():
     """
     Current rank on COMM_WORLD
diff --git a/pyuvsim/source.py b/pyuvsim/source.py
index 5a64b3c..1a3bc81 100644
--- a/pyuvsim/source.py
+++ b/pyuvsim/source.py
@@ -20,18 +20,45 @@ from . import utils as simutils
 from . import spherical_coordinates_basis_transformation as scbt
 
 
-def _skymodel_basesize():
+def estimate_skymodel_memory_usage(Ncomponents, Nfreqs):
     """
-    Estimate the memory footprint of a SkyModel with a single source.
+    Estimate the memory footprint of a SkyModel by summing the sizes
+    of the data types that go into SkyModel.
 
-    Sum the sizes of the data types that go into SkyModel
+    This aims to anticipate the full memory required to handle a SkyModel
+    class in simulation, accounting for its attributes as well as temporary
+    and intermediate data.
+
+    Parameters
+    ----------
+
+    Ncomponents : int
+        Number of source components.
+    Nfreqs : int
+        Number of frequencies per source component.
+        (size of SkyModel.freq_array)
+
+    Returns
+    -------
+    mem_est : float
+        Estimate of memory usage in bytes
     """
-    attrs = [
-        '', Quantity(1.0, 'Hz'), [0.0] * 4,
-        Angle(np.pi, 'rad'), Angle(np.pi, 'rad'),
-        [1.5] * 4, [0.3] * 2, [0.0] * 3
-    ]
-    return np.sum([sys.getsizeof(a) for a in attrs])
+
+    base_float = [1.5]    # A float
+    base_bool = [True]
+    base_str = ["source_name"]
+
+    Ncomp_attrs = {'ra': base_float, 'dec': base_float,
+                   'alt_az': 2 * base_float, 'rise_lst': base_float, 'set_lst': base_float,
+                   'pos_lmn': 3 * base_float, 'name': base_str, 'horizon_mask': base_bool}
+    Ncomp_Nfreq_attrs = {'stokes': 4 * base_float,
+                         'coherency_radec': 4 * base_float,
+                         'coherency_local': 4 * base_float}
+
+    mem_est = np.sum([sys.getsizeof(v) * Ncomponents for k, v in six.iteritems(Ncomp_attrs)])
+    mem_est += np.sum([sys.getsizeof(v) * Ncomponents * Nfreqs
+                       for k, v in six.iteritems(Ncomp_Nfreq_attrs)])
+    return mem_est
 
 
 class SkyModel(object):
@@ -68,8 +95,6 @@ class SkyModel(object):
 
     Ncomponents = None  # Number of point source components represented here.
 
-    _basesize = _skymodel_basesize()
-
     _Ncomp_attrs = ['ra', 'dec', 'coherency_radec', 'coherency_local',
                     'stokes', 'alt_az', 'rise_lst', 'set_lst', 'freq',
                     'pos_lmn', 'name', 'horizon_mask']
@@ -389,7 +414,7 @@ class SkyModel(object):
         """
         Estimate the SkyModel memory footprint in bytes
         """
-        return self.Ncomponents * self._basesize
+        return estimate_skymodel_memory_usage(self.Ncomponents, self.Nfreqs)
 
 
 def read_healpix_hdf5(hdf5_filename):
diff --git a/pyuvsim/utils.py b/pyuvsim/utils.py
index 069b4ca..a57795b 100644
--- a/pyuvsim/utils.py
+++ b/pyuvsim/utils.py
@@ -36,7 +36,14 @@ def get_version_string():
 
 class progsteps:
     """
-        Similar to a progress bar, this prints a percentage of task completion.
+    Similar to a progress bar, this prints a percentage of task completion.
+
+    Parameters
+    ----------
+
+    maxval : int
+        Maximum value to count to.
+
     """
 
     def __init__(self, maxval=None):
@@ -51,15 +58,25 @@ class progsteps:
         self.curval = -1
 
     def update(self, count):
+        """
+        Update the progress bar.
+
+        Parameters
+        ----------
+        count : int
+            Current value of counter
+        """
         if count >= self.curval + self.step:
             doprint = False
             if not self.curval == count:
                 doprint = True
                 self.curval = count
             if doprint:
-                print("{:0.2f}% completed. {:0.3f} minutes elapsed \n".format(
-                    (count / self.maxval) * 100., (pytime.time() - self.t0) / 60.))
-            sys.stdout.flush()
+                ostr = "{:0.2f}% completed. {:0.3f} minutes elapsed".format(
+                    (count / self.maxval) * 100., (pytime.time() - self.t0) / 60.)
+                ostr += '\n'
+                print(ostr)
+                sys.stdout.flush()
 
     def finish(self):
         self.update(self.maxval)
diff --git a/pyuvsim/uvsim.py b/pyuvsim/uvsim.py
index aa3bf79..9aa104c 100644
--- a/pyuvsim/uvsim.py
+++ b/pyuvsim/uvsim.py
@@ -23,7 +23,7 @@ from . import simsetup
 from . import utils as simutils
 from .antenna import Antenna
 from .baseline import Baseline
-from .source import SkyModel
+from .source import estimate_skymodel_memory_usage
 from .telescope import Telescope
 
 __all__ = ['UVTask', 'UVEngine', 'uvdata_to_task_iter', 'run_uvsim', 'run_uvdata_uvsim',
@@ -178,13 +178,13 @@ def _make_task_inds(Nbls, Ntimes, Nfreqs, Nsrcs, rank, Npus):
         Ntasks_local = Nbltf
     else:
         task_inds, Ntasks_local = simutils.iter_array_split(rank, Nbltf, Npus)
-        src_inds = range(Nsrcs)
+        src_inds = slice(None)
         Nsrcs_local = Nsrcs
 
     return task_inds, src_inds, Ntasks_local, Nsrcs_local
 
 
-def uvdata_to_task_iter(task_ids, input_uv, catalog, beam_list, beam_dict):
+def uvdata_to_task_iter(task_ids, input_uv, catalog, beam_list, beam_dict, Nsky_parts=1):
     """
     Generate local tasks, reusing quantities where possible.
 
@@ -212,17 +212,12 @@ def uvdata_to_task_iter(task_ids, input_uv, catalog, beam_list, beam_dict):
         raise TypeError("catalog must be a record array")
 
     # Splitting the catalog for memory's sake.
-    # Total memory footprint will be \approx Npus (on node) * (size of skymodel)
     Nsrcs_total = len(catalog)
-    mem_avail = simutils.get_avail_memory()
-    Npus_node = mpi.Npus_node
-    skymodel_mem_footprint = SkyModel._basesize * Nsrcs_total * Npus_node
-    if skymodel_mem_footprint > mem_avail:
-        Nsky_parts = np.ceil(skymodel_mem_footprint / float(mem_avail))
-        partsize = int(np.floor(Nsrcs_total / Nsky_parts))
-        src_iter = [range(s, s + partsize) for s in range(0, Nsrcs_total, partsize)]
+    if Nsky_parts > 1:
+        Nsky_parts = int(Nsky_parts)
+        src_iter = [simutils.iter_array_split(s, Nsrcs_total, Nsky_parts)[0]
+                    for s in range(Nsky_parts)]
     else:
-        Nsky_parts = 1
         src_iter = [slice(None)]
 
     # Build the antenna list.
@@ -251,7 +246,6 @@ def uvdata_to_task_iter(task_ids, input_uv, catalog, beam_list, beam_dict):
 
     freq_array = input_uv.freq_array * units.Hz
     time_array = Time(input_uv.time_array, scale='utc', format='jd', location=telescope.location)
-
     for src_i in src_iter:
         sky = simsetup.array_to_skymodel(catalog[src_i])
         if sky.spectral_type == 'values':
@@ -282,6 +276,7 @@ def uvdata_to_task_iter(task_ids, input_uv, catalog, beam_list, beam_dict):
             task.uvdata_index = (blti, 0, freq_i)    # 0 = spectral window index
 
             yield task
+        del sky
 
 
 def serial_gather(uvtask_list, uv_out):
@@ -340,6 +335,7 @@ def run_uvdata_uvsim(input_uv, beam_list, beam_dict=None, catalog=None):
         print('Ntimes:', input_uv.Ntimes)
         print('Nfreqs:', input_uv.Nfreqs)
         print('Nsrcs:', len(catalog))
+        sys.stdout.flush()
         uv_container = simsetup._complete_uvdata(input_uv, inplace=False)
 
     Nbls = input_uv.Nbls
@@ -356,20 +352,36 @@ def run_uvdata_uvsim(input_uv, beam_list, beam_dict=None, catalog=None):
     for bm in beam_models:
         bm.interpolation_function = 'az_za_simple'
 
+    # Estimating required memory to decide how to split source array.
+
+    Nsrcs_total = len(catalog)
+    Nsrcfreqs = np.atleast_1d(catalog['frequency'][0]).size
+    mem_avail = (simutils.get_avail_memory()
+                 - mpi.get_max_node_rss(return_per_node=True) * 2**30)
+
+    Npus_node = mpi.node_comm.Get_size()
+    skymodel_mem_footprint = estimate_skymodel_memory_usage(Nsrcs_total, Nsrcfreqs) * Npus_node
+
+    # Allow up to 50% of available memory for SkyModel data.
+    skymodel_mem_max = 0.5 * mem_avail
+
+    Nsky_parts = np.ceil(skymodel_mem_footprint / float(skymodel_mem_max))
+
+    Ntasks_tot = Ntimes * Nbls * Nfreqs * Nsky_parts
+
     local_task_iter = uvdata_to_task_iter(
-        task_inds, input_uv, catalog[src_inds], beam_models, beam_dict
+        task_inds, input_uv, catalog[src_inds], beam_models, beam_dict, Nsky_parts=Nsky_parts
     )
 
     summed_task_dict = {}
+    Ntasks_tot = comm.reduce(Ntasks_tot, op=mpi.MPI.MAX, root=0)
     if rank == 0:
-        tot = Ntimes * Nbls * Nfreqs
-        print("Tasks: ", Ntimes * Nbls * Nfreqs)
+        print("Tasks: ", Ntasks_tot)
         sys.stdout.flush()
-        pbar = simutils.progsteps(maxval=tot)
+        pbar = simutils.progsteps(maxval=Ntasks_tot)
 
     engine = UVEngine()
     count = mpi.Counter()
-
     for task in local_task_iter:
         engine.set_task(task)
         if task.uvdata_index not in summed_task_dict.keys():
@@ -378,9 +390,11 @@ def run_uvdata_uvsim(input_uv, beam_list, beam_dict=None, catalog=None):
             summed_task_dict[task.uvdata_index].visibility_vector = engine.make_visibility()
         else:
             summed_task_dict[task.uvdata_index].visibility_vector += engine.make_visibility()
+
         count.next()
         if rank == 0:
             pbar.update(count.current_value())
+
     comm.Barrier()
     count.free()
     if rank == 0:


Test patch:
diff --git a/pyuvsim/tests/test_mpi_uvsim.py b/pyuvsim/tests/test_mpi_uvsim.py
index 5ae5760..fa3767c 100644
--- a/pyuvsim/tests/test_mpi_uvsim.py
+++ b/pyuvsim/tests/test_mpi_uvsim.py
@@ -9,6 +9,8 @@ import os
 import mpi4py
 import numpy as np
 import yaml
+import resource
+import time
 
 mpi4py.rc.initialize = False  # noqa
 import pytest
@@ -143,6 +145,25 @@ def test_shared_mem():
     pytest.raises(ValueError, sA.itemset, 0, 3.0)
 
 
+def test_mem_usage():
+    # Check that the mpi-enabled memory check is consistent
+    # with a local memory check.
+
+    # Also check that making a variable of a given size
+    # increases memory usage by the expected amount.
+
+    mpi.start_mpi()
+
+    memory_usage_GiB = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 2**10 / 2**30
+    assert memory_usage_GiB == mpi.get_max_node_rss()
+    incsize = 50 * 2**20    # 50 MiB
+    arr = bytearray(incsize)
+    time.sleep(1)
+    change = mpi.get_max_node_rss() - memory_usage_GiB
+    assert np.isclose(change, incsize / 2**30, atol=5e-2)
+    del arr
+
+
 @pytest.mark.skip
 def test_mpi_counter():
     # This test should be run in parallel to check likely bug sources.
diff --git a/pyuvsim/tests/test_uvsim.py b/pyuvsim/tests/test_uvsim.py
index d3aa580..4ffea32 100644
--- a/pyuvsim/tests/test_uvsim.py
+++ b/pyuvsim/tests/test_uvsim.py
@@ -603,6 +603,7 @@ def test_task_coverage():
             task_inds, src_inds, Ntasks_local, Nsrcs_local = pyuvsim.uvsim._make_task_inds(
                 Nbls, Ntimes, Nfreqs, Nsrcs, rank, Npus
             )
+            src_inds = range(Nsrcs)[src_inds]   # Turn slice into iterator
             tasks = itertools.product(task_inds, src_inds)
             tasks_all.append(tasks)
         tasks_all = itertools.chain(*tasks_all)
@@ -652,7 +653,7 @@ def test_source_splitting():
     )
 
     # Spoof environmental parameters.
-    # Choose an absurdlly large number of Npus per node and very small available memory
+    # Choose an absurdly large number of tasks per node and very small available memory
     # to trigger the source splitting condition in uvdata_to_task_iter
     # The alternative would be to make a very large source catalog, but that's not ideal in a test.
     Npus_node = 2000
@@ -667,13 +668,6 @@ def test_source_splitting():
     Nsrcs = len(sources)
     Ntasks = Nblts * Nfreqs
     beam_dict = None
-    taskiter = pyuvsim.uvdata_to_task_iter(
-        np.arange(Ntasks), hera_uv, sources, beam_list, beam_dict
-    )
-    uvtask_list = uvtest.checkWarnings(
-        list, [taskiter], message='The default for the `center` keyword has changed.',
-        category=DeprecationWarning
-    )
 
     skymodel = pyuvsim.array_to_skymodel(sources)
     skymodel_mem_footprint = skymodel.get_size() * Npus_node
@@ -682,14 +676,20 @@ def test_source_splitting():
     Nsky_parts = np.ceil(skymodel_mem_footprint / float(mem_avail))
     partsize = int(np.floor(Nsrcs / Nsky_parts))
 
-    assert partsize * pyuvsim.SkyModel._basesize * Npus_node < mem_avail
-    print(skymodel_mem_footprint / 1e6,
-          partsize * pyuvsim.SkyModel._basesize * Npus_node / 1e6,
-          mem_avail / 1e6)
+    taskiter = pyuvsim.uvdata_to_task_iter(
+        np.arange(Ntasks), hera_uv, sources, beam_list, beam_dict, Nsky_parts=Nsky_parts
+    )
+
+    uvtask_list = uvtest.checkWarnings(
+        list, [taskiter], message='The default for the `center` keyword has changed.',
+        category=DeprecationWarning
+    )
+
+    assert pyuvsim.estimate_skymodel_memory_usage(partsize, 1) * Npus_node < mem_avail
 
     # Normally, the number of tasks is Nbls * Ntimes * Nfreqs (partsize = Nsrcs)
     # If the source list is split within the task iterator, then it will be larger.
-    assert len(uvtask_list) == Ntasks * Nsrcs / partsize
+    assert len(uvtask_list) == Ntasks * Nsky_parts
 
     # Reset spoofed parameters.
     del os.environ['SLURM_MEM_PER_NODE']