support max_concurrency in upload_blob and download_blob operat…

…ions (#420)

* spec: expose max_concurrency param in blob_upload/blob_download

* prefer batch_size in batched async methods

adlfs/spec.py

@@ -26,7 +26,7 @@
 from azure.storage.blob._shared.base_client import create_configuration
 from azure.storage.blob.aio import BlobServiceClient as AIOBlobServiceClient
 from azure.storage.blob.aio._list_blobs_helper import BlobPrefix
-from fsspec.asyn import AsyncFileSystem, get_loop, sync, sync_wrapper
+from fsspec.asyn import AsyncFileSystem, _get_batch_size, get_loop, sync, sync_wrapper
 from fsspec.spec import AbstractBufferedFile
 from fsspec.utils import infer_storage_options
 
@@ -165,6 +165,10 @@ class AzureBlobFileSystem(AsyncFileSystem):
  False throws if retrieving container properties fails, which might happen if your
  authentication is only valid at the storage container level, and not the
  storage account level.
+ max_concurrency:
+ The number of concurrent connections to use when uploading or downloading a blob.
+ If None it will be inferred from fsspec.asyn._get_batch_size().
+
  Pass on to fsspec:
 
  skip_instance_cache: to control reuse of instances
@@ -227,6 +231,7 @@ def __init__(
  default_cache_type: str = "bytes",
  version_aware: bool = False,
  assume_container_exists: Optional[bool] = None,
+ max_concurrency: Optional[int] = None,
  **kwargs,
  ):
  super_kwargs = {
@@ -292,6 +297,12 @@ def __init__(
  if self.credential is not None:
  weakref.finalize(self, sync, self.loop, close_credential, self)
 
+ if max_concurrency is None:
+ batch_size = _get_batch_size()
+ if batch_size > 0:
+ max_concurrency = batch_size
+ self.max_concurrency = max_concurrency
+
  @classmethod
  def _strip_protocol(cls, path: str):
  """
@@ -1426,7 +1437,9 @@ async def _dir_exists(self, container, path):
  except ResourceNotFoundError:
  return False
 
- async def _pipe_file(self, path, value, overwrite=True, **kwargs):
+ async def _pipe_file(
+ self, path, value, overwrite=True, max_concurrency=None, **kwargs
+ ):
  """Set the bytes of given file"""
  container_name, path, _ = self.split_path(path)
  async with self.service_client.get_blob_client(
@@ -1436,14 +1449,23 @@ async def _pipe_file(self, path, value, overwrite=True, **kwargs):
  data=value,
  overwrite=overwrite,
  metadata={"is_directory": "false"},
+ max_concurrency=max_concurrency or self.max_concurrency,
  **kwargs,
  )
  self.invalidate_cache(self._parent(path))
  return result
 
  pipe_file = sync_wrapper(_pipe_file)
 
- async def _cat_file(self, path, start=None, end=None, **kwargs):
+ async def _pipe(self, *args, batch_size=None, max_concurrency=None, **kwargs):
+ max_concurrency = max_concurrency or 1
+ return await super()._pipe(
+ *args, batch_size=batch_size, max_concurrency=max_concurrency, **kwargs
+ )
+
+ async def _cat_file(
+ self, path, start=None, end=None, max_concurrency=None, **kwargs
+ ):
  path = self._strip_protocol(path)
  if end is not None:
  start = start or 0 # download_blob requires start if length is provided.
@@ -1456,7 +1478,10 @@ async def _cat_file(self, path, start=None, end=None, **kwargs):
  ) as bc:
  try:
  stream = await bc.download_blob(
- offset=start, length=length, version_id=version_id
+ offset=start,
+ length=length,
+ version_id=version_id,
+ max_concurrency=max_concurrency or self.max_concurrency,
  )
  except ResourceNotFoundError as e:
  raise FileNotFoundError from e
@@ -1497,6 +1522,12 @@ def cat(self, path, recursive=False, on_error="raise", **kwargs):
  else:
  return self.cat_file(paths[0])
 
+ async def _cat_ranges(self, *args, batch_size=None, max_concurrency=None, **kwargs):
+ max_concurrency = max_concurrency or 1
+ return await super()._cat_ranges(
+ *args, batch_size=batch_size, max_concurrency=max_concurrency, **kwargs
+ )
+
  def url(self, path, expires=3600, **kwargs):
  return sync(self.loop, self._url, path, expires, **kwargs)
 
@@ -1593,7 +1624,14 @@ async def _expand_path(
  return list(sorted(out))
 
  async def _put_file(
- self, lpath, rpath, delimiter="/", overwrite=False, callback=None, **kwargws
+ self,
+ lpath,
+ rpath,
+ delimiter="/",
+ overwrite=False,
+ callback=None,
+ max_concurrency=None,
+ **kwargws,
  ):
  """
  Copy single file to remote
@@ -1621,6 +1659,7 @@ async def _put_file(
  raw_response_hook=make_callback(
  "upload_stream_current", callback
  ),
+ max_concurrency=max_concurrency or self.max_concurrency,
  )
  self.invalidate_cache()
  except ResourceExistsError:
@@ -1633,6 +1672,12 @@ async def _put_file(
 
  put_file = sync_wrapper(_put_file)
 
+ async def _put(self, *args, batch_size=None, max_concurrency=None, **kwargs):
+ max_concurrency = max_concurrency or 1
+ return await super()._put(
+ *args, batch_size=batch_size, max_concurrency=max_concurrency, **kwargs
+ )
+
  async def _cp_file(self, path1, path2, **kwargs):
  """Copy the file at path1 to path2"""
  container1, path1, version_id = self.split_path(path1, delimiter="/")
@@ -1668,7 +1713,14 @@ def download(self, rpath, lpath, recursive=False, **kwargs):
  return self.get(rpath, lpath, recursive=recursive, **kwargs)
 
  async def _get_file(
- self, rpath, lpath, recursive=False, delimiter="/", callback=None, **kwargs
+ self,
+ rpath,
+ lpath,
+ recursive=False,
+ delimiter="/",
+ callback=None,
+ max_concurrency=None,
+ **kwargs,
  ):
  """Copy single file remote to local"""
  if os.path.isdir(lpath):
@@ -1683,6 +1735,7 @@ async def _get_file(
  "download_stream_current", callback
  ),
  version_id=version_id,
+ max_concurrency=max_concurrency or self.max_concurrency,
  )
  with open(lpath, "wb") as my_blob:
  await stream.readinto(my_blob)
@@ -1691,6 +1744,12 @@ async def _get_file(
 
  get_file = sync_wrapper(_get_file)
 
+ async def _get(self, *args, batch_size=None, max_concurrency=None, **kwargs):
+ max_concurrency = max_concurrency or 1
+ return await super()._get(
+ *args, batch_size=batch_size, max_concurrency=max_concurrency, **kwargs
+ )
+
  def getxattr(self, path, attr):
  meta = self.info(path).get("metadata", {})
  return meta[attr]

adlfs/tests/test_spec.py

@@ -392,7 +392,9 @@ def test_info_missing(storage, path):
 
 def test_time_info(storage):
  fs = AzureBlobFileSystem(
- account_name=storage.account_name, connection_string=CONN_STR
+ account_name=storage.account_name,
+ connection_string=CONN_STR,
+ max_concurrency=1,
  )
 
  creation_time = fs.created("data/root/d/file_with_metadata.txt")
@@ -1486,7 +1488,10 @@ async def test_cat_file_versioned(storage, mocker):
 
  await fs._cat_file(f"data/root/a/file.txt?versionid={DEFAULT_VERSION_ID}")
  download_blob.assert_called_once_with(
- offset=None, length=None, version_id=DEFAULT_VERSION_ID
+ offset=None,
+ length=None,
+ version_id=DEFAULT_VERSION_ID,
+ max_concurrency=fs.max_concurrency,
  )
 
  download_blob.reset_mock()
@@ -1744,7 +1749,9 @@ async def test_get_file_versioned(storage, mocker, tmp_path):
  f"data/root/a/file.txt?versionid={DEFAULT_VERSION_ID}", tmp_path / "file.txt"
  )
  download_blob.assert_called_once_with(
- raw_response_hook=mocker.ANY, version_id=DEFAULT_VERSION_ID
+ raw_response_hook=mocker.ANY,
+ version_id=DEFAULT_VERSION_ID,
+ max_concurrency=fs.max_concurrency,
  )
  download_blob.reset_mock()
  download_blob.side_effect = ResourceNotFoundError