Fix memory leaks

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

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)

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"))

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())
 

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