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

rslearn/main.py

@@ -1,39 +1,56 @@
 """Entrypoint for the rslearn command-line interface."""
 
 import argparse
-import logging
 import multiprocessing
 import random
 import sys
+import time
 from collections.abc import Callable
-from datetime import datetime, timezone
-from pathlib import Path
+from datetime import UTC, datetime, timedelta
+from typing import Any, TypeVar
 
 import tqdm
-import wandb
-from lightning.pytorch.cli import LightningCLI
 from rasterio.crs import CRS
 from upath import UPath
 
 from rslearn.config import LayerConfig
 from rslearn.const import WGS84_EPSG
-from rslearn.data_sources import Item, data_source_from_config
-from rslearn.dataset import Dataset, Window
+from rslearn.data_sources import Item
+from rslearn.dataset import Dataset, Window, WindowLayerData
 from rslearn.dataset.add_windows import add_windows_from_box, add_windows_from_file
-from rslearn.dataset.manage import materialize_dataset_windows, prepare_dataset_windows
-from rslearn.tile_stores import get_tile_store_for_layer
-from rslearn.train.data_module import RslearnDataModule
-from rslearn.train.lightning_module import RslearnLightningModule
+from rslearn.dataset.handler_summaries import (
+    ErrorOutcome,
+    IngestCounts,
+    IngestDatasetJobsSummary,
+    LayerIngestSummary,
+    MaterializeDatasetWindowsSummary,
+    PrepareDatasetWindowsSummary,
+    UnknownIngestCounts,
+)
+from rslearn.dataset.manage import (
+    AttemptsCounter,
+    materialize_dataset_windows,
+    prepare_dataset_windows,
+    retry,
+)
+from rslearn.dataset.storage.file import FileWindowStorage
+from rslearn.log_utils import get_logger
+from rslearn.tile_stores import get_tile_store_with_layer
 from rslearn.utils import Projection, STGeometry
 
-logging.basicConfig()
+logger = get_logger(__name__)
+
 handler_registry = {}
 
+ItemType = TypeVar("ItemType", bound="Item")
+
+MULTIPROCESSING_CONTEXT = "forkserver"
 
-def register_handler(category, command):
+
+def register_handler(category: Any, command: str) -> Callable:
     """Register a new handler for a command."""
 
-    def decorator(f):
+    def decorator(f: Callable) -> Callable:
         handler_registry[(category, command)] = f
         return f
 
@@ -47,7 +64,7 @@ def parse_time(time_str: str) -> datetime:
     """
     ts = datetime.fromisoformat(time_str)
     if not ts.tzinfo:
-        ts = ts.replace(tzinfo=timezone.utc)
+        ts = ts.replace(tzinfo=UTC)
     return ts
 
 
@@ -60,8 +77,13 @@ def parse_time_range(
     return (parse_time(start), parse_time(end))
 
 
+def parse_disabled_layers(disabled_layers: str) -> list[str]:
+    """Parse the disabled layers string."""
+    return disabled_layers.split(",") if disabled_layers else []
+
+
 @register_handler("dataset", "add_windows")
-def add_windows():
+def add_windows() -> None:
     """Handler for the rslearn dataset add_windows command."""
     parser = argparse.ArgumentParser(
         prog="rslearn dataset add_windows",
@@ -156,7 +178,13 @@ def add_windows():
     )
     args = parser.parse_args(args=sys.argv[3:])
 
-    def parse_projection(crs_str, resolution, x_res, y_res, default_crs=None):
+    def parse_projection(
+        crs_str: str | None,
+        resolution: float | None,
+        x_res: float,
+        y_res: float,
+        default_crs: CRS | None = None,
+    ) -> Projection | None:
         if not crs_str:
             if default_crs:
                 crs = default_crs
@@ -197,7 +225,8 @@ def add_windows():
         box = [float(value) for value in args.box.split(",")]
 
         windows = add_windows_from_box(
-            box=box,
+            # TODO: we should have an object for box
+            box=box,  # type: ignore
             src_projection=parse_projection(
                 args.src_crs, args.src_resolution, args.src_x_res, args.src_y_res
             ),
@@ -210,10 +239,10 @@ def add_windows():
     else:
         raise Exception("one of box or fname must be specified")
 
-    print(f"created {len(windows)} windows")
+    logger.info(f"created {len(windows)} windows")
 
 
-def add_apply_on_windows_args(parser: argparse.ArgumentParser):
+def add_apply_on_windows_args(parser: argparse.ArgumentParser) -> None:
     """Add arguments for handlers that use the apply_on_windows helper.
 
     Args:
@@ -223,10 +252,14 @@ def add_apply_on_windows_args(parser: argparse.ArgumentParser):
         "--root", type=str, required=True, help="Dataset root directory"
     )
     parser.add_argument(
-        "--group", type=str, default=None, help="Only prepare windows in this group"
+        "--group",
+        type=str,
+        nargs="*",
+        default=None,
+        help="Only prepare windows in these groups",
     )
     parser.add_argument(
-        "--window", type=str, default=None, help="Only prepare this window"
+        "--window", type=str, nargs="*", default=None, help="Only prepare these windows"
     )
     parser.add_argument(
         "--workers",
@@ -234,6 +267,12 @@ def add_apply_on_windows_args(parser: argparse.ArgumentParser):
         default=0,
         help="Number of worker processes (default 0 to use main process only)",
     )
+    parser.add_argument(
+        "--load-workers",
+        type=int,
+        default=None,
+        help="Number of workers for loading windows (defaults to --workers)",
+    )
     parser.add_argument(
         "--batch-size",
         type=int,
@@ -255,25 +294,31 @@ def add_apply_on_windows_args(parser: argparse.ArgumentParser):
 
 
 def apply_on_windows(
-    f: Callable[[list[Window]], None],
+    f: Callable[[list[Window]], Any],
     dataset: Dataset,
-    group: str | None = None,
-    window: str | None = None,
+    group: str | list[str] | None = None,
+    names: list[str] | None = None,
     workers: int = 0,
+    load_workers: int | None = None,
     batch_size: int = 1,
     jobs_per_process: int | None = None,
     use_initial_job: bool = True,
-):
+) -> None:
     """A helper to apply a function on windows in a dataset.
 
     Args:
         f: the function to apply on lists of windows.
         dataset: the dataset.
         group: optional, only apply on windows in this group.
-        window: optional, only apply on windows with this name.
+        names: optional, only apply on windows with these names.
         workers: the number of parallel workers to use, default 0 (main thread only).
+        load_workers: optional different number of workers to use for loading the
+            windows. If set, workers controls the number of workers to process the
+            jobs, while load_workers controls the number of workers to use for reading
+            windows from the rslearn dataset. Workers is only passed if the window
+            storage is FileWindowStorage.
         batch_size: if workers > 0, the maximum number of windows to pass to the
-            function. If workers == 0, all windows are always passed.
+            function.
         jobs_per_process: optional, terminate processes after they have handled this
             many jobs. This is useful if there is a memory leak in a dependency.
         use_initial_job: if workers > 0, by default, an initial job is run on the first
@@ -284,30 +329,33 @@ def apply_on_windows(
     if hasattr(f, "set_dataset"):
         f.set_dataset(dataset)
 
-    groups = None
-    names = None
-    if group:
+    # Handle group. It can be None (load all groups) or list of groups. But it can also
+    # just be group name, in which case we must convert to list.
+    groups: list[str] | None
+    if isinstance(group, str):
         groups = [group]
-    if window:
-        names = [window]
-    windows = dataset.load_windows(
-        groups=groups, names=names, workers=workers, show_progress=True
-    )
-    print(f"found {len(windows)} windows")
+    else:
+        groups = group
+
+    # Load the windows. We pass workers and show_progress if it is FileWindowStorage.
+    kwargs: dict[str, Any] = {}
+    if isinstance(dataset.storage, FileWindowStorage):
+        if load_workers is None:
+            load_workers = workers
+        kwargs["workers"] = load_workers
+        kwargs["show_progress"] = True
+    windows = dataset.load_windows(groups=groups, names=names, **kwargs)
+    logger.info(f"found {len(windows)} windows")
 
     if hasattr(f, "get_jobs"):
-        jobs = f.get_jobs(windows, workers)
-        print(f"got {len(jobs)} jobs")
+        jobs = f.get_jobs(windows, load_workers)
+        logger.info(f"got {len(jobs)} jobs")
     else:
         jobs = windows
 
-    if workers == 0:
-        f(jobs)
-        return
-
     random.shuffle(jobs)
 
-    if use_initial_job:
+    if use_initial_job and len(jobs) > 0:
         # Apply directly on first window to get any initialization out of the way.
         f([jobs[0]])
         jobs = jobs[1:]
@@ -316,41 +364,59 @@ def apply_on_windows(
     for i in range(0, len(jobs), batch_size):
         batches.append(jobs[i : i + batch_size])
 
-    p = multiprocessing.Pool(processes=workers, maxtasksperchild=jobs_per_process)
-    outputs = p.imap_unordered(f, batches)
-    for _ in tqdm.tqdm(outputs, total=len(batches)):
-        pass
-    p.close()
+    num_batches = len(batches)
+    if workers == 0:
+        # Process batches sequentially but with same error handling as parallel
+        for batch in tqdm.tqdm(batches, total=num_batches):
+            f(batch)
+    else:
+        # Process batches in parallel
+        p = multiprocessing.Pool(processes=workers, maxtasksperchild=jobs_per_process)
+        outputs = p.imap_unordered(f, batches)
+        for _ in tqdm.tqdm(outputs, total=num_batches):
+            pass
+        p.close()
 
 
-def apply_on_windows_args(f: Callable[[list[Window]], None], args: argparse.Namespace):
+def apply_on_windows_args(f: Callable[..., Any], args: argparse.Namespace) -> None:
     """Call apply_on_windows with arguments passed via command-line interface."""
-    dataset = Dataset(UPath(args.root))
+    dataset = Dataset(UPath(args.root), disabled_layers=args.disabled_layers)
     apply_on_windows(
-        f,
-        dataset,
-        args.group,
-        args.window,
-        args.workers,
-        args.batch_size,
-        args.jobs_per_process,
-        args.use_initial_job,
+        f=f,
+        dataset=dataset,
+        group=args.group,
+        names=args.window,
+        workers=args.workers,
+        load_workers=args.load_workers,
+        batch_size=args.batch_size,
+        jobs_per_process=args.jobs_per_process,
+        use_initial_job=args.use_initial_job,
     )
 
 
 class PrepareHandler:
     """apply_on_windows handler for the rslearn dataset prepare command."""
 
-    def __init__(self, force: bool):
+    def __init__(
+        self,
+        force: bool,
+        retry_max_attempts: int = 0,
+        retry_backoff: timedelta = timedelta(minutes=1),
+    ) -> None:
         """Initialize a new PrepareHandler.
 
         Args:
             force: force prepare
