thds.adls 3.1.20250226213817__py3-none-any.whl → 3.1.20250227181712__py3-none-any.whl

thds/adls/_progress.py

@@ -138,11 +138,13 @@ class Tracker:
         self._reporter(list(self._progresses.values()))
         return self, key
 
-    def __call__(self, key: str, written: int):
+    def __call__(self, key: str, written: int = 0, total_written: int = 0):
         assert written >= 0, "cannot write negative bytes: {written}"
         try:
             start, total, n = self._progresses[key]
-            self._progresses[key] = ProgressState(start, total, n + written)
+            if not total_written:
+                total_written = n + written
+            self._progresses[key] = ProgressState(start, total, total_written)
             self._reporter(list(self._progresses.values()))
             if self._progresses[key].n >= total:
                 del self._progresses[key]
@@ -155,6 +157,10 @@ _GLOBAL_UP_TRACKER = Tracker(TqdmReporter("thds.adls uploading"))
 T = ty.TypeVar("T", bound=ty.IO)
 
 
+def get_global_download_tracker() -> Tracker:
+    return _GLOBAL_DN_TRACKER
+
+
 def _proxy_io(io_type: str, stream: T, key: str, total_len: int) -> T:
     assert io_type in ("read", "write"), io_type
 

thds/adls/azcopy/__init__.py

@@ -0,0 +1 @@
+from . import download  # noqa

thds/adls/azcopy/download.py

