I/O hangs when getting/setting dataset distributed over last axis

[jrs65]

This one is a little nebulous. I had originally thought that the issue was when the distributed axis is shorter than the number o processes and so one has nothing to do. That may still be true, but some of the log messages seemed to indicate that it was actually doing the IO split along that axis?? Regardless, reducing the number of MPI processes being used seemed to fix the issue.

Anyway this clearly needs a little more debugging to figure out what exactly is happening, but we shouldn't crash when this happens.

[sjforeman]

I'd like to revive this issue, because Haochen and I have recently run into something that seems related, and it would nice to resolve it.

Here is a config that can be used to reproduce the issue on cedar:

cluster:
  name: df_day_{lsd}

  directory: {out_dir}

  time: 10
  system: cedar
  account: rpp-chime
  nodes: {nodes}
  ompnum: 4
  pernode: 12
  mem: 192000M

  venv: {venv}


# Pipeline task configuration
pipeline:

  logging:
    root: DEBUG
    peewee: INFO
    matplotlib: INFO
    h5py: INFO

  tasks:
    - type: draco.core.task.SetMPILogging
      params:
        level_rank0: DEBUG
        level_all: DEBUG

    - type: draco.core.io.LoadFilesFromParams
      out: ringmap
      params:
        files: ["/project/rpp-chime/chime/chime_processed/daily/rev_06/{lsd}/ringmap_intercyl_lsd_{lsd}.zarr.zip"]
        distributed: true

    - type: draco.analysis.dayenu.DayenuDelayFilterMap
      in: ringmap
      out: ringmap_filtered
      params:
        tauw: 0.5
        single_mask: true
        atten_threshold: 0.20
        save: true
        output_name: "ringmap_intercyl_postfilter.h5"

You can pick your favourite LSD, but I've mostly been working with 3245. With nodes=2, this runs without issue. (Expanding the walltime would let it finish, but 10 minutes is sufficient to see if it crashes.) With nodes=4 or higher, we get an ugly crash, usually some variation of

slurmstepd: error:  mpi/pmix_v3: _errhandler: cdr1585 [1]: pmixp_client_v2.c:211: Error handler invoked: status = -25, source = [slurm.pmix.61907741.0:19]
srun: error: cdr1585: task 19: Killed
srun: launch/slurm: _step_signal: Terminating StepId=61907741.0

The crash occurs in this line of mpiarray.MPIArray.to_hdf5():

caput/caput/mpiarray.py

Line 1143 in 0beefd6

dset[islice] = self[fslice]

I've attempted some dumb debugging by adding log statements as shown in this git diff:

         # Read using collective MPI-IO if specified
+        logger.debug("RIGHT BEFORE with dset.collective if use_collective else DummyContext():")
+        logger.debug(f"use_collective = {use_collective}, faster_without_collective = {faster_without_collective}, fh.is_mpi = {fh.is_mpi}, no_null_slices = {no_null_slices}")
+        logger.debug(f"split_axis = {split_axis}, dataset = {dataset}, partitions = {partitions}")
         with dset.collective if use_collective else DummyContext():
             # Loop over partitions of the IO and perform them
             for part in partitions:
                 islice, fslice = _partition_sel(
                     sel, split_axis, self.global_shape[split_axis], part
                 )
+                logger.debug(f"islice, fslice, self.shape, dset.shape, self[fslice].shape: {islice}, {fslice}, {self.shape}, {dset.shape}, {self[fslice].shape}")
                 dset[islice] = self[fslice]
+                logger.debug("COMPLETE: dset[islice] = self[fslice]")

         if fh.opened:
             fh.close()

These statements reveal that there is no MPI task that completes its write successfully. Here's what I get from task 0, right before the crash:

   115.2s [MPI  0/48] - DEBUG    caput.mpiarray: RIGHT BEFORE with dset.collective if use_collective else DummyContext():
   115.2s [MPI  0/48] - DEBUG    caput.mpiarray: use_collective = True, faster_without_collective = False, fh.is_mpi = True, no_null_slices = True
   115.2s [MPI  0/48] - DEBUG    caput.mpiarray: split_axis = 0, dataset = /dirty_beam, partitions = [slice(0, 1, None)]
   115.2s [MPI  0/48] - DEBUG    caput.mpiarray: islice, fslice, self.shape, dset.shape, self[fslice].shape: (slice(0, 1, 1), slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(0, 11, None)), (slice(0, 1, None), slice(None, None, None), slice(None, None, None), slice(None, None, None), slice(None, None, None)), (1, 4, 1024, 4096, 11), (1, 4, 1024, 4096, 512), (1, 4, 1024, 4096, 11)

