memory_issue_fix

.circleci/config.yml

@@ -2,14 +2,13 @@ version: 2
 jobs:
   build:
     docker:
-      - image: continuumio/miniconda3
+      - image: cimg/python:3.8
     steps:
       - setup_remote_docker
       - checkout
       - run:
           name: Install buildstock
           command: |
-            conda install -c conda-forge -y -q scipy numpy "pandas>=1.0.0,!=1.0.4" "pyarrow>=0.14.1" dask joblib pyyaml
             pip install .[dev] --progress-bar off
       - run:
           name: Run PyTest
@@ -35,7 +34,6 @@ jobs:
           name: Build documentation
           when: always
           command: |
-            apt-get install make
             cd docs
             make html
             mkdir /tmp/docs

buildstockbatch/base.py

@@ -784,4 +784,8 @@ def process_results(self, skip_combine=False, force_upload=False):
             if 'athena' in aws_conf:
                 postprocessing.create_athena_tables(aws_conf, os.path.basename(self.output_dir), s3_bucket, s3_prefix)
 
-        postprocessing.remove_intermediate_files(fs, self.results_dir)
+        if not self.cfg['eagle'].get('postprocessing', {}).get('keep_intermediate_files', False):
+            logger.info("Removing intermediate files.")
+            postprocessing.remove_intermediate_files(fs, self.results_dir)
+        else:
+            logger.info("Skipped removing intermediate files.")

buildstockbatch/eagle.py

@@ -494,7 +494,8 @@ def queue_post_processing(self, after_jobids=[], upload_only=False, hipri=False)
         # Configuration values
         account = self.cfg['eagle']['account']
         walltime = self.cfg['eagle'].get('postprocessing', {}).get('time', '1:30:00')
-
+        memory = self.cfg['eagle'].get('postprocessing', {}).get('memory', 85248)
+        print(f"Submitting job to {memory}MB memory nodes.")
         # Throw an error if the files already exist.
         if not upload_only:
             for subdir in ('parquet', 'results_csvs'):
@@ -518,19 +519,20 @@ def queue_post_processing(self, after_jobids=[], upload_only=False, hipri=False)
         env['MY_CONDA_ENV'] = os.environ['CONDA_PREFIX']
         env['OUT_DIR'] = self.output_dir
         env['UPLOADONLY'] = str(upload_only)