+            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).
         """
         self.force = force
-        self.dataset = None
+        self.dataset: Dataset | None = None
+        self.retry_max_attempts = retry_max_attempts
+        self.retry_backoff = retry_backoff
 
-    def set_dataset(self, dataset: Dataset):
+    def set_dataset(self, dataset: Dataset) -> None:
         """Captures the dataset from apply_on_windows_args.
 
         Args:
@@ -358,13 +424,22 @@ class PrepareHandler:
         """
         self.dataset = dataset
 
-    def __call__(self, windows: list[Window]):
+    def __call__(self, windows: list[Window]) -> PrepareDatasetWindowsSummary:
         """Prepares the windows from apply_on_windows."""
-        prepare_dataset_windows(self.dataset, windows, self.force)
+        logger.info(f"Running prepare on {len(windows)} windows")
+        if self.dataset is None:
+            raise ValueError("dataset not set")
+        return prepare_dataset_windows(
+            self.dataset,
+            windows,
+            self.force,
+            retry_max_attempts=self.retry_max_attempts,
+            retry_backoff=self.retry_backoff,
+        )
 
 
 @register_handler("dataset", "prepare")
-def dataset_prepare():
+def dataset_prepare() -> None:
     """Handler for the rslearn dataset prepare command."""
     parser = argparse.ArgumentParser(
         prog="rslearn dataset prepare",
@@ -377,14 +452,38 @@ def dataset_prepare():
         action=argparse.BooleanOptionalAction,
         help="Prepare windows even if they were previously prepared",
     )
+    parser.add_argument(
+        "--disabled-layers",
+        type=parse_disabled_layers,
+        default="",
+        help="List of layers to disable e.g 'layer1,layer2'",
+    )
+    parser.add_argument(
+        "--retry-max-attempts",
+        type=int,
+        default=0,
+        help="Retry for this many attempts",
+    )
+    parser.add_argument(
+        "--retry-backoff-seconds",
+        type=int,
+        default=0,
+        help="Backoff time (seconds) between retries",
+    )
     add_apply_on_windows_args(parser)
     args = parser.parse_args(args=sys.argv[3:])
 
-    fn = PrepareHandler(args.force)
+    fn = PrepareHandler(
+        args.force,
+        retry_max_attempts=args.retry_max_attempts,
+        retry_backoff=timedelta(seconds=args.retry_backoff_seconds),
+    )
     apply_on_windows_args(fn, args)
 
 
-def _load_window_layer_datas(window: Window):
+def _load_window_layer_datas(
+    window: Window,
+) -> tuple[Window, dict[str, WindowLayerData]]:
     # Helper for IngestHandler to use with multiprocessing.
     return window, window.load_layer_datas()
 
@@ -392,11 +491,19 @@ def _load_window_layer_datas(window: Window):
 class IngestHandler:
     """apply_on_windows handler for the rslearn dataset ingest command."""
 
-    def __init__(self):
+    def __init__(
+        self,
+        ignore_errors: bool = False,
+        retry_max_attempts: int = 0,
+        retry_backoff: timedelta = timedelta(minutes=1),
+    ) -> None:
         """Initialize a new IngestHandler."""
-        self.dataset = None
+        self.dataset: Dataset | None = None
+        self.ignore_errors = ignore_errors
+        self.retry_max_attempts = retry_max_attempts
+        self.retry_backoff = retry_backoff
 
-    def set_dataset(self, dataset: Dataset):
+    def set_dataset(self, dataset: Dataset) -> None:
         """Captures the dataset from apply_on_windows_args.
 
         Args:
@@ -404,21 +511,32 @@ class IngestHandler:
         """
         self.dataset = dataset
 
-    def __call__(self, jobs: list[tuple[str, LayerConfig, Item, list[STGeometry]]]):
+    def __call__(
+        self, jobs: list[tuple[str, LayerConfig, Item, list[STGeometry]]]
+    ) -> IngestDatasetJobsSummary:
         """Ingest the specified items.
 
         The items are computed from list of windows via IngestHandler.get_jobs.
 
         Args:
-            jobs: list of (layer_name, item, geometries) tuples to ingest.
+            jobs: list of (layer_name, layer_cfg, item, geometries) tuples to ingest.
+
+        Returns:
+            summary of the ingest jobs operation fit for telemetry purposes.
         """
+        start_time = time.monotonic()
+        layer_summaries: list[LayerIngestSummary] = []
+
+        logger.info(f"Running ingest for {len(jobs)} jobs")
         import gc
 
+        if self.dataset is None:
+            raise ValueError("dataset not set")
         tile_store = self.dataset.get_tile_store()
 
         # Group jobs by layer name.
-        jobs_by_layer = {}
-        configs_by_layer = {}
+        jobs_by_layer: dict = {}
+        configs_by_layer: dict = {}
         for layer_name, layer_cfg, item, geometries in jobs:
             if layer_name not in jobs_by_layer:
                 jobs_by_layer[layer_name] = []
@@ -426,24 +544,81 @@ class IngestHandler:
             configs_by_layer[layer_name] = layer_cfg
 
         for layer_name, items_and_geometries in jobs_by_layer.items():
-            cur_tile_store = get_tile_store_for_layer(tile_store, layer_name, layer_cfg)
+            layer_tile_store = get_tile_store_with_layer(
+                tile_store, layer_name, layer_cfg
+            )
             layer_cfg = self.dataset.layers[layer_name]
-            data_source = data_source_from_config(layer_cfg, self.dataset.path)
+            data_source = layer_cfg.instantiate_data_source(self.dataset.path)
 
+            attempts_counter = AttemptsCounter()
+            ingest_counts: IngestCounts | UnknownIngestCounts
             try:
-                data_source.ingest(
-                    tile_store=cur_tile_store,
-                    items=[item for item, _ in items_and_geometries],
-                    geometries=[geometries for _, geometries in items_and_geometries],
+                retry(
+                    lambda: data_source.ingest(
+                        tile_store=layer_tile_store,
+                        items=[item for item, _ in items_and_geometries],
+                        geometries=[
+                            geometries for _, geometries in items_and_geometries
+                        ],
+                    ),
+                    retry_max_attempts=self.retry_max_attempts,
+                    retry_backoff=self.retry_backoff,
+                    attempts_counter=attempts_counter,
+                )
+                ingest_counts = IngestCounts(
+                    items_ingested=len(items_and_geometries),
+                    geometries_ingested=sum(
+                        len(geometries) for _, geometries in items_and_geometries
+                    ),
                 )
             except Exception as e:
-                print(
+                if not self.ignore_errors:
+                    raise
+
+                ingest_counts = UnknownIngestCounts(
+                    items_attempted=len(items_and_geometries),
+                    geometries_attempted=sum(
+                        len(geometries) for _, geometries in items_and_geometries
+                    ),
+                )
+                logger.error(
                     "warning: got error while ingesting "
                     + f"{len(items_and_geometries)} items: {e}"
                 )
 
+            layer_summaries.append(
+                LayerIngestSummary(
+                    layer_name=layer_name,
+                    data_source_name=getattr(layer_cfg.data_source, "name", "N/A"),
+                    duration_seconds=time.monotonic() - start_time,
+                    ingest_counts=ingest_counts,
+                    ingest_attempts=attempts_counter.value,
+                )
+            )
+
         gc.collect()
 
+        return IngestDatasetJobsSummary(
+            duration_seconds=time.monotonic() - start_time,
+            num_jobs=len(jobs),
+            layer_summaries=layer_summaries,
+        )
+
+    def _load_layer_data_for_windows(
+        self, windows: list[Window], workers: int
+    ) -> list[tuple[Window, dict[str, WindowLayerData]]]:
+        if workers == 0:
+            return [(_load_window_layer_datas(window)) for window in windows]
+        p = multiprocessing.Pool(workers)
+        outputs = p.imap_unordered(_load_window_layer_datas, windows)
+        windows_and_layer_datas = []
+        for window, layer_datas in tqdm.tqdm(
+            outputs, total=len(windows), desc="Loading window layer datas"
+        ):
+            windows_and_layer_datas.append((window, layer_datas))
+        p.close()
+        return windows_and_layer_datas
+
     def get_jobs(
         self, windows: list[Window], workers: int
     ) -> list[tuple[str, LayerConfig, Item, list[STGeometry]]]:
@@ -455,17 +630,12 @@ class IngestHandler:
         This makes sure that jobs are grouped by item rather than by window, which
         makes sense because there's no reason to ingest the same item twice.
         """
