rslearn 0.0.7__tar.gz → 0.0.9__tar.gz

{rslearn-0.0.7/rslearn.egg-info → rslearn-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.7
+Version: 0.0.9
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License
@@ -212,7 +212,6 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: boto3>=1.39
-Requires-Dist: class_registry>=2.1
 Requires-Dist: fiona>=1.10
 Requires-Dist: fsspec>=2025.9.0
 Requires-Dist: jsonargparse>=4.35.0
@@ -233,7 +232,7 @@ Requires-Dist: cdsapi>=0.7.6; extra == "extra"
 Requires-Dist: earthdaily[platform]>=1.0.7; extra == "extra"
 Requires-Dist: earthengine-api>=1.6.3; extra == "extra"
 Requires-Dist: einops>=0.8; extra == "extra"
-Requires-Dist: gcsfs>=2025.9.0; extra == "extra"
+Requires-Dist: fsspec[gcs,s3]; extra == "extra"
 Requires-Dist: google-cloud-bigquery>=3.35; extra == "extra"
 Requires-Dist: google-cloud-storage>=2.18; extra == "extra"
 Requires-Dist: huggingface_hub>=0.34.4; extra == "extra"
@@ -244,7 +243,6 @@ Requires-Dist: planetary_computer>=1.0; extra == "extra"
 Requires-Dist: pycocotools>=2.0; extra == "extra"
 Requires-Dist: pystac_client>=0.9; extra == "extra"
 Requires-Dist: rtree>=1.4; extra == "extra"
-Requires-Dist: s3fs>=2025.9.0; extra == "extra"
 Requires-Dist: satlaspretrain_models>=0.3; extra == "extra"
 Requires-Dist: scipy>=1.16; extra == "extra"
 Requires-Dist: terratorch>=1.0.2; extra == "extra"
@@ -285,6 +283,7 @@ Quick links:
 - [Examples](docs/Examples.md) contains more examples, including customizing different
   stages of rslearn with additional code.
 - [DatasetConfig](docs/DatasetConfig.md) documents the dataset configuration file.
+- [ModelConfig](docs/ModelConfig.md) documents the model configuration file.
 
 
 Setup

{rslearn-0.0.7 → rslearn-0.0.9}/README.md

@@ -21,6 +21,7 @@ Quick links:
 - [Examples](docs/Examples.md) contains more examples, including customizing different
   stages of rslearn with additional code.
 - [DatasetConfig](docs/DatasetConfig.md) documents the dataset configuration file.
+- [ModelConfig](docs/ModelConfig.md) documents the model configuration file.
 
 
 Setup

{rslearn-0.0.7 → rslearn-0.0.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rslearn"
-version = "0.0.7"
+version = "0.0.9"
 description = "A library for developing remote sensing datasets and models"
 authors = [
     { name = "OlmoEarth Team" },
@@ -10,9 +10,8 @@ license = {file = "LICENSE"}
 requires-python = ">=3.11"
 dependencies = [
     "boto3>=1.39",
-    "class_registry>=2.1",
     "fiona>=1.10",
-    "fsspec>=2025.9.0",
+    "fsspec>=2025.9.0", # this is used both directly and indirectly (via universal_pathlib) in our code
     "jsonargparse>=4.35.0",
     "lightning>=2.5.1.post0",
     "Pillow>=11.3",
@@ -35,7 +34,9 @@ extra = [
     "earthdaily[platform]>=1.0.7",
     "earthengine-api>=1.6.3",
     "einops>=0.8",
-    "gcsfs>=2025.9.0",
+    # https://github.com/fsspec/universal_pathlib?tab=readme-ov-file#adding-universal_pathlib-to-your-project
+    # https://github.com/fsspec/filesystem_spec?tab=readme-ov-file#install
+    "fsspec[gcs, s3]", # for both direct use via fsspec and indirect use via universal_pathlib, docs suggest enabling specific backends like this
     "google-cloud-bigquery>=3.35",
     "google-cloud-storage>=2.18",
     "huggingface_hub>=0.34.4",
@@ -46,7 +47,6 @@ extra = [
     "pycocotools>=2.0",
     "pystac_client>=0.9",
     "rtree>=1.4",
-    "s3fs>=2025.9.0",
     "satlaspretrain_models>=0.3",
     "scipy>=1.16",
     "terratorch>=1.0.2",

{rslearn-0.0.7 → rslearn-0.0.9}/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-0.0.7 → rslearn-0.0.9}/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-0.0.9/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