rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

rslearn/dataset/handler_summaries.py

@@ -0,0 +1,131 @@
+"""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
+    windows_rejected: 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

@@ -1,19 +1,89 @@
 """Functions to manage datasets."""
 
-import rslearn.data_sources
-from rslearn.config import LayerConfig, LayerType
+import random
+import time
+from collections.abc import Callable
+from datetime import timedelta
+from typing import Any
+
+from rslearn.config import (
+    LayerConfig,
+    LayerType,
+)
 from rslearn.data_sources import DataSource, Item
-from rslearn.tile_stores import TileStore, get_tile_store_for_layer
-from rslearn.utils import logger
+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
 
 from .dataset import Dataset
-from .materialize import Materializers
+from .materialize import Materializer, RasterMaterializer, VectorMaterializer
 from .window import Window, WindowLayerData
 
+logger = get_logger(__name__)
+
+
+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
+    runs successfully without raising an Exception.
+
+    Args:
+        fn: the function to call.
+        retry_max_attempts: retry this many times (plus the original attempt) before
+            giving up (and raising Exception).
+        retry_backoff: the base backoff time used to compute how long to wait between
+            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:
+            logger.warning(f"Retrying after catching error in retry loop: {e}")
+            sleep_base_seconds = retry_backoff.total_seconds() * (attempt_idx + 1)
+            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()
+
 
 def prepare_dataset_windows(
-    dataset: Dataset, windows: list[Window], force: bool = False
-) -> None:
+    dataset: Dataset,
+    windows: list[Window],
+    force: bool = False,
+    retry_max_attempts: int = 0,
+    retry_backoff: timedelta = timedelta(minutes=1),
+) -> PrepareDatasetWindowsSummary:
     """Prepare windows in a dataset.
 
     Preparing a window involves looking up items corresponding to the window in each of
@@ -24,28 +94,73 @@ def prepare_dataset_windows(
         windows: the windows to prepare
         force: whether to prepare windows even if they were previously prepared
             (default false)
+        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),
+                    windows_rejected=0,
+                    get_items_attempts=0,
+                )
+            )
             continue
+        data_source_cfg = layer_cfg.data_source
+        min_matches = data_source_cfg.query_config.min_matches
 
         # Get windows that need to be prepared for this layer.
+        # Also track which windows are skipped vs previously rejected.
         needed_windows = []
+        windows_skipped = 0
+        windows_rejected = 0
         for window in windows:
             layer_datas = window.load_layer_datas()
             if layer_name in layer_datas and not force:
+                # Window already has layer data - check if it was previously rejected
+                layer_data = layer_datas[layer_name]
+                if len(layer_data.serialized_item_groups) == 0 and min_matches > 0:
+                    # Previously rejected due to min_matches
+                    windows_rejected += 1
+                else:
+                    # Successfully prepared previously
+                    windows_skipped += 1
                 continue
             needed_windows.append(window)
-        print(f"Preparing {len(needed_windows)} windows for layer {layer_name}")
+        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.class_path,
+                    duration_seconds=time.monotonic() - layer_start_time,
+                    windows_prepared=0,
+                    windows_skipped=windows_skipped,
+                    windows_rejected=windows_rejected,
+                    get_items_attempts=0,
+                )
+            )
             continue
 
         # Create data source after checking for at least one window so it can be fast
         # if there are no windows to prepare.
-        data_source = rslearn.data_sources.data_source_from_config(
-            layer_cfg, dataset.path
-        )
+        data_source = layer_cfg.instantiate_data_source(dataset.path)
 
         # Get STGeometry for each window.
         geometries = []
