Initial messy commit

vast_pipeline/daskmanager/config.py

@@ -19,10 +19,10 @@
 # pause and terminate conditions on the workers which can happen
 # routinely durin the IO steps in the pipeline, when a subset
 # of workers read a set of FITS files into memory.
-dask_config['distributed']['worker']['memory']['spill'] = False
 dask_config['distributed']['worker']['memory']['target'] = False
-dask_config['distributed']['worker']['memory']['terminate'] = False
+dask_config['distributed']['worker']['memory']['spill'] = 0.9
 dask_config['distributed']['worker']['memory']['pause'] = False
+dask_config['distributed']['worker']['memory']['terminate'] = False
 
 # Further distributed configuration
 # (see: https://distributed.dask.org/en/stable/worker-memory.html)

vast_pipeline/daskmanager/manager.py

@@ -4,7 +4,7 @@
 import random
 import time
 
-from dask.distributed import Client, LocalCluster
+from dask.distributed import Client, LocalCluster, Semaphore
 from django.conf import settings as s
 from . import config # noqa: F401
 
@@ -32,6 +32,9 @@ def _start_cluster():
     logger.info('Connected to local Dask Cluster')
     return client
 
+def get_semaphore(name, num_workers=5):
+    return Semaphore(name=name, max_leases=num_workers)
+
 class Singleton(type):
     _instances = {}
 
@@ -91,3 +94,17 @@ def shutdown(self):
         logger.debug("Running close...")
         self.client.close()
         logger.info("Dask Cluster shut down.")
+
+    def log_cluster_memory(self):
+        workers = self.client.scheduler_info()['workers']
+        logger.info("Logging memory usage for %d workers...", len(workers))
+        for addr, info in workers.items():
+            memory_limit = info['memory_limit'] / 1e9
+
+            mem_metrics = info['metrics']
+            managed = mem_metrics['managed_bytes'] / 1e9
+            spilled_memory = mem_metrics['spilled_bytes']['memory'] / 1e9
+            spilled_disk = mem_metrics['spilled_bytes']['disk'] / 1e9
+            memory_used = mem_metrics['memory'] / 1e9
+
+            logger.info(f"Worker {addr}: {memory_used:.2f}GB (managed: {managed:.2f}GB, spilled disk: {spilled_disk:.2f}GB, spilled memory: {spilled_memory:.2f}GB) of {memory_limit:.2f}GB.")

vast_pipeline/pipeline/finalise.py

