rslearn 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

rslearn/data_sources/local_files.py

@@ -2,12 +2,12 @@
 
 import functools
 import json
+from collections.abc import Callable
 from typing import Any, Generic, TypeVar
 
 import fiona
 import shapely
 import shapely.geometry
-from class_registry import ClassRegistry
 from rasterio.crs import CRS
 from upath import UPath
 
@@ -23,7 +23,24 @@ from rslearn.utils.geometry import Projection, STGeometry, get_global_geometry
 from .data_source import DataSource, Item, QueryConfig
 
 logger = get_logger("__name__")
-Importers = ClassRegistry()
+_ImporterT = TypeVar("_ImporterT", bound="Importer")
+
+
+class _ImporterRegistry(dict[str, type["Importer"]]):
+    """Registry for Importer classes."""
+
+    def register(self, name: str) -> Callable[[type[_ImporterT]], type[_ImporterT]]:
+        """Decorator to register an importer class."""
+
+        def decorator(cls: type[_ImporterT]) -> type[_ImporterT]:
+            self[name] = cls
+            return cls
+
+        return decorator
+
+
+Importers = _ImporterRegistry()
+
 
 ItemType = TypeVar("ItemType", bound=Item)
 LayerConfigType = TypeVar("LayerConfigType", bound=LayerConfig)
@@ -425,7 +442,7 @@ class LocalFiles(DataSource):
         """
         self.config = config
 
-        self.importer = Importers[config.layer_type.value]
+        self.importer = Importers[config.layer_type.value]()
         self.src_dir = src_dir
 
     @staticmethod

rslearn/data_sources/planetary_computer.py

@@ -83,6 +83,10 @@ class PlanetaryComputer(DataSource, TileStore):
 
     STAC_ENDPOINT = "https://planetarycomputer.microsoft.com/api/stac/v1"
 
+    # Default threshold for recreating the STAC client to prevent memory leaks
+    # from the pystac Catalog's resolved objects cache growing unbounded
+    DEFAULT_MAX_ITEMS_PER_CLIENT = 1000
+
     def __init__(
         self,
         collection_name: str,
@@ -93,6 +97,7 @@ class PlanetaryComputer(DataSource, TileStore):
         timeout: timedelta = timedelta(seconds=10),
         skip_items_missing_assets: bool = False,
         cache_dir: UPath | None = None,
+        max_items_per_client: int | None = None,
     ):
         """Initialize a new PlanetaryComputer instance.
 
@@ -109,6 +114,9 @@ class PlanetaryComputer(DataSource, TileStore):
             cache_dir: optional directory to cache items by name, including asset URLs.
                 If not set, there will be no cache and instead STAC requests will be
                 needed each time.
+            max_items_per_client: number of STAC items to process before recreating
+                the client to prevent memory leaks from the resolved objects cache.
+                Defaults to DEFAULT_MAX_ITEMS_PER_CLIENT.
         """
         self.collection_name = collection_name
         self.asset_bands = asset_bands
@@ -118,12 +126,15 @@ class PlanetaryComputer(DataSource, TileStore):
         self.timeout = timeout
         self.skip_items_missing_assets = skip_items_missing_assets
         self.cache_dir = cache_dir
+        self.max_items_per_client = (
+            max_items_per_client or self.DEFAULT_MAX_ITEMS_PER_CLIENT
+        )
 
         if self.cache_dir is not None:
             self.cache_dir.mkdir(parents=True, exist_ok=True)
 
         self.client: pystac_client.Client | None = None
-        self.collection: pystac_client.CollectionClient | None = None
+        self._client_item_count = 0
 
     @staticmethod
     def from_config(config: RasterLayerConfig, ds_path: UPath) -> "PlanetaryComputer":
@@ -142,7 +153,12 @@ class PlanetaryComputer(DataSource, TileStore):
         if "cache_dir" in d:
             kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
 
-        simple_optionals = ["query", "sort_by", "sort_ascending"]
+        simple_optionals = [
+            "query",
+            "sort_by",
+            "sort_ascending",
+            "max_items_per_client",
+        ]
         for k in simple_optionals:
             if k in d:
                 kwargs[k] = d[k]
@@ -151,20 +167,40 @@ class PlanetaryComputer(DataSource, TileStore):
 
     def _load_client(
         self,
-    ) -> tuple[pystac_client.Client, pystac_client.CollectionClient]:
+    ) -> pystac_client.Client:
         """Lazily load pystac client.
 
         We don't load it when creating the data source because it takes time and caller
         may not be calling get_items. Additionally, loading it during the get_items
         call enables leveraging the retry loop functionality in
         prepare_dataset_windows.