+        if self.dataset is None:
+            raise ValueError("dataset not set")
         # TODO: avoid duplicating ingest_dataset_windows...
 
         # Load layer datas of each window.
-        p = multiprocessing.Pool(workers)
-        outputs = p.imap_unordered(_load_window_layer_datas, windows)
-        windows_and_layer_datas = []
-        for window, layer_datas in tqdm.tqdm(
-            outputs, total=len(windows), desc="Loading window layer datas"
-        ):
-            windows_and_layer_datas.append((window, layer_datas))
-        p.close()
+        windows_and_layer_datas = self._load_layer_data_for_windows(windows, workers)
 
         jobs: list[tuple[str, LayerConfig, Item, list[STGeometry]]] = []
         for layer_name, layer_cfg in self.dataset.layers.items():
@@ -474,9 +644,9 @@ class IngestHandler:
             if not layer_cfg.data_source.ingest:
                 continue
 
-            data_source = data_source_from_config(layer_cfg, self.dataset.path)
+            data_source = layer_cfg.instantiate_data_source(self.dataset.path)
 
-            geometries_by_item = {}
+            geometries_by_item: dict = {}
             for window, layer_datas in windows_and_layer_datas:
                 if layer_name not in layer_datas:
                     continue
@@ -484,7 +654,9 @@ class IngestHandler:
                 layer_data = layer_datas[layer_name]
                 for group in layer_data.serialized_item_groups:
                     for serialized_item in group:
-                        item = data_source.deserialize_item(serialized_item)
+                        item = data_source.deserialize_item(  # type: ignore
+                            serialized_item
+                        )
                         if item not in geometries_by_item:
                             geometries_by_item[item] = []
                         geometries_by_item[item].append(geometry)
@@ -492,32 +664,69 @@ class IngestHandler:
             for item, geometries in geometries_by_item.items():
                 jobs.append((layer_name, layer_cfg, item, geometries))
 
-        print(f"computed {len(jobs)} ingest jobs from {len(windows)} windows")
+        logger.info(f"computed {len(jobs)} ingest jobs from {len(windows)} windows")
         return jobs
 
 
 @register_handler("dataset", "ingest")
-def dataset_ingest():
+def dataset_ingest() -> None:
     """Handler for the rslearn dataset ingest command."""
     parser = argparse.ArgumentParser(
         prog="rslearn dataset ingest",
         description="rslearn dataset ingest: ingest items in retrieved data sources",
     )
