rslearn 0.0.17__tar.gz → 0.0.18__tar.gz

{rslearn-0.0.17/rslearn.egg-info → rslearn-0.0.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.17
+Version: 0.0.18
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License

{rslearn-0.0.17 → rslearn-0.0.18}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rslearn"
-version = "0.0.17"
+version = "0.0.18"
 description = "A library for developing remote sensing datasets and models"
 authors = [
     { name = "OlmoEarth Team" },

{rslearn-0.0.17 → rslearn-0.0.18}/rslearn/config/__init__.py

@@ -10,6 +10,7 @@ from .dataset import (
     LayerType,
     QueryConfig,
     SpaceMode,
+    StorageConfig,
     TimeMode,
 )
 
@@ -23,5 +24,6 @@ __all__ = [
     "LayerType",
     "QueryConfig",
     "SpaceMode",
+    "StorageConfig",
     "TimeMode",
 ]

{rslearn-0.0.17 → rslearn-0.0.18}/rslearn/config/dataset.py

@@ -31,6 +31,7 @@ from rslearn.utils.vector_format import VectorFormat
 
 if TYPE_CHECKING:
     from rslearn.data_sources.data_source import DataSource
+    from rslearn.dataset.storage.storage import WindowStorageFactory
 
 logger = get_logger("__name__")
 
@@ -132,7 +133,11 @@ class BandSetConfig(BaseModel):
     bands.
     """
 
-    dtype: DType = Field(description="Pixel value type to store the data under")
+    model_config = ConfigDict(extra="forbid")
+
+    dtype: DType = Field(
+        description="Pixel value type to store the data under. This is used during dataset materialize and model predict."
+    )
     bands: list[str] = Field(
         default_factory=lambda: [],
         description="List of band names in this BandSetConfig. One of bands or num_bands must be set.",
@@ -329,7 +334,7 @@ class TimeMode(StrEnum):
 class QueryConfig(BaseModel):
     """A configuration for querying items in a data source."""
 
-    model_config = ConfigDict(frozen=True)
+    model_config = ConfigDict(frozen=True, extra="forbid")
 
     space_mode: SpaceMode = Field(
         default=SpaceMode.MOSAIC,
@@ -363,7 +368,7 @@ class QueryConfig(BaseModel):
 class DataSourceConfig(BaseModel):
     """Configuration for a DataSource in a dataset layer."""
 
-    model_config = ConfigDict(frozen=True)
+    model_config = ConfigDict(frozen=True, extra="forbid")
 
     class_path: str = Field(description="Class path for the data source.")
     init_args: dict[str, Any] = Field(
@@ -469,7 +474,7 @@ class CompositingMethod(StrEnum):
 class LayerConfig(BaseModel):
     """Configuration of a layer in a dataset."""
 
-    model_config = ConfigDict(frozen=True)
+    model_config = ConfigDict(frozen=True, extra="forbid")
 
     type: LayerType = Field(description="The LayerType (raster or vector).")
     data_source: DataSourceConfig | None = Field(
@@ -592,11 +597,51 @@ class LayerConfig(BaseModel):
         return vector_format
 
 
+class StorageConfig(BaseModel):
+    """Configuration for the WindowStorageFactory (window metadata storage backend)."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    class_path: str = Field(
+        default="rslearn.dataset.storage.file.FileWindowStorageFactory",
+        description="Class path for the WindowStorageFactory.",
+    )
+    init_args: dict[str, Any] = Field(
+        default_factory=lambda: {},
+        description="jsonargparse init args for the WindowStorageFactory.",
+    )
+
+    def instantiate_window_storage_factory(self) -> "WindowStorageFactory":
+        """Instantiate the WindowStorageFactory specified by this config."""
+        from rslearn.dataset.storage.storage import WindowStorageFactory
+        from rslearn.utils.jsonargparse import init_jsonargparse
+
+        init_jsonargparse()
+        parser = jsonargparse.ArgumentParser()
+        parser.add_argument("--wsf", type=WindowStorageFactory)
+        cfg = parser.parse_object(
+            {
+                "wsf": dict(
+                    class_path=self.class_path,
+                    init_args=self.init_args,
+                )
+            }
+        )
+        wsf = parser.instantiate_classes(cfg).wsf
+        return wsf
+
+
 class DatasetConfig(BaseModel):
     """Overall dataset configuration."""
 
+    model_config = ConfigDict(extra="forbid")
+
     layers: dict[str, LayerConfig] = Field(description="Layers in the dataset.")
     tile_store: dict[str, Any] = Field(
         default={"class_path": "rslearn.tile_stores.default.DefaultTileStore"},
         description="jsonargparse configuration for the TileStore.",
     )
+    storage: StorageConfig = Field(
+        default_factory=lambda: StorageConfig(),
+        description="jsonargparse configuration for the WindowStorageFactory.",
+    )

{rslearn-0.0.17 → rslearn-0.0.18}/rslearn/dataset/add_windows.py

@@ -131,7 +131,7 @@ def add_windows_from_geometries(
                     f"_{time_range[0].isoformat()}_{time_range[1].isoformat()}"
                 )
         window = Window(
-            path=dataset.path / "windows" / group / cur_window_name,
+            storage=dataset.storage,
             group=group,
             name=cur_window_name,
             projection=cur_projection,

{rslearn-0.0.17 → rslearn-0.0.18}/rslearn/dataset/dataset.py

@@ -1,9 +1,8 @@
 """rslearn dataset class."""
 
 import json
-import multiprocessing
+from typing import Any
 
-import tqdm
 from upath import UPath
 
 from rslearn.config import DatasetConfig
@@ -11,7 +10,6 @@ from rslearn.log_utils import get_logger
 from rslearn.template_params import substitute_env_vars_in_string
 from rslearn.tile_stores import TileStore, load_tile_store
 
-from .index import DatasetIndex
 from .window import Window
 
 logger = get_logger(__name__)
@@ -68,80 +66,26 @@ class Dataset:
                 self.layers[layer_name] = layer_config
 
             self.tile_store_config = config.tile_store
-
-    def _get_index(self) -> DatasetIndex | None:
-        index_fname = self.path / DatasetIndex.FNAME
-        if not index_fname.exists():
-            return None
-        return DatasetIndex.load_index(self.path)
+            self.storage = (
+                config.storage.instantiate_window_storage_factory().get_storage(
+                    self.path
+                )
+            )
 
     def load_windows(
         self,
         groups: list[str] | None = None,
         names: list[str] | None = None,
-        show_progress: bool = False,
-        workers: int = 0,
-        no_index: bool = False,
+        **kwargs: Any,
     ) -> list[Window]:
         """Load the windows in the dataset.
 
         Args:
             groups: an optional list of groups to filter loading
             names: an optional list of window names to filter loading
-            show_progress: whether to show tqdm progress bar
-            workers: number of parallel workers, default 0 (use main thread only to load windows)
-            no_index: don't use the dataset index even if it exists.
+            kwargs: optional keyword arguments to pass to WindowStorage.get_windows.
         """
-        # Load from index if it exists.
-        # We never use the index if names is set since loading the index will likely be
-        # slower than loading a few windows.
-        if not no_index and names is None:
-            dataset_index = self._get_index()
-            if dataset_index is not None:
-                return dataset_index.get_windows(groups=groups, names=names)
-
-        # Avoid directory does not exist errors later.
-        if not (self.path / "windows").exists():
-            return []
-
-        window_dirs = []
-        if not groups:
-            groups = []
-            for p in (self.path / "windows").iterdir():
-                groups.append(p.name)
-        for group in groups:
-            group_dir = self.path / "windows" / group
-            if not group_dir.exists():
-                logger.warning(
-                    f"Skipping group directory {group_dir} since it does not exist"
-                )
-                continue
-            if names:
-                cur_names = names
-            else:
-                cur_names = []
-                for p in group_dir.iterdir():
-                    cur_names.append(p.name)
-
-            for window_name in cur_names:
-                window_dir = group_dir / window_name
-                window_dirs.append(window_dir)
-
-        if workers == 0:
-            windows = [Window.load(window_dir) for window_dir in window_dirs]
-        else:
-            p = multiprocessing.Pool(workers)
-            outputs = p.imap_unordered(Window.load, window_dirs)
-            if show_progress:
-                outputs = tqdm.tqdm(
-                    outputs, total=len(window_dirs), desc="Loading windows"
-                )
-            windows = []
-            for window in outputs:
-                windows.append(window)
-            p.close()
-
-        return windows
+        return self.storage.get_windows(groups=groups, names=names, **kwargs)
 
     def get_tile_store(self) -> TileStore:
         """Get the tile store associated with this dataset.

{rslearn-0.0.17 → rslearn-0.0.18}/rslearn/dataset/materialize.py

@@ -161,7 +161,7 @@ def build_first_valid_composite(
     nodata_vals: list[Any],
     bands: list[str],
     bounds: PixelBounds,
-    band_dtype: Any,
+    band_dtype: npt.DTypeLike,
     tile_store: TileStoreWithLayer,
     projection: Projection,
     remapper: Remapper | None,
@@ -233,7 +233,7 @@ def read_and_stack_raster_windows(
     projection: Projection,
     nodata_vals: list[Any],
     remapper: Remapper | None,
-    band_dtype: Any,
+    band_dtype: npt.DTypeLike,
     resampling_method: Resampling = Resampling.bilinear,
 ) -> npt.NDArray[np.generic]:
     """Create a stack of extent aligned raster windows.
@@ -326,7 +326,7 @@ def build_mean_composite(
     nodata_vals: list[Any],
     bands: list[str],
     bounds: PixelBounds,
-    band_dtype: Any,
+    band_dtype: npt.DTypeLike,
     tile_store: TileStoreWithLayer,
     projection: Projection,
     remapper: Remapper | None,
@@ -383,7 +383,7 @@ def build_median_composite(
     nodata_vals: list[Any],
     bands: list[str],
     bounds: PixelBounds,
-    band_dtype: Any,
+    band_dtype: npt.DTypeLike,
     tile_store: TileStoreWithLayer,
     projection: Projection,
     remapper: Remapper | None,
@@ -471,7 +471,7 @@ def build_composite(
         nodata_vals=nodata_vals,
         bands=band_cfg.bands,
         bounds=bounds,
-        band_dtype=band_cfg.dtype.value,
+        band_dtype=band_cfg.dtype.get_numpy_dtype(),
         tile_store=tile_store,
         projection=projection,
         resampling_method=layer_cfg.resampling_method.get_rasterio_resampling(),

rslearn-0.0.18/rslearn/dataset/storage/__init__.py

@@ -0,0 +1 @@
+"""Storage backends for rslearn window metadata."""

rslearn-0.0.18/rslearn/dataset/storage/file.py

@@ -0,0 +1,202 @@
+"""The default file-based window storage backend."""
+
+import json
+import multiprocessing
+
+import tqdm
+from typing_extensions import override
+from upath import UPath
+
+from rslearn.dataset.window import (
+    LAYERS_DIRECTORY_NAME,
+    Window,
+    WindowLayerData,
+    get_layer_and_group_from_dir_name,
+    get_window_layer_dir,
+)
+from rslearn.log_utils import get_logger
+from rslearn.utils.fsspec import open_atomic
+from rslearn.utils.mp import star_imap_unordered
+
+from .storage import WindowStorage, WindowStorageFactory
+
+logger = get_logger(__name__)
+
+
+def load_window(storage: "FileWindowStorage", window_dir: UPath) -> Window:
+    """Load the window from its directory by reading metadata.json.
+
+    Args:
+        storage: the underlying FileWindowStorage.
+        window_dir: the path where the window is stored.
+
+    Returns:
+        the window object.
+    """
+    metadata_fname = window_dir / "metadata.json"
+    with metadata_fname.open() as f:
+        metadata = json.load(f)
+    return Window.from_metadata(storage, metadata)
+
+
+class FileWindowStorage(WindowStorage):
+    """The default file-backed window storage."""
+
+    def __init__(self, path: UPath):
+        """Create a new FileWindowStorage.
+
+        Args:
+            path: the path to the dataset.
+        """
+        self.path = path
+
+    @override
+    def get_window_root(self, group: str, name: str) -> UPath:
+        return Window.get_window_root(self.path, group, name)
+
+    @override
+    def get_windows(
+        self,
+        groups: list[str] | None = None,
+        names: list[str] | None = None,
+        show_progress: bool = False,
+        workers: int = 0,
+    ) -> list["Window"]:
+        """Load the windows in the dataset.
+
+        Args:
+            groups: an optional list of groups to filter loading
+            names: an optional list of window names to filter loading
+            show_progress: whether to show tqdm progress bar
+            workers: number of parallel workers, default 0 (use main thread only to load windows)
+        """
+        # Avoid directory does not exist errors later.
+        if not (self.path / "windows").exists():
+            return []
+
+        window_dirs = []
+        if not groups:
+            groups = []
+            for p in (self.path / "windows").iterdir():
+                groups.append(p.name)
+        for group in groups:
+            group_dir = self.path / "windows" / group
+            if not group_dir.exists():
+                logger.warning(
+                    f"Skipping group directory {group_dir} since it does not exist"
+                )
+                continue
+            if names:
+                cur_names = names
+            else:
+                cur_names = []
+                for p in group_dir.iterdir():
+                    cur_names.append(p.name)
+
+            for window_name in cur_names:
+                window_dir = group_dir / window_name
+                window_dirs.append(window_dir)
+
+        if workers == 0:
+            windows = [load_window(self, window_dir) for window_dir in window_dirs]
+        else:
+            p = multiprocessing.Pool(workers)
+            outputs = star_imap_unordered(
+                p,
+                load_window,
+                [
+                    dict(storage=self, window_dir=window_dir)
+                    for window_dir in window_dirs
+                ],
+            )
+            if show_progress:
+                outputs = tqdm.tqdm(
+                    outputs, total=len(window_dirs), desc="Loading windows"
+                )
+            windows = []
+            for window in outputs:
+                windows.append(window)
+            p.close()
+
+        return windows
+
+    @override
+    def create_or_update_window(self, window: Window) -> None:
+        window_path = self.get_window_root(window.group, window.name)
+        window_path.mkdir(parents=True, exist_ok=True)
+        metadata_path = window_path / "metadata.json"
+        logger.debug(f"Saving window metadata to {metadata_path}")
+        with open_atomic(metadata_path, "w") as f:
+            json.dump(window.get_metadata(), f)
+
+    @override
+    def get_layer_datas(self, group: str, name: str) -> dict[str, "WindowLayerData"]:
+        window_path = self.get_window_root(group, name)
+        items_fname = window_path / "items.json"
+        if not items_fname.exists():
+            return {}
+
+        with items_fname.open() as f:
+            layer_datas = [
+                WindowLayerData.deserialize(layer_data) for layer_data in json.load(f)
+            ]
+
+        return {layer_data.layer_name: layer_data for layer_data in layer_datas}
+
+    @override
+    def save_layer_datas(
+        self, group: str, name: str, layer_datas: dict[str, "WindowLayerData"]
+    ) -> None:
+        window_path = self.get_window_root(group, name)
+        json_data = [layer_data.serialize() for layer_data in layer_datas.values()]
+        items_fname = window_path / "items.json"
+        logger.info(f"Saving window items to {items_fname}")
+        with open_atomic(items_fname, "w") as f:
+            json.dump(json_data, f)
+
+    @override
+    def list_completed_layers(self, group: str, name: str) -> list[tuple[str, int]]:
+        window_path = self.get_window_root(group, name)
+        layers_directory = window_path / LAYERS_DIRECTORY_NAME
+        if not layers_directory.exists():
+            return []
+
+        completed_layers = []
+        for layer_dir in layers_directory.iterdir():
+            layer_name, group_idx = get_layer_and_group_from_dir_name(layer_dir.name)
+            if not self.is_layer_completed(group, name, layer_name, group_idx):
+                continue
+            completed_layers.append((layer_name, group_idx))
+
+        return completed_layers
+
+    @override
+    def is_layer_completed(
+        self, group: str, name: str, layer_name: str, group_idx: int = 0
+    ) -> bool:
+        window_path = self.get_window_root(group, name)
+        layer_dir = get_window_layer_dir(
+            window_path,
+            layer_name,
+            group_idx,
+        )
+        return (layer_dir / "completed").exists()
+
+    @override
+    def mark_layer_completed(
+        self, group: str, name: str, layer_name: str, group_idx: int = 0
+    ) -> None:
+        window_path = self.get_window_root(group, name)
+        layer_dir = get_window_layer_dir(window_path, layer_name, group_idx)
+        # We assume the directory exists because the layer should be materialized before
+        # being marked completed.
+        (layer_dir / "completed").touch()
+
+
+class FileWindowStorageFactory(WindowStorageFactory):
+    """Factory class for FileWindowStorage."""
+
+    @override
+    def get_storage(self, ds_path: UPath) -> FileWindowStorage:
+        """Get a FileWindowStorage for the given dataset path."""
+        return FileWindowStorage(ds_path)

rslearn-0.0.18/rslearn/dataset/storage/storage.py

@@ -0,0 +1,140 @@
+"""Abstract classes for window metadata storage."""
+
+import abc
+from typing import TYPE_CHECKING
+
+from upath import UPath
+
+if TYPE_CHECKING:
+    from rslearn.dataset.window import Window, WindowLayerData
+
+
+class WindowStorage(abc.ABC):
+    """An abstract class for the storage backend for window metadata.
+
+    This is instantiated by a WindowStorageFactory for a specific rslearn dataset.
+
+    Window metadata includes the location and time range of windows (metadata.json),
+    the window layer datas (items.json), and the completed (materialized) layers. It
+    excludes the actual materialized data. All operations involving window metadata go
+    through the WindowStorage, including enumerating windows, creating new windows, and
+    updating window layer datas during `rslearn dataset prepare` or the completed
+    layers during `rslearn dataset materialize`.
+    """
+
+    @abc.abstractmethod
+    def get_window_root(self, group: str, name: str) -> UPath:
+        """Get the path where the window should be stored."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_windows(
+        self,
+        groups: list[str] | None = None,
+        names: list[str] | None = None,
+    ) -> list["Window"]:
+        """Load the windows in the dataset.
+
+        Args:
+            groups: an optional list of groups to filter loading
+            names: an optional list of window names to filter loading
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def create_or_update_window(self, window: "Window") -> None:
+        """Create or update the window.
+
+        An existing window is only updated if there is one with the same name and group.
+
+        If there is a window with the same name but a different group, the behavior is
+        undefined.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_layer_datas(self, group: str, name: str) -> dict[str, "WindowLayerData"]:
+        """Get the window layer datas for the specified window.
+
+        Args:
+            group: the window group.
+            name: the window name.
+
+        Returns:
+            a dict mapping from the layer name to the layer data for that layer, if one
+                was previously saved.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def save_layer_datas(
+        self, group: str, name: str, layer_datas: dict[str, "WindowLayerData"]
+    ) -> None:
+        """Set the window layer datas for the specified window."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def list_completed_layers(self, group: str, name: str) -> list[tuple[str, int]]:
+        """List the layers available for this window that are completed.
+
+        Args:
+            group: the window group.
+            name: the window name.
+
+        Returns:
+            a list of (layer_name, group_idx) completed layers.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def is_layer_completed(
+        self, group: str, name: str, layer_name: str, group_idx: int = 0
+    ) -> bool:
+        """Check whether the specified layer is completed in the given window.
+
+        Completed means there is data in the layer and the data has been written
+        (materialized).
+
+        Args:
+            group: the window group.
+            name: the window name.
+            layer_name: the layer name.
+            group_idx: the index of the group within the layer.
+
+        Returns:
+            whether the layer is completed.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def mark_layer_completed(
+        self, group: str, name: str, layer_name: str, group_idx: int = 0
+    ) -> None:
+        """Mark the specified layer completed for the given window.
+
+        This must be done after the contents of the layer have been written. If a layer
+        has multiple groups, the caller should wait until the contents of all groups
+        have been written before marking them completed; this is because, when
+        materializing a window, we skip materialization if the first group
+        (group_idx=0) is marked completed.
+
+        Args:
+            group: the window group.
+            name: the window name.
+            layer_name: the layer name.
+            group_idx: the index of the group within the layer.
+        """
+        raise NotImplementedError
+
+
+class WindowStorageFactory(abc.ABC):
+    """An abstract class for a configurable storage backend for window metadata.
+
+    The dataset config includes a StorageConfig that configures a WindowStorageFactory,
+    which in turn creates a WindowStorage given a dataset path.
+    """
+
+    @abc.abstractmethod
+    def get_storage(self, ds_path: UPath) -> WindowStorage:
+        """Get a WindowStorage for the given dataset path."""
+        raise NotImplementedError