-        """
-        if self.client is not None:
-            return self.client, self.collection
 
+        Note: We periodically recreate the client to prevent memory leaks from the
+        pystac Catalog's resolved objects cache, which grows unbounded as STAC items
+        are deserialized and cached. The cache cannot be cleared or disabled.
+        """
+        if self.client is None:
+            logger.info("Creating initial STAC client")
+            self.client = pystac_client.Client.open(self.STAC_ENDPOINT)
+            return self.client
+
+        if self._client_item_count < self.max_items_per_client:
+            return self.client
+
+        # Recreate client to clear the resolved objects cache
+        current_client = self.client
+        logger.debug(
+            "Recreating STAC client after processing %d items (threshold: %d)",
+            self._client_item_count,
+            self.max_items_per_client,
+        )
+        client_root = current_client.get_root()
+        client_root.clear_links()
+        client_root.clear_items()
+        client_root.clear_children()
+        self._client_item_count = 0
         self.client = pystac_client.Client.open(self.STAC_ENDPOINT)
-        self.collection = self.client.get_collection(self.collection_name)
-        return self.client, self.collection
+        return self.client
 
     def _stac_item_to_item(self, stac_item: pystac.Item) -> PlanetaryComputerItem:
         shp = shapely.geometry.shape(stac_item.geometry)
@@ -210,10 +246,26 @@ class PlanetaryComputer(DataSource, TileStore):
 
         # No cache or not in cache, so we need to make the STAC request.
         logger.debug("Getting STAC item {name}")
-        _, collection = self._load_client()
-        stac_item = collection.get_item(name)
+        client = self._load_client()
+
+        search_result = client.search(ids=[name], collections=[self.collection_name])
+        stac_items = list(search_result.items())
+
+        if not stac_items:
+            raise ValueError(
+                f"Item {name} not found in collection {self.collection_name}"
+            )
+        if len(stac_items) > 1:
+            raise ValueError(
+                f"Multiple items found for ID {name} in collection {self.collection_name}"
+            )
+
+        stac_item = stac_items[0]
         item = self._stac_item_to_item(stac_item)
 
+        # Track items processed for client recreation threshold (after deserialization)
+        self._client_item_count += 1
+
         # Finally we cache it if cache_dir is set.
         if cache_fname is not None:
             with cache_fname.open("w") as f:
@@ -233,7 +285,7 @@ class PlanetaryComputer(DataSource, TileStore):
         Returns:
             List of groups of items that should be retrieved for each geometry.
         """
-        client, _ = self._load_client()
+        client = self._load_client()
 
         groups = []
         for geometry in geometries:
@@ -247,7 +299,9 @@ class PlanetaryComputer(DataSource, TileStore):
                 datetime=wgs84_geometry.time_range,
                 query=self.query,
             )
-            stac_items = [item for item in result.item_collection()]
+            stac_items = [item for item in result.items()]
+            # Track items processed for client recreation threshold (after deserialization)
+            self._client_item_count += len(stac_items)
             logger.debug("STAC search yielded %d items", len(stac_items))
 
             if self.skip_items_missing_assets:
@@ -580,7 +634,13 @@ class Sentinel2(PlanetaryComputer):
         if "cache_dir" in d:
             kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
 
-        simple_optionals = ["harmonize", "query", "sort_by", "sort_ascending"]
+        simple_optionals = [
+            "harmonize",
+            "query",
+            "sort_by",
+            "sort_ascending",
+            "max_items_per_client",
+        ]
         for k in simple_optionals:
             if k in d:
                 kwargs[k] = d[k]
@@ -756,7 +816,12 @@ class Sentinel1(PlanetaryComputer):
         if "cache_dir" in d:
             kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
 
-        simple_optionals = ["query", "sort_by", "sort_ascending"]
+        simple_optionals = [
+            "query",
+            "sort_by",
+            "sort_ascending",
+            "max_items_per_client",
+        ]
         for k in simple_optionals:
             if k in d:
                 kwargs[k] = d[k]

rslearn/dataset/handler_summaries.py

@@ -0,0 +1,130 @@
+"""This module contains dataclasses for summarizing the results of dataset operations.
+
+They can be used by callers to emit telemetry / logs, or discarded.
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class LayerPrepareSummary:
+    """Results for preparing a single layer."""
+
+    # Identity
+    layer_name: str
+    data_source_name: str
+
+    # Timing
+    duration_seconds: float
+
+    # Counts
+    windows_prepared: int
+    windows_skipped: int
+    get_items_attempts: int
+
+
+@dataclass
+class PrepareDatasetWindowsSummary:
+    """Results from prepare_dataset_windows operation for telemetry purposes."""
+
+    # Timing
+    duration_seconds: float
+
+    # Counts
+    total_windows_requested: int
+
+    # Per-layer summaries
+    layer_summaries: list[LayerPrepareSummary]
+
+
+@dataclass
+class IngestCounts:
+    """Known ingestion counts."""
+
+    items_ingested: int
+    geometries_ingested: int
+
+
+@dataclass
+class UnknownIngestCounts:
+    """Indicates ingestion counts are unknown due to partial failure."""
+
+    items_attempted: int
+    geometries_attempted: int
+
+
+@dataclass
+class LayerIngestSummary:
+    """Results for ingesting a single layer."""
+
+    # Identity
+    layer_name: str
+    data_source_name: str
+
+    # Timing
+    duration_seconds: float
+
+    # Counts - either known or unknown
+    ingest_counts: IngestCounts | UnknownIngestCounts
+    ingest_attempts: int
+
+
+@dataclass
+class IngestDatasetJobsSummary:
+    """Results from ingesting a set of jobs; for telemetry purposes."""
+
+    # Timing
+    duration_seconds: float
+
+    # Counts
+    num_jobs: int
+
+    # Per-layer summaries
+    layer_summaries: list[LayerIngestSummary]
+
+
+@dataclass
+class MaterializeWindowLayerSummary:
+    """Results for materializing a single window layer."""
+
+    skipped: bool
+    materialize_attempts: int
+
+
+@dataclass
+class MaterializeWindowLayersSummary:
+    """Results for materialize a given layer for all windows in a materialize call."""
+
+    # Identity
+    layer_name: str
+    data_source_name: str
+
+    # Timing
+    duration_seconds: float
+
+    # Counts
+    total_windows_requested: int
+    num_windows_materialized: int
+    materialize_attempts: int
+
+
+@dataclass
+class MaterializeDatasetWindowsSummary:
+    """Results from materialize_dataset_windows operation for telemetry purposes."""
+
+    # Timing
+    duration_seconds: float
+
+    # Counts
+    total_windows_requested: int
+
+    # Per-layer summaries
+    layer_summaries: list[MaterializeWindowLayersSummary]
+
+
+@dataclass
+class ErrorOutcome:
+    """TBD what goes in here, if anything."""
+
+    # Timing
+    duration_seconds: float

rslearn/dataset/manage.py

@@ -13,6 +13,13 @@ from rslearn.config import (
     RasterLayerConfig,
 )
 from rslearn.data_sources import DataSource, Item
+from rslearn.dataset.handler_summaries import (
+    LayerPrepareSummary,
+    MaterializeDatasetWindowsSummary,
+    MaterializeWindowLayersSummary,
+    MaterializeWindowLayerSummary,
+    PrepareDatasetWindowsSummary,
+)
 from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStore, get_tile_store_with_layer
 
@@ -23,7 +30,24 @@ from .window import Window, WindowLayerData
 logger = get_logger(__name__)
 
 
-def retry(fn: Callable, retry_max_attempts: int, retry_backoff: timedelta) -> Any:
+class AttemptsCounter:
+    """A simple counter for tracking attempts (including initial attempt and retries)."""
+
+    def __init__(self) -> None:
+        """Initialize counter with value 0."""
+        self.value = 0
+
+    def increment(self) -> None:
+        """Increment the counter by 1."""
+        self.value += 1
+
+
+def retry(
+    fn: Callable,
+    retry_max_attempts: int,
+    retry_backoff: timedelta,
+    attempts_counter: AttemptsCounter | None = None,
+) -> Any:
     """Retry the function multiple times in case of error.
 
     The function is retried until either the attempts are exhausted, or the function
@@ -37,8 +61,11 @@ def retry(fn: Callable, retry_max_attempts: int, retry_backoff: timedelta) -> An
             retries. The actual time is (retry_backoff * attempts) * r, where r is a
             random number between 1 and 2, and attempts is the number of attempts tried
             so far.
+        attempts_counter: an optional counter to increment for each attempt
     """
     for attempt_idx in range(retry_max_attempts):
+        if attempts_counter:
+            attempts_counter.increment()
         try:
             return fn()
         except Exception as e:
@@ -47,6 +74,8 @@ def retry(fn: Callable, retry_max_attempts: int, retry_backoff: timedelta) -> An
             time.sleep(sleep_base_seconds * (1 + random.random()))
 
     # Last attempt. This time we don't catch the exception.
+    if attempts_counter:
+        attempts_counter.increment()
     return fn()
 
 