@@ -53,13 +168,13 @@ def prepare_dataset_windows(
             geometry = window.get_geometry()
 
             # Apply temporal modifiers.
-            time_offset = layer_cfg.data_source.time_offset
+            time_offset = data_source_cfg.time_offset
             if geometry.time_range and time_offset:
                 geometry.time_range = (
                     geometry.time_range[0] + time_offset,
                     geometry.time_range[1] + time_offset,
                 )
-            duration = layer_cfg.data_source.duration
+            duration = data_source_cfg.duration
             if geometry.time_range and duration:
                 geometry.time_range = (
                     geometry.time_range[0],
@@ -68,7 +183,15 @@ def prepare_dataset_windows(
 
             geometries.append(geometry)
 
-        results = data_source.get_items(geometries, layer_cfg.data_source.query_config)
+        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,
+        )
+
+        windows_prepared = 0
         for window, result in zip(needed_windows, results):
             layer_datas = window.load_layer_datas()
             layer_datas[layer_name] = WindowLayerData(
@@ -79,8 +202,39 @@ def prepare_dataset_windows(
             )
             window.save_layer_datas(layer_datas)
 
+            # If result is empty and min_matches > 0, window was rejected due to min_matches
+            if len(result) == 0 and min_matches > 0:
+                windows_rejected += 1
+            else:
+                windows_prepared += 1
 
-def ingest_dataset_windows(dataset: Dataset, windows: list[Window]) -> None:
+        layer_summaries.append(
+            LayerPrepareSummary(
+                layer_name=layer_name,
+                data_source_name=data_source_cfg.class_path,
+                duration_seconds=time.monotonic() - layer_start_time,
+                windows_prepared=windows_prepared,
+                windows_skipped=windows_skipped,
+                windows_rejected=windows_rejected,
+                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,
+    windows: list[Window],
+    retry_max_attempts: int = 0,
+    retry_backoff: timedelta = timedelta(minutes=1),
+) -> None:
     """Ingest items for retrieved layers in a dataset.
 
     The items associated with the specified windows are downloaded and divided into
@@ -89,6 +243,9 @@ def ingest_dataset_windows(dataset: Dataset, windows: list[Window]) -> None:
     Args:
         dataset: the dataset
         windows: the windows to ingest
+        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).
     """
     tile_store = dataset.get_tile_store()
     for layer_name, layer_cfg in dataset.layers.items():
@@ -97,11 +254,9 @@ def ingest_dataset_windows(dataset: Dataset, windows: list[Window]) -> None:
         if not layer_cfg.data_source.ingest:
             continue
 
-        data_source = rslearn.data_sources.data_source_from_config(
-            layer_cfg, dataset.path
-        )
+        data_source = layer_cfg.instantiate_data_source(dataset.path)
 
-        geometries_by_item = {}
+        geometries_by_item: dict = {}
         for window in windows:
             layer_datas = window.load_layer_datas()
             if layer_name not in layer_datas:
@@ -116,12 +271,20 @@ def ingest_dataset_windows(dataset: Dataset, windows: list[Window]) -> None:
                     geometries_by_item[item].append(geometry)
 
         print(f"Ingesting {len(geometries_by_item)} items in layer {layer_name}")
-        cur_tile_store = get_tile_store_for_layer(tile_store, layer_name, layer_cfg)
         geometries_and_items = list(geometries_by_item.items())
-        data_source.ingest(
-            tile_store=cur_tile_store,
-            items=[item for item, _ in geometries_and_items],
-            geometries=[geometries for _, geometries in geometries_and_items],
+
+        # Use retry loop for the actual data source ingest call.
+        def ingest() -> None:
+            data_source.ingest(
+                tile_store=get_tile_store_with_layer(tile_store, layer_name, layer_cfg),
+                items=[item for item, _ in geometries_and_items],
+                geometries=[geometries for _, geometries in geometries_and_items],
+            )
+
+        retry(
+            fn=ingest,
+            retry_max_attempts=retry_max_attempts,
+            retry_backoff=retry_backoff,
         )
 
 
@@ -145,47 +308,31 @@ def is_window_ingested(
             continue
         if layer_name not in layer_datas:
             return False
+
+        layer_tile_store = get_tile_store_with_layer(tile_store, layer_name, layer_cfg)
+
         layer_data = layer_datas[layer_name]
         for group in layer_data.serialized_item_groups:
             for serialized_item in group:
                 item = Item.deserialize(serialized_item)
 
-                if layer_cfg.layer_type == LayerType.RASTER:
+                if layer_cfg.type == LayerType.RASTER:
                     for band_set in layer_cfg.band_sets:
-                        projection, _ = band_set.get_final_projection_and_bounds(
-                            window.projection, window.bounds
-                        )
-                        cur_tile_store = get_tile_store_for_layer(
-                            tile_store, layer_name, layer_cfg
-                        )
-                        layer_prefix = (item.name,)
                         # Make sure that layers exist containing each configured band.
                         # And that those layers are marked completed.
-                        suffixes = cur_tile_store.list_layers(layer_prefix)
-                        needed_suffixes = []
-                        needed_bands = {band for band in band_set.bands}
-                        for suffix in suffixes:
-                            cur_bands = suffix.split("_")
+                        available_bands = layer_tile_store.get_raster_bands(item.name)
+                        wanted_bands = {band for band in band_set.bands}
+                        for cur_bands in available_bands:
                             is_needed = False
                             for band in cur_bands:
-                                if band in needed_bands:
+                                if band in wanted_bands:
                                     is_needed = True
-                                    needed_bands.remove(band)
+                                    wanted_bands.remove(band)
                             if not is_needed:
                                 continue
-                            needed_suffixes.append(suffix)
-                        if len(needed_bands) > 0:
+                        if len(wanted_bands) > 0:
                             return False
 
-                        for suffix in needed_suffixes:
-                            layer_id = (item.name, suffix, str(projection))
-                            ts_layer = get_tile_store_for_layer(
-                                tile_store, layer_name, layer_cfg
-                            ).get_layer(layer_id)
-                            if not ts_layer:
-                                return False
-                            if not ts_layer.get_metadata().properties.get("completed"):
-                                return False
     return True
 
 
@@ -196,7 +343,9 @@ def materialize_window(
     tile_store: TileStore,
     layer_name: str,
     layer_cfg: LayerConfig,
-) -> None:
+    retry_max_attempts: int = 0,
+    retry_backoff: timedelta = timedelta(minutes=1),
+) -> MaterializeWindowLayerSummary:
     """Materialize a window.
 
     Args:
@@ -206,11 +355,19 @@ def materialize_window(
         tile_store: tile store of the dataset to materialize from
         layer_name: the layer name
         layer_cfg: the layer config
+        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.
-    completed_fname = window.path / "layers" / layer_name / "completed"
-    if completed_fname.exists():
-        return
+    if window.is_layer_completed(layer_name):
+        return MaterializeWindowLayerSummary(
+            skipped=True,
+            materialize_attempts=0,
+        )
 
     layer_datas = window.load_layer_datas()
     if layer_name not in layer_datas:
@@ -219,7 +376,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:
@@ -229,6 +390,10 @@ def materialize_window(
             item_group.append(item)
         item_groups.append(item_group)
 
+    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(
@@ -236,30 +401,62 @@ 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: Materializer
+        if layer_cfg.type == LayerType.RASTER:
+            materializer = RasterMaterializer()
+        elif layer_cfg.type == LayerType.VECTOR:
+            materializer = VectorMaterializer()
         else:
-            materializer = Materializers[layer_cfg.layer_type.value]
-        materializer.materialize(tile_store, window, layer_name, layer_cfg, item_groups)
+            raise ValueError(f"unknown layer type {layer_cfg.type}")
+
+        retry(
+            fn=lambda: materializer.materialize(
+                get_tile_store_with_layer(tile_store, layer_name, layer_cfg),
+                window,
+                layer_name,
+                layer_cfg,
+                item_groups,
+            ),
+            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"
         )
-        try:
-            data_source.materialize(window, item_groups, layer_name, layer_cfg)
-        except Exception as e:
-            print(f"error materializing window {window.name}: {e}")
+        retry(
+            fn=lambda: data_source.materialize(
+                window, item_groups, layer_name, layer_cfg
+            ),
+            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]) -> None:
+
+def materialize_dataset_windows(
+    dataset: Dataset,
+    windows: list[Window],
+    retry_max_attempts: int = 0,
+    retry_backoff: timedelta = timedelta(minutes=1),
+) -> MaterializeDatasetWindowsSummary:
     """Materialize items for retrieved layers in a dataset.
 
     The portions of items corresponding to dataset windows are extracted from the tile
@@ -268,17 +465,59 @@ def materialize_dataset_windows(dataset: Dataset, windows: list[Window]) -> None
     Args:
         dataset: the dataset
         windows: the windows to materialize
+        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.class_path
+            data_source = layer_cfg.instantiate_data_source(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
 
-        data_source = rslearn.data_sources.data_source_from_config(
-            layer_cfg, dataset.path
+        layer_summaries.append(
+            MaterializeWindowLayersSummary(
+                layer_name=layer_name,
+                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,
+            )
         )
 
-        for window in windows:
-            materialize_window(
-                window, dataset, data_source, tile_store, layer_name, layer_cfg
-            )
+    return MaterializeDatasetWindowsSummary(
+        duration_seconds=time.monotonic() - start_time,
+        total_windows_requested=len(windows),
+        layer_summaries=layer_summaries,
+    )