@@ -0,0 +1,185 @@
+# azcopy is, for whatever reason, quite a bit faster than the Python SDK for downloads.
+# This code allows us to have a thin "try azcopy if it's present, fall back to the Python
+# SDK if it's not" layer.
+#
+# However, azcopy is also quite a bit 'dumber' when it comes to progress reporting (its
+# reported progress totals seem to way underestimate progress and then rubber-band at the
+# very end of the download), so for local users who don't have huge bandwidth, it's likely
+# a better user experience to disable this globally.
+import asyncio
+import json
+import os
+import subprocess
+import typing as ty
+import urllib.parse
+from contextlib import contextmanager
+from pathlib import Path
+
+from azure.storage.filedatalake import DataLakeFileClient
+
+from thds.core import cache, config, log
+
+from .. import _progress, conf, uri
+
+DONT_USE_AZCOPY = config.item("dont_use", default=False, parse=config.tobool)
+
+_AZCOPY_LOGIN_WORKLOAD_IDENTITY = "azcopy login --login-type workload".split()
+_AZCOPY_LOGIN_LOCAL_STATUS = "azcopy login status".split()
+# device login is an interactive process involving a web browser,
+# which is not acceptable for large scale automation.
+# So instead of logging in, we check to see if you _are_ logged in,
+# and if you are, we try using azcopy in the future.
+
+logger = log.getLogger(__name__)
+
+
+class DownloadRequest(ty.NamedTuple):
+    """Use one or the other, but not both, to write the results."""
+
+    writer: ty.IO[bytes]
+    temp_path: Path
+
+
+@cache.locking  # only run this once per process.
+def _good_azcopy_login() -> bool:
+    if DONT_USE_AZCOPY():
+        return False
+
+    try:
+        subprocess.run(_AZCOPY_LOGIN_WORKLOAD_IDENTITY, check=True, capture_output=True)
+        logger.info("Will use azcopy for downloads in this process...")
+        return True
+
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        pass
+    try:
+        subprocess.run(_AZCOPY_LOGIN_LOCAL_STATUS, check=True)
+        logger.info("Will use azcopy for downloads in this process...", dl=None)
+        return True
+    except FileNotFoundError:
+        logger.info("azcopy is not installed or not on your PATH, so we cannot speed up downloads")
+    except subprocess.CalledProcessError as cpe:
+        logger.warning(
+            "You are not logged in with azcopy, so we cannot speed up downloads."
+            f" Run `azcopy login` to fix this. Return code was {cpe.returncode}"
+        )
+    return False
+
+
+def _azcopy_download_command(dl_file_client: DataLakeFileClient, path: Path) -> ty.List[str]:
+    return ["azcopy", "copy", dl_file_client.url, str(path), "--output-type=json"]
+
+
+class AzCopyMessage(ty.TypedDict):
+    TotalBytesEnumerated: str
+    TotalBytesTransferred: str
+
+
+class AzCopyJsonLine(ty.TypedDict):
+    MessageType: str
+    MessageContent: AzCopyMessage
+
+
+def _parse_azcopy_json_output(line: str) -> AzCopyJsonLine:
+    outer_msg = json.loads(line)
+    return AzCopyJsonLine(
+        MessageType=outer_msg["MessageType"],
+        MessageContent=json.loads(outer_msg["MessageContent"]),
+    )
+
+
+@contextmanager
+def _track_azcopy_progress(http_url: str) -> ty.Iterator[ty.Callable[[str], None]]:
+    """Context manager that tracks progress from AzCopy JSON lines. This works for both async and sync impls."""
+    tracker = _progress.get_global_download_tracker()
+    adls_uri = urllib.parse.unquote(str(uri.parse_uri(http_url)))
+
+    def track(line: str):
+        if not line:
+            return
+
+        try:
+            prog = _parse_azcopy_json_output(line)
+            if prog["MessageType"] == "Progress":
+                tracker(adls_uri, total_written=int(prog["MessageContent"]["TotalBytesTransferred"]))
+        except json.JSONDecodeError:
+            pass
+
+    yield track
+
+
+def _restrict_mem() -> dict:
+    return dict(os.environ, AZCOPY_BUFFER_GB="0.3")
+
+
+def sync_fastpath(
+    dl_file_client: DataLakeFileClient,
+    download_request: DownloadRequest,
+) -> None:
+    if _good_azcopy_login():
+        try:
+            # Run the copy
+            process = subprocess.Popen(
+                _azcopy_download_command(dl_file_client, download_request.temp_path),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                env=_restrict_mem(),
+            )
+            assert process.stdout
+            with _track_azcopy_progress(dl_file_client.url) as track:
+                for line in process.stdout:
+                    track(line)
+            return  # success
+
+        except (subprocess.SubprocessError, FileNotFoundError):
+            logger.warning("Falling back to Python SDK for download")
+
+    dl_file_client.download_file(
+        max_concurrency=conf.DOWNLOAD_FILE_MAX_CONCURRENCY(),
+        connection_timeout=conf.CONNECTION_TIMEOUT(),
+    ).readinto(download_request.writer)
+
+
+async def async_fastpath(
+    dl_file_client: DataLakeFileClient,
+    download_request: DownloadRequest,
+) -> None:
+    # technically it would be 'better' to do this login in an async subproces,
+    # but it involves a lot of boilerplate, _and_ we have no nice way to cache
+    # the value, which is going to be computed one per process and never again.
+    # So we'll just block the async loop for a couple of seconds one time...
+    if _good_azcopy_login():
+        try:
+            # Run the copy
+            copy_proc = await asyncio.create_subprocess_exec(
+                *_azcopy_download_command(dl_file_client, download_request.temp_path),
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                env=_restrict_mem(),
+            )
+            assert copy_proc.stdout
+
+            # Feed lines to the tracker asynchronously
+            with _track_azcopy_progress(dl_file_client.url) as track:
+                while True:
+                    line = await copy_proc.stdout.readline()
+                    if not line:  # EOF
+                        break
+                    track(line.decode().strip())
+
+            # Wait for process completion
+            exit_code = await copy_proc.wait()
+            if exit_code != 0:
+                raise subprocess.SubprocessError()
+
+            return  # success
+
+        except (subprocess.SubprocessError, FileNotFoundError):
+            logger.warning("Falling back to Python SDK for download")
+
+    reader = await dl_file_client.download_file(
+        max_concurrency=conf.DOWNLOAD_FILE_MAX_CONCURRENCY(),
+        connection_timeout=conf.CONNECTION_TIMEOUT(),
+    )
+    await reader.readinto(download_request.writer)

thds/adls/download.py

@@ -17,8 +17,8 @@ from thds.core import log, scope, tmp
 from thds.core.hashing import b64
 from thds.core.types import StrOrPath
 
+from . import azcopy
 from ._progress import report_download_progress
-from .conf import CONNECTION_TIMEOUT, DOWNLOAD_FILE_MAX_CONCURRENCY
 from .download_lock import download_lock
 from .errors import MD5MismatchError, translate_azure_error
 from .etag import match_etag
