rslearn 0.0.3__tar.gz → 0.0.4__tar.gz

{rslearn-0.0.3/rslearn.egg-info → rslearn-0.0.4}/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.3
+Version: 0.0.4
 Summary: A library for developing remote sensing datasets and models
-Author-email: Favyen Bastani <favyenb@allenai.org>, Yawen Zhang <yawenz@allenai.org>, Patrick Beukema <patrickb@allenai.org>, Henry Herzog <henryh@allenai.org>, Piper Wolters <piperw@allenai.org>
+Author: OlmoEarth Team
 License:                                  Apache License
                                    Version 2.0, January 2004
                                 http://www.apache.org/licenses/
@@ -227,7 +227,7 @@ Requires-Dist: universal_pathlib>=0.2.6
 Provides-Extra: extra
 Requires-Dist: accelerate>=1.10; extra == "extra"
 Requires-Dist: cdsapi>=0.7.6; extra == "extra"
-Requires-Dist: earthdaily[platform]>=1.0.0; 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.3.0; extra == "extra"

{rslearn-0.0.3 → rslearn-0.0.4}/pyproject.toml

@@ -1,13 +1,9 @@
 [project]
 name = "rslearn"
-version = "0.0.3"
+version = "0.0.4"
 description = "A library for developing remote sensing datasets and models"
 authors = [
-    {name = "Favyen Bastani", email = "favyenb@allenai.org"},
-    {name = "Yawen Zhang", email = "yawenz@allenai.org"},
-    {name = "Patrick Beukema", email = "patrickb@allenai.org"},
-    {name = "Henry Herzog", email = "henryh@allenai.org"},
-    {name = "Piper Wolters", email = "piperw@allenai.org"},
+    { name = "OlmoEarth Team" },
 ]
 readme = "README.md"
 license = {file = "LICENSE"}
