Merge pull request #33 from coecms/beta

Combined branch for recent pull requests

Author: aidanheerdegen

.travis.yml

@@ -5,7 +5,7 @@ before_install:
     - bash Miniconda3-latest-Linux-x86_64.sh -b -p ~/conda
     - source ~/conda/bin/activate
     - conda create --yes -c conda-forge -n mppnccombine-fast-build gcc openmpi 'hdf5>=1.10.2' libnetcdf nomkl
-    - conda create --yes -c conda-forge -n mppnccombine-fast-test openmpi pytest xarray nomkl
+    - conda create --yes -c conda-forge -n mppnccombine-fast-test openmpi pytest xarray nomkl netcdf4
 
 script:
     - make CONDA_BUILD=1 mppnccombine-fast

Makefile

@@ -25,7 +25,7 @@ endif
 
 all: mppnccombine-fast
 
-test: mppnccombine-fast
+test check: mppnccombine-fast
 	${TEST_ENV} pytest -v --capture=no test.py
 
 mppnccombine-fast: async.o error.o read_chunked.o

async.c

@@ -47,6 +47,7 @@
 #define TAG_VAR_INFO       22
 
 #define MAX_VARIABLES 100
+#define MAX_DIMS 10
 
 typedef struct {
     hid_t file_id;
@@ -55,6 +56,7 @@ typedef struct {
         hid_t var_id;
         char varname[NC_MAX_NAME+1];
         size_t refcount;
+        hsize_t chunk_shape[MAX_DIMS]; // Set by get_var_info
     } vars[MAX_VARIABLES];
 
     int total_vars;
@@ -78,6 +80,34 @@ varid_t open_variable_async(
     return out;
 }
 
+// Get a dataspace, making sure it can fit the field
+hid_t get_resized_space(async_state_t * state, int idx, int ndims, hsize_t offset[], hsize_t shape[]) {
+    hid_t data_space = H5Dget_space(state->vars[idx].var_id);
+    H5ERR(data_space);
+
+    // Check if we need to extend the dataset
+    hsize_t file_size[ndims];
+    hsize_t max_size[ndims];
+    H5ERR(H5Sget_simple_extent_dims(data_space, file_size, max_size));
+    bool needs_resize = false;
+    for (int d=0;d<ndims;++d) {
+        if (offset[d] + shape[d] > file_size[d]) {
+            needs_resize = true;
+            file_size[d] = offset[d] + shape[d];
+            if ((max_size[d] != H5S_UNLIMITED) && (file_size[d] > max_size[d])) {
+                file_size[d] = max_size[d];
+            }
+        }
+    }
+    if (needs_resize) {
+        H5ERR(H5Dset_extent(state->vars[idx].var_id, file_size));
+        // Re-open
+        data_space = H5Dget_space(state->vars[idx].var_id);
+        H5ERR(data_space);
+    }
+    return data_space;
+}
+
 void close_variable_async(
     varid_t varid,
     int async_writer_rank
@@ -172,6 +202,7 @@ void receive_variable_info_async(
     H5ERR(plist);
 
     H5ERR(H5Pget_chunk(plist, ndims, varinfo + 0));
+    H5ERR(H5Pget_chunk(plist, ndims, state->vars[idx].chunk_shape));
 
     int nfilters = H5Pget_nfilters(plist);
     H5ERR(nfilters);
@@ -317,8 +348,8 @@ static size_t receive_write_uncompressed_async(
     hid_t mem_space = H5Screate_simple(ndims, shape, NULL);
     H5ERR(mem_space);
 
-    hid_t data_space = H5Dget_space(state->vars[idx].var_id);
-    H5ERR(data_space);
+    hid_t data_space = get_resized_space(state, idx, ndims, offset, shape);
+
     H5ERR(H5Sselect_hyperslab(data_space, H5S_SELECT_SET,
                               offset, NULL, shape, NULL));
 
@@ -419,6 +450,9 @@ static size_t receive_write_chunk_async(
     hsize_t offset_[ndims];
     for (size_t d=0; d<ndims; ++d) {offset_[d] = offset[d];}
 
+    hid_t data_space = get_resized_space(state, idx, ndims, offset_, state->vars[idx].chunk_shape);
+    H5ERR(H5Sclose(data_space));
+
     int err = (H5DOwrite_chunk(state->vars[idx].var_id, H5P_DEFAULT, filter_mask,
                           offset_, buffer_size, buffer));
     if (err < 0) {
@@ -509,9 +543,8 @@ static void receive_close_variable_async(
 void close_async(
     int async_writer_rank
     ) {
-    int buffer = 0;
     log_message(LOG_DEBUG, "SEND close file");
-    MPI_Send(&buffer, 1, MPI_INT, async_writer_rank, TAG_CLOSE, MPI_COMM_WORLD);
+    MPI_Send(NULL, 0, MPI_INT, async_writer_rank, TAG_CLOSE, MPI_COMM_WORLD);
 }
 
 // Async runner to accept writes
@@ -545,19 +578,28 @@ size_t run_async_writer(
                 receive_variable_info_async(&state, status);
                 break;
             case (TAG_WRITE_CHUNK):
-                receive_write_chunk_async(&state, status);
+                total_size += receive_write_chunk_async(&state, status);
                 break;
             case (TAG_WRITE_FILTER):
-                receive_write_uncompressed_async(&state, status);
+                total_size += receive_write_uncompressed_async(&state, status);
                 break;
             case (TAG_CLOSE_VARIABLE):
                 receive_close_variable_async(&state, status);
                 break;
             case (TAG_CLOSE):
                 --open_senders;
+                MPI_Recv(NULL, 0, MPI_INT, status.MPI_SOURCE, TAG_CLOSE, MPI_COMM_WORLD,
+                         MPI_STATUS_IGNORE);
                 break;
         }
     }
+
+    for (int v=0; v<state.total_vars; ++v) {
+        H5ERR(H5Dclose(state.vars[v].var_id));
+    }
+
     H5ERR(H5Fclose(state.file_id));
     log_message(LOG_DEBUG, "DONE run_async_writer");
+
+    return total_size;
 }