@@ -34,12 +34,14 @@ def _atomic_download_and_move(
     fqn: AdlsFqn,
     dest: StrOrPath,
     properties: ty.Optional[FileProperties] = None,
-) -> ty.Iterator[ty.IO[bytes]]:
+) -> ty.Iterator[azcopy.download.DownloadRequest]:
     with tmp.temppath_same_fs(dest) as dpath:
         with open(dpath, "wb") as f:
             known_size = (properties.size or 0) if properties else 0
             logger.debug("Downloading %s", fqn)
-            yield report_download_progress(f, str(dest), known_size)
+            yield azcopy.download.DownloadRequest(
+                report_download_progress(f, str(fqn), known_size), dpath
+            )
         try:
             os.rename(dpath, dest)  # will succeed even if dest is read-only
         except OSError as oserr:
@@ -139,7 +141,7 @@ class _IoRequest(enum.Enum):
     FILE_PROPERTIES = "file_properties"
 
 
-IoRequest = ty.Union[_IoRequest, ty.IO[bytes]]
+IoRequest = ty.Union[_IoRequest, azcopy.download.DownloadRequest]
 IoResponse = ty.Union[FileProperties, None]
 
 
@@ -332,12 +334,12 @@ def download_or_use_verified(
                     # only fetch these if they haven't already been requested
                     file_properties = dl_file_client.get_file_properties()
                 co_request = co.send(file_properties)
-            else:  # needs file object
-                dl_file_client.download_file(
-                    max_concurrency=DOWNLOAD_FILE_MAX_CONCURRENCY(),
-                    connection_timeout=CONNECTION_TIMEOUT(),
-                ).readinto(co_request)
+            elif isinstance(co_request, azcopy.download.DownloadRequest):
+                # coroutine is requesting download
+                azcopy.download.sync_fastpath(dl_file_client, co_request)
                 co_request = co.send(None)
+            else:
+                raise ValueError(f"Unexpected coroutine request: {co_request}")
     except StopIteration as si:
         if cs := _set_md5_if_missing(file_properties, si.value.md5b64):
             try:
@@ -370,13 +372,13 @@ async def async_download_or_use_verified(
                     # only fetch these if they haven't already been requested
                     file_properties = await dl_file_client.get_file_properties()
                 co_request = co.send(file_properties)
-            else:  # needs file object
-                reader = await dl_file_client.download_file(
-                    max_concurrency=DOWNLOAD_FILE_MAX_CONCURRENCY(),
-                    connection_timeout=CONNECTION_TIMEOUT(),
-                )
-                await reader.readinto(co_request)
+            elif isinstance(co_request, azcopy.download.DownloadRequest):
+                # coroutine is requesting download
+                await azcopy.download.async_fastpath(dl_file_client, co_request)
                 co_request = co.send(None)
+            else:
+                raise ValueError(f"Unexpected coroutine request: {co_request}")
+
     except StopIteration as si:
         if cs := _set_md5_if_missing(file_properties, si.value.md5b64):
             try:

thds/adls/meta.json

@@ -1,8 +1,8 @@
 {
-  "git_commit": "923bef905389724895eddd2b2d00e87a533d4c52",
+  "git_commit": "f4211224c05a5de6426bfd3efe4d62e3c0a7d64d",
   "git_branch": "main",
   "git_is_clean": true,
-  "pyproject_version": "3.1.20250226213817",
+  "pyproject_version": "3.1.20250227181712",
   "thds_user": "runner",
   "misc": {}
 }

thds/adls/tools/download.py

@@ -3,26 +3,45 @@ from pathlib import Path
 
 from thds.adls.cached_up_down import download_directory, download_to_cache
 from thds.adls.file_properties import get_file_properties, is_directory
+from thds.adls.impl import ADLSFileSystem
 from thds.adls.uri import resolve_uri
 from thds.core.link import link
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("uri", type=resolve_uri, help="A fully qualified path to an ADLS location")
+    parser.add_argument(
+        "adls_fqn",
+        type=resolve_uri,
+        help="A fully qualified path to an ADLS location. Accepts adls://, https:// and abfss:// URIs.",
+    )
     parser.add_argument(
         "--copy-to",
         "-c",
         type=Path,
         help="This will create a link to the cached download at the specified location",
     )
+    parser.add_argument(
+        "--use-async",
+        action="store_true",
+        help="Only useful if you want to exercise the async code; does not make things faster or better.",
+    )
 
     args = parser.parse_args()
 
-    if is_directory(get_file_properties(args.uri)):
-        cache_path = download_directory(args.uri)
+    is_dir = is_directory(get_file_properties(args.adls_fqn))
+
+    if args.use_async:
+        fs = ADLSFileSystem(args.adls_fqn.sa, args.adls_fqn.container)
+        if is_dir:
+            cache_path = fs.fetch_directory(args.adls_fqn.path)[0]
+        else:
+            cache_path = fs.fetch_file(args.adls_fqn.path)
     else:
-        cache_path = download_to_cache(args.uri)
+        if is_dir:
+            cache_path = download_directory(args.adls_fqn)
+        else:
+            cache_path = download_to_cache(args.adls_fqn)
 
     if args.copy_to:
         link(cache_path, args.copy_to)

{thds_adls-3.1.20250226213817.dist-info → thds_adls-3.1.20250227181712.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: thds.adls
-Version: 3.1.20250226213817
+Version: 3.1.20250227181712
 Summary: ADLS tools
 Author: Trilliant Health
 Description-Content-Type: text/markdown

{thds_adls-3.1.20250226213817.dist-info → thds_adls-3.1.20250227181712.dist-info}/RECORD

@@ -1,5 +1,5 @@
 thds/adls/__init__.py,sha256=er14MoCC9PlJMxWVS4G1hAeMJItaJj4EAsrTZlvlb0M,797
-thds/adls/_progress.py,sha256=SVOiVRvfvOfUjp-y-UUgA6n0FhAnpKesxltaQ16kmOk,6394
+thds/adls/_progress.py,sha256=ZzCHn_G7nHakioNFxdvoJZRr-jN6ymsp5JXf-iReROM,6580
 thds/adls/_upload.py,sha256=q6Sk0CRnNcAjUOUPiBj4CfO4tJD196SQY0lT25CTSE4,4364
 thds/adls/abfss.py,sha256=ZRJOLjDuXmS4bIbQAQpQxWWWeu74N9NKEKCNfXQek80,726
 thds/adls/cached_up_down.py,sha256=qTMrsHHrDq3YZ0gaLMFIPnGNraz0ctG9yoFd3HCVit4,1896
@@ -7,7 +7,7 @@ thds/adls/conf.py,sha256=q1SPrgb46NpobVzwt_Oyv71-BvsIbZLq9nRWS3LZjz0,1990
 thds/adls/copy.py,sha256=jUWbGvTpb4B3yRGS0nhGSbDzqRPzUqYgH0z1lFRJB3k,6365
 thds/adls/dbfs.py,sha256=pPAjbIZRKJsaXKQljDMUgqS_zy1yKeEZHGMueXbuv3g,2219
 thds/adls/defaults.py,sha256=VxyaGz1gCcz2rGTR9acNRybbxGOYhTldvD_-SdGLBh0,652
-thds/adls/download.py,sha256=liVfnW4PJRjvBgPKsMqoo_F-9NT8syKsFMDaVSL8A4E,16345
+thds/adls/download.py,sha256=Vfd9DEHUmSKPKpKxf9WjkqAzZFZwIgjo4fgM5aIvW0g,16455
 thds/adls/download_lock.py,sha256=_JZj-kjCUfHk9FvrmEuYpJYknmbam5eReFhGNDgzdLQ,2520
 thds/adls/errors.py,sha256=B_rMsQvQnNmP_sf-x8kmGsv2vIeOh4G9kVbdNVyk350,1469
 thds/adls/etag.py,sha256=ct7jpHhNFcKzbekn5rZ3m6DhjK48A7qOZGwDiHkc-pc,242
@@ -16,7 +16,7 @@ thds/adls/fqn.py,sha256=0zHmHhBWN7GEfKRB3fBC1NVhaiIHHifBdCRanyT01X8,5822
 thds/adls/global_client.py,sha256=f4VJw5y_Yh__8gQUcdSYTh1aU6iEPlauMchVirSAwDQ,3716
 thds/adls/impl.py,sha256=x1nSqc8W4NeuX8-JGOp2MRkK8ff6GnelTWedGxPs-qY,42494
 thds/adls/md5.py,sha256=qOX4_7WUj1QkbH_IwREcQNHvvZccOj-HpHZBfsKn1gY,1846
-thds/adls/meta.json,sha256=lX2VX33yybCNVd5IFfQ04M5BMyXhW_eDcmT6V0UEFfQ,195
+thds/adls/meta.json,sha256=MhZqsLGrHq1cqZB6OydUf5yDbfo_GySrRRtjdpImBQc,195
 thds/adls/named_roots.py,sha256=7SLbAoQQpV_mrFZaUPjYoS-F9dxQxN5Hg4M3YPirF_w,751
 thds/adls/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 thds/adls/ro_cache.py,sha256=F0uXol0t95mcRuBukNg3A7wt7XXQxpD5Sy09d9sl8f0,4825
@@ -25,15 +25,17 @@ thds/adls/shared_credential.py,sha256=-x42aXoIM001KW59oS8PpuXQd4-F2vg-1gB6OMHlpk
 thds/adls/source.py,sha256=1JYliDqafqpLewOQ2XMKp4T2WD7b5Rl7XYViyLYtzX8,2437
 thds/adls/source_tree.py,sha256=CiyCwgItF186aOoYiICkWt2oMLGMD6v526HRi8XWHtM,2219
 thds/adls/uri.py,sha256=pDH956p_VEHnjLLUnjWY6sGgRqksp9gdpc9KOW4gEP0,1205
+thds/adls/azcopy/__init__.py,sha256=nTNbgz2GcEiGeswYbAgy4oPhivnzl_5crF3HqCdWWiw,31
+thds/adls/azcopy/download.py,sha256=_I0smR2f9pu0xIiQC9C97Xb-GHMGwn5dqbEhz41H-qo,6548
 thds/adls/resource/__init__.py,sha256=IZ7_aRf1b3jEp7wXOxqHop0gV2gUcf9SOLeEEjIWlCU,1669
 thds/adls/resource/core.py,sha256=BVM91xsZ_B_CoGTc9DDD3FnGy8g6X-9eFpa86ZCzuZI,2717
 thds/adls/resource/file_pointers.py,sha256=PLru_3lwut_ZvrX5Keu-wJkPOt5o7UGf-OOT4ixaXME,2049
 thds/adls/resource/up_down.py,sha256=3uNlTvm2gVhSyYdQTBwsGecOgwtINQfINckR-awwV0Y,9907
-thds/adls/tools/download.py,sha256=ZZ8t1g9MRRy-aENGjr10E4aAqyVkKq2OqpyH70pCD_8,965
+thds/adls/tools/download.py,sha256=vvBO8lSDl9oPugv75qpCkoemT9pOM9BV6yeExlkyG08,1594
 thds/adls/tools/ls.py,sha256=OgEaIfTK359twlZIj-A0AW_nv81Z6zi0b9Tw6OJJfWA,1083
 thds/adls/tools/upload.py,sha256=eMk4pdug1aCMPDDWpIE3Zoq77i5APp9Uuh-sVCCDNJE,493
-thds_adls-3.1.20250226213817.dist-info/METADATA,sha256=g0fHOJuklbyOLusIxUSeAa2FIRcLAsEwxCmkzF-ChoY,397
-thds_adls-3.1.20250226213817.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-thds_adls-3.1.20250226213817.dist-info/entry_points.txt,sha256=uTqreT1AIwqJboMfLv5w6sviM8mNbAkln765gIjzoA4,152
-thds_adls-3.1.20250226213817.dist-info/top_level.txt,sha256=LTZaE5SkWJwv9bwOlMbIhiS-JWQEEIcjVYnJrt-CriY,5
-thds_adls-3.1.20250226213817.dist-info/RECORD,,
+thds_adls-3.1.20250227181712.dist-info/METADATA,sha256=O8W1_KMHFvwQZ--B8xPvFDMLL774NvfbhpxvEKkACWE,397
+thds_adls-3.1.20250227181712.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+thds_adls-3.1.20250227181712.dist-info/entry_points.txt,sha256=uTqreT1AIwqJboMfLv5w6sviM8mNbAkln765gIjzoA4,152
+thds_adls-3.1.20250227181712.dist-info/top_level.txt,sha256=LTZaE5SkWJwv9bwOlMbIhiS-JWQEEIcjVYnJrt-CriY,5
+thds_adls-3.1.20250227181712.dist-info/RECORD,,