+        env['MEMORY'] = str(memory)
         here = os.path.dirname(os.path.abspath(__file__))
         eagle_post_sh = os.path.join(here, 'eagle_postprocessing.sh')
 
         args = [
             'sbatch',
             '--account={}'.format(account),
             '--time={}'.format(walltime),
-            '--export=PROJECTFILE,MY_CONDA_ENV,OUT_DIR,UPLOADONLY',
+            '--export=PROJECTFILE,MY_CONDA_ENV,OUT_DIR,UPLOADONLY,MEMORY',
             '--job-name=bstkpost',
             '--output=postprocessing.out',
             '--nodes=1',
             ':',
-            '--mem=180000',
+            '--mem={}'.format(memory),
             '--output=dask_workers.out',
             '--nodes={}'.format(self.cfg['eagle'].get('postprocessing', {}).get('n_workers', 2)),
             eagle_post_sh
@@ -599,7 +601,6 @@ def user_cli(argv=sys.argv[1:]):
     '''
     This is the user entry point for running buildstockbatch on Eagle
     '''
-
     # set up logging, currently based on within-this-file hard-coded config
     logging.config.dictConfig(logging_config)
 

buildstockbatch/eagle_postprocessing.sh

@@ -1,5 +1,4 @@
 #!/bin/bash
-
 echo "begin eagle_postprocessing.sh"
 
 echo "Job ID: $SLURM_JOB_ID"
@@ -11,6 +10,7 @@ source activate "$MY_CONDA_ENV"
 export POSTPROCESS=1
 
 echo "UPLOADONLY: ${UPLOADONLY}"
+echo "MEMORY: ${MEMORY}"
 
 SCHEDULER_FILE=$OUT_DIR/dask_scheduler.json
 
@@ -22,6 +22,6 @@ echo $SLURM_JOB_NODELIST_PACK_GROUP_1
 pdsh -w $SLURM_JOB_NODELIST_PACK_GROUP_1 "free -h"
 
 $MY_CONDA_ENV/bin/dask-scheduler --scheduler-file $SCHEDULER_FILE &> $OUT_DIR/dask_scheduler.out &
-pdsh -w $SLURM_JOB_NODELIST_PACK_GROUP_1 "$MY_CONDA_ENV/bin/dask-worker --scheduler-file $SCHEDULER_FILE --local-directory /tmp/scratch/dask --nprocs 9" &> $OUT_DIR/dask_workers.out &
+pdsh -w $SLURM_JOB_NODELIST_PACK_GROUP_1 "$MY_CONDA_ENV/bin/dask-worker --scheduler-file $SCHEDULER_FILE --local-directory /tmp/scratch/dask --nprocs 1 --nthreads 1 --memory-limit ${MEMORY}MB" &> $OUT_DIR/dask_workers.out &
 
 time python -u -m buildstockbatch.eagle "$PROJECTFILE"

buildstockbatch/postprocessing.py

@@ -9,10 +9,10 @@
 :copyright: (c) 2018 by The Alliance for Sustainable Energy
 :license: BSD-3
 """
-
+import traceback
 import boto3
 import dask.bag as db
-import dask.dataframe as dd
+from dask.distributed import performance_report
 import dask
 import datetime as dt
 from fsspec.implementations.local import LocalFileSystem
@@ -34,7 +34,7 @@
 
 logger = logging.getLogger(__name__)
 
-MAX_PARQUET_MEMORY = 1e9  # maximum size of the parquet file in memory when combining multiple parquets
+MAX_PARQUET_MEMORY = 4000  # maximum size (MB) of the parquet file in memory when combining multiple parquets
 
 
 def read_data_point_out_json(fs, reporting_measures, filename):
@@ -218,17 +218,28 @@ def read_results_json(fs, filename):
 
 
 def read_enduse_timeseries_parquet(fs, filename, all_cols):
-    with fs.open(filename, 'rb') as f:
-        df = pd.read_parquet(f, engine='pyarrow')
-    building_id = int(re.search(r'bldg(\d+).parquet', filename).group(1))
-    df['building_id'] = building_id
-    for col in set(all_cols).difference(df.columns.values):
-        df[col] = np.nan
-    return df[all_cols]
-
-
-def read_and_concat_enduse_timeseries_parquet(fs, filenames, all_cols):
-    return pd.concat(read_enduse_timeseries_parquet(fs, filename, all_cols) for filename in filenames)
+    try:
+        with fs.open(filename, 'rb') as f:
+            df = pd.read_parquet(f, engine='pyarrow')
+        building_id = int(re.search(r'bldg(\d+).parquet', filename).group(1))
+        df['building_id'] = building_id
+        for col in set(all_cols).difference(df.columns.values):
+            df[col] = np.nan
+        return df[all_cols]
+    except Exception:
+        return pd.DataFrame(columns=[all_cols]+['building_id'])
+
+
+def read_and_concat_enduse_timeseries_parquet(fs, all_cols, output_dir, filenames, group_id):
+    try:
+        dfs = [read_enduse_timeseries_parquet(fs, filename, all_cols) for filename in filenames]
+        grouped_df = pd.concat(dfs)
+        grouped_df.set_index('building_id', inplace=True)
+        grouped_df.to_parquet(output_dir + f'group{group_id}.parquet')
+        return ("success", group_id, len(grouped_df))
+        del grouped_df
+    except Exception:
+        return ("fail", group_id, f"Exp: {traceback.format_exc()}")
 
 
 def combine_results(fs, results_dir, cfg, do_timeseries=True):
@@ -330,32 +341,39 @@ def combine_results(fs, results_dir, cfg, do_timeseries=True):
             get_ts_mem_usage_d = dask.delayed(lambda x: read_ts_parquet(x).memory_usage(deep=True).sum())
             sample_size = min(len(ts_filenames), 36 * 3)
             mean_mem = np.mean(dask.compute(map(get_ts_mem_usage_d, random.sample(ts_filenames, sample_size)))[0])
-            total_mem = mean_mem * len(ts_filenames)
+            total_mem = mean_mem * len(ts_filenames) / 1e6  # total_mem in MB
 
             # Determine how many files should be in each partition and group the files
-            npartitions = math.ceil(total_mem / MAX_PARQUET_MEMORY)  # 1 GB per partition
+            parquet_memory = int(cfg['eagle'].get('postprocessing', {}).get('parquet_memory_mb', MAX_PARQUET_MEMORY))
+            logger.info(f"Max parquet memory: {parquet_memory} MB")
+            npartitions = math.ceil(total_mem / parquet_memory)
             npartitions = min(len(ts_filenames), npartitions)  # cannot have less than one file per partition
             ts_files_in_each_partition = np.array_split(ts_filenames, npartitions)
+            N = len(ts_files_in_each_partition[0])
+            group_ids = list(range(npartitions))
 
-            # Read the timeseries into a dask dataframe
-            read_and_concat_ts_pq_d = dask.delayed(
-                partial(read_and_concat_enduse_timeseries_parquet, fs, all_cols=all_ts_cols_sorted)
-            )
-            ts_df = dd.from_delayed(map(read_and_concat_ts_pq_d, ts_files_in_each_partition))
-            ts_df = ts_df.set_index('building_id', sorted=True)
-
-            # Write out new dask timeseries dataframe.
+            logger.info(f"Combining about {N} parquets together. Creating {npartitions} groups.")
             if isinstance(fs, LocalFileSystem):
-                ts_out_loc = f"{ts_dir}/upgrade={upgrade_id}"
+                ts_out_loc = f"{ts_dir}/upgrade={upgrade_id}/"
             else:
                 assert isinstance(fs, S3FileSystem)
-                ts_out_loc = f"s3://{ts_dir}/upgrade={upgrade_id}"
-            logger.info(f'Writing {ts_out_loc}')
-            ts_df.to_parquet(
-                ts_out_loc,
-                engine='pyarrow',
-                flavor='spark'
+                ts_out_loc = f"s3://{ts_dir}/upgrade={upgrade_id}/"
+
+            fs.makedirs(ts_out_loc)
+            logger.info(f'Created directory {ts_out_loc} for writing.')
+            # Read the timeseries into a dask dataframe
+            read_and_concat_ts_pq_d = dask.delayed(
+                # fs, all_cols, output_dir, filenames, group_id
+                partial(read_and_concat_enduse_timeseries_parquet, fs, all_ts_cols_sorted, ts_out_loc)
             )
+            group_ids = list(range(npartitions))
+            with performance_report(filename='dask_combine_report.html'):
+                results = dask.compute(map(read_and_concat_ts_pq_d, ts_files_in_each_partition, group_ids))[0]
+
+            logger.info(f"Finished combining and saving. First 10 results: {results[:10]}")
+            failed_results = [r for r in results if r[0] == 'fail']
+            if failed_results:
+                logger.info(f"{len(failed_results)} groupings have failed. First 10 failed results: {results[:10]}")
 
 
 def remove_intermediate_files(fs, results_dir):

buildstockbatch/schemas/v0.2.yaml

@@ -50,6 +50,9 @@ hpc-spec:
 hpc-postprocessing-spec:
   time: int(required=True)
   n_workers: int(required=False)
+  node_memory_mb: enum(85248, 180224, 751616, required=False)
+  parquet_memory_mb: int(required=False)
+  keep_intermediate_files: bool(required=False)
 
 sampling-spec:
   time: int(required=True)

buildstockbatch/test/test_base.py

@@ -12,7 +12,7 @@
 import shutil
 import tempfile
 from unittest.mock import patch, MagicMock
-
+import buildstockbatch
 from buildstockbatch.base import BuildStockBatchBase
 from buildstockbatch.postprocessing import write_dataframe_as_parquet
 
@@ -21,6 +21,8 @@
 
 OUTPUT_FOLDER_NAME = 'output'
 
+buildstockbatch.postprocessing.performance_report = MagicMock()
+
 
 def test_reference_scenario(basic_residential_project_file):
     # verify that the reference_scenario get's added to the upgrade file
@@ -304,13 +306,13 @@ def test_upload_files(mocked_s3, basic_residential_project_file):
     assert (source_file_path, s3_file_path) in files_uploaded
     files_uploaded.remove((source_file_path, s3_file_path))
 
-    s3_file_path = s3_path + 'timeseries/upgrade=0/part.0.parquet'
-    source_file_path = os.path.join(source_path, 'timeseries', 'upgrade=0', 'part.0.parquet')
+    s3_file_path = s3_path + 'timeseries/upgrade=0/group0.parquet'
+    source_file_path = os.path.join(source_path, 'timeseries', 'upgrade=0', 'group0.parquet')
     assert (source_file_path, s3_file_path) in files_uploaded
     files_uploaded.remove((source_file_path, s3_file_path))
 
-    s3_file_path = s3_path + 'timeseries/upgrade=1/part.0.parquet'
-    source_file_path = os.path.join(source_path, 'timeseries', 'upgrade=1', 'part.0.parquet')
+    s3_file_path = s3_path + 'timeseries/upgrade=1/group0.parquet'
+    source_file_path = os.path.join(source_path, 'timeseries', 'upgrade=1', 'group0.parquet')
     assert (source_file_path, s3_file_path) in files_uploaded
     files_uploaded.remove((source_file_path, s3_file_path))
 

buildstockbatch/test/test_postprocessing.py

@@ -7,10 +7,11 @@
 import tarfile
 import pytest
 import shutil
-
 from buildstockbatch import postprocessing
 from buildstockbatch.base import BuildStockBatchBase
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
+
+postprocessing.performance_report = MagicMock()
 
 
 def test_report_additional_results_csv_columns(basic_residential_project_file):