diff --git a/CHANGELOG.md b/CHANGELOG.md index 68d8a148a..457def5da 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/vast_pipeline/pipeline/finalise.py b/vast_pipeline/pipeline/finalise.py index 34a8b49a9..59322a14e 100644 --- a/vast_pipeline/pipeline/finalise.py +++ b/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', diff --git a/vast_pipeline/pipeline/loading.py b/vast_pipeline/pipeline/loading.py index cd9b6d9e5..437160ba8 100644 --- a/vast_pipeline/pipeline/loading.py +++ b/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: