rslearn 0.0.3__tar.gz → 0.0.5__tar.gz

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

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.3
+Version: 0.0.5
 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/
@@ -205,6 +205,9 @@ License:                                  Apache License
            See the License for the specific language governing permissions and
            limitations under the License.
         
+Project-URL: homepage, https://github.com/allenai/rslearn
+Project-URL: issues, https://github.com/allenai/rslearn/issues
+Project-URL: repository, https://github.com/allenai/rslearn
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
@@ -227,7 +230,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.5}/pyproject.toml

@@ -1,18 +1,13 @@
 [project]
 name = "rslearn"
-version = "0.0.3"
+version = "0.0.5"
 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"}
 requires-python = ">=3.11"
-
 dependencies = [
     "boto3>=1.39",
     "class_registry>=2.1",
@@ -39,7 +34,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",
@@ -71,6 +66,11 @@ dev = [
     "pytest-xdist",
 ]
 
+[project.urls]
+homepage = "https://github.com/allenai/rslearn"
+issues = "https://github.com/allenai/rslearn/issues"
+repository = "https://github.com/allenai/rslearn"
+
 [build-system]
 requires = ["setuptools>=61"]
 build-backend = "setuptools.build_meta"

rslearn-0.0.5/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.5}/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):
@@ -319,7 +319,6 @@ class Copernicus(DataSource):
                 then we attempt to read the username/password from COPERNICUS_USERNAME
                 and COPERNICUS_PASSWORD (this is useful since access tokens are only
                 valid for an hour).
-            password: set API username/password instead of access token.
             query_filter: filter string to include when searching for items. This will
                 be appended to other name, geographic, and sensing time filters where
                 applicable. For example, "Collection/Name eq 'SENTINEL-2'". See the API
@@ -368,6 +367,7 @@ class Copernicus(DataSource):
             "order_by",
             "sort_by",
             "sort_desc",
+            "timeout",
         ]
         for k in simple_optionals:
             if k in d:
@@ -709,6 +709,8 @@ class Sentinel2(Copernicus):
         "B12": ["B12"],
         "B8A": ["B8A"],
         "TCI": ["R", "G", "B"],
+        # L1C-only products.
+        "B10": ["B10"],
         # L2A-only products.
         "AOT": ["AOT"],
         "WVP": ["WVP"],
@@ -809,17 +811,16 @@ class Sentinel2(Copernicus):
 
         kwargs: dict[str, Any] = dict(
             assets=list(needed_assets),
+            product_type=Sentinel2ProductType[d["product_type"]],
         )
 
-        if "product_type" in d:
-            kwargs["product_type"] = Sentinel2ProductType(d["product_type"])
-
         simple_optionals = [
             "harmonize",
             "access_token",
             "order_by",
             "sort_by",
             "sort_desc",
+            "timeout",
         ]
         for k in simple_optionals:
             if k in d:
@@ -965,6 +966,7 @@ class Sentinel1(Copernicus):
             "order_by",
             "sort_by",
             "sort_desc",
+            "timeout",
         ]
         for k in simple_optionals:
             if k in d:

{rslearn-0.0.3 → rslearn-0.0.5}/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.5/rslearn/data_sources/eurocrops.py

@@ -0,0 +1,246 @@
+"""Data source for vector EuroCrops crop type data."""
+
+import glob
+import os
+import tempfile
+import zipfile
+from datetime import UTC, datetime, timedelta
+from typing import Any
+
+import fiona
+import requests
+from rasterio.crs import CRS
+from upath import UPath
+
+from rslearn.config import QueryConfig, VectorLayerConfig
+from rslearn.const import WGS84_PROJECTION
+from rslearn.data_sources import DataSource, Item
+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.feature import Feature
+from rslearn.utils.geometry import Projection, STGeometry, get_global_geometry
+
+logger = get_logger(__name__)
+
+
+class EuroCropsItem(Item):
+    """An item in the EuroCrops data source.
+
+    For simplicity, we have just one item per year, so each item combines all of the
+    country-level files for that year.
+    """
+
+    def __init__(self, name: str, geometry: STGeometry, zip_fnames: list[str]):
+        """Creates a new EuroCropsItem.
+
+        Args:
+            name: unique name of the item. It is just the year that this item
+                corresponds to.
+            geometry: the spatial and temporal extent of the item
+            zip_fnames: the filenames of the zip files that contain country-level crop
+                type data for this year.
+        """
+        super().__init__(name, geometry)
+        self.zip_fnames = zip_fnames
+
+    def serialize(self) -> dict:
+        """Serializes the item to a JSON-encodable dictionary."""
+        d = super().serialize()
+        d["zip_fnames"] = self.zip_fnames
+        return d
+
+    @staticmethod
+    def deserialize(d: dict) -> "EuroCropsItem":
+        """Deserializes an item from a JSON-decoded dictionary."""
+        item = super(EuroCropsItem, EuroCropsItem).deserialize(d)
+        return EuroCropsItem(
+            name=item.name, geometry=item.geometry, zip_fnames=d["zip_fnames"]
+        )
+
+
+class EuroCrops(DataSource[EuroCropsItem]):
+    """A data source for EuroCrops vector data (v11).
+
+    See https://zenodo.org/records/14094196 for details.
+
+    While the source data is split into country-level files, this data source uses one
+    item per year for simplicity. So each item corresponds to all of the country-level
+    files for that year.
+
+    Note that the RO_ny.zip file is not used.
+    """
+
+    BASE_URL = "https://zenodo.org/records/14094196/files/"
+    FILENAMES_BY_YEAR = {
+        2018: [
+            "FR_2018.zip",
+        ],
+        2019: [
+            "DK_2019.zip",
+        ],
+        2020: [
+            "ES_NA_2020.zip",
+            "FI_2020.zip",
+            "HR_2020.zip",
+            "NL_2020.zip",
+        ],
+        2021: [
+            "AT_2021.zip",
+            "BE_VLG_2021.zip",
+            "BE_WAL_2021.zip",
+            "EE_2021.zip",
+            "LT_2021.zip",
+            "LV_2021.zip",
+            "PT_2021.zip",
+            "SE_2021.zip",
+            "SI_2021.zip",
+            "SK_2021.zip",
+        ],
+        2023: [
+            "CZ_2023.zip",
+            "DE_BB_2023.zip",
+            "DE_LS_2021.zip",
+            "DE_NRW_2021.zip",
+            "ES_2023.zip",
+            "IE_2023.zip",
+        ],
+    }
+    TIMEOUT = timedelta(seconds=10)
+
+    @staticmethod
+    def from_config(config: VectorLayerConfig, ds_path: UPath) -> "EuroCrops":
+        """Creates a new EuroCrops instance from a configuration dictionary."""
+        if config.data_source is None:
+            raise ValueError("data_source is required")
+        return EuroCrops()
+
+    def _get_all_items(self) -> list[EuroCropsItem]:
+        """Get a list of all available items in the data source."""
+        items: list[EuroCropsItem] = []
+        for year, fnames in self.FILENAMES_BY_YEAR.items():
+            items.append(
+                EuroCropsItem(
+                    str(year),
+                    get_global_geometry(
+                        time_range=(
+                            datetime(year, 1, 1, tzinfo=UTC),
+                            datetime(year + 1, 1, 1, tzinfo=UTC),
+                        ),
+                    ),
+                    fnames,
+                )
+            )
+        return items
+
+    def get_items(
+        self, geometries: list[STGeometry], query_config: QueryConfig
+    ) -> list[list[list[EuroCropsItem]]]:
+        """Get a list of items in the data source intersecting the given geometries.
+
+        Args:
+            geometries: the spatiotemporal geometries
+            query_config: the query configuration
+
+        Returns:
+            List of groups of items that should be retrieved for each geometry.
+        """
+        wgs84_geometries = [
+            geometry.to_projection(WGS84_PROJECTION) for geometry in geometries
+        ]
+        all_items = self._get_all_items()
+        groups = []
+        for geometry in wgs84_geometries:
+            cur_groups = match_candidate_items_to_window(
+                geometry, all_items, query_config
+            )
+            groups.append(cur_groups)
+        return groups
+
+    def deserialize_item(self, serialized_item: Any) -> EuroCropsItem:
+        """Deserializes an item from JSON-decoded data."""
+        return EuroCropsItem.deserialize(serialized_item)
+
+    def _extract_features(self, fname: str) -> list[Feature]:
+        """Download the given zip file, extract shapefile, and return list of features."""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Download the zip file.
+            url = self.BASE_URL + fname
+            logger.debug(f"Downloading zip file from {url}")
+            response = requests.get(
+                url,
+                stream=True,
+                timeout=self.TIMEOUT.total_seconds(),
+                allow_redirects=False,
+            )
+            response.raise_for_status()
+            zip_fname = os.path.join(tmp_dir, "data.zip")
+            with open(zip_fname, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+
+            # Extract all of the files and look for shapefile filename.
+            logger.debug(f"Extracting zip file {fname}")
+            with zipfile.ZipFile(zip_fname) as zip_f:
+                zip_f.extractall(path=tmp_dir)
+
+            # The shapefiles or geopackage files can appear at any level in the hierarchy.
+            # Most zip files contain one but some contain multiple (one per region).
+            shp_fnames = glob.glob(
+                "**/*.shp", root_dir=tmp_dir, recursive=True
+            ) + glob.glob("**/*.gpkg", root_dir=tmp_dir, recursive=True)
+            if len(shp_fnames) == 0:
+                tmp_dir_fnames = os.listdir(tmp_dir)
+                raise ValueError(
+                    f"expected {fname} to contain .shp file but none found (matches={shp_fnames}, ls={tmp_dir_fnames})"
+                )
+
+            # Load the features from the shapefile(s).
+            features = []
+            for shp_fname in shp_fnames:
+                logger.debug(f"Loading feature list from {shp_fname}")
+                with fiona.open(os.path.join(tmp_dir, shp_fname)) as src:
+                    crs = CRS.from_wkt(src.crs.to_wkt())
+                    # Normal GeoJSON should have coordinates in CRS coordinates, i.e. it
+                    # should be 1 projection unit/pixel.
+                    projection = Projection(crs, 1, 1)
+
+                    for feat in src:
+                        features.append(
+                            Feature.from_geojson(
+                                projection,
+                                {
+                                    "type": "Feature",
+                                    "geometry": dict(feat.geometry),
+                                    "properties": dict(feat.properties),
+                                },
+                            )
+                        )
+
+            return features
+
+    def ingest(
+        self,
+        tile_store: TileStoreWithLayer,
+        items: list[EuroCropsItem],
+        geometries: list[list[STGeometry]],
+    ) -> None:
+        """Ingest items into the given tile store.
+
+        Args:
+            tile_store: the tile store to ingest into
+            items: the items to ingest
+            geometries: a list of geometries needed for each item
+        """
+        for item in items:
+            if tile_store.is_vector_ready(item.name):
+                continue
+
+            # Get features across all shapefiles.
+            features: list[Feature] = []
+            for fname in item.zip_fnames:
+                logger.debug(f"Getting features from {fname} for item {item.name}")
+                features.extend(self._extract_features(fname))
+
+            logger.debug(f"Writing features for {item.name} to the tile store")
+            tile_store.write_vector(item.name, features)

{rslearn-0.0.3 → rslearn-0.0.5}/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.5}/rslearn/data_sources/local_files.py

