Add further debug logging

vast_pipeline/pipeline/finalise.py

@@ -133,6 +133,8 @@ def final_operations(
         sources_df.source.unique().shape[0]
     )
     srcs_df = parallel_groupby(sources_df)
+    mem_usage = srcs_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"Initial srcs_df memory: {mem_usage}MB")
     logger.info('Groupby-apply time: %.2f seconds', timer.reset())
 
     # add new sources
@@ -145,6 +147,8 @@ def final_operations(
         how="left",
     )
     srcs_df["new_high_sigma"] = srcs_df["new_high_sigma"].fillna(0.0)
+    mem_usage = srcs_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"srcs_df memory after adding new sources: {mem_usage}MB")
 
     # calculate nearest neighbour
     srcs_skycoord = SkyCoord(
@@ -159,12 +163,16 @@ def final_operations(
 
     # add the separation distance in degrees
     srcs_df['n_neighbour_dist'] = d2d.deg
+    mem_usage = srcs_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"srcs_df memory after nearest-neighbour: {mem_usage}MB")
 
     # create measurement pairs, aka 2-epoch metrics
     if calculate_pairs:
         timer.reset()
         measurement_pairs_df = calculate_measurement_pair_metrics(sources_df)
         logger.info('Measurement pair metrics time: %.2f seconds', timer.reset())
+        mem_usage = measurement_pairs_df.memory_usage(deep=True).sum() / 1e6
+        logger.debug(f"measurement_pairs_df memory: {mem_usage}MB")
 
         # calculate measurement pair metric aggregates for sources by finding the row indices
         # of the aggregate max of the abs(m) metric for each flux type.
@@ -189,6 +197,8 @@ def final_operations(
             "m_abs_significant_max_int": 0.0,
         })
         logger.info("Measurement pair aggregate metrics time: %.2f seconds", timer.reset())
+        mem_usage = srcs_df.memory_usage(deep=True).sum() / 1e6
+        logger.debug(f"srcs_df memory after calculate_pairs: {mem_usage}MB")
     else:
         logger.info(
             "Skipping measurement pair metric calculation as specified in the run configuration."
@@ -201,18 +211,31 @@ def final_operations(
         # upload new ones first (new id's are fetched)
         src_done_mask = srcs_df.index.isin(done_source_ids)
         srcs_df_upload = srcs_df.loc[~src_done_mask].copy()
+        mem_usage = srcs_df_upload.memory_usage(deep=True).sum() / 1e6
+        logger.debug(f"srcs_df_upload initial memory: {mem_usage}MB")
+
         srcs_df_upload = make_upload_sources(srcs_df_upload, p_run, add_mode)
+        mem_usage = srcs_df_upload.memory_usage(deep=True).sum() / 1e6
+        logger.debug(f"srcs_df_upload memory after upload: {mem_usage}MB")
         # And now update
         srcs_df_update = srcs_df.loc[src_done_mask].copy()
+        mem_usage = srcs_df_update.memory_usage(deep=True).sum() / 1e6
+        logger.debug(f"srcs_df_update memory: {mem_usage}MB")
         logger.info(
             f"Updating {srcs_df_update.shape[0]} sources with new metrics.")
+
         srcs_df = update_sources(srcs_df_update, batch_size=1000)
+        mem_usage = srcs_df_update.memory_usage(deep=True).sum() / 1e6
+        logger.debug(f"srcs_df_update memory after update: {mem_usage}MB")
         # Add back together
         if not srcs_df_upload.empty:
             srcs_df = pd.concat([srcs_df, srcs_df_upload])
     else:
         srcs_df = make_upload_sources(srcs_df, p_run, add_mode)
 
+    mem_usage = srcs_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"srcs_df memory after upload_sources: {mem_usage}MB")
+
     # gather the related df, upload to db and save to parquet file
     # the df will look like
     #
@@ -230,11 +253,15 @@ def final_operations(
         .explode("related_list")
         .rename(columns={"id": "from_source_id", "related_list": "to_source_id"})
     )
+    mem_usage = related_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"related_df memory: {mem_usage}MB")
 
     # for the column 'from_source_id', replace relation source ids with db id
     related_df["to_source_id"] = related_df["to_source_id"].map(srcs_df["id"].to_dict())
     # drop relationships with the same source
     related_df = related_df[related_df["from_source_id"] != related_df["to_source_id"]]
+    mem_usage = related_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"related_df memory after calcs: {mem_usage}MB")
 
     # write symmetrical relations to parquet
     related_df.to_parquet(
@@ -256,7 +283,12 @@ def final_operations(
         )
         logger.debug(f'Add mode: #{related_df.shape[0]} relations to upload.')
 
+        mem_usage = related_df.memory_usage(deep=True).sum() / 1e6
+        logger.debug(f"related_df memory after partitioning: {mem_usage}MB")
+
     make_upload_related_sources(related_df)
+    mem_usage = related_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"related_df memory after upload: {mem_usage}MB")
 
     del related_df
 
@@ -272,6 +304,8 @@ def final_operations(
         sources_df.drop('related', axis=1)
         .merge(srcs_df.rename(columns={'id': 'source_id'}), on='source')
     )
+    mem_usage = sources_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"sources_df memory after srcs_df merge: {mem_usage}MB")
 
     if add_mode:
         # Load old associations so the already uploaded ones can be removed

vast_pipeline/pipeline/loading.py

@@ -165,6 +165,9 @@ def make_upload_sources(
     Returns:
         The input dataframe with the 'id' column added.
     '''
+    logger.debug("Uploading sources...")
+    mem_usage = sources_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"sources_df memory usage: {mem_usage}MB")
     # create sources in DB
     with transaction.atomic():
         if (add_mode is False and
@@ -204,6 +207,8 @@ def make_upload_related_sources(related_df: pd.DataFrame) -> None:
         None.
     """
     logger.info('Populate "related" field of sources...')
+    mem_usage = related_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"related_df memory usage: {mem_usage}MB")
     bulk_upload_model(RelatedSource, related_models_generator(related_df))
 
 
@@ -220,6 +225,8 @@ def make_upload_associations(associations_df: pd.DataFrame) -> None:
         None.
     """
     logger.info('Upload associations...')
+    mem_usage = associations_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"associations_df memory usage: {mem_usage}MB")
     bulk_upload_model(
         Association, association_models_generator(associations_df)
     )
@@ -237,6 +244,9 @@ def make_upload_measurements(measurements_df: pd.DataFrame) -> pd.DataFrame:
     Returns:
         Original DataFrame with the database ID attached to each row.
     """
+    logger.info("Upload measurements...")
+    mem_usage = measurements_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"measurements_df memory usage: {mem_usage}MB")
     meas_dj_ids = bulk_upload_model(
         Measurement,
         measurement_models_generator(measurements_df),

vast_pipeline/pipeline/main.py

@@ -232,7 +232,8 @@ def process_pipeline(self, p_run: Run) -> None:
                 self.previous_parquets,
                 done_images_df
             )
-
+        mem_usage = sources_df.memory_usage(deep=True).sum() / 1e6
+        logger.debug(f"Step 2: sources_df memory usage: {mem_usage}MB")
         # Obtain the number of selavy measurements for the run
         # n_selavy_measurements = sources_df.
         nr_selavy_measurements = sources_df['id'].unique().shape[0]
@@ -285,6 +286,8 @@ def process_pipeline(self, p_run: Run) -> None:
                 done_images_df,
                 done_source_ids
             )
+            mem_usage = sources_df.memory_usage(deep=True).sum() / 1e6
+            logger.debug(f"Step 5: sources_df memory usage: {mem_usage}MB")
 
         del missing_sources_df