@@ -56,7 +85,7 @@ def prepare_dataset_windows(
     force: bool = False,
     retry_max_attempts: int = 0,
     retry_backoff: timedelta = timedelta(minutes=1),
-) -> None:
+) -> PrepareDatasetWindowsSummary:
     """Prepare windows in a dataset.
 
     Preparing a window involves looking up items corresponding to the window in each of
@@ -70,10 +99,28 @@ def prepare_dataset_windows(
         retry_max_attempts: set greater than zero to retry for this many attempts in
             case of error.
         retry_backoff: how long to wait before retrying (see retry).
+
+    Returns:
+        a summary of the prepare operation, fit for telemetry purposes
     """
+    start_time = time.monotonic()
+    layer_summaries: list[LayerPrepareSummary] = []
+
     # Iterate over retrieved layers, and prepare each one.
     for layer_name, layer_cfg in dataset.layers.items():
+        layer_start_time = time.monotonic()
+
         if not layer_cfg.data_source:
+            layer_summaries.append(
+                LayerPrepareSummary(
+                    layer_name=layer_name,
+                    data_source_name="N/A",
+                    duration_seconds=time.monotonic() - layer_start_time,
+                    windows_prepared=0,
+                    windows_skipped=len(windows),
+                    get_items_attempts=0,
+                )
+            )
             continue
         data_source_cfg = layer_cfg.data_source
 
