• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import os
2from typing import List, Type, Union
3
4from .filesystem import FileSystemReader, FileSystemWriter
5from .storage import StorageReader, StorageWriter
6
7
8def _storage_setup(
9    storage: Union[StorageReader, StorageWriter, None],
10    checkpoint_id: Union[str, os.PathLike, None],
11    reader: bool = False,
12) -> Union[None, StorageReader, StorageWriter]:
13    if storage:
14        if checkpoint_id is not None:
15            storage.reset(checkpoint_id)
16        return storage
17
18    if not checkpoint_id:
19        raise RuntimeError(
20            "`checkpoint_id` must be specificed if "
21            "storage_reader/storage_writer is None."
22        )
23
24    targets: List[Type[Union[StorageReader, StorageWriter]]] = []
25    if reader:
26        targets = [
27            FileSystemReader,
28        ]
29    else:
30        targets = [
31            FileSystemWriter,
32        ]
33    try:
34        from ._fsspec_filesystem import FsspecReader, FsspecWriter
35
36        targets.append(FsspecReader if reader else FsspecWriter)
37    except Exception:
38        pass
39
40    for target in targets:
41        if target.validate_checkpoint_id(checkpoint_id):
42            storage = target(checkpoint_id)  # type: ignore[call-arg]
43            storage.reset(checkpoint_id)
44            return storage
45
46    raise RuntimeError(
47        "Cannot detect which StorageReader or StorageWriter to use. "
48        "Please specify the storage_reader/storage_writer."
49    )
50