rslearn.train.dataset_index

dataset_index

Dataset index for caching window lists to speed up ModelDataset initialization.

DatasetIndex

Manages indexed window lists for faster ModelDataset initialization.

Note: The index does NOT automatically detect when windows are added or removed from the dataset. Use refresh=True after modifying dataset windows.

Source code in rslearn/train/dataset_index.py

class DatasetIndex:
    """Manages indexed window lists for faster ModelDataset initialization.

    Note: The index does NOT automatically detect when windows are added or removed
    from the dataset. Use refresh=True after modifying dataset windows.
    """

    def __init__(
        self,
        storage: "WindowStorage",
        data_storage_factory: WindowDataStorageFactory,
        dataset_path: UPath,
        groups: list[str] | None,
        names: list[str] | None,
        tags: dict[str, Any] | None,
        num_samples: int | None,
        skip_targets: bool,
        inputs: dict[str, Any],
    ) -> None:
        """Initialize DatasetIndex with specific configuration.

        Args:
            storage: WindowStorage for deserializing windows.
            data_storage_factory: factory to bind data storage to deserialized windows.
            dataset_path: Path to the dataset directory.
            groups: list of window groups to include.
            names: list of window names to include.
            tags: tags to filter windows by.
            num_samples: limit on number of samples.
            skip_targets: whether targets are skipped.
            inputs: dict mapping input names to DataInput objects.
        """
        self.storage = storage
        self.data_storage_factory = data_storage_factory
        self.dataset_path = dataset_path
        self.index_dir = dataset_path / INDEX_DIR_NAME

        # Compute index key from configuration
        inputs_data = {}
        for name, inp in inputs.items():
            inputs_data[name] = {
                "layers": inp.layers,
                "required": inp.required,
                "load_all_layers": inp.load_all_layers,
                "is_target": inp.is_target,
            }

        key_data = {
            "groups": groups,
            "names": names,
            "tags": tags,
            "num_samples": num_samples,
            "skip_targets": skip_targets,
            "inputs": inputs_data,
        }
        self.index_key = hashlib.sha256(
            json.dumps(key_data, sort_keys=True).encode()
        ).hexdigest()

    def _get_config_hash(self) -> str:
        """Get hash of config.json for quick validation.

        Returns:
            A 16-character hex string hash of the config, or empty string if no config.
        """
        config_path = self.dataset_path / "config.json"
        if config_path.exists():
            with config_path.open() as f:
                return hashlib.sha256(f.read().encode()).hexdigest()[:16]
        return ""

    def load_windows(self, refresh: bool = False) -> list[Window] | None:
        """Load indexed window list if valid, else return None.

        Args:
            refresh: If True, ignore existing index and return None.

        Returns:
            List of Window objects if index is valid, None otherwise.
        """
        if refresh:
            logger.info("refresh=True, rebuilding index")
            return None

        index_file = self.index_dir / f"{self.index_key}.json"
        if not index_file.exists():
            logger.info(f"No index found at {index_file}, will build")
            return None

        try:
            with index_file.open() as f:
                index_data = json.load(f)
        except (OSError, json.JSONDecodeError):
            logger.warning(f"Corrupted index file at {index_file}, will rebuild")
            return None

        # Check index version
        if index_data.get("version") != INDEX_VERSION:
            logger.info(
                f"Index version mismatch (got {index_data.get('version')}, "
                f"expected {INDEX_VERSION}), will rebuild"
            )
            return None

        # Quick validation: check config hash
        if index_data.get("config_hash") != self._get_config_hash():
            logger.info("Config hash mismatch, index invalidated")
            return None

        # Deserialize windows
        windows = [Window.from_metadata(self.storage, w) for w in index_data["windows"]]
        for window in windows:
            window._data = self.data_storage_factory.create(window)
        return windows

    def save_windows(self, windows: list[Window]) -> None:
        """Save processed windows to index with atomic write.

        Args:
            windows: List of Window objects to index.
        """
        self.index_dir.mkdir(parents=True, exist_ok=True)
        index_file = self.index_dir / f"{self.index_key}.json"

        # Serialize windows
        serialized_windows = [w.get_metadata() for w in windows]

        index_data = {
            "version": INDEX_VERSION,
            "config_hash": self._get_config_hash(),
            "created_at": datetime.now().isoformat(),
            "num_windows": len(windows),
            "windows": serialized_windows,
        }
        with open_atomic(index_file, "w") as f:
            json.dump(index_data, f)
        logger.info(f"Saved {len(windows)} windows to index at {index_file}")

load_windows

load_windows(refresh: bool = False) -> list[Window] | None

Load indexed window list if valid, else return None.

Parameters:

Name	Type	Description	Default
refresh	bool	If True, ignore existing index and return None.	False

Returns:

Type	Description
list[Window] | None	List of Window objects if index is valid, None otherwise.

Source code in rslearn/train/dataset_index.py

def load_windows(self, refresh: bool = False) -> list[Window] | None:
    """Load indexed window list if valid, else return None.

    Args:
        refresh: If True, ignore existing index and return None.

    Returns:
        List of Window objects if index is valid, None otherwise.
    """
    if refresh:
        logger.info("refresh=True, rebuilding index")
        return None

    index_file = self.index_dir / f"{self.index_key}.json"
    if not index_file.exists():
        logger.info(f"No index found at {index_file}, will build")
        return None

    try:
        with index_file.open() as f:
            index_data = json.load(f)
    except (OSError, json.JSONDecodeError):
        logger.warning(f"Corrupted index file at {index_file}, will rebuild")
        return None

    # Check index version
    if index_data.get("version") != INDEX_VERSION:
        logger.info(
            f"Index version mismatch (got {index_data.get('version')}, "
            f"expected {INDEX_VERSION}), will rebuild"
        )
        return None

    # Quick validation: check config hash
    if index_data.get("config_hash") != self._get_config_hash():
        logger.info("Config hash mismatch, index invalidated")
        return None

    # Deserialize windows
    windows = [Window.from_metadata(self.storage, w) for w in index_data["windows"]]
    for window in windows:
        window._data = self.data_storage_factory.create(window)
    return windows

save_windows

save_windows(windows: list[Window]) -> None

Save processed windows to index with atomic write.

Parameters:

Name	Type	Description	Default
windows	list[Window]	List of Window objects to index.	required

Source code in rslearn/train/dataset_index.py

def save_windows(self, windows: list[Window]) -> None:
    """Save processed windows to index with atomic write.

    Args:
        windows: List of Window objects to index.
    """
    self.index_dir.mkdir(parents=True, exist_ok=True)
    index_file = self.index_dir / f"{self.index_key}.json"

    # Serialize windows
    serialized_windows = [w.get_metadata() for w in windows]

    index_data = {
        "version": INDEX_VERSION,
        "config_hash": self._get_config_hash(),
        "created_at": datetime.now().isoformat(),
        "num_windows": len(windows),
        "windows": serialized_windows,
    }
    with open_atomic(index_file, "w") as f:
        json.dump(index_data, f)
    logger.info(f"Saved {len(windows)} windows to index at {index_file}")