@@ -232,6 +232,17 @@ class RasterImporter(Importer):
                 projection = Projection(crs, x_resolution, y_resolution)
                 geometry = STGeometry(projection, shp, None)
 
+            if geometry.is_too_large():
+                geometry = get_global_geometry(time_range=None)
+                logger.warning(
+                    "Global geometry detected: this geometry will be matched against all "
+                    "windows in the rslearn dataset. When using settings like "
+                    "max_matches=1 and space_mode=MOSAIC, this may cause windows outside "
+                    "the geometry’s valid bounds to be materialized from the global raster "
+                    "instead of a more appropriate source. Consider using COMPOSITE mode, "
+                    "or increasing max_matches if this behavior is unintended."
+                )
+
             if spec.name:
                 item_name = spec.name
             else:

{rslearn-0.0.3 → rslearn-0.0.5}/rslearn/data_sources/openstreetmap.py

@@ -1,4 +1,4 @@
-"""Data source for raster data on public Cloud Storage buckets."""
+"""Data source for OpenStreetMap vector features."""
 
 import json
 import shutil
@@ -392,7 +392,7 @@ class OpenStreetMap(DataSource[OsmItem]):
         bounds_fname: UPath,
         categories: dict[str, Filter],
     ):
-        """Initialize a new Sentinel2 instance.
+        """Initialize a new OpenStreetMap instance.
 
         Args:
             config: the configuration of this layer.
@@ -508,8 +508,6 @@ class OpenStreetMap(DataSource[OsmItem]):
             items: the items to ingest
             geometries: a list of geometries needed for each item
         """
-        item_names = [item.name for item in items]
-        item_names.sort()
         for cur_item, cur_geometries in zip(items, geometries):
             if tile_store.is_vector_ready(cur_item.name):
                 continue

{rslearn-0.0.3 → rslearn-0.0.5}/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.5}/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,
     )