Modify UCObjectStore.list_objects to lists all files recursively

composer/utils/object_store/uc_object_store.py

@@ -219,10 +219,6 @@ def get_object_size(self, object_name: str) -> int:
     def list_objects(self, prefix: Optional[str]) -> List[str]:
         """List all objects in the object store with the given prefix.
 
-         .. note::
-
-            This function removes the directories from the returned list.
-
         Args:
             prefix (str): The prefix to search for.
 
@@ -234,14 +230,30 @@ def list_objects(self, prefix: Optional[str]) -> List[str]:
 
         from databricks.sdk.core import DatabricksError
         try:
-            data = json.dumps({'path': self._get_object_path(prefix)})
             # NOTE: This API is in preview and should not be directly used outside of this instance
-            resp = self.client.api_client.do(method='GET',
-                                             path=self._UC_VOLUME_LIST_API_ENDPOINT,
-                                             data=data,
-                                             headers={'Source': 'mosaicml/composer'})
-            assert isinstance(resp, dict)
-            return [f['path'] for f in resp.get('files', []) if not f['is_dir']]
+            logging.warn('UCObjectStore.list_objects is experimental.')
+            max_recursion_depth = 4
+
+            def get_uc_files(dir_path: str, recursion_depth: int = 0) -> list[str]:
+                if recursion_depth == max_recursion_depth:
+                    raise Exception(
+                        f'Objects at {dir_path} cannot be downloaded. Please reduce the' +
+                        ' level of folder nesting from {prefix} in UC Volumes to under {max_recursion_depth}.')
+                resp = self.client.api_client.do(method='GET',
+                                                 path=self._UC_VOLUME_LIST_API_ENDPOINT,
+                                                 data=json.dumps({'path': self._get_object_path(dir_path)}),
+                                                 headers={'Source': 'mosaicml/composer'})
+                assert isinstance(resp, dict)
+                files = []
+                for f in resp.get('files', []):
+                    fpath = f['path']
+                    if f['is_dir']:
+                        files.extend(get_uc_files(fpath, recursion_depth=recursion_depth + 1))
+                    else:
+                        files.append(fpath)
+                return files
+
+            return get_uc_files(prefix)
         except DatabricksError as e:
             _wrap_errors(self.get_uri(prefix), e)
         return []