Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Minor cleanup of UC get_object_size #2989

Merged
merged 3 commits into from
Feb 9, 2024
Merged

Minor cleanup of UC get_object_size #2989

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5293f32
lstrip and unit test
dakinggg Feb 9, 2024
2107887
merge
dakinggg Feb 9, 2024
78c22a5
path join
dakinggg Feb 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions composer/utils/object_store/uc_object_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,9 +211,9 @@ def get_object_size(self, object_name: str) -> int:
# Note: The UC team is working on changes to fix the files.get_status API, but it currently
# does not work. Once fixed, we will call the files API endpoint. We currently only use this
# function in Composer and LLM-foundry to check the UC object's existence.
self.client.api_client.do(method='HEAD',
path=f'{self._UC_VOLUME_FILES_API_ENDPOINT}/{self._get_object_path(object_name)}',
headers={'Source': 'mosaicml/composer'})
object_path = self._get_object_path(object_name).lstrip('/')
path = os.path.join(self._UC_VOLUME_FILES_API_ENDPOINT, object_path)
self.client.api_client.do(method='HEAD', path=path, headers={'Source': 'mosaicml/composer'})
return 1000000 # Dummy value, as we don't have a way to get the size of the file
except DatabricksError as e:
# If the code reaches here, the file was not found
Expand Down
8 changes: 8 additions & 0 deletions tests/utils/object_store/test_uc_object_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,14 @@ def test_get_object_size(ws_client, uc_object_store, result: str):
raise NotImplementedError(f'Test for result={result} is not implemented.')


def test_get_object_size_full_path(ws_client, uc_object_store):
ws_client.api_client.do.return_value = {}
assert uc_object_store.get_object_size('Volumes/catalog/schema/volume/train.txt') == 1000000
ws_client.api_client.do.assert_called_with(method='HEAD',
path=f'/api/2.0/fs/files/Volumes/catalog/schema/volume/train.txt',
headers={'Source': 'mosaicml/composer'})


def test_get_uri(uc_object_store):
assert uc_object_store.get_uri('train.txt') == 'dbfs:/Volumes/catalog/schema/volume/train.txt'
assert uc_object_store.get_uri('Volumes/catalog/schema/volume/checkpoint/model.bin'
Expand Down
Loading