askap-vast · ddobie · May 21, 2025 · May 20, 2025 · May 20, 2025 · May 20, 2025
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,6 +20,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 #### Changed
 
+- V2: Compute associations prior to upload, and then upload in serial [#846](https://github.com/askap-vast/vast-pipeline/pull/846)
+- V2: Persist srcs_df prior to measurement pairs df output file computation [#846](https://github.com/askap-vast/vast-pipeline/pull/846)
+- V2: Change measurement and association DB indices to UUIDs [#844](https://github.com/askap-vast/vast-pipeline/pull/844)
 - V2: Limit associations upload to using num_io_workers [#833](https://github.com/askap-vast/vast-pipeline/pull/833)
 - V2: Allow user specification of dask dashboard paramters and add some further logging to dask setup [#829](https://github.com/askap-vast/vast-pipeline/pull/829)
 - V2: Replace all source, measurement, run etc. IDs with UUID indices. [#803](https://github.com/askap-vast/vast-pipeline/pull/803)
@@ -35,6 +38,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)
@@ -58,6 +64,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 #### List of PRs
 
+- [#846](https://github.com/askap-vast/vast-pipeline/pull/846): fix: V2: Compute associations prior to upload and persist srcs_df in final pairs calculation
+- [#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
 - [#829](https://github.com/askap-vast/vast-pipeline/pull/829): feat: V2: Allow user specification of dask dashboard paramters and add some further logging to dask setup

diff --git a/vast_pipeline/image/main.py b/vast_pipeline/image/main.py
@@ -6,6 +6,7 @@
 import logging
 import numpy as np
 import pandas as pd
+import uuid
 
 from django.conf import settings
 from astropy.io import fits
@@ -18,9 +19,7 @@
 
 from vast_pipeline import models
 from vast_pipeline.survey.translators import tr_selavy
-from vast_pipeline.utils.utils import (
-    generate_shortuuid, UUID_LEN_MEAS,
-)
+
 from vast_pipeline.image.utils import get_fits_header
 
 
@@ -341,7 +340,7 @@ def read_selavy(self, dj_image: models.Image) -> pd.DataFrame:
                 df[key["name"]] = df[key["name"]].astype(key["dtype"])
 
         # Add id column
-        df["id"] = df.apply(lambda _: generate_shortuuid(UUID_LEN_MEAS), axis=1)
+        df["id"] = df.apply(lambda _: str(uuid.uuid4()), axis=1)
 
         # do checks and fill in missing field for uploading sources
         # in DB (see fields in models.py -> Source model)

diff --git a/vast_pipeline/migrations/0003_alter_association_id_alter_measurement_id.py b/vast_pipeline/migrations/0003_alter_association_id_alter_measurement_id.py
@@ -0,0 +1,24 @@
+# Generated by Django 5.1.6 on 2025-03-30 23:19
+
+import uuid
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('vast_pipeline', '0002_q3c'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='association',
+            name='id',
+            field=models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False),
+        ),
+        migrations.AlterField(
+            model_name='measurement',
+            name='id',
+            field=models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False),
+        ),
+    ]
diff --git a/vast_pipeline/models.py b/vast_pipeline/models.py
@@ -1,4 +1,5 @@
 import numpy as np
+import uuid
 
 from dataclasses import dataclass
 from itertools import combinations
@@ -417,7 +418,7 @@ class Measurement(CommentableModel):
     Essentially a source single measurement in time.
     """
 
-    id = ShortUUIDField(primary_key=True, editable=False, length=UUID_LEN_MEAS, alphabet=UUID_ALPHABET)
+    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
     image = models.ForeignKey(
         Image, null=True, on_delete=models.CASCADE, to_field="id"
     )  # first image seen in
@@ -713,7 +714,7 @@ class Association(models.Model):
     some parameters
     """
 
-    id = ShortUUIDField(primary_key=True, editable=False, length=UUID_LEN_MEAS, alphabet=UUID_ALPHABET)
+    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
     source = models.ForeignKey(Source, on_delete=models.CASCADE, to_field="id")
     meas = models.ForeignKey(Measurement, on_delete=models.CASCADE, to_field="id")
 

diff --git 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)
@@ -331,8 +343,8 @@ def final_operations(
 
     # upload associations into DB
     if not __TESTING__:
-        assoc_df = associations_df_upload.loc[:, ["id", "source", "d2d", "dr"]]
-        copy_upload_associations(assoc_df, io_workers)
+        associations_df_upload = associations_df_upload.loc[:, ["id", "source", "d2d", "dr"]]
+        copy_upload_associations(associations_df_upload)
 
     # write associations to parquet file
     associations_df[['source', 'id', 'd2d', 'dr']] \
@@ -347,7 +359,7 @@ def final_operations(
         timer.reset()
         # ingest to dask data frames
         srcs_df.index.name = "source_id"
-        srcs_df = dd.from_pandas(srcs_df, npartitions=n_partitions)
+        srcs_df = dd.from_pandas(srcs_df, npartitions=n_partitions).persist()
         columns = ['id_a', 'id_b', 'flux_int_a', 'flux_int_err_a', 'flux_peak_a',
        'flux_peak_err_a', 'image_name_a', 'flux_int_b', 'flux_int_err_b',
        'flux_peak_b', 'flux_peak_err_b', 'image_name_b', 'vs_peak', 'vs_int',

diff --git a/vast_pipeline/pipeline/forced_extraction.py b/vast_pipeline/pipeline/forced_extraction.py
@@ -1,6 +1,9 @@
 import os
 import logging
 import datetime
+import gc
+import uuid
+
 import numpy as np
 import pandas as pd
 import dask.dataframe as dd
@@ -20,11 +23,7 @@
 from vast_pipeline.pipeline.loading import copy_upload_measurements
 
 from forced_phot import ForcedPhot
-from ..utils.utils import (
-    StopWatch,
-    generate_shortuuid,
-    UUID_LEN_MEAS
-)
+from ..utils.utils import StopWatch
 from vast_pipeline.image.utils import open_fits
 
 # NOTE: We check here to see if we're in a testing environment.
@@ -486,7 +485,7 @@ def _update_forced_measurements(df: pd.DataFrame) -> pd.DataFrame:
             The forced extraction dataframe updated with defaults.
         """
         df["name"] = df["name"] + f"_f_{p_run_id}"
-        df["id"] = df.apply(lambda _: generate_shortuuid(UUID_LEN_MEAS), axis=1)
+        df["id"] = df.apply(lambda _: str(uuid.uuid4()), axis=1)
         default_pos_err = settings.POS_DEFAULT_MIN_ERROR / 3600.0
         df["ra_err"] = default_pos_err
         df["dec_err"] = default_pos_err
@@ -682,7 +681,7 @@ 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()
 
     if add_mode:
         # If we are adding images to the run we assume that monitoring was
@@ -693,7 +692,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"])],
@@ -715,12 +714,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,
@@ -744,18 +745,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
@@ -1,5 +1,7 @@
 import os
 import logging
+import gc
+
 import numpy as np
 import pandas as pd
 import dask.dataframe as dd
@@ -9,6 +11,7 @@
 from itertools import islice
 from django.db import transaction, connection, models
 from contextlib import closing
+from uuid import uuid4
 
 from vast_pipeline.image.main import SelavyImage
 from vast_pipeline.pipeline.model_generator import (
@@ -36,7 +39,6 @@
     deg2hms,
     deg2dms,
     generate_shortuuid,
-    UUID_LEN_MEAS,
     UUID_LEN_SOURCE
 )
 
@@ -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())
 
@@ -403,19 +406,15 @@ def make_upload_related_sources(related_df: pd.DataFrame) -> None:
 
 def copy_upload_associations(
     associations_df: dd.DataFrame,
-    io_workers: List[str],
     batch_size: int = 10_000,
 ) -> None:
     """Upload associations using django-postgres-copy in-memory csv method.
 
     Args:
         associations_df: The associations dataframe to upload.
-        io_workers:
-            List of dask worker addresses to use for the compute.
-            This is likely the output of `DaskManager.get_n_random_workers()`.
         batch_size: The batch size. Defaults to 10_000.
     """
-    logger.info("Upload associations...")
+    logger.info("Uploading associations in batches of %d", batch_size)
     columns_to_upload = ["source"]
     for fld in Association._meta.get_fields():
         if getattr(fld, "attname", None) and fld.attname in associations_df.columns:
@@ -430,19 +429,16 @@ def copy_upload_associations(
         "d2d": "d2d",
         "dr": "dr"
     }
-
-    def upload(df, Association, mapping, batch_size):
-        df["db_id"] = df.apply(lambda _: generate_shortuuid(UUID_LEN_MEAS), axis=1)
-        copy_upload_model(df, Association, mapping=mapping, batch_size=batch_size)
-
-    associations_df = associations_df[columns_to_upload].map_partitions(upload,
-                                                                        Association,
-                                                                        mapping,
-                                                                        batch_size,
-                                                                        enforce_metadata=False,
-                                                                        meta={})
-
-    associations_df.compute(workers=io_workers)
+    timer = StopWatch()
+    associations_df = associations_df.compute()
+    logger.debug("Time to compute associations_df: %.1f s", timer.reset())
+
+    associations_df["db_id"] = associations_df.apply(lambda _: str(uuid4()), axis=1)
+    logger.debug("Time to add db_id: %.1f s", timer.reset())
+
+    copy_upload_model(associations_df, Association, mapping=mapping, batch_size=batch_size)
+    logger.debug("Time to upload associations: %.1f s", timer.reset())
+    logger.info("Associations upload complete")
 
 
 def make_upload_associations(associations_df: pd.DataFrame) -> None:

diff --git a/vast_pipeline/plots.py b/vast_pipeline/plots.py
@@ -84,6 +84,7 @@ def plot_lightcurve(
 
     # lightcurve required cols: taustart_ts, flux, flux_err_upper, flux_err_lower, forced
     lightcurve = pd.DataFrame(measurements_qs)
+    lightcurve['id'] = lightcurve['id'].astype(str)
 
     # remap method values to labels to make a better legend
     lightcurve["method"] = lightcurve.forced.map({True: "Forced", False: "Selavy"})
@@ -167,6 +168,13 @@ def plot_lightcurve(
             )
             .reset_index()
         )
+
+        candidate_measurement_pairs_df = candidate_measurement_pairs_df.astype(
+            {'measurement_a_id': 'str',
+             'measurement_b_id': 'str',
+            }
+        )
+
         g = nx.Graph()
         for _row in candidate_measurement_pairs_df.itertuples(index=False):
             g.add_edge(_row.measurement_a_id, _row.measurement_b_id)

diff --git a/vast_pipeline/serializers.py b/vast_pipeline/serializers.py
@@ -56,7 +56,7 @@ class Meta:
 
 
 class MeasurementSerializer(serializers.ModelSerializer):
-    id = serializers.CharField(read_only=True)
+    id = serializers.UUIDField(read_only=True)
     frequency = serializers.SerializerMethodField(read_only=True)
 
     def get_frequency(self, obj):

diff --git a/vast_pipeline/urls.py b/vast_pipeline/urls.py
@@ -38,6 +38,11 @@
         name="image_detail",
     ),
     path("measurements/", views.MeasurementIndex, name="measurement_index"),
+    path(
+        "measurements/<uuid:id>/",
+        views.MeasurementDetail,
+        name="measurement_detail"
+    ),
     re_path(
         fr"^measurements/(?P<id>[\w]{{{UUID_LEN_MEAS}}})(?:/(?P<action>[\w]+))?/$",
         views.MeasurementDetail,