+    parser.add_argument(
+        "--disabled-layers",
+        type=parse_disabled_layers,
+        default="",
+        help="List of layers to disable e.g 'layer1,layer2'",
+    )
+    parser.add_argument(
+        "--ignore-errors",
+        type=bool,
+        default=False,
+        help="Ignore ingestion errors in individual jobs",
+        action=argparse.BooleanOptionalAction,
+    )
+    parser.add_argument(
+        "--retry-max-attempts",
+        type=int,
+        default=0,
+        help="Retry for this many attempts",
+    )
+    parser.add_argument(
+        "--retry-backoff-seconds",
+        type=int,
+        default=0,
+        help="Backoff time (seconds) between retries",
+    )
     add_apply_on_windows_args(parser)
     args = parser.parse_args(args=sys.argv[3:])
 
-    fn = IngestHandler()
+    fn = IngestHandler(
+        ignore_errors=args.ignore_errors,
+        retry_max_attempts=args.retry_max_attempts,
+        retry_backoff=timedelta(seconds=args.retry_backoff_seconds),
+    )
     apply_on_windows_args(fn, args)
 
 
 class MaterializeHandler:
     """apply_on_windows handler for the rslearn dataset materialize command."""
 
-    def __init__(self):
+    def __init__(
+        self,
+        ignore_errors: bool = False,
+        retry_max_attempts: int = 0,
+        retry_backoff: timedelta = timedelta(minutes=1),
+    ) -> None:
         """Initialize a MaterializeHandler."""
-        self.dataset = None
+        self.dataset: Dataset | None = None
+        self.ignore_errors = ignore_errors
+        self.retry_max_attempts = retry_max_attempts
+        self.retry_backoff = retry_backoff
 
-    def set_dataset(self, dataset: Dataset):
+    def set_dataset(self, dataset: Dataset) -> None:
         """Captures the dataset from apply_on_windows_args.
 
         Args:
@@ -525,13 +734,31 @@ class MaterializeHandler:
         """
         self.dataset = dataset
 
-    def __call__(self, windows: list[Window]):
+    def __call__(
+        self, windows: list[Window]
+    ) -> MaterializeDatasetWindowsSummary | ErrorOutcome:
         """Materializes the windows from apply_on_windows."""
-        materialize_dataset_windows(self.dataset, windows)
+        logger.info(f"Running Materialize with {len(windows)} windows")
+        start_time = time.monotonic()
+        if self.dataset is None:
+            raise ValueError("dataset not set")
+        try:
+            return materialize_dataset_windows(
+                self.dataset,
+                windows,
+                retry_max_attempts=self.retry_max_attempts,
+                retry_backoff=self.retry_backoff,
+            )
+        except Exception as e:
+            if not self.ignore_errors:
+                logger.error(f"Error materializing windows: {e}")
+                raise
+            logger.warning(f"Ignoring error while materializing windows: {e}")
+            return ErrorOutcome(duration_seconds=time.monotonic() - start_time)
 
 
 @register_handler("dataset", "materialize")
-def dataset_materialize():
+def dataset_materialize() -> None:
     """Handler for the rslearn dataset materialize command."""
     parser = argparse.ArgumentParser(
         prog="rslearn dataset materialize",
@@ -540,110 +767,87 @@ def dataset_materialize():
             + "materialize data from retrieved data sources"
         ),
     )
+    parser.add_argument(
+        "--disabled-layers",
+        type=parse_disabled_layers,
+        default="",
+        help="List of layers to disable e.g 'layer1,layer2'",
+    )
+    parser.add_argument(
+        "--ignore-errors",
+        type=bool,
+        default=False,
+        help="Ignore errors in individual jobs",
+        action=argparse.BooleanOptionalAction,
+    )
+    parser.add_argument(
+        "--retry-max-attempts",
+        type=int,
+        default=0,
+        help="Retry for this many attempts",
+    )
+    parser.add_argument(
+        "--retry-backoff-seconds",
+        type=int,
+        default=0,
+        help="Backoff time (seconds) between retries",
+    )
     add_apply_on_windows_args(parser)
     args = parser.parse_args(args=sys.argv[3:])
