1
2
3
4
5
6
7
8
9
10
11
12
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
|
import os
from typing import List, Type, Union
from .filesystem import FileSystemReader, FileSystemWriter
from .storage import StorageReader, StorageWriter
def _storage_setup(
storage: Union[StorageReader, StorageWriter, None],
checkpoint_id: Union[str, os.PathLike, None],
reader: bool = False,
) -> Union[None, StorageReader, StorageWriter]:
if storage:
if checkpoint_id is not None:
storage.reset(checkpoint_id)
return storage
if not checkpoint_id:
raise RuntimeError(
"`checkpoint_id` must be specificed if "
"storage_reader/storage_writer is None."
)
targets: List[Type[Union[StorageReader, StorageWriter]]] = []
if reader:
targets = [
FileSystemReader,
]
else:
targets = [
FileSystemWriter,
]
try:
from ._fsspec_filesystem import FsspecReader, FsspecWriter
targets.append(FsspecReader if reader else FsspecWriter)
except Exception:
pass
for target in targets:
if target.validate_checkpoint_id(checkpoint_id):
storage = target(checkpoint_id) # type: ignore[call-arg]
storage.reset(checkpoint_id)
return storage
raise RuntimeError(
"Cannot detect which StorageReader or StorageWriter to use. "
"Please specify the storage_reader/storage_writer."
)
|
