diff --git a/CHANGELOG.md b/CHANGELOG.md index 740742e99..68d8a148a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), #### Fixed +- V2: Fixed duplicate source ID error by switching from sources_df.repartition() to sources_df.shuffle() in pipeline.finalise.final_operations [#845](https://github.com/askap-vast/vast-pipeline/pull/845) +- V2: Tentative fix for image upload memory leak via garbage collect [#845](https://github.com/askap-vast/vast-pipeline/pull/845) +- V2: Fixed dd.concat memory blow-up in pipeline.forced_extraction.forced_extraction by persisting both dataframes prior [#845](https://github.com/askap-vast/vast-pipeline/pull/845) - V2: Pair metrics working with V2 dask LocalCluster changes [#817](https://github.com/askap-vast/vast-pipeline/pull/817) - V2: Fixed missing JS9 overlays on source webpage [#809](https://github.com/askap-vast/vast-pipeline/pull/809) - V2: Fix bug when deleting source tags from database which are referenced by multiple sources [#803](https://github.com/askap-vast/vast-pipeline/pull/803) @@ -59,6 +62,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), #### List of PRs +- [#845](https://github.com/askap-vast/vast-pipeline/pull/845): fix: V2: Fix memory leaks and duplicate source ID issue - [#844](https://github.com/askap-vast/vast-pipeline/pull/844): fix: V2: Change measurements and association DB indices to UUIDs - [#843](https://github.com/askap-vast/vast-pipeline/pull/843): feat: V2: Enable specification of Dask worker memory limits - [#833](https://github.com/askap-vast/vast-pipeline/pull/833): feat: V2: Limit associations upload to using num_io_workers diff --git a/vast_pipeline/pipeline/finalise.py b/vast_pipeline/pipeline/finalise.py index afec07aea..34a8b49a9 100644 --- a/vast_pipeline/pipeline/finalise.py +++ b/vast_pipeline/pipeline/finalise.py @@ -1,9 +1,11 @@ import os import logging import warnings + import pandas as pd import pyarrow as pa import dask.dataframe as dd +import dask.config as dc from astropy import units as u from astropy.coordinates import SkyCoord @@ -12,7 +14,10 @@ from vast_pipeline.models import Run from vast_pipeline.utils.utils import ( - StopWatch, optimise_numeric, delete_file_or_dir + StopWatch, + optimise_numeric, + delete_file_or_dir, + calculate_n_partitions ) from vast_pipeline.pipeline.loading import ( update_sources, @@ -96,8 +101,15 @@ def final_operations( logger.info("Calculating statistics for sources...") log_total_memory_usage() - sources_df = sources_df.set_index("source") \ - .repartition(partition_size=f"{upload_chunk_size_mb}MB") + npartitions = calculate_n_partitions( + sources_df, + partition_size_mb=upload_chunk_size_mb + ) + + with dc.set({"dataframe.shuffle.method": "p2p"}): + sources_df = sources_df.set_index("source") \ + .shuffle(npartitions=npartitions, on_index=True) + srcs_df = parallel_groupby(sources_df) mem_usage = get_df_memory_usage(srcs_df) diff --git a/vast_pipeline/pipeline/forced_extraction.py b/vast_pipeline/pipeline/forced_extraction.py index 809fab5c0..fb7dd9f26 100644 --- a/vast_pipeline/pipeline/forced_extraction.py +++ b/vast_pipeline/pipeline/forced_extraction.py @@ -1,6 +1,7 @@ import os import logging import datetime +import gc import uuid import numpy as np @@ -680,7 +681,8 @@ def forced_extraction( # Explode out the img_diff column. extr_df = extr_df.explode("img_diff").reset_index() - total_to_extract = extr_df.shape[0] + total_to_extract = extr_df.shape[0].compute() + logger.info("%d forced measurements to extract", total_to_extract) if add_mode: # If we are adding images to the run we assume that monitoring was @@ -691,7 +693,7 @@ def forced_extraction( # images. # 3. A new relation has been created and they need the forced # measuremnts filled in (actually covered by 2.) - total_to_extract = extr_df.shape[0].compute() + logger.info("Running extraction in add mode...") extr_df = dd.concat( [ extr_df[~extr_df["img_diff"].isin(done_images_df["name"])], @@ -713,12 +715,14 @@ def forced_extraction( min_sigma, edge_buffer, cluster_threshold, allow_nan, add_mode, p_run.path, io_workers ) + logger.info("Completed parallel extraction step.") # Dask needs type metadata for map_partitions sources_meta = dd.utils.make_meta(sources_df).drop(['epoch', 'interim_ns', 'interim_ew'], axis=1) # Get expected database measurements schema columns = read_schema(images_df.iloc[0]["measurements_path"]).names + logger.debug("Building save and upload...") extr_df = extr_df.map_partitions(save_and_upload_forced_df, p_run_path=p_run.path, p_run_id=p_run.id, @@ -742,18 +746,27 @@ def forced_extraction( else: extr_df["epoch"] = sources_df['epoch'].compute().iloc[0] + logger.info("Persisting extr_df...") + extr_df = extr_df.persist() + wait(extr_df) + logger.info("Persisted extr_df.") + sources_df = dd.concat( - [sources_df, extr_df] + [sources_df, extr_df], + interleave_partitions=True ) # Wait for the forced extraction step to complete # NOTE: Ideally we would have some optimised way of sorting sources_df # by source id at this point to avoid needing to `set_index` on it # during the finalise step. + logger.info("Persisting sources_df...") sources_df = sources_df.persist() wait(sources_df) + logger.info("Persisted sources_df") del extr_df + gc.collect() # get the number of forced extractions for the run forced_parquets = glob(os.path.join(p_run.path, "forced_measurements*.parquet")) diff --git a/vast_pipeline/pipeline/loading.py b/vast_pipeline/pipeline/loading.py index 0d5ef6e19..cd9b6d9e5 100644 --- a/vast_pipeline/pipeline/loading.py +++ b/vast_pipeline/pipeline/loading.py @@ -1,5 +1,7 @@ import os import logging +import gc + import numpy as np import pandas as pd import dask.dataframe as dd @@ -196,7 +198,8 @@ def make_upload_images( os.makedirs(base_folder) measurements.to_parquet(img.measurements_path, index=False) - del measurements, image, band, img + del measurements, image, band + gc.collect() logger.info("Total images upload/loading time: %.2f seconds", timer.reset_init()) diff --git a/vast_pipeline/utils/utils.py b/vast_pipeline/utils/utils.py index f8b605fad..c9fef6f4f 100644 --- a/vast_pipeline/utils/utils.py +++ b/vast_pipeline/utils/utils.py @@ -393,8 +393,8 @@ def timeStamped(fname, fmt="%Y-%m-%d-%H-%M-%S_{fname}"): def calculate_n_partitions( - df: pd.DataFrame, - n_cpu: int, + df: Union[pd.DataFrame, dd.DataFrame], + n_cpu: Optional[int] = None, partition_size_mb: Optional[int] = 15 ) -> int: """ @@ -427,15 +427,17 @@ def calculate_n_partitions( n_partitions = int(np.ceil(mem_usage_mb / partition_size_mb)) # n_partitions should be >= n_cpu for optimal parallel processing - if n_partitions < n_cpu: - n_partitions = n_cpu + if n_cpu is not None: + if n_partitions < n_cpu: + n_partitions = n_cpu - partition_size_mb = int(np.ceil(mem_usage_mb / n_partitions)) + partition_size_mb = mem_usage_mb / n_partitions logger.debug( - "Using %d partitions of %dMB", + "Using %d partitions of %.1fMB", n_partitions, - partition_size_mb) + partition_size_mb + ) return n_partitions