Add further debug logging

vast_pipeline/_version.py

@@ -1 +1 @@
-__version__ = '1.0.0dev'
+__version__ = '1.0.1dev'

vast_pipeline/image/main.py

@@ -431,11 +431,11 @@ def read_selavy(self, dj_image: models.Image) -> pd.DataFrame:
 
             logger.debug("Condon errors done.")
 
-        logger.debug("Calculating positional errors...")
         # TODO: avoid extra column given that it is a single value
         df['ew_sys_err'] = self.config["ra_uncertainty"] / 3600.
         df['ns_sys_err'] = self.config["dec_uncertainty"] / 3600.
 
+
         df['error_radius'] = calc_error_radius(
             df['ra'].values,
             df['ra_err'].values,
@@ -446,17 +446,15 @@ def read_selavy(self, dj_image: models.Image) -> pd.DataFrame:
         df['uncertainty_ew'] = np.hypot(
             df['ew_sys_err'].values, df['error_radius'].values
         )
-
+        
         df['uncertainty_ns'] = np.hypot(
             df['ns_sys_err'].values, df['error_radius'].values
         )
 
         # weight calculations to use later
         df['weight_ew'] = 1. / df['uncertainty_ew'].values**2
         df['weight_ns'] = 1. / df['uncertainty_ns'].values**2
-
-        logger.debug('Positional errors done.')
-
+
         # Initialise the forced column as False
         df['forced'] = False
 

vast_pipeline/pipeline/association.py

@@ -934,7 +934,7 @@ def advanced_association(
             association.
     '''
     # read the needed sources fields
-    # Step 1: get matches within semimajor axis of image.
+    # Step 1: get matches within semimajor axis of image.    
     idx_skyc1, idx_skyc2, d2d, d3d = skyc2.search_around_sky(
         skyc1, bw_max
     )
@@ -1298,7 +1298,7 @@ def association(
         )
 
         sources_df = sources_df.drop(['ra_wrap'], axis=1)
-
+        
         tmp_srcs_df = (
             sources_df.loc[
                 (sources_df['source'] != -1) & (sources_df['forced'] == False),
@@ -1334,7 +1334,9 @@ def association(
                     'weight_ns': 'uncertainty_ns'
             })
         )
-
+        nan_indices = weighted_df.query("ra.isnull()", engine='python').index
+        nan_sources = weighted_df.iloc[nan_indices].source.values
+
         # correct the RA wrapping
         ra_wrap_mask = weighted_df.ra >= 360.
         weighted_df.loc[

vast_pipeline/pipeline/finalise.py

@@ -14,7 +14,7 @@
     update_sources
 )
 from vast_pipeline.pipeline.pairs import calculate_measurement_pair_metrics
-from vast_pipeline.pipeline.utils import parallel_groupby
+from vast_pipeline.pipeline.utils import parallel_groupby, get_memory_usage
 
 
 logger = logging.getLogger(__name__)
@@ -132,8 +132,15 @@ def final_operations(
         'Calculating statistics for %i sources...',
         sources_df.source.unique().shape[0]
     )
+    mem_usage = sources_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"sources_df memory: {mem_usage}MB")
+
     srcs_df = parallel_groupby(sources_df)
+
     logger.info('Groupby-apply time: %.2f seconds', timer.reset())
+    mem_usage = srcs_df.memory_usage(deep=True).sum() / 1e6
+    logger.debug(f"Initial srcs_df memory: {mem_usage}MB")
+    logger.debug(get_memory_usage())
 
     # add new sources
     srcs_df["new"] = srcs_df.index.isin(new_sources_df.index)
@@ -145,6 +152,9 @@ 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")
+    logger.debug(get_memory_usage())
 
     # calculate nearest neighbour
     srcs_skycoord = SkyCoord(
@@ -159,12 +169,19 @@ 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")
+    logger.debug(get_memory_usage())
 
     # create measurement pairs, aka 2-epoch metrics
     if calculate_pairs:
+        sources_df.to_parquet('calcalate_measurement_pair_metrics_input_df.parquet')
         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")
+        logger.debug(get_memory_usage())
 
         # 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 +206,9 @@ 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")
+        logger.debug(get_memory_usage())
     else:
         logger.info(
             "Skipping measurement pair metric calculation as specified in the run configuration."
@@ -201,18 +221,36 @@ 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")
+        logger.debug(get_memory_usage())
+
         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")
+        logger.debug(get_memory_usage())
         # 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.debug(get_memory_usage())
         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")
+        logger.debug(get_memory_usage())
         # 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")
+    logger.debug(get_memory_usage())
+
     # gather the related df, upload to db and save to parquet file
     # the df will look like
     #
@@ -230,11 +268,17 @@ 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")
+    logger.debug(get_memory_usage())
 
     # 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")
+    logger.debug(get_memory_usage())
 
     # write symmetrical relations to parquet
     related_df.to_parquet(
@@ -256,7 +300,14 @@ 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")
+        logger.debug(get_memory_usage())
+
     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")
+    logger.debug(get_memory_usage())
 
     del related_df
 
@@ -272,6 +323,9 @@ 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")
+    logger.debug(get_memory_usage())
 
     if add_mode:
         # Load old associations so the already uploaded ones can be removed

vast_pipeline/pipeline/forced_extraction.py

@@ -166,13 +166,15 @@ def extract_from_image(
         Dictionary with input dataframe with added columns (flux_int,
             flux_int_err, chi_squared_fit) and image name.
     """
+    timer = StopWatch()
     # create the skycoord obj to pass to the forced extraction
     # see usage https://github.com/dlakaplan/forced_phot
     P_islands = SkyCoord(
         df['wavg_ra'].values,
         df['wavg_dec'].values,
         unit=(u.deg, u.deg)
     )
+
     # load the image, background and noisemaps into memory
     # a dedicated function may seem unneccesary, but will be useful if we
     # split the load to a separate thread.
@@ -183,12 +185,14 @@ def extract_from_image(
                                            )
     FP = ForcedPhot(*forcedphot_input)
 
+
     flux, flux_err, chisq, DOF, cluster_id = FP.measure(
         P_islands,
         cluster_threshold=cluster_threshold,
         allow_nan=allow_nan,
         edge_buffer=edge_buffer
     )
+    logger.debug(f"Time to measure FP for {image}: {timer.reset()}s")
     df['flux_int'] = flux * 1.e3
     df['flux_int_err'] = flux_err * 1.e3
     df['chi_squared_fit'] = chisq
@@ -238,7 +242,7 @@ def finalise_forced_dfs(
     df['component_id'] = df['island_id'].str.replace(
         'island', 'component'
     ) + 'a'
-    img_prefix = image.split('.')[0] + '_'
+    img_prefix = ""#image.split('.')[0] + '_'
     df['name'] = img_prefix + df['component_id']
     # assign all the other columns
     # convert fluxes to mJy
@@ -390,7 +394,8 @@ def image_data_func(image_name: str) -> Dict[str, Any]:
     )
     del col_to_drop
 
-    n_cpu = cpu_count() - 1
+    #n_cpu = cpu_count() - 1 # this doesn't work because cpu_count returns the number of CPUs in the machine, not the container.
+    n_cpu = 10
     bags = db.from_sequence(list_to_map, npartitions=len(list_to_map))
     forced_dfs = (
         bags.map(lambda x: extract_from_image(
@@ -399,7 +404,7 @@ def image_data_func(image_name: str) -> Dict[str, Any]:
             allow_nan=allow_nan,
             **x
         ))
-        .compute()
+        .compute(scheduler='processes', num_workers=n_cpu)
     )
     del bags
     # create intermediates dfs combining the mapping data and the forced
@@ -477,7 +482,7 @@ def get_fname(n): return os.path.join(
         'forced_measurements_' + n.replace('.', '_') + '.parquet'
     )
     dfs = list(map(lambda x: (df[df['image'] == x], get_fname(x)), images))
-    n_cpu = cpu_count() - 1
+    n_cpu = 10 #cpu_count() - 1 # temporarily hardcode n_cpu
 
     # writing parquets using Dask bag
     bags = db.from_sequence(dfs)
@@ -610,7 +615,7 @@ def forced_extraction(
     )
 
     # make measurement names unique for db constraint
-    extr_df['name'] = extr_df['name'] + f'_f_run{p_run.id:06d}'
+    extr_df['name'] = extr_df['name'] + f'_f_run{p_run.id:03d}'
 
     # select sensible flux values and set the columns with fix values
     values = {
@@ -678,6 +683,10 @@ def forced_extraction(
     extr_df = extr_df[col_order + remaining]
 
     # upload the measurements, a column 'id' is returned with the DB id
+    long_names = extr_df.loc[extr_df['name'].str.len() > 63]
+    long_comps = extr_df.loc[extr_df['component_id'].str.len() > 63]
+    long_isls = extr_df.loc[extr_df['island_id'].str.len() > 63]
+
     extr_df = make_upload_measurements(extr_df)
 
     extr_df = extr_df.rename(columns={'source_tmp_id': 'source'})

vast_pipeline/pipeline/loading.py

@@ -18,7 +18,9 @@
     Association, Band, Measurement, SkyRegion, Source, RelatedSource,
     Run, Image
 )
-from vast_pipeline.pipeline.utils import get_create_img, get_create_img_band
+from vast_pipeline.pipeline.utils import (
+    get_create_img, get_create_img_band, get_memory_usage
+)
 from vast_pipeline.utils.utils import StopWatch
 
 
@@ -31,6 +33,7 @@ def bulk_upload_model(
     generator: Iterable[Generator[models.Model, None, None]],
     batch_size: int=10_000,
     return_ids: bool=False,
+    log_mem_usage: bool=False,
 ) -> List[int]:
     '''
     Bulk upload a list of generator objects of django models to db.
@@ -55,6 +58,8 @@ def bulk_upload_model(
     bulk_ids = []
     while True:
         items = list(islice(generator, batch_size))
+        if log_mem_usage:
+            logger.debug(get_memory_usage())
         if not items:
             break
         out_bulk = djmodel.objects.bulk_create(items)
@@ -168,6 +173,10 @@ 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")
+    logger.debug(get_memory_usage())
     # create sources in DB
     with transaction.atomic():
         if (add_mode is False and
@@ -207,6 +216,9 @@ 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")
+    logger.debug(get_memory_usage())
     bulk_upload_model(RelatedSource, related_models_generator(related_df))
 
 
@@ -223,6 +235,7 @@ def make_upload_associations(associations_df: pd.DataFrame) -> None:
         None.
     """
     logger.info('Upload associations...')
+
     assoc_chunk_size = 100000
     for i in range(0,len(associations_df),assoc_chunk_size):
         bulk_upload_model(
@@ -243,6 +256,10 @@ 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")
+    logger.debug(get_memory_usage())
     meas_dj_ids = bulk_upload_model(
         Measurement,
         measurement_models_generator(measurements_df),