@@ -39,7 +35,7 @@ dependencies = [
 extra = [
     "accelerate>=1.10",
     "cdsapi>=0.7.6",
-    "earthdaily[platform]>=1.0.0",
+    "earthdaily[platform]>=1.0.7",
     "earthengine-api>=1.6.3",
     "einops>=0.8",
     "gcsfs==2025.3.0",

rslearn-0.0.4/rslearn/arg_parser.py

@@ -0,0 +1,59 @@
+"""Custom Lightning ArgumentParser with environment variable substitution support."""
+
+import os
+import re
+from typing import Any
+
+from jsonargparse import Namespace
+from lightning.pytorch.cli import LightningArgumentParser
+
+
+def substitute_env_vars_in_string(content: str) -> str:
+    """Substitute environment variables in a string.
+
+    Replaces ${VAR_NAME} patterns with os.getenv(VAR_NAME, "") values.
+    This works on raw string content before YAML parsing.
+
+    Args:
+        content: The string content containing template variables
+
+    Returns:
+        The string with environment variables substituted
+    """
+    pattern = r"\$\{([^}]+)\}"
+
+    def replace_variable(match_obj: re.Match[str]) -> str:
+        var_name = match_obj.group(1)
+        env_value = os.getenv(var_name, "")
+        return env_value if env_value is not None else ""
+
+    return re.sub(pattern, replace_variable, content)
+
+
+class RslearnArgumentParser(LightningArgumentParser):
+    """Custom ArgumentParser that substitutes environment variables in config files.
+
+    This parser extends LightningArgumentParser to automatically substitute
+    ${VAR_NAME} patterns with environment variable values before parsing
+    configuration content. This allows config files to use environment
+    variables while still passing Lightning's validation.
+    """
+
+    def parse_string(
+        self,
+        cfg_str: str,
+        cfg_path: str | os.PathLike = "",
+        ext_vars: dict | None = None,
+        env: bool | None = None,
+        defaults: bool = True,
+        with_meta: bool | None = None,
+        **kwargs: Any,
+    ) -> Namespace:
+        """Pre-processes string for environment variable substitution before parsing."""
+        # Substitute environment variables in the config string before parsing
+        substituted_cfg_str = substitute_env_vars_in_string(cfg_str)
+
+        # Call the parent method with the substituted config
+        return super().parse_string(
+            substituted_cfg_str, cfg_path, ext_vars, env, defaults, with_meta, **kwargs
+        )

{rslearn-0.0.3 → rslearn-0.0.4}/rslearn/data_sources/copernicus.py

@@ -34,7 +34,7 @@ from rslearn.utils.geometry import (
     FloatBounds,
     STGeometry,
     flatten_shape,
-    split_shape_at_prime_meridian,
+    split_shape_at_antimeridian,
 )
 from rslearn.utils.grid_index import GridIndex
 from rslearn.utils.raster_format import get_raster_projection_and_bounds
@@ -160,7 +160,7 @@ def get_sentinel2_tile_index() -> dict[str, list[FloatBounds]]:
         # issues where the tile bounds go from -180 to 180 longitude and thus match
         # with anything at the same latitude.
         union_shp = shapely.unary_union(shapes)
-        split_shapes = flatten_shape(split_shape_at_prime_meridian(union_shp))
+        split_shapes = flatten_shape(split_shape_at_antimeridian(union_shp))
         bounds_list: list[FloatBounds] = []
         for shp in split_shapes:
             bounds_list.append(shp.bounds)
@@ -222,10 +222,10 @@ def get_sentinel2_tiles(geometry: STGeometry, cache_dir: UPath) -> list[str]:
     """
     tile_index = load_sentinel2_tile_index(cache_dir)
     wgs84_geometry = geometry.to_projection(WGS84_PROJECTION)
-    # If the shape is a collection, it could be cutting across prime meridian.
+    # If the shape is a collection, it could be cutting across antimeridian.
     # So we query each component shape separately and collect the results to avoid
     # issues.
-    # We assume the caller has already applied split_at_prime_meridian.
+    # We assume the caller has already applied split_at_antimeridian.
     results = set()
     for shp in flatten_shape(wgs84_geometry.shp):
         for result in tile_index.query(shp.bounds):

{rslearn-0.0.3 → rslearn-0.0.4}/rslearn/data_sources/earthdaily.py

@@ -82,6 +82,8 @@ class EarthDaily(DataSource, TileStore):
         timeout: timedelta = timedelta(seconds=10),
         skip_items_missing_assets: bool = False,
         cache_dir: UPath | None = None,
+        max_retries: int = 3,
+        retry_backoff_factor: float = 5.0,
         service_name: Literal["platform"] = "platform",
     ):
         """Initialize a new EarthDaily instance.
@@ -99,6 +101,11 @@ class EarthDaily(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_retries: the maximum number of retry attempts for HTTP requests that fail
+                due to transient errors (e.g., 429, 500, 502, 503, 504 status codes).
+            retry_backoff_factor: backoff factor for exponential retry delays between HTTP
+                request attempts.  The delay between retries is calculated using the formula:
+                `(retry_backoff_factor * (2 ** (retry_count - 1)))` seconds.
             service_name: the service name, only "platform" is supported, the other
                 services "legacy" and "internal" are not supported.
         """
@@ -110,6 +117,8 @@ class EarthDaily(DataSource, TileStore):
         self.timeout = timeout
         self.skip_items_missing_assets = skip_items_missing_assets
         self.cache_dir = cache_dir
+        self.max_retries = max_retries
+        self.retry_backoff_factor = retry_backoff_factor
         self.service_name = service_name
 
         if cache_dir is not None:
@@ -139,6 +148,12 @@ class EarthDaily(DataSource, TileStore):
         if "cache_dir" in d:
             kwargs["cache_dir"] = join_upath(ds_path, d["cache_dir"])
 
+        if "max_retries" in d:
+            kwargs["max_retries"] = d["max_retries"]
+
+        if "retry_backoff_factor" in d:
+            kwargs["retry_backoff_factor"] = d["retry_backoff_factor"]
+
         simple_optionals = ["query", "sort_by", "sort_ascending"]
         for k in simple_optionals:
             if k in d:
@@ -159,7 +174,12 @@ class EarthDaily(DataSource, TileStore):
         if self.eds_client is not None:
             return self.eds_client, self.client, self.collection
 
-        self.eds_client = EDSClient(EDSConfig())
+        self.eds_client = EDSClient(
+            EDSConfig(
+                max_retries=self.max_retries,
+                retry_backoff_factor=self.retry_backoff_factor,
+            )
+        )
 
         if self.service_name == "platform":
             self.client = self.eds_client.platform.pystac_client

{rslearn-0.0.3 → rslearn-0.0.4}/rslearn/data_sources/gcp_public_data.py

@@ -26,7 +26,7 @@ from rslearn.data_sources.utils import match_candidate_items_to_window
 from rslearn.log_utils import get_logger
 from rslearn.tile_stores import TileStoreWithLayer
 from rslearn.utils.fsspec import join_upath, open_atomic
-from rslearn.utils.geometry import STGeometry, flatten_shape, split_at_prime_meridian
+from rslearn.utils.geometry import STGeometry, flatten_shape, split_at_antimeridian
 from rslearn.utils.raster_format import get_raster_projection_and_bounds
 
 from .copernicus import get_harmonize_callback, get_sentinel2_tiles
@@ -358,7 +358,7 @@ class Sentinel2(DataSource):
             shp = shapely.box(*bounds)
             sensing_time = row["sensing_time"]
             geometry = STGeometry(WGS84_PROJECTION, shp, (sensing_time, sensing_time))
-            geometry = split_at_prime_meridian(geometry)
+            geometry = split_at_antimeridian(geometry)
 
             cloud_cover = float(row["cloud_cover"])
 
@@ -511,7 +511,7 @@ class Sentinel2(DataSource):
 
         time_range = (product_xml.start_time, product_xml.start_time)
         geometry = STGeometry(WGS84_PROJECTION, product_xml.shp, time_range)
-        geometry = split_at_prime_meridian(geometry)
+        geometry = split_at_antimeridian(geometry)
 
         # Sometimes the geometry is not valid.
         # We just apply make_valid on it to correct issues.

{rslearn-0.0.3 → rslearn-0.0.4}/rslearn/data_sources/utils.py

@@ -256,23 +256,7 @@ def match_candidate_items_to_window(
             if item_geom.is_global():
                 item_geom = geometry
             else:
-                # Windows are usually smaller than items.
-                # So we first clip the item to the window bounds in the item's
-                # projection, then re-project the item to the window's projection.
-                buffered_window_geom = STGeometry(
-                    geometry.projection,
-                    geometry.shp.buffer(1),
-                    geometry.time_range,
-                )
-                window_shp_in_item_proj = buffered_window_geom.to_projection(
-                    item_geom.projection
-                ).shp
-                clipped_item_geom = STGeometry(
-                    item_geom.projection,
-                    item_geom.shp.intersection(window_shp_in_item_proj),
-                    item_geom.time_range,
-                )
-                item_geom = clipped_item_geom.to_projection(geometry.projection)
+                item_geom = item_geom.to_projection(geometry.projection)
         item_shps.append(item_geom.shp)
 
     if query_config.space_mode == SpaceMode.CONTAINS:

{rslearn-0.0.3 → rslearn-0.0.4}/rslearn/main.py

@@ -13,6 +13,7 @@ from lightning.pytorch.cli import LightningArgumentParser, LightningCLI
 from rasterio.crs import CRS
 from upath import UPath
 
+from rslearn.arg_parser import RslearnArgumentParser
 from rslearn.config import LayerConfig
 from rslearn.const import WGS84_EPSG
 from rslearn.data_sources import Item, data_source_from_config
@@ -779,7 +780,7 @@ def dataset_build_index() -> None:
 
 
 class RslearnLightningCLI(LightningCLI):
-    """LightningCLI that links data.tasks to model.tasks."""
+    """LightningCLI that links data.tasks to model.tasks and supports environment variables."""
 
     def add_arguments_to_parser(self, parser: LightningArgumentParser) -> None:
         """Link data.tasks to model.tasks.
@@ -787,6 +788,7 @@ class RslearnLightningCLI(LightningCLI):
         Args:
             parser: the argument parser
         """
+        # Link data.tasks to model.tasks
         parser.link_arguments(
             "data.init_args.task", "model.init_args.task", apply_on="instantiate"
         )
@@ -815,6 +817,12 @@ class RslearnLightningCLI(LightningCLI):
         # sampler as needed.
         c.trainer.use_distributed_sampler = False
 
+        # For predict, make sure that return_predictions is False.
+        # Otherwise all the predictions would be stored in memory which can lead to
+        # high memory consumption.
+        if subcommand == "predict":
+            c.return_predictions = False
+
 
 def model_handler() -> None:
     """Handler for any rslearn model X commands."""
@@ -825,6 +833,7 @@ def model_handler() -> None:
         subclass_mode_model=True,
         subclass_mode_data=True,
         save_config_kwargs={"overwrite": True},
+        parser_class=RslearnArgumentParser,
     )
 
 

rslearn-0.0.4/rslearn/models/trunk.py

@@ -0,0 +1,136 @@
+"""Trunk module for decoder."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+import torch
+
+from rslearn.log_utils import get_logger
+from rslearn.models.task_embedding import BaseTaskEmbedding
+
+logger = get_logger(__name__)
+
+
+class DecoderTrunkLayer(torch.nn.Module, ABC):
+    """Trunk layer for decoder."""
+
+    def __init__(self) -> None:
+        """Initialize the DecoderTrunkLayer module."""
+        super().__init__()
+
+    @abstractmethod
+    def forward(
+        self, x: torch.Tensor, task_embedding: torch.Tensor | None = None
+    ) -> dict[str, torch.Tensor]:
+        """Forward pass.
+
+        Args:
+            x: input tensor of shape (batch_size, seq_len, dim)
+            task_embedding: task embedding tensor of shape (batch_size, dim), or None
+
+        Returns:
+            dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
+            and optionally other keys.
+        """
+
+    @abstractmethod
+    def apply_auxiliary_losses(
+        self, trunk_out: dict[str, Any], outs: dict[str, Any]
+    ) -> None:
+        """Apply auxiliary losses in-place.
+
+        Args:
+            trunk_out: The output of the trunk.
+            outs: The output of the decoders, with key "loss_dict" containing the losses.
+        """
+
+
+class DecoderTrunk(torch.nn.Module):
+    """Trunk module for decoder, including arbitrary layers plus an optional task embedding."""
+
+    def __init__(
+        self,
+        task_embedding: BaseTaskEmbedding | None = None,
+        layers: list[DecoderTrunkLayer] | None = None,
+    ) -> None:
+        """Initialize the DecoderTrunk module.
+
+        Args:
+            task_embedding: Task-specific embedding module, or None if not using task embedding.
+            layers: List of other shared layers. The first one should expect a
+                B x T x C tensor, and the last should output a B x T x C tensor.
+                All layers must output a dict with key "outputs" (output tensor of shape
+                (B, T, C)) and optionally other keys.
+        """
+        super().__init__()
+        self.layers = torch.nn.ModuleList(layers or [])
+        self.task_embedding = task_embedding
+
+        # If we have multiple instances of the same layer class, output keys will get overwritten
+        if layers is not None:
+            types = [type(layer) for layer in layers]
+            if len(set(types)) != len(types):
+                logger.warning(
+                    "Multiple instances of the same layer class found in trunk. "
+                    "Only the keys from the last instance will be used"
+                )
+
+    def register_tasks(self, task_names: list[str]) -> None:
+        """Register tasks.
+
+        Args:
+            task_names: list of task names
+        """
+        if self.task_embedding is not None:
+            self.task_embedding.register_tasks(task_names)
+
+    def forward(
+        self,
+        features: list[torch.tensor],
+        inputs: list[dict[str, Any]],
+    ) -> dict[str, Any]:
+        """Forward pass.
+
+        Args:
+            features: The encoder features, a 1-list of B x C x H x W features.
+            inputs: The original inputs to the encoder.
+
+        Returns:
+            dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
+            and optionally other keys from the other layers.
+        """
+        embeds = None
+        if self.task_embedding is not None:
+            embeds = self.task_embedding.compute_embeds(features, inputs)
+            features = self.task_embedding(features, inputs, embeds=embeds)
+
+        if not self.layers:
+            return {"outputs": features}
+
+        assert len(features) == 1, "DecoderTrunk only supports one feature map"
+        x = torch.einsum("bchw->bhwc", features[0])
+        x = torch.flatten(x, start_dim=1, end_dim=2)  # B x T x C, T = HW
+        out = {}
+        for layer in self.layers:
+            layer_out = layer(x, task_embedding=embeds)
+            x = layer_out.pop("outputs")  # unspecified shape
+            out.update(layer_out)
+        x = torch.einsum("btc->bct", x)  # B x C x T
+        x = x.view(*features[0].shape)  # B x C x H x W
+
+        out["outputs"] = [x]
+        return out
+
+    def apply_auxiliary_losses(
+        self, trunk_out: dict[str, Any], outs: dict[str, Any]
+    ) -> None:
+        """Apply auxiliary losses in-place.
+
+        Each layer handles its own auxiliary losses, assuming the loss key is `loss_dict`.
+
+        Args:
+            trunk_out: The output of the trunk.
+            outs: The output of the decoders, with key "loss_dict" containing the losses.
+        """
+        for layer in self.layers:
+            layer.apply_auxiliary_losses(trunk_out, outs)

rslearn-0.0.4/rslearn/train/callbacks/adapters.py

@@ -0,0 +1,53 @@
+"""Callback to activate/deactivate adapter layers."""
+
+from typing import Any
+
+from lightning.pytorch import LightningModule
+from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.trainer import Trainer
+
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class ActivateLayers(Callback):
+    """Activates adapter layers on a given epoch.
+
+    By default, at every epoch, every adapter layer is deactivated.
+    To activate an adapter layer, add a selector with the name of the adapter layer
+    and the epoch at which to activate it. Once an adapter layer is activated, it
+    remains active until the end of training.
+    """
+
+    def __init__(self, selectors: list[dict[str, Any]]) -> None:
+        """Initialize the callback.
+
+        Args:
+            selectors: List of selectors to activate.
+                Each selector is a dictionary with the following keys:
+                - "name": Substring selector of modules to activate (str).
+                - "at_epoch": The epoch at which to activate (int).
+        """
+        self.selectors = selectors
+
+    def on_train_epoch_start(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+    ) -> None:
+        """Activate adapter layers on a given epoch.
+
+        Adapter layers are activated/deactivated by setting the `active` attribute.
+
+        Args:
+            trainer: The trainer object.
+            pl_module: The LightningModule object.
+        """
+        status = {}
+        for name, module in pl_module.named_modules():
+            for selector in self.selectors:
+                if selector["name"] in name:
+                    module.active = trainer.current_epoch >= selector["at_epoch"]
+                    status[selector["name"]] = "active" if module.active else "inactive"
+        logger.info(f"Updated adapter status: {status}")