-
-    fn = MaterializeHandler()
-    apply_on_windows_args(fn, args)
-
-
-class RslearnLightningCLI(LightningCLI):
-    """LightningCLI that links data.tasks to model.tasks."""
-
-    def add_arguments_to_parser(self, parser) -> None:
-        """Link data.tasks to model.tasks.
-
-        Args:
-            parser: the argument parser
-        """
-        parser.link_arguments(
-            "data.init_args.task", "model.init_args.task", apply_on="instantiate"
-        )
-        parser.add_argument(
-            "--wandb_run_id",
-            default="",
-            type=str,
-            help="W&B run ID to load checkpoint from",
-        )
-        parser.add_argument(
-            "--wandb_resume",
-            default=False,
-            type=bool,
-            help="Whether to resume from specified wandb_run_id",
-        )
-
-    def before_instantiate_classes(self):
-        """Called before Lightning class initialization.
-
-        Sets up wandb_run_id / wandb_resume arguments.
-        """
-        subcommand = self.config.subcommand
-        c = self.config[subcommand]
-
-        if c.wandb_run_id:
-            api = wandb.Api()
-            artifact_id = (
-                f"{c.trainer.logger.init_args.project}/model-{c.wandb_run_id}:latest"
-            )
-            print(f"restoring from artifact {artifact_id} on wandb")
-            artifact = api.artifact(artifact_id, type="model")
-            artifact_dir = artifact.download()
-            c.ckpt_path = str(Path(artifact_dir) / "model.ckpt")
-
-        if c.wandb_resume:
-            c.trainer.logger.init_args.id = c.wandb_run_id
-
-        # If there is a RslearnPredictionWriter, set its path.
-        prediction_writer_callback = None
-        if "callbacks" in c.trainer:
-            for existing_callback in c.trainer.callbacks:
-                if (
-                    existing_callback.class_path
-                    == "rslearn.train.prediction_writer.RslearnWriter"
-                ):
-                    prediction_writer_callback = existing_callback
-        if prediction_writer_callback:
-            prediction_writer_callback.init_args.path = c.data.init_args.path
-
-
-def model_handler():
-    """Handler for any rslearn model X commands."""
-    RslearnLightningCLI(
-        model_class=RslearnLightningModule,
-        datamodule_class=RslearnDataModule,
-        args=sys.argv[2:],
-        subclass_mode_model=True,
-        subclass_mode_data=True,
-        save_config_kwargs={"overwrite": True},
+    fn = MaterializeHandler(
+        ignore_errors=args.ignore_errors,
+        retry_max_attempts=args.retry_max_attempts,
+        retry_backoff=timedelta(seconds=args.retry_backoff_seconds),
     )
+    apply_on_windows_args(fn, args)
 
 
 @register_handler("model", "fit")
-def model_fit():
+def model_fit() -> None:
     """Handler for rslearn model fit."""
+    from .lightning_cli import model_handler
+
     model_handler()
 
 
 @register_handler("model", "validate")
-def model_validate():
+def model_validate() -> None:
     """Handler for rslearn model validate."""
+    from .lightning_cli import model_handler
+
     model_handler()
 
 
 @register_handler("model", "test")
-def model_test():
+def model_test() -> None:
     """Handler for rslearn model test."""
+    from .lightning_cli import model_handler
+
     model_handler()
 
 
 @register_handler("model", "predict")
-def model_predict():
+def model_predict() -> None:
     """Handler for rslearn model predict."""
+    from .lightning_cli import model_handler
+
     model_handler()
 
 
-def main():
+def main() -> None:
     """CLI entrypoint."""
+    try:
+        multiprocessing.set_start_method(MULTIPROCESSING_CONTEXT)
+    except RuntimeError as e:
+        logger.error(
+            f"Multiprocessing context already set to {multiprocessing.get_context()}: "
+            + f"ignoring {e}"
+        )
+    except Exception as e:
+        logger.error(f"Failed to set multiprocessing context: {e}")
+        raise
+    finally:
+        logger.info(f"Using multiprocessing context: {multiprocessing.get_context()}")
     parser = argparse.ArgumentParser(description="rslearn")
     parser.add_argument(
         "category", help="Command category: dataset, annotate, or model"
@@ -653,12 +857,11 @@ def main():
 
     handler = handler_registry.get((args.category, args.command))
     if handler is None:
-        print(f"Unknown command: {args.category} {args.command}", file=sys.stderr)
+        logger.error(f"Unknown command: {args.category} {args.command}")
         sys.exit(1)
 
     handler()
 
 
 if __name__ == "__main__":
-    multiprocessing.set_start_method("forkserver")
     main()