In this case, MPIArray._partition_io() determined that there's no need to further partition the I/O, beyond the way the array is already distributed along the last axis. The size of the slice being written is 1.375GB, so it's not running into any known 2GB-write issues in h5py. If it makes a difference, the "0" axis (the beam axis of the ringmap container) has 1 element.

If we use 2 nodes (24 tasks), the analogous output instead looks like:

   154.3s [MPI  1/24] - DEBUG    caput.mpiarray: RIGHT BEFORE with dset.collective if use_collective else DummyContext():
   154.3s [MPI  1/24] - DEBUG    caput.mpiarray: use_collective = True, faster_without_collective = False, fh.is_mpi = True, no_null_slices = True
   154.3s [MPI  1/24] - DEBUG    caput.mpiarray: 1, /dirty_beam, [slice(0, 2, None), slice(2, 4, None)]
   154.3s [MPI  1/24] - DEBUG    caput.mpiarray: islice, fslice, self.shape, dset.shape, self[fslice].shape: (slice(None, None, None), slice(0, 2, 1), slice(None, None, None), slice(None, None, None), slice(22, 44, None)), (slice(None, None, None), slice(0, 2, None), slice(None, None, None), slice(None, None, None), slice(None, None, None)), (1, 4, 1024, 4096, 22), (1, 4, 1024, 4096, 512), (1, 2, 1024, 4096, 22)

Here, the output is partitioned into 2 slices along the "1" axis (i.e. the pol axis, which has 4 elements). The slice being written is 1.375GB here as well, but now the write operation completes successfully.

I've found that restricting the slice size to 1GB, using threshold=0.99 in MPIArray._partition_io(), seems to solve the problem, but I'm not sure why, unless the h5py output issues with 2GB chunks actually also apply to chunks between 1 and 2GB. Any expert opinions would be appreciated!

[ljgray]

I think that the slice into dset in the mpiarray line is causing the issue. Just slicing into it without trying to set anything (just logging dset[islice] for example) hangs but doesn't actually crash. In this test case, the distributed axis is significantly longer than the number of processes, so each process has something to do.

One thing I noticed is that the fixes mentioned by both @jrs65 and @sjforeman involve reducing the number of processes, which would increase the size of the local array on each process. In this case, by reducing the number of processes, _partition_io will return multiple parts rather than just one, and those parts will not be across the 0th axis (which here only has length 1). I'm wondering if the issue is somehow with slicing along the zeroth axis

[ljgray]

This seems relevant, although I can't access the JIRA page linked to see if the issue was ever resolved in hdf5
h5py/h5py#1176

[sjforeman]

Thanks @ljgray! I'm experimenting with possible fixes related to the slicing issues you've described

[ljgray]

I should note that our situation isn't exactly the same as the one in the h5py issue, as they're trying to take non-consecutive elements, but it's always possible that we're seeing the effects of that issue in another way

[jrs65]

I think that issue is a red herring as it only applies to chunked data (which this is not).

That said, I can't really see any obvious issues with what's going on.

Can you try forcing collective IO to be off? That is an obvious contender as I think the h5py implementation has some issues (e.g. the eliding of zero length writes, which we already need to work around).

[ljgray]

I tried forcing collective IO off and the job just timed out after 90 minutes. Interestingly, a few ranks did complete their write quite quickly but most never did

[jrs65]

Can you go one step further and try forcing use_mpi = False so it doesn't open the file with MPI-IO?

[jrs65]

Github really needs to add the 🤔 emoji. That's pretty much the reaction I want to add to every issue comment.

[sjforeman]

In the case where we're just slicing a length-1 axis with slice(0,1,1), I'm trying to test whether simply using slice(None) would solve things

[jrs65]

Can you go one step further and try forcing use_mpi = False so it doesn't open the file with MPI-IO?

Actually that doesn't make any sense. If it can't use MPI-IO it needs to use _to_hdf5_serial as it can't coordinate the writes, so opening with MPI-IO is essential to this.

Sorry, bad suggestion!

[jrs65]

In the case where we're just slicing a length-1 axis with slice(0,1,1), I'm trying to test whether simply using slice(None) would solve things

Good test.

I think broadly it would be good to try and reduce this to a simpler test case. Can you get it down to something that only needs 4 MPI processes and where the dataset is smaller and with fewer dimensions?

[ljgray]

Yeah it would be beneficial to be able to iterate more quickly, but it might be tricky to simplify when we don't really know exactly what the issue is.

@sjforeman just note that using a slice(None) will fail in _partition_sel, I believe in _reslice, so there will have to be some modifications for that. I tried doing that test yesterday but didn't get around to actually making that change.