@@ -85,7 +132,18 @@ def prepare_dataset_windows(
                 continue
             needed_windows.append(window)
         logger.info(f"Preparing {len(needed_windows)} windows for layer {layer_name}")
+
         if len(needed_windows) == 0:
+            layer_summaries.append(
+                LayerPrepareSummary(
+                    layer_name=layer_name,
+                    data_source_name=data_source_cfg.name,
+                    duration_seconds=time.monotonic() - layer_start_time,
+                    windows_prepared=0,
+                    windows_skipped=len(windows),
+                    get_items_attempts=0,
+                )
+            )
             continue
 
         # Create data source after checking for at least one window so it can be fast
@@ -115,10 +173,12 @@ def prepare_dataset_windows(
 
             geometries.append(geometry)
 
+        attempts_counter = AttemptsCounter()
         results = retry(
             fn=lambda: data_source.get_items(geometries, data_source_cfg.query_config),
             retry_max_attempts=retry_max_attempts,
             retry_backoff=retry_backoff,
+            attempts_counter=attempts_counter,
         )
 
         for window, result in zip(needed_windows, results):
@@ -131,6 +191,25 @@ def prepare_dataset_windows(
             )
             window.save_layer_datas(layer_datas)
 
+        layer_summaries.append(
+            LayerPrepareSummary(
+                layer_name=layer_name,
+                data_source_name=data_source_cfg.name,
+                duration_seconds=time.monotonic() - layer_start_time,
+                windows_prepared=len(needed_windows),  # we assume all have succeeded
+                windows_skipped=len(windows) - len(needed_windows),
+                get_items_attempts=attempts_counter.value,
+            )
+        )
+
+    summary = PrepareDatasetWindowsSummary(
+        duration_seconds=time.monotonic() - start_time,
+        total_windows_requested=len(windows),
+        layer_summaries=layer_summaries,
+    )
+
+    return summary
+
 
 def ingest_dataset_windows(
     dataset: Dataset,
@@ -251,7 +330,7 @@ def materialize_window(
     layer_cfg: LayerConfig,
     retry_max_attempts: int = 0,
     retry_backoff: timedelta = timedelta(minutes=1),
-) -> None:
+) -> MaterializeWindowLayerSummary:
     """Materialize a window.
 
     Args:
@@ -264,10 +343,16 @@ def materialize_window(
         retry_max_attempts: set greater than zero to retry for this many attempts in
             case of error.
         retry_backoff: how long to wait before retrying (see retry).
+
+    Returns:
+        a summary of the materialize operation, fit for telemetry purposes
     """
     # Check if layer is materialized already.
     if window.is_layer_completed(layer_name):
-        return
+        return MaterializeWindowLayerSummary(
+            skipped=True,
+            materialize_attempts=0,
+        )
 
     layer_datas = window.load_layer_datas()
     if layer_name not in layer_datas:
@@ -276,7 +361,11 @@ def materialize_window(
             layer_name,
             window.name,
         )
-        return
+        return MaterializeWindowLayerSummary(
+            skipped=True,
+            materialize_attempts=0,
+        )
+
     layer_data = layer_datas[layer_name]
     item_groups = []
     for serialized_group in layer_data.serialized_item_groups:
@@ -288,6 +377,8 @@ def materialize_window(
 
     if layer_cfg.data_source is None:
         raise ValueError("data_source is required")
+
+    attempts_counter = AttemptsCounter()
     if layer_cfg.data_source.ingest:
         if not is_window_ingested(dataset, window, check_layer_name=layer_name):
             logger.info(
@@ -295,16 +386,19 @@ def materialize_window(
                 layer_name,
                 window.name,
             )
-            return
+            return MaterializeWindowLayerSummary(
+                skipped=True,
+                materialize_attempts=0,
+            )
 
-        print(
+        logger.info(
             f"Materializing {len(item_groups)} item groups in layer {layer_name} from tile store"
         )
 
         if dataset.materializer_name:
-            materializer = Materializers[dataset.materializer_name]
+            materializer = Materializers[dataset.materializer_name]()
         else:
-            materializer = Materializers[layer_cfg.layer_type.value]
+            materializer = Materializers[layer_cfg.layer_type.value]()
 
         retry(
             fn=lambda: materializer.materialize(
@@ -316,11 +410,12 @@ def materialize_window(
             ),
             retry_max_attempts=retry_max_attempts,
             retry_backoff=retry_backoff,
+            attempts_counter=attempts_counter,
         )
 
     else:
         # This window is meant to be materialized directly from the data source.
-        print(
+        logger.info(
             f"Materializing {len(item_groups)} item groups in layer {layer_name} via data source"
         )
         retry(
@@ -329,15 +424,21 @@ def materialize_window(
             ),
             retry_max_attempts=retry_max_attempts,
             retry_backoff=retry_backoff,
+            attempts_counter=attempts_counter,
         )
 
+    return MaterializeWindowLayerSummary(
+        skipped=False,
+        materialize_attempts=attempts_counter.value,
+    )
+
 
 def materialize_dataset_windows(
     dataset: Dataset,
     windows: list[Window],
     retry_max_attempts: int = 0,
     retry_backoff: timedelta = timedelta(minutes=1),
-) -> None:
+) -> MaterializeDatasetWindowsSummary:
     """Materialize items for retrieved layers in a dataset.
 
     The portions of items corresponding to dataset windows are extracted from the tile
@@ -349,24 +450,58 @@ def materialize_dataset_windows(
         retry_max_attempts: set greater than zero to retry for this many attempts in
             case of error.
         retry_backoff: how long to wait before retrying (see retry).
+
+    Returns:
+        a summary of the materialize operation, fit for telemetry purposes
     """
+    start_time = time.monotonic()
+
+    layer_summaries: list[MaterializeWindowLayersSummary] = []
+
     tile_store = dataset.get_tile_store()
     for layer_name, layer_cfg in dataset.layers.items():
+        layer_start_time = time.monotonic()
+
+        total_materialize_attempts = 0
+        total_skipped = 0
+        data_source_name = "N/A"
+
         if not layer_cfg.data_source:
-            continue
+            total_skipped = len(windows)
+        else:
+            data_source_name = layer_cfg.data_source.name
+            data_source = rslearn.data_sources.data_source_from_config(
+                layer_cfg, dataset.path
+            )
 
-        data_source = rslearn.data_sources.data_source_from_config(
-            layer_cfg, dataset.path
-        )
+            for window in windows:
+                window_summary = materialize_window(
+                    window=window,
+                    dataset=dataset,
+                    data_source=data_source,
+                    tile_store=tile_store,
+                    layer_name=layer_name,
+                    layer_cfg=layer_cfg,
+                    retry_max_attempts=retry_max_attempts,
+                    retry_backoff=retry_backoff,
+                )
+                total_materialize_attempts += window_summary.materialize_attempts
+                if window_summary.skipped:
+                    total_skipped += 1
 
-        for window in windows:
-            materialize_window(
-                window=window,
-                dataset=dataset,
-                data_source=data_source,
-                tile_store=tile_store,
+        layer_summaries.append(
+            MaterializeWindowLayersSummary(
                 layer_name=layer_name,
-                layer_cfg=layer_cfg,
-                retry_max_attempts=retry_max_attempts,
-                retry_backoff=retry_backoff,
+                data_source_name=data_source_name,
+                duration_seconds=time.monotonic() - layer_start_time,
+                total_windows_requested=len(windows),
+                num_windows_materialized=len(windows) - total_skipped,
+                materialize_attempts=total_materialize_attempts,
             )
+        )
+
+    return MaterializeDatasetWindowsSummary(
+        duration_seconds=time.monotonic() - start_time,
+        total_windows_requested=len(windows),
+        layer_summaries=layer_summaries,
+    )