@@ -143,8 +143,10 @@ def final_operations(
     # add the separation distance in degrees
     srcs_df["n_neighbour_dist"] = d2d.deg
 
+    print(new_sources_df)
+    print(new_sources_df.columns)
     # add new sources
-    srcs_df["new"] = srcs_df.index.isin(new_sources_df.index.compute().values)
+    srcs_df["new"] = srcs_df.index.isin(new_sources_df.index.values)
 
     mem_usage = get_df_memory_usage(srcs_df)
     logger.debug(f"srcs_df memory after nearest-neighbour: {mem_usage}MB")

vast_pipeline/pipeline/forced_extraction.py

@@ -25,6 +25,7 @@
 from forced_phot import ForcedPhot
 from ..utils.utils import StopWatch
 from vast_pipeline.image.utils import open_fits
+from vast_pipeline.daskmanager.manager import get_semaphore
 
 # NOTE: We check here to see if we're in a testing environment.
 # This is done since the django does all its testing inside an
@@ -127,10 +128,11 @@ def _forcedphot_preload(image: str,
     Returns:
         A tuple containing the HDU lists
     """
-
-    image_hdul = open_fits(image, memmap=memmap)
-    background_hdul = open_fits(background, memmap=memmap)
-    noise_hdul = open_fits(noise, memmap=memmap)
+    sem = get_semaphore('io_throttle')
+    with sem:
+        image_hdul = open_fits(image, memmap=memmap)
+        background_hdul = open_fits(background, memmap=memmap)
+        noise_hdul = open_fits(noise, memmap=memmap)
 
     return image_hdul, background_hdul, noise_hdul
 
@@ -207,6 +209,8 @@ def extract_from_image(
     )
     logger.debug("%s - Time to measure FP: %.3fs", image, FP_timer.reset())
 
+    del FP
+
     num_fits = np.sum(flux>0.0)
 
     logger.debug("%s: Obtained %d measurements "
@@ -430,7 +434,7 @@ def parallel_extraction(
     # Persist at this point uning the number of io workers.
     # df_out will contain the forced extraction measurments per image.
     # df_out should be sorted and partitioned by image at this point.
-    df_out = dd.from_delayed(func_d).persist(workers=io_workers)
+    df_out = dd.from_delayed(func_d).persist()#workers=io_workers)
 
     del out, func_d, df_per_image, measurements_parquet_data
 

vast_pipeline/pipeline/loading.py

@@ -41,6 +41,7 @@
     generate_shortuuid,
     UUID_LEN_SOURCE
 )
+from vast_pipeline.daskmanager.manager import get_semaphore
 
 
 logger = logging.getLogger(__name__)
@@ -64,6 +65,13 @@ def in_memory_csv(df: pd.DataFrame) -> StringIO:
     return mem_csv
 
 
+def _csv_upload(djmodel, mem_csv, mapping):
+    with closing(mem_csv) as csv_io:
+        num_copied = djmodel.copies.from_csv(
+            csv_io, drop_constraints=False, drop_indexes=False, mapping=mapping
+        )
+    return num_copied
+
 def copy_upload_model(
         df: pd.DataFrame,
         djmodel: models.Model,
@@ -87,14 +95,15 @@ def copy_upload_model(
 
     while start_index < total_rows:
         end_index = min(start_index + batch_size, total_rows)
+        logging.debug("Calculated end_index")
         batch = df.iloc[start_index:end_index]
-
+        logging.debug("Got batch")
         mem_csv = in_memory_csv(batch)
-        with closing(mem_csv) as csv_io:
-            num_copied = djmodel.copies.from_csv(
-                csv_io, drop_constraints=False, drop_indexes=False, mapping=mapping
-            )
-            logging.info(f"Copied {num_copied} {djmodel.__name__} objects to database.")
+        logging.debug("Got mem_csv")
+
+        with get_semaphore('csv_upload', num_workers=10):
+            num_copied = _csv_upload(djmodel, mem_csv, mapping)
+        logging.info(f"Copied {num_copied} {djmodel.__name__} objects to database.")
 
         start_index = end_index
 
@@ -406,15 +415,16 @@ def make_upload_related_sources(related_df: pd.DataFrame) -> None:
 
 def copy_upload_associations(
     associations_df: dd.DataFrame,
-    batch_size: int = 10_000,
+    batch_size: int = 1000,
 ) -> None:
     """Upload associations using django-postgres-copy in-memory csv method.
 
     Args:
         associations_df: The associations dataframe to upload.
         batch_size: The batch size. Defaults to 10_000.
     """
-    logger.info("Uploading associations in batches of %d", batch_size)
+    n_associations = len(associations_df)
+    logger.info("Uploading %d associations in batches of %d", n_associations, 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:
@@ -429,6 +439,23 @@ def copy_upload_associations(
         "d2d": "d2d",
         "dr": "dr"
     }
+
+    def upload(df, Association, mapping, batch_size):
+        logger.info("Uploading %d associations in batches of %d", len(df), batch_size)
+        df["db_id"] = df.apply(lambda _: str(uuid4()), axis=1)
+        copy_upload_model(df, Association, mapping=mapping, batch_size=batch_size)
+
+    logger.debug(f"associations partitions: {associations_df.npartitions}")
+    associations_df = associations_df[columns_to_upload].map_partitions(upload,
+                                                                        Association,
+                                                                        mapping,
+                                                                        batch_size,
+                                                                        enforce_metadata=False,
+                                                                        meta={})
+    timer = StopWatch()
+    associations_df.compute()
+    logger.info("Uploaded %d associations in %.1f s", n_associations, timer.reset())
+    """
     timer = StopWatch()
     associations_df = associations_df.compute()
     logger.debug("Time to compute associations_df: %.1f s", timer.reset())
@@ -439,6 +466,7 @@ def copy_upload_associations(
     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:

vast_pipeline/pipeline/main.py

@@ -160,6 +160,8 @@ def process_pipeline(self, p_run: Run) -> None:
             images, skyregions, bands, self.config["run"]["path"]
         )
 
+        self.dm.log_cluster_memory()
+
         # STEP #2: measurements association
         # order images by time
         logger.info("Running step #2: measurements association...")
@@ -179,6 +181,8 @@ def process_pipeline(self, p_run: Run) -> None:
         duplicate_limit = Angle(
             self.config["source_association"]["epoch_duplicate_radius"] * u.arcsec
         )
+
+        self.dm.log_cluster_memory()
 
         # 2.1 Check if sky regions to be associated can be
         # split into connected point groups
@@ -198,6 +202,8 @@ def process_pipeline(self, p_run: Run) -> None:
         else:
             done_images_df = None
             done_source_ids = None
+
+        self.dm.log_cluster_memory()
 
         # 2.2 Associate with other measurements
         if self.config["source_association"]["parallel"] and n_skyregion_groups > 1:
@@ -251,6 +257,7 @@ def process_pipeline(self, p_run: Run) -> None:
         mem_usage = get_df_memory_usage(sources_df)
         logger.debug(f"Step 2: sources_df memory usage: {mem_usage}MB")
         log_total_memory_usage()
+        self.dm.log_cluster_memory()
 
         # Obtain the number of selavy measurements for the run
         # n_selavy_measurements = sources_df.
@@ -281,6 +288,7 @@ def process_pipeline(self, p_run: Run) -> None:
         )
         del images_df
         del unforced_df
+        self.dm.log_cluster_memory()
 
         # Make missing sources into Dask dataframe
         # NOTE: This would not be necessary if the get_src_skyregion_merged_df
@@ -293,6 +301,7 @@ def process_pipeline(self, p_run: Run) -> None:
             npartitions=npartitions
         )
         wait(missing_sources_df)
+        self.dm.log_cluster_memory()
 
         # STEP #4 New source analysis
         logger.info("Running step #4: new source analysis...")
@@ -304,6 +313,7 @@ def process_pipeline(self, p_run: Run) -> None:
             p_run,
             self.dm.get_n_random_workers(self.config['processing']['num_workers_io']),
         )
+        self.dm.log_cluster_memory()
 
         # Drop column no longer required in missing_sources_df.
         missing_sources_df = missing_sources_df.drop(["in_primary"], axis=1)
@@ -329,10 +339,12 @@ def process_pipeline(self, p_run: Run) -> None:
             mem_usage = get_df_memory_usage(sources_df)
             logger.debug(f"Step 5: sources_df memory usage: {mem_usage}MB")
             log_total_memory_usage()
+            self.dm.log_cluster_memory()
 
         del missing_sources_df
 
         log_total_memory_usage()
+        self.dm.log_cluster_memory()
 
         # STEP #6: finalise the df getting unique sources, calculating
         # metrics and upload data to database
@@ -351,6 +363,7 @@ def process_pipeline(self, p_run: Run) -> None:
         )
 
         log_total_memory_usage()
+        self.dm.log_cluster_memory()
 
         # calculate number processed images
         nr_img_processed = len(images)