Tweak associations upload

CHANGELOG.md

@@ -20,6 +20,8 @@ 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)
@@ -62,6 +64,7 @@ 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

vast_pipeline/pipeline/finalise.py

@@ -343,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']] \
@@ -359,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',

vast_pipeline/pipeline/loading.py

@@ -406,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:
@@ -433,19 +429,16 @@ def copy_upload_associations(
         "d2d": "d2d",
         "dr": "dr"
     }
-
-    def upload(df, Association, mapping, batch_size):
-        df["db_id"] = df.apply(lambda _: str(uuid4()), 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: