memory usage logging

SkyTruth · Jan 2, 2025 · 667e450 · 667e450
1 parent 12a35bf
commit 667e450
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/cerulean_cloud/cloud_run_orchestrator/clients.py b/cerulean_cloud/cloud_run_orchestrator/clients.py
@@ -15,6 +15,7 @@
 
 import httpx
 import numpy as np
+import psutil
 import rasterio
 import skimage
 from rasterio.enums import Resampling
@@ -83,6 +84,17 @@ def __init__(
             layers, self.sceneid, tileset_envelope_bounds, image_hw_pixels
         )
 
+        # Log memory usage and aux_datasets size
+        self.logger.info(
+            structured_log(
+                "MEMORY USAGE 1",
+                severity="INFO",
+                scene_id=self.sceneid,
+                aux_datasets_size_mb=sys.getsizeof(self.aux_datasets) / (1024**2),
+                memory_usage_mb=psutil.Process().memory_info().rss / (1024**2),
+            )
+        )
+
     async def fetch_and_process_image(
         self, tile_bounds, rescale=(0, 255), num_channels=1
     ):
@@ -303,6 +315,15 @@ async def run_parallel_inference(self, tileset):
             # If asyncio.gather tasks fail, raise ValueError
             raise ValueError(f"Failed to gather inference: {e}")
         finally:
+            self.logger.info(
+                structured_log(
+                    "MEMORY USAGE 2",
+                    severity="INFO",
+                    scene_id=self.sceneid,
+                    aux_datasets_size_mb=sys.getsizeof(self.aux_datasets) / (1024**2),
+                    memory_usage_mb=psutil.Process().memory_info().rss / (1024**2),
+                )
+            )
             # Cleanup: close and delete aux_datasets
             if self.aux_datasets:
                 del self.aux_datasets