Skip to content

rslearn.dataset.storage.migrate

migrate

Utilities for migrating between WindowStorage backends.

migrate_window_storage 

migrate_window_storage(source: WindowStorage, target: WindowStorage, fail_if_target_nonempty: bool = True, source_get_windows_kwargs: dict[str, Any] | None = None) -> int

Migrate all window metadata from source to target storage.

Parameters:

Name Type Description Default
source WindowStorage

source storage to read windows from.

 required
target WindowStorage

target storage to write windows to.

 required
fail_if_target_nonempty bool

whether to fail if target already has windows.

 True
source_get_windows_kwargs dict[str, Any] | None

optional keyword args to pass to source.get_windows, e.g. {"workers": 8, "show_progress": True} for FileWindowStorage.

 None

Returns:

Type Description
int

number of migrated windows.
Source code in rslearn/dataset/storage/migrate.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def migrate_window_storage(
    source: WindowStorage,
    target: WindowStorage,
    fail_if_target_nonempty: bool = True,
    source_get_windows_kwargs: dict[str, Any] | None = None,
) -> int:
    """Migrate all window metadata from source to target storage.

    Args:
        source: source storage to read windows from.
        target: target storage to write windows to.
        fail_if_target_nonempty: whether to fail if target already has windows.
        source_get_windows_kwargs: optional keyword args to pass to
            source.get_windows, e.g. {"workers": 8, "show_progress": True}
            for FileWindowStorage.

    Returns:
        number of migrated windows.
    """
    if fail_if_target_nonempty and len(target.get_windows()) > 0:
        raise ValueError(
            "target window storage is not empty; rerun with --no-fail-if-target-nonempty to bypass this check"
        )

    if source_get_windows_kwargs is None:
        source_get_windows_kwargs = {}
    windows = source.get_windows(**source_get_windows_kwargs)
    total = len(windows)
    logger.info(f"Found {total} windows in source storage")
    if total == 0:
        return 0

    for window in tqdm.tqdm(windows, total=total, desc="Migrating windows"):
        target.create_or_update_window(window)

        layer_datas = source.get_layer_datas(window.group, window.name)
        if layer_datas:
            target.save_layer_datas(window.group, window.name, layer_datas)

        for layer_name, group_idx in source.list_completed_layers(
            window.group, window.name
        ):
            target.mark_layer_completed(
                window.group, window.name, layer_name, group_idx
            )

    return total