Skip to content

Commit

Permalink
Add back comments
Browse files Browse the repository at this point in the history
  • Loading branch information
irenedea committed Oct 24, 2024
1 parent 5feab8b commit a501786
Showing 1 changed file with 14 additions and 2 deletions.
16 changes: 14 additions & 2 deletions composer/utils/file_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -412,7 +412,9 @@ def parse_uri(uri: str) -> tuple[str, str, str]:


def maybe_create_object_store_from_uri(uri: str) -> Optional[ObjectStore]:
"""Automatically creates an ObjectStore from supported URI formats.
"""Automatically creates an :class:`composer.utils.ObjectStore` from supported URI formats.
Currently supported backends are ``s3://``, ``oci://``, and local paths (in which case ``None`` will be returned)
Args:
uri (str): The path to (maybe) create an :class:`composer.utils.ObjectStore` from.
Expand All @@ -439,14 +441,21 @@ def maybe_create_object_store_from_uri(uri: str) -> Optional[ObjectStore]:
store = None
if dist.get_global_rank() == 0:
store = MLFlowObjectStore(path)

# The path may have had placeholders, so update it with the experiment/run IDs initialized by the store
path = store.get_dbfs_path(path)

# Broadcast the rank 0 updated path to all ranks for their own object stores
path_list = [path]
dist.broadcast_object_list(path_list, src=0)
path = path_list[0]

# Create the object store for all other ranks
if dist.get_global_rank() != 0:
store = MLFlowObjectStore(path)
return store
else:
# validate if the path conforms to the requirements for UC volume paths
UCObjectStore.validate_path(path)
return UCObjectStore(path=path)

Expand All @@ -455,7 +464,10 @@ def maybe_create_object_store_from_uri(uri: str) -> Optional[ObjectStore]:
return BACKEND_TO_OBJECT_STORE_FACTORY[backend](bucket_name, path)

# If backend is unknown, raise NotImplementedError
raise NotImplementedError(f'There is no implementation for the cloud backend {backend} via URI.')
raise NotImplementedError(
f'There is no implementation for the cloud backend {backend} via URI. Please use '
'one of the supported object stores',
)


def maybe_create_remote_uploader_downloader_from_uri(
Expand Down

0 comments on commit a501786

Please sign in to comment.