From 380f920f5ce389a52ec8e5efa1c56852a35bdb54 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Thu, 22 Oct 2026 19:29:43 +0000 Subject: [PATCH 01/15] Add ADLS Gen2 connector with bounded transfers and staged exports --- .github/workflows/linting.yml | 1 + .github/workflows/testing.yml | 5 +- docs/README.md | 2 +- docs/connector-test-matrix.md | 4 +- packages/kumo-adls/LICENSE | 204 ++++++ packages/kumo-adls/README.md | 155 ++++ packages/kumo-adls/docs/performance.md | 137 ++++ packages/kumo-adls/pyproject.toml | 52 ++ packages/kumo-adls/src/kumo_adls/__init__.py | 16 + packages/kumo-adls/src/kumo_adls/backend.py | 163 +++++ packages/kumo-adls/src/kumo_adls/errors.py | 61 ++ .../kumo-adls/src/kumo_adls/range_cache.py | 62 ++ .../kumo-adls/src/kumo_adls/sdk_filesystem.py | 402 +++++++++++ packages/kumo-adls/src/kumo_adls/session.py | 130 ++++ packages/kumo-adls/src/kumo_adls/sink.py | 35 + .../kumo-adls/src/kumo_adls/source_access.py | 16 + packages/kumo-adls/src/kumo_adls/transfer.py | 59 ++ packages/kumo-adls/tests/conftest.py | 43 ++ packages/kumo-adls/tests/fakes.py | 132 ++++ .../tests/live/test_live_filesystem.py | 90 +++ packages/kumo-adls/tests/test_backend.py | 306 ++++++++ packages/kumo-adls/tests/test_filesystem.py | 205 ++++++ .../kumo-adls/tests/test_identity_listing.py | 306 ++++++++ .../kumo-adls/tests/test_package_boundary.py | 23 + packages/kumo-adls/tests/test_range_cache.py | 199 ++++++ packages/kumo-adls/tests/test_sdk_transfer.py | 432 ++++++++++++ packages/kumo-adls/tests/test_session.py | 141 ++++ packages/kumo-adls/tests/test_transfer.py | 534 ++++++++++++++ packages/kumo-api/kumoapi/adls.py | 11 + .../src/kumo_core/object_store_sink.py | 50 +- .../src/kumo_core/sources/object_store.py | 7 +- .../tests/sources/package_boundary.py | 3 +- packages/kumo-engine/README.md | 2 + packages/kumo-engine/kumo/engine/__init__.py | 4 + packages/kumo-engine/kumo/engine/_portable.py | 8 + .../kumo/engine/connector/__init__.py | 4 + .../kumo-engine/kumo/engine/connector/adls.py | 81 +++ .../kumo/engine/connector/identity.py | 17 + packages/kumo-engine/pyproject.toml | 4 + .../kumo-engine/test/test_adls_connector.py | 108 +++ .../kumo-engine/test/test_optional_imports.py | 3 + pyproject.toml | 2 + scripts/README.md | 8 + .../benchmarks/object_store/compare_live.py | 575 +++++++++++++++ .../object_store/compare_streams.py | 663 ++++++++++++++++++ .../object_store/matched-transfer.md | 108 +++ scripts/check_python_license_metadata.py | 2 + scripts/ci/check_dependency_boundaries.sh | 3 +- scripts/gen_third_party.py | 1 + scripts/test_public_connectors.py | 2 +- test/scripts/test_engine_wheel.py | 17 +- tests/public/Dockerfile.dockerignore | 4 + tests/public/pytest.ini | 1 + uv.lock | 41 +- 54 files changed, 5616 insertions(+), 28 deletions(-) create mode 100644 packages/kumo-adls/LICENSE create mode 100644 packages/kumo-adls/README.md create mode 100644 packages/kumo-adls/docs/performance.md create mode 100644 packages/kumo-adls/pyproject.toml create mode 100644 packages/kumo-adls/src/kumo_adls/__init__.py create mode 100644 packages/kumo-adls/src/kumo_adls/backend.py create mode 100644 packages/kumo-adls/src/kumo_adls/errors.py create mode 100644 packages/kumo-adls/src/kumo_adls/range_cache.py create mode 100644 packages/kumo-adls/src/kumo_adls/sdk_filesystem.py create mode 100644 packages/kumo-adls/src/kumo_adls/session.py create mode 100644 packages/kumo-adls/src/kumo_adls/sink.py create mode 100644 packages/kumo-adls/src/kumo_adls/source_access.py create mode 100644 packages/kumo-adls/src/kumo_adls/transfer.py create mode 100644 packages/kumo-adls/tests/conftest.py create mode 100644 packages/kumo-adls/tests/fakes.py create mode 100644 packages/kumo-adls/tests/live/test_live_filesystem.py create mode 100644 packages/kumo-adls/tests/test_backend.py create mode 100644 packages/kumo-adls/tests/test_filesystem.py create mode 100644 packages/kumo-adls/tests/test_identity_listing.py create mode 100644 packages/kumo-adls/tests/test_package_boundary.py create mode 100644 packages/kumo-adls/tests/test_range_cache.py create mode 100644 packages/kumo-adls/tests/test_sdk_transfer.py create mode 100644 packages/kumo-adls/tests/test_session.py create mode 100644 packages/kumo-adls/tests/test_transfer.py create mode 100644 packages/kumo-engine/kumo/engine/connector/adls.py create mode 100644 packages/kumo-engine/test/test_adls_connector.py create mode 100644 scripts/benchmarks/object_store/compare_live.py create mode 100644 scripts/benchmarks/object_store/compare_streams.py create mode 100644 scripts/benchmarks/object_store/matched-transfer.md diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml index c96961cc1..33abb304d 100644 --- a/.github/workflows/linting.yml +++ b/.github/workflows/linting.yml @@ -149,6 +149,7 @@ jobs: packages/kumo-core packages/kumo-databricks packages/kumo-gcs + packages/kumo-adls packages/kumo-pql packages/kumo-s3 packages/kumo-sdk diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index cd20b2b7e..c628fe5ca 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -543,7 +543,7 @@ jobs: run: | status=0 for package in kumo-pql kumo-core kumo-snowflake kumo-bigquery \ - kumo-databricks kumo-s3 kumo-gcs; do + kumo-databricks kumo-s3 kumo-gcs kumo-adls; do echo "::group::pytest packages/$package" if [ "$package" = kumo-core ]; then (cd "packages/$package" && python -m pytest -rs --cov=kumo_core --cov-report=term-missing --cov-report=xml) || status=1 @@ -574,6 +574,9 @@ jobs: # namespace package the root app owns. Exercise the wheel in isolation # before/after install so a stale or partial install can't leave # `kumo.engine` resolvable when it shouldn't be. + - name: Install ADLS provider for engine integration tests + run: uv pip install --no-deps -e packages/kumo-adls + - name: Build and test the isolated kumo-engine wheel run: | python -c "import importlib.util; assert importlib.util.find_spec('kumo.engine') is None" diff --git a/docs/README.md b/docs/README.md index 0c54622cb..e628c202b 100644 --- a/docs/README.md +++ b/docs/README.md @@ -37,4 +37,4 @@ guides for their interfaces and development workflows. | Query language | [kumo-pql](../packages/kumo-pql/README.md) and [existing instructions](../packages/kumo-pql/AGENTS.md) | | UI component library | [catalog](../packages/kumo-ui-catalog/README.md) and [existing component guide](../packages/kumo-ui-catalog/CLAUDE.md) | | Native graph engine | [DiskGraph](../packages/diskgraph/README.md); maintained directly in this repository | -| Source providers | [S3](../packages/kumo-s3/README.md), [GCS](../packages/kumo-gcs/README.md), [BigQuery](../packages/kumo-bigquery/README.md), [Databricks](../packages/kumo-databricks/README.md), [Snowflake](../packages/kumo-snowflake/README.md) | +| Source providers | [S3](../packages/kumo-s3/README.md), [GCS](../packages/kumo-gcs/README.md), [ADLS](../packages/kumo-adls/README.md), [BigQuery](../packages/kumo-bigquery/README.md), [Databricks](../packages/kumo-databricks/README.md), [Snowflake](../packages/kumo-snowflake/README.md) | diff --git a/docs/connector-test-matrix.md b/docs/connector-test-matrix.md index 2e6404e04..f0cc3d673 100644 --- a/docs/connector-test-matrix.md +++ b/docs/connector-test-matrix.md @@ -23,7 +23,7 @@ during a transfer. | `parquet` | [`packages/kumo-core/tests/sources/parquet`](../packages/kumo-core/tests/sources/parquet) | Covers reads, writes, columns, row counts, versions, partitions, cleanup after failure, and limits on file paths. | | `s3` | [`packages/kumo-s3/tests`](../packages/kumo-s3/tests) | Covers credentials, errors, sessions, regions, service addresses, connection probes, imports, and downloads that preserve files on failure. | | `gcs` | [`packages/kumo-gcs/tests`](../packages/kumo-gcs/tests) | Covers downloads, object versions, credential renewal, concurrent transfers, cleanup, billing configuration, and imports. | -| `adls` | [`test/io_utils/fsutil/adls`](../test/io_utils/fsutil/adls) | Covers file listing, limits on Parquet preview reads, missing paths, read-only mode, and storage addresses. | +| `adls` | [`packages/kumo-adls/tests`](../packages/kumo-adls/tests), [`test/io_utils/fsutil/adls`](../test/io_utils/fsutil/adls) | Provider identity binding, CSV/Parquet ingestion, rooted output, version checks, URI fencing, and filesystem preview behavior. | | `bigquery` | [`packages/kumo-bigquery/tests`](../packages/kumo-bigquery/tests) | Covers queries, streamed and staged reads, versions, row counts, errors, cleanup, and rules that enable cloud tests. | | `databricks` | [`packages/kumo-databricks/tests`](../packages/kumo-databricks/tests) | Covers catalogs, versions, queries, reads, write helpers, and limits on cloud file downloads called CloudFetch. | | `snowflake` | [`packages/kumo-snowflake/tests`](../packages/kumo-snowflake/tests) | Covers catalogs, types, queries, staged reads, concurrent downloads, retries, cleanup, write helpers, and hiding sensitive values. | @@ -35,7 +35,7 @@ need separate tests. The selected tests also have these limits: - Snowflake tests use fake download bytes without decoding them as Parquet. It includes local output-writer and Snowflake Container Services (SPCS) tests. Native ingestion runs in `shared-sql`. -- ADLS tests the code that presents cloud storage as a filesystem. It does not test connector identity or writes. +- ADLS provider tests use offline Azure responses for identities, reads, and writes. Filesystem tests cover preview and addressing; live credentials and Azure permissions remain separate checks. - Parquet tests do not prove detection of corrupt Parquet data pages or every timestamp conversion. - BigQuery tests use fake clients to test rules that enable real cloud tests. Pytest loads the cloud test file, but does not run its cloud tests. - The BigQuery, Snowflake, and Databricks groups test their live-test gates with fake credentials and blocked provider imports. diff --git a/packages/kumo-adls/LICENSE b/packages/kumo-adls/LICENSE new file mode 100644 index 000000000..d1593aaa1 --- /dev/null +++ b/packages/kumo-adls/LICENSE @@ -0,0 +1,204 @@ +Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/packages/kumo-adls/README.md b/packages/kumo-adls/README.md new file mode 100644 index 000000000..391fa4ae8 --- /dev/null +++ b/packages/kumo-adls/README.md @@ -0,0 +1,155 @@ +# kumo-adls + +Azure Data Lake Storage Gen2 source and sink provider for `kumo-engine`. +Install through `kumo-engine[adls]`. The package registers source, sink, and +identity plugins using the same entry-point groups as other providers. + +```python +from kumo.engine import ADLSConnector, Table + +lake = ADLSConnector( + root_dir="abfss://curated@companylake.dfs.core.windows.net/orders/", +) +orders = Table.from_source_table(lake["orders.parquet"], primary_key="id") +``` + +A rooted connector discovers and reads CSV/Parquet tables beneath its +filesystem/prefix and writes predictions, embeddings, and model bundles there. +A rootless connector reads exact canonical `abfss://` URIs; discovery and output +are unavailable. SQL and append output are unsupported. + +Without an identity, requests use Azure's ambient credential chain. An explicit +`ADLSIdentity(tenant_id=..., client_id=..., cloud=...)` uses workload federation +with `AZURE_FEDERATED_TOKEN_FILE`. Identity secrets, storage keys, and SAS tokens +are not connector arguments. The selected identity cloud must match the URI's +public, US Government, or China endpoint. See the shared +[ADLS storage contract](../../docs/cloud-storage/adls.md). + +The provider adapts Azure SDK operations to shared Arrow source and sink +machinery in `kumo-core`. ETags guard range/full-file reads and source versions; +account-qualified paths distinguish accounts and clouds. Azure Identity refreshes +tokens, and identity cleanup closes cached clients. Sessions are cached for the +process lifetime, without automatic eviction; identity cleanup or +`kumo_adls.session.close()` releases them. Construction runs under a global lock +to prevent duplicate clients, so concurrent first use of different accounts can +serialize. Arrow/fsspec wrappers are lightweight and created per access. + +Identity probing checks token acquisition only. It does not verify storage +existence, RBAC/ACL permissions, or connectivity; list or read the configured +source to verify those. + +Bulk downloads use one conditional SDK stream per file, with 16 MiB initial +and subsequent responses and 16 MiB HTTP transport blocks. Files download +concurrently; each file holds one transfer slot and uses no per-file executor +or queued chunk futures. The SDK writes each response to the local temporary +file before requesting the next chunk. The initial response remains buffered +for the stream's lifetime; this is not zero-copy network-to-disk streaming. +Prepared uploads retain up to four 4 MiB chunks. Large projected Arrow reads +can use up to four SDK workers, bounded by the shared transfer slots. A 64 MiB +range uses up to three workers; ranges above 64 MiB can use four. The initial +response plus four 16 MiB chunks can account for about 80 MiB before transport +buffers, SDK copies, and the complete requested range returned by `readall()`. +Projected reads are not covered by the 64 MiB bulk-copy reservation. + +Transfers across files and accounts share a CPU-derived concurrency ceiling +(at most 128, lowered by `KUMO_SOURCE_COPY_WORKERS`). The HTTP connection pool +matches that ceiling. Metadata bypasses transfer slots and shares the HTTP +transport. Uploads reserve 64 MiB per active file through the shared memory +ledger; source bulk copies use the core copy reservation. These reservations +are accounting estimates, not caps on total process RSS. Larger responses +reduce requests and fragment assembly while increasing memory per active file. +Plan roughly 48–64 MiB of payload buffers per active bulk download, with +additional headroom for SDK copies and other allocations. At 128 active files, +that estimate alone is 6–8 GiB; the reservation does not guarantee that this +fits a host's memory limit. Actual memory depends on allocation lifetimes. + +For memory-constrained hosts, set `KUMO_ADLS_TRANSFER_WORKERS=4` before +importing the provider to limit ADLS transfer requests and its HTTP pool +without lowering other providers' copy limits. The value must be a positive +integer and cannot raise the shared copy/CPU ceiling. The limit is process-wide +across ADLS accounts and readers; it is read once at import. Even with one slot, +files use the same serial SDK stream. Arrow buffers, cached bytes, SDK copies, +loaded libraries, and Python overhead still consume memory. Reduced ceilings +are covered by correctness tests; a fixed-RAM host requires separate validation. +An invalid value fails at the provider's first import. Set these variables in +the launching environment and restart an existing process after changing them. +Catchable allocation failures inside ADLS SDK operations retain their original +exception and traceback, with a note suggesting these controls if the failure +occurred during a transfer. They may not help with listing or session setup. +An operating-system OOM kill cannot provide a Python error message. +Reducing the shared copy ceiling +also limits waiting file workers and their buffers; the ADLS-only limit does not. + +Projected Arrow reads disable fsspec read-ahead, avoiding unrelated column data. +Successful small ranges can reuse immutable bytes under the opened file's raw +ETag. The process-wide LRU holds at most 64 MiB of payload/1,024 entries, admits +ranges up to 128 KiB, and isolates filesystem/principal namespaces. Opening a +file still performs fresh properties checks; metadata is never cached here. +Existing opened streams may reuse their pinned version, while ingest's final +identity checks still reject source changes. Filesystem cleanup releases its +cached ranges; Python container/key overhead is additional. Each source access +creates a new filesystem namespace. This cache avoids repeated footer reads +within one reader, not across separate reads or connector operations. The +64 MiB limit bounds aggregate retained payload across concurrent readers; it +is not a persistent connector cache or a preallocated buffer. +Recursive discovery uses Azure's paginated recursive listing. Version ETag +checks use fresh bounded listings when complete ETags are available, falling +back to concurrent HEADs for missing metadata, narrow permissions, and exact +single files. ADLS uses the shared object-store `prepare`: discover files once, +stamp the version from that same selected list, and read those files. A supplied +version is stamped without refreshing it, matching S3/GCS. Shared per-file +identity guards check selected objects before and after reading; files added +after discovery belong to a subsequent read. Identity checks may still perform +listings or HEADs, so one discovery does not mean one metadata request. +ETag quote and directory-path normalization can invalidate existing version +hashes once, requiring restaging. +Prepared export files and model files use seekable parallel uploads; Arrow +writers keep the buffered streaming path. Failed prepared uploads delete only +the object matching their creation ETag, preserving another writer's output. +Prepared uploads create the destination before uploading; concurrent readers +can observe an empty or incomplete object. Publishing an individual file is +not atomic, so consumers must wait for export completion. +Streaming writes are not transactional: Arrow closes its wrapper without +forwarding serialization exceptions to the fsspec context manager, so failed +serialization can leave partial output. Direct fsspec contexts can discard +on caller errors; upload failures also trigger best-effort cleanup. + +Downloaded files remain owner-private, matching the S3 staging path. Local +engine ingestion uses the same OS user; sharing scratch files across users +requires explicitly managed permissions. Temporary files are closed before +promotion with `os.replace`, including on failure cleanup paths. This avoids +requiring POSIX open-file replacement; it does not establish full Windows +engine support. + +These settings remove known serial paths and request amplification; they are +not evidence of optimal live Azure throughput. Validate full logical tables, +client placement, resource limits, retries, and service throttling before making +production performance claims. +See the [performance methodology](docs/performance.md) for complete-table +validation, native-host comparisons, per-reader cache scope, and resource +tradeoffs. The [matched-request stream method](../../scripts/benchmarks/object_store/matched-transfer.md) +separately controls HEADs, range GETs, bytes, and worker limits for ADLS and +native S3; it uses fixed historical diagnostic geometry rather than the +production 16 MiB download configuration or full connector ingestion. +For live comparisons on Linux using your own storage locations, the +[bulk-reader transfer harness](../../scripts/benchmarks/object_store/compare_live.py) accepts complete-table +manifests and verifies staged files after each production bulk-reader trial. + +Run offline provider and existing filesystem tests from the repository root: + +```bash +python3 scripts/test_public.py --group adls +``` + +Live access, token renewal, RBAC/ACL propagation, and sovereign endpoints require +validation against an explicitly configured Azure environment. Offline tests do +not establish live cloud support. + +Run the read-only live filesystem checks with an existing HNS directory: + +```bash +KUMO_ADLS_LIVE_ROOT='abfss://filesystem@account.dfs.core.windows.net/existing-directory/' \ +python -m pytest packages/kumo-adls/tests/live --confcutdir=packages/kumo-adls/tests +``` + +The live checks use ambient credentials and do not create or delete cloud paths. diff --git a/packages/kumo-adls/docs/performance.md b/packages/kumo-adls/docs/performance.md new file mode 100644 index 000000000..4f7b5cd2c --- /dev/null +++ b/packages/kumo-adls/docs/performance.md @@ -0,0 +1,137 @@ +# Object-store transfer performance methodology + +Measure complete, verified workloads before attributing performance to a +connector. This guide describes engineering tradeoffs and reproducible controls; +it contains no deployment-specific results or resource identifiers. + +## Implementation and tradeoffs + +- Bound queued tasks as well as active threads. A worker limit alone can still + retain one future per chunk of a large object. Lazy submission bounds that + overhead independently of object size. +- Budget concurrency across files. Larger chunks or fewer workers per file + do not necessarily reduce aggregate memory when the request ceiling stays + unchanged. Capture CPU and RSS alongside elapsed time. +- Avoid redundant buffering when the SDK already holds a response. Bulk ADLS + transfers use one serial SDK downloader per file, with files copied in + parallel. Matching response and transport-block sizes reduces fragment + assembly. This is not zero-copy or direct network-to-disk streaming. +- Cache immutable bytes under a version identity with an explicit lifetime. + The small-range cache avoids repeated footer reads within one reader; it + does not cache across separate reader/source-access instances or replace fresh + metadata checks. +- Separate experiments. Do not pool measurements from different revisions, + layouts, client placements, or request policies into a cumulative speedup. + +The [filesystem adapter](../src/kumo_adls/sdk_filesystem.py), +[session setup](../src/kumo_adls/session.py), and +[transfer limits](../src/kumo_adls/transfer.py) own these behaviors. + +## Resource controls + +Bulk ADLS downloads hold one transfer slot per file and use serial conditional +16 MiB responses, with no per-file executor or chunk-future backlog. The SDK +retains its initial response while processing later chunks. Transport blocks +are also 16 MiB; prepared uploads retain up to four 4 MiB chunks. The SDK still +buffers responses; a 64 MiB shared copy reservation is not a cap on total RSS. +Arrow buffers, cached bytes, metadata, loaded libraries, and Python overhead +consume additional memory. +Budget roughly 48–64 MiB of bulk payload buffers per active file, plus SDK-copy +headroom; 128 active files can therefore imply 6–8 GiB of payload alone. These +are planning estimates, not measured peaks or strict allocation bounds. + +Projected reads use `readall()` and materialize the complete requested range. +A 64 MiB range uses up to three SDK workers after its initial 16 MiB response; +ranges above 64 MiB can use four. The initial response plus four 16 MiB chunks +can account for about 80 MiB before transport blocks, SDK copies and the output +range itself. Parallel SDK range reads can also queue futures for the requested +range. Neither the 80 MiB payload estimate nor the 64 MiB bulk-copy reservation +is a memory bound for projected reads. + +`KUMO_SOURCE_COPY_WORKERS` lowers the shared copy ceiling. Set the positive +startup value `KUMO_ADLS_TRANSFER_WORKERS` before importing the provider to +lower only ADLS transfer slots and HTTP pooling, capped by the shared CPU/copy +ceiling. Bulk downloads use the same serial SDK path at every ceiling. +Metadata requests bypass transfer slots but share the HTTP pool. These controls bound requests, +not total RSS; a reduced ceiling alone is not a fixed-RAM-quota benchmark. +Both settings belong in the launching environment; restart a process after +changing the ADLS limit, which is fixed at import. Invalid ADLS values fail at +first provider import. Catchable allocation failures inside ADLS SDK operations +retain their original exception/traceback and add a note suggesting both controls +if the failure occurred during transfer; they may not help with listing or +session setup. An OS OOM kill cannot show this guidance. For lower memory across whole reads, +reduce shared file concurrency as well as ADLS transfer slots. + +The process-wide small-range LRU caps payload at 64 MiB and 1,024 entries, with +Python key/container overhead additional. It admits complete successful ranges +up to 128 KiB, keyed by filesystem namespace, full path, raw ETag, and exact +range. Each source access creates a new namespace and cleanup releases it. +Opens and final source checks remain fresh. Prepared exports do not use this +reader cache. + +## Live comparison recipe + +Use identical complete tables and physical files on both stores. Prepare +trusted manifests containing every relative path, size, SHA256, footer row +count, and pinned provider ETag. Fixture copying and equivalence validation +must finish before timing. Do not substitute unrelated objects of equal +compressed size for equivalent work. + +The [live bulk-reader harness](../../../scripts/benchmarks/object_store/compare_live.py) only reads cloud +objects and writes local results/scratch. Use authenticated, existing sources +and a Linux environment with provider dependencies and Arrow S3 support. +For native deployment comparisons, run S3 on the AWS client and ADLS on the +Azure client in the respective storage regions. Record actual CPU/RAM, OS, +network, disk and storage-class differences; matching vCPU counts does not +make the hosts identical. + +Example single-provider phases, executed on their respective hosts: + +```bash +# AWS host +python scripts/benchmarks/object_store/compare_live.py --providers s3 \ + --manifest /data/manifests/table.json --placement /data/aws-placement.json \ + --output /data/transfer-scratch --results /data/s3-results \ + --workers 32 --trials 3 --verification-workers 8 --s3-region us-west-2 + +# Azure host, using its configured ambient identity +python scripts/benchmarks/object_store/compare_live.py --providers adls \ + --manifest /data/manifests/table.json --placement /data/azure-placement.json \ + --output /data/transfer-scratch --results /data/adls-results \ + --workers 32 --trials 3 --verification-workers 8 +``` + +Results directories must be new and scratch must fit the complete table plus +headroom. Run one timed workload per host. Preserve every individual trial, +rejection, manifest/harness/imported-source hash, dependency version and +placement record. Cloud-side cache state remains unknown in a fresh process. + +## Timing, acceptance, and limits + +Setup, authentication preflight, provenance and initial source identity checks +occur outside timing. The preparation timer covers the production reader's +`prepare(transform=False)` metadata guards and raw copying. A nested download +phase includes fetch's own metadata/checks/retries; it is not pure wire time, +and the remainder is not pure metadata latency. This workload does not measure +Graph construction, typed decoding, projection, conversion, training, or export. + +Capture process CPU and RSS/HWM before untimed output verification. Loaded +provider stacks and setup allocations can affect the high-water mark; kernel +page cache and remote service memory are excluded. Local buffered writes, +close and atomic promotion do not establish durable-disk throughput without +an explicit durability requirement. + +Require unchanged exact remote membership/identities and every output file's +SHA256, size and footer rows before accepting a trial. Verification is bounded +and untimed; its scheduling/cache advice must remain in provenance. Request, +retry, token-refresh counts and effective worker overlap are uninstrumented +in the live harness, so do not infer them from configured ceilings. + +The [matched-request harness](../../../scripts/benchmarks/object_store/matched-transfer.md) separately isolates stream +client stacks under loopback HTTP. It bypasses production bulk paths and does +not replace native cloud measurements. Its fixed 8 MiB windows, 4 MiB SDK +chunks and 1 MiB transport blocks model different geometry from production +ADLS downloads (16 MiB for all three). Historical raw artifacts are audit +provenance, not current deployment results. Projections/conversions, CSV, +fixed-RAM limits, sustained refresh/throttling, DataFrame writes, model bundles, +and exports require their own representative evidence. diff --git a/packages/kumo-adls/pyproject.toml b/packages/kumo-adls/pyproject.toml new file mode 100644 index 000000000..726d67e17 --- /dev/null +++ b/packages/kumo-adls/pyproject.toml @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=77", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "kumo-adls" +version = "0.1.0.dev1" +description = "Azure Data Lake Storage Gen2 source and sink integration for Kumo" +readme = "README.md" +requires-python = ">=3.11" +license = "Apache-2.0" +license-files = ["LICENSE"] +dependencies = [ + "kumo-core==0.1.0.dev1", + "kumo-api>=3.0.0", + "pyarrow>=14", + "fsspec", + "requests>=2.31,<3", + "azure-core>=1.31,<2", + "azure-identity>=1.25,<2", + "azure-storage-file-datalake>=12.22,<13", +] + +[project.optional-dependencies] +test = ["pytest>=7.4"] + +[project.entry-points."kumoai_studio.sources"] +adls = "kumo_adls.backend:ADLSSourceBackend" + +[project.entry-points."kumoai_studio.sinks"] +adls = "kumo_adls.sink:ADLSSinkBackend" + +[project.entry-points."kumoai_studio.identities"] +adls = "kumo_adls.session" + +[tool.uv.sources] +kumo-core = { workspace = true } + +[tool.setuptools] +package-dir = {"" = "src"} + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +addopts = "-q" +testpaths = ["tests"] +pythonpath = ["src", "../kumo-api", "../kumo-core/src", "../kumo-core/tests"] diff --git a/packages/kumo-adls/src/kumo_adls/__init__.py b/packages/kumo-adls/src/kumo_adls/__init__.py new file mode 100644 index 000000000..2a7c0132f --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Azure Data Lake Storage plugins; provider clients load on first use.""" + +from kumo_core.sources.lazy_exports import module_dir, module_getattr + +ADLS_NAMESPACE = 'azure-adls' +_LAZY_EXPORTS = { + 'ADLSSourceBackend': 'backend', + 'ADLSSinkBackend': 'sink', + 'probe': 'session', +} +__all__ = ['ADLS_NAMESPACE', *sorted(_LAZY_EXPORTS)] +__getattr__ = module_getattr(__name__, _LAZY_EXPORTS) +__dir__ = module_dir(__all__) diff --git a/packages/kumo-adls/src/kumo_adls/backend.py b/packages/kumo-adls/src/kumo_adls/backend.py new file mode 100644 index 000000000..6cd4a6235 --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/backend.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +from typing import Optional +from urllib.parse import quote, unquote, urlsplit + +from kumo_core.sources.object_store import ObjectStoreSourceReader +from kumo_core.sources.object_store_backend import ObjectStoreSourceBackend +from kumo_core.sources.object_store_sessions import ( + ObjectStorePaths, + ObjectStoreRoot, + ObjectStoreSessions, +) +from kumo_core.sources.principals import IdentityRef +from kumo_core.sources.types import ReadMode +from kumoapi.adls import ( + normalize_adls_source_relative_path, + normalize_adls_source_uri, +) +from kumoapi.data_source import ( + ADLSConnectorResourceConfig, + ADLSRootedConnectorTarget, + ConnectorType, + DataSourceType, +) +from kumoapi.source_table import ( + ADLSRelativeSource, + ADLSSourceTable, + ADLSURISource, +) +from pyarrow.fs import FSSpecHandler, PyFileSystem + +from . import ADLS_NAMESPACE +from .sdk_filesystem import ADLSFileSystem +from .source_access import source_access + + +def canonical_uri(path: str) -> str: + authority, _, key = path.partition('/') + return normalize_adls_source_uri( + f'abfss://{authority}/{quote(key, safe="/-._~!$&()*+,;=:@")}') + + +class ADLSSourceReader(ObjectStoreSourceReader): + source_type = ConnectorType.ADLS + + def _identities(self, found, worker_ceiling): + if isinstance(self.filesystem, PyFileSystem): + handler = self.filesystem.handler + adapter = handler.fs if isinstance(handler, FSSpecHandler) else None + if (isinstance(adapter, ADLSFileSystem) and getattr( + self.object_identity, '__self__', None) is adapter + and getattr(self.object_identity, '__func__', + None) is ADLSFileSystem.object_identity): + identities = adapter.object_identities( + [entry.path for entry in found], root=self.root) + if identities is not None: + return identities + return super()._identities(found, worker_ceiling) + + def _data_version(self, path, entries): + # The shared single-file version is only an mtime. Include authority + # and ETags even for one object, so accounts and same-second rewrites + # cannot alias the same cached source. + ordered = sorted(entries, key=lambda entry: entry.path) + metadata = list( + zip((entry.path for entry in ordered), + self._identities(ordered, None))) + # Discovery and preparation spell the same directory with/without '/'. + return hashlib.sha256( + json.dumps((path.rstrip('/'), metadata)).encode()).hexdigest() + + +class ADLSSessions(ObjectStoreSessions): + namespace = ADLS_NAMESPACE + rooted = ADLSRootedConnectorTarget + resource_config = ADLSConnectorResourceConfig + paths = ObjectStorePaths( + scheme='abfss', + provider='ADLS', + source_uri=normalize_adls_source_uri, + destination_uri=normalize_adls_source_uri, + destination_path=normalize_adls_source_relative_path, + ) + + def _root_for(self, connector): + target = self._target_config(connector).target + if not isinstance(target, self.rooted): + return None + authority = (f'{target.filesystem}@{target.storage_account}.' + f'{target.endpoint_suffix}') + return ObjectStoreRoot(authority, target.prefix) + + @staticmethod + def _contained_key(root, uri, *, purpose): + parsed = urlsplit(normalize_adls_source_uri(uri)) + key = unquote(parsed.path).lstrip('/') + prefix = root.prefix.rstrip('/') + outside_prefix = (bool(prefix) and key != prefix + and not key.startswith(prefix + '/')) + if parsed.netloc != root.bucket or outside_prefix: + raise ValueError(f'ADLS {purpose} is outside the connector root') + return key + + def _placement(self, connector, table_name): + root = self._root_for(connector) + if root is None: + return super()._placement(connector, table_name) + if '://' in table_name: + key = self._contained_key(root, table_name, purpose='source') + prefix = root.prefix.rstrip('/') + table_name = key[len(prefix):].lstrip('/') if prefix else key + table_name = normalize_adls_source_relative_path(table_name) + return root.bucket, root.path, table_name, None + + def _source_access(self, connector, bucket, region): + del region + _, host = bucket.split('@', 1) + account, suffix = host.split('.', 1) + return source_access(self._principal(connector), account, suffix) + + def _over(self, connector, bucket, root, region): + access = self._source_access(connector, bucket, region) + return ADLSSourceReader(access.filesystem, root=root, + source_type=ConnectorType.ADLS, + fetch=access.fetch, + object_identity=access.object_identity) + + def _filesystem(self, connector, bucket, region): + return self._source_access(connector, bucket, region).filesystem + + def _resolved(self, connector, coordinate): + root = self._root_for(connector) + if root is None: + raise NotImplementedError( + 'Rootless ADLS connectors are source-only') + if '://' in coordinate: + key = self._contained_key(root, coordinate, purpose='destination') + else: + relative = normalize_adls_source_relative_path(coordinate) + key = '/'.join(part for part in (root.prefix, relative) if part) + return root.bucket, f'{root.bucket}/{key}', None + + +class ADLSSourceBackend(ADLSSessions, ObjectStoreSourceBackend): + source_type = DataSourceType.ADLS + relative_source = ADLSRelativeSource + uri_source = ADLSURISource + source_table = ADLSSourceTable + + def __init__(self, identity: Optional[IdentityRef] = None, + storage_namespace: Optional[str] = None, + mode: Optional[ReadMode] = None): + super().__init__(identity, storage_namespace) + self.mode = mode + + def get_download_urls(self, connector, table_name, project_id=None): + del project_id + reader, name = self._reader(connector, table_name) + return [canonical_uri(path) for path in reader.object_paths(name)] diff --git a/packages/kumo-adls/src/kumo_adls/errors.py b/packages/kumo-adls/src/kumo_adls/errors.py new file mode 100644 index 000000000..4f02f87b7 --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/errors.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from contextlib import contextmanager + +from azure.core.exceptions import ( + AzureError, + ClientAuthenticationError, + HttpResponseError, + ResourceModifiedError, + ResourceNotFoundError, + ServiceRequestError, + ServiceResponseError, +) +from kumo_core.sources.errors import ( + PermanentSourceError, + ReconnectRequiredError, + SourceAccessDeniedError, + TransientSourceError, +) + + +@contextmanager +def translated(): + """Classify Azure failures without exposing request metadata or tokens.""" + try: + yield + except MemoryError as error: + error.add_note( + 'If this occurred during a data transfer, lowering ' + 'KUMO_ADLS_TRANSFER_WORKERS and KUMO_SOURCE_COPY_WORKERS may help. ' + 'Set them before importing the provider (restart an existing ' + 'process). These limits do not cap total memory or address all ' + 'allocation failures; projected reads also buffer their requested ' + 'range.') + raise + except ClientAuthenticationError: + raise ReconnectRequiredError('ADLS authentication failed') from None + except ResourceNotFoundError: + raise FileNotFoundError('ADLS path was not found') from None + except ResourceModifiedError: + raise TransientSourceError( + 'ADLS object changed; retry the read') from None + except (ServiceRequestError, ServiceResponseError): + raise TransientSourceError( + 'ADLS transport failed; retry the request') from None + except HttpResponseError as error: + status = error.status_code + if status == 401: + kind = ReconnectRequiredError + elif status == 403: + kind = SourceAccessDeniedError + elif status in (408, 412, 429) or (status is not None + and status >= 500): + kind = TransientSourceError + else: + kind = PermanentSourceError + raise kind('ADLS request failed') from None + except AzureError: + raise PermanentSourceError('ADLS request failed') from None diff --git a/packages/kumo-adls/src/kumo_adls/range_cache.py b/packages/kumo-adls/src/kumo_adls/range_cache.py new file mode 100644 index 000000000..d5221e3c5 --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/range_cache.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import threading +from collections import OrderedDict + +MAX_BYTES = 64 << 20 +MAX_ENTRIES = 1024 +MAX_ENTRY_BYTES = 128 << 10 + + +class RangeBytes: + """Process-bounded payload LRU; keys never retain a filesystem or client. + + Budgets cover payload bytes and entry count. Python key/container overhead + is additional; filesystem namespaces prevent sharing across principals. + """ + def __init__(self, max_bytes=MAX_BYTES, max_entries=MAX_ENTRIES, + max_entry_bytes=MAX_ENTRY_BYTES): + self.max_bytes = max_bytes + self.max_entries = max_entries + self.max_entry_bytes = max_entry_bytes + self._entries = OrderedDict() + self._bytes = 0 + self._lock = threading.Lock() + + def get(self, key): + with self._lock: + value = self._entries.get(key) + if value is not None: + self._entries.move_to_end(key) + return value + + def put(self, key, value): + # Admit immutable, complete successful ranges only. Empty/error/short + # responses must never turn into a later successful cache hit. + if (not isinstance(value, bytes) or not value + or len(value) > self.max_entry_bytes + or len(value) > self.max_bytes + or len(value) != key[-1] - key[-2]): + return + with self._lock: + previous = self._entries.pop(key, None) + if previous is not None: + self._bytes -= len(previous) + self._entries[key] = value + self._bytes += len(value) + while (self._bytes > self.max_bytes + or len(self._entries) > self.max_entries): + _, removed = self._entries.popitem(last=False) + self._bytes -= len(removed) + + def clear_namespace(self, namespace): + with self._lock: + for key in list(self._entries): + if key[0] is namespace: + self._bytes -= len(self._entries.pop(key)) + + +# One global allowance, not a new 64 MiB allocation per connector/filesystem. +RANGE_BYTES = RangeBytes() diff --git a/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py b/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py new file mode 100644 index 000000000..a939b1680 --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import posixpath +import threading +import weakref +from collections import OrderedDict +from contextlib import suppress +from tempfile import NamedTemporaryFile + +from azure.core import MatchConditions +from azure.core.exceptions import ResourceExistsError +from fsspec.spec import AbstractBufferedFile, AbstractFileSystem +from kumo_core.sources.cpu import COPY_WORKER_BYTES, memory_share +from kumo_core.sources.errors import ( + SourceAccessDeniedError, + TransientSourceError, +) +from kumo_core.sources.layout import MAX_SCANNED_ENTRIES +from kumoapi.adls import normalize_adls_filesystem + +from .errors import translated +from .range_cache import MAX_ENTRY_BYTES, RANGE_BYTES +from .transfer import CHUNK_BYTES, concurrency + + +def _identity(etag, size): + # List Paths may omit the quotes present on HEAD response ETags. Tokens + # remain opaque and case-sensitive; downloads keep their original ETag. + if isinstance(etag, str) and len(etag) >= 2 and etag[0] == etag[-1] == '"': + etag = etag[1:-1] + if not etag or not isinstance(etag, + str) or type(size) is not int or size < 0: + return None + return etag, str(size) + + +def _is_directory(properties): + # Azure preserves metadata spelling; HNS can return Hdi_isfolder too. + return any(name.casefold() == 'hdi_isfolder' and value.casefold() == 'true' + for name, value in (properties.metadata or {}).items()) + + +class _ReadFile(AbstractBufferedFile): + def __init__(self, filesystem, path, client, properties, **kwargs): + self.client = client + self.etag = properties.etag + super().__init__(filesystem, path, mode='rb', size=properties.size, + **kwargs) + + def _fetch_range(self, start, end): + # fsspec/Arrow may request past EOF (CSV sniffing and readall). The + # fresh open properties pin size together with the conditional ETag. + end = min(end, self.size) + if end <= start: + return b'' + length = end - start + key = None + if length <= MAX_ENTRY_BYTES and isinstance(self.etag, + str) and self.etag: + key = (self.fs._range_namespace, self.path, self.etag, start, end) + cached = RANGE_BYTES.get(key) + if cached is not None: + return cached + with concurrency(length) as workers, translated(): + data = self.client.download_file( + offset=start, length=length, etag=self.etag, + match_condition=MatchConditions.IfNotModified, + max_concurrency=workers).readall() + if len(data) != length: + raise TransientSourceError('ADLS download range was truncated') + if key is not None: + RANGE_BYTES.put(key, data) + return data + + +class _WriteFile(AbstractBufferedFile): + """Abort upload failures and exceptions in a direct fsspec context. + + Arrow's PythonFile wrapper calls close without forwarding caller errors. + That path can flush partial output; this is not a transactional writer. + """ + def __init__(self, filesystem, path, client, **kwargs): + self.client = client + super().__init__(filesystem, path, mode='wb', **kwargs) + + def _initiate_upload(self): + with translated(): + self.client.create_file() + + def _upload_chunk(self, final=False): + data = self.buffer.getvalue() + try: + with translated(): + if data: + self.client.append_data(data, offset=self.offset, + length=len(data)) + if final: + self.client.flush_data(self.offset + len(data)) + except Exception: + self.discard() + raise + + def discard(self): + # Prevent close/__del__ from retrying a failed append or flush. + if self.closed: + return + self.closed = True + if self.offset is None: + return + try: + self.client.delete_file() + except Exception: + # Cleanup must preserve the original upload failure. + pass + + def __exit__(self, exc_type, exc_value, traceback): + if exc_type is not None: + self.discard() + else: + self.close() + + +class ADLSFileSystem(AbstractFileSystem): + """SDK-backed Arrow access with account-qualified virtual paths. + + Virtual paths retain the full ADLS authority so shared versioning and URI + reconstruction distinguish accounts. Azure clients receive filesystem/path. + """ + protocol = 'kumo-adls' + cachable = False + + def __init__(self, client, account: str, suffix: str): + super().__init__() + self._range_namespace = object() + weakref.finalize(self, RANGE_BYTES.clear_namespace, + self._range_namespace) + self.client = client + self.host = f'{account}.{suffix}' + self._filesystems = OrderedDict() + self._filesystems_lock = threading.Lock() + + def _address(self, path): + path = path.strip('/') + authority, _, key = path.partition('/') + filesystem, separator, host = authority.partition('@') + if not separator or host != self.host: + raise ValueError( + 'ADLS path does not match the bound storage account') + normalize_adls_filesystem(filesystem) + with self._filesystems_lock: + if filesystem not in self._filesystems: + self._filesystems[ + filesystem] = self.client.get_file_system_client(filesystem) + if len(self._filesystems) > 32: + self._filesystems.popitem(last=False) + else: + self._filesystems.move_to_end(filesystem) + client = self._filesystems[filesystem] + return client, key, authority + + @staticmethod + def _entry(authority, item): + directory = item.is_directory + return { + 'name': f'{authority}/{item.name}'.rstrip('/'), + 'size': 0 if directory else item.content_length, + 'type': 'directory' if directory else 'file', + 'mtime': item.last_modified + } + + def info(self, path, **kwargs): + filesystem, key, authority = self._address(path) + with translated(): + if not key: + filesystem.get_file_system_properties() + return {'name': authority, 'size': 0, 'type': 'directory'} + properties = filesystem.get_file_client(key).get_file_properties() + directory = _is_directory(properties) + return { + 'name': f'{authority}/{key}', + 'size': 0 if directory else properties.size, + 'type': 'directory' if directory else 'file', + 'mtime': properties.last_modified, + } + + def ls(self, path, detail=True, **kwargs): + info = self.info(path) + if info['type'] == 'file': + entries = [info] + else: + filesystem, key, authority = self._address(path) + with translated(): + entries = [ + self._entry(authority, + item) for item in filesystem.get_paths( + path=key or None, recursive=False) + ] + return entries if detail else [entry['name'] for entry in entries] + + def find(self, path, maxdepth=None, withdirs=False, detail=False, **kwargs): + if maxdepth is not None and maxdepth < 1: + raise ValueError('maxdepth must be at least 1') + info = self.info(path) + if info['type'] == 'file': + entries = [info] + else: + filesystem, key, authority = self._address(path) + with translated(): + entries = [] + for item in filesystem.get_paths(path=key or None, + recursive=maxdepth != 1): + relative = item.name[len(key.rstrip('/')) + + 1:] if key else item.name + if maxdepth is not None and len( + relative.split('/')) > maxdepth: + continue + if withdirs or not item.is_directory: + entries.append(self._entry(authority, item)) + if withdirs: + entries.append(info) + found = {entry['name']: entry for entry in entries} + return found if detail else sorted(found) + + def mkdir(self, path, create_parents=True, **kwargs): + filesystem, key, _ = self._address(path) + if not key: + # A connector never creates a storage account or filesystem. + with translated(): + filesystem.get_file_system_properties() + return + parts = key.split('/') + keys = ['/'.join(parts[:i]) for i in range(1, len(parts) + 1)] + for name in keys if create_parents else [key]: + with translated(): + try: + filesystem.get_directory_client(name).create_directory() + except ResourceExistsError: + pass + + def makedirs(self, path, exist_ok=False): + if not exist_ok and self.exists(path): + raise FileExistsError(path) + self.mkdir(path) + + def rmdir(self, path): + filesystem, key, _ = self._address(path) + if not key: + raise ValueError('Deleting an ADLS filesystem is not supported') + with translated(): + filesystem.get_directory_client(key).delete_directory() + + def rm_file(self, path): + filesystem, key, _ = self._address(path) + with translated(): + filesystem.get_file_client(key).delete_file() + + def _open(self, path, mode='rb', block_size=None, cache_options=None, + **kwargs): + filesystem, key, _ = self._address(path) + client = filesystem.get_file_client(key) + if mode == 'rb': + with translated(): + properties = client.get_file_properties() + return _ReadFile(self, path, client, properties, + block_size=block_size, cache_options=cache_options, + cache_type=kwargs.get('cache_type', 'none')) + if mode == 'wb': + return _WriteFile(self, path, client, block_size=block_size) + raise NotImplementedError( + 'ADLS supports binary reads and overwrite writes') + + def object_identities(self, paths, *, root=''): + """Fresh, bounded listing identities, or None to use per-file HEADs. + + No metadata outlives this call. Listing is an optimization: narrower + file grants and incomplete list metadata retain the HEAD path. + """ + if len(paths) < 2 or len(set(paths)) != len(paths): + return None + groups = {} + for path in paths: + authority, _, key = path.partition('/') + groups.setdefault(authority, []).append((path, key)) + identities = {} + # Count unrelated entries too, so sibling-heavy roots cannot turn a + # metadata optimization into an unbounded recursive scan. + limit = min(MAX_SCANNED_ENTRIES, max(1024, 4 * len(paths))) + scanned = 0 + for authority, selected in groups.items(): + if len(selected) < 2: + return None + prefix = posixpath.commonpath( + [posixpath.dirname(key) for _, key in selected]) + if not prefix: + return None + root_authority, _, root_key = root.partition('/') + outside_prefix = (bool(root_key) and prefix != root_key and + not prefix.startswith(root_key.rstrip('/') + '/')) + outside_root = (bool(root) + and (authority != root_authority or outside_prefix)) + if outside_root: + return None + filesystem, _, _ = self._address(f'{authority}/{prefix}') + wanted = {key: path for path, key in selected} + try: + with translated(): + for item in filesystem.get_paths(path=prefix, + recursive=True): + scanned += 1 + if scanned > limit: + return None + path = wanted.get(item.name) + if path is None: + continue + identity = _identity(item.etag, item.content_length) + if item.is_directory or identity is None or path in identities: + return None + identities[path] = identity + except (SourceAccessDeniedError, FileNotFoundError, + NotImplementedError): + return None + if len(identities) != len(paths): + return None + return [identities[path] for path in paths] + + def object_identity(self, path): + filesystem, key, _ = self._address(path) + with translated(): + properties = filesystem.get_file_client(key).get_file_properties() + if _is_directory(properties) or not properties.etag: + raise ValueError('ADLS source must be a versioned file') + identity = _identity(properties.etag, properties.size) + if identity is None: + raise ValueError('ADLS source must have a complete object identity') + return identity + + def fetch(self, source, destination): + filesystem, key, _ = self._address(source) + client = filesystem.get_file_client(key) + with translated(): + properties = client.get_file_properties() + target = NamedTemporaryFile( + dir=os.path.dirname(os.path.abspath(destination)), + prefix='.kumo-adls-', delete=False) + staging = target.name + try: + with target: + with concurrency(properties.size, ranged=True): + written = client.download_file( + **({ + 'length': properties.size, + 'offset': 0 + } if properties.size else {}), etag=properties.etag, + match_condition=MatchConditions.IfNotModified, + max_concurrency=1).readinto(target) + # The SDK can report the requested size after a short + # write, so independently verify the stream position. + position = target.tell() + if written != properties.size or position != properties.size: + raise TransientSourceError( + 'ADLS download was truncated') + target.flush() + if os.path.getsize(staging) != properties.size: + raise TransientSourceError( + 'ADLS download was truncated') + # Close before promotion or cleanup, including on Windows. + os.replace(staging, destination) + finally: + with suppress(OSError): + os.unlink(staging) + return None + + def put_file(self, source, destination, **kwargs): + filesystem, key, _ = self._address(destination) + client = filesystem.get_file_client(key) + with memory_share(COPY_WORKER_BYTES, lambda ceiling: 1), \ + open(source, 'rb') as stream: + size = os.fstat(stream.fileno()).st_size + created_etag = None + try: + with concurrency(size, upload=True) as workers, translated(): + created_etag = client.create_file().get('etag') + if not created_etag: + raise TransientSourceError( + 'ADLS upload returned no object identity') + if size: + client.upload_data( + stream, length=size, overwrite=False, + etag=created_etag, + match_condition=MatchConditions.IfNotModified, + chunk_size=CHUNK_BYTES, max_concurrency=workers) + except Exception: + # Never delete a replacement created by another writer. + if created_etag: + with suppress(Exception): + client.delete_file( + etag=created_etag, + match_condition=MatchConditions.IfNotModified) + raise diff --git a/packages/kumo-adls/src/kumo_adls/session.py b/packages/kumo-adls/src/kumo_adls/session.py new file mode 100644 index 000000000..f348b0839 --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/session.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import atexit +import threading + +from azure.core.pipeline.transport import RequestsTransport +from azure.identity import ( + AzureAuthorityHosts, + DefaultAzureCredential, + WorkloadIdentityCredential, +) +from azure.storage.filedatalake import DataLakeServiceClient +from kumo_core.sources.principals import Principal +from kumoapi.data_source import ADLSIdentityConfig, AzureCloud +from requests import Session +from requests.adapters import HTTPAdapter + +from .errors import translated +from .transfer import ( + BULK_DOWNLOAD_CHUNK_BYTES, + CONNECTIONS, + INITIAL_GET_BYTES, + NETWORK_BLOCK_BYTES, +) + +# Storage uses the same generic token audience in every cloud. +# DFS endpoints and Entra authorities remain cloud-specific. +CLOUDS = { + AzureCloud.PUBLIC: ('dfs.core.windows.net', 'https://storage.azure.com', + AzureAuthorityHosts.AZURE_PUBLIC_CLOUD), + AzureCloud.US_GOVERNMENT: + ('dfs.core.usgovcloudapi.net', 'https://storage.azure.com', + AzureAuthorityHosts.AZURE_GOVERNMENT), + AzureCloud.CHINA: ('dfs.core.chinacloudapi.cn', 'https://storage.azure.com', + AzureAuthorityHosts.AZURE_CHINA), +} + + +def credential_for(principal: Principal | None, endpoint_suffix: str): + cloud = next( + (cloud + for cloud, values in CLOUDS.items() if values[0] == endpoint_suffix), + None) + if cloud is None: + raise ValueError('Unsupported ADLS endpoint suffix') + _, audience, authority = CLOUDS[cloud] + if principal is None: + return DefaultAzureCredential(authority=authority), audience + config = principal.config + if not isinstance(config, ADLSIdentityConfig): + raise ValueError('The selected identity is not an ADLS identity') + if config.cloud != cloud: + raise ValueError( + 'ADLS endpoint does not match the selected identity cloud') + # Explicit identities must never fall back to another ambient principal. + return WorkloadIdentityCredential(tenant_id=config.tenant_id, + client_id=config.client_id, + authority=authority), audience + + +_LOCK = threading.RLock() +_SESSIONS = {} + + +def session_for(principal: Principal | None, account: str, suffix: str): + # Config participates in the key so a rotated identity cannot reuse a + # client signed as its predecessor. Tokens refresh inside Azure Identity. + key = (principal.name if principal else None, + repr(principal.config) if principal else None, account, suffix) + with _LOCK: + if key not in _SESSIONS: + with translated(): + credential, audience = credential_for(principal, suffix) + transport_session = Session() + adapter = HTTPAdapter(pool_connections=CONNECTIONS, + pool_maxsize=CONNECTIONS, pool_block=True) + transport_session.mount('https://', adapter) + transport_session.mount('http://', adapter) + transport = RequestsTransport( + session=transport_session, session_owner=True, + connection_data_block_size=NETWORK_BLOCK_BYTES) + try: + client = DataLakeServiceClient( + account_url=f'https://{account}.{suffix}', + credential=credential, audience=audience, + max_single_get_size=INITIAL_GET_BYTES, + max_chunk_get_size=BULK_DOWNLOAD_CHUNK_BYTES, + connection_data_block_size=NETWORK_BLOCK_BYTES, + transport=transport) + except Exception: + transport.close() + credential.close() + raise + _SESSIONS[key] = client, credential + return _SESSIONS[key][0] + + +def close(identity_name: str | None = None) -> None: + """Drop a principal's clients, or every session during process shutdown.""" + with _LOCK: + keys = [ + key for key in _SESSIONS + if identity_name is None or key[0] == identity_name + ] + sessions = [_SESSIONS.pop(key) for key in keys] + for client, credential in sessions: + for resource in (client, credential): + try: + resource.close() + except Exception: + # Teardown must not prevent closing the remaining sessions + # or mask the caller's original authentication failure. + pass + + +atexit.register(close) + + +def probe(principal: Principal) -> None: + """Validate token acquisition; this does not attest storage permissions.""" + config = principal.config + if not isinstance(config, ADLSIdentityConfig): + raise ValueError('The selected identity is not an ADLS identity') + with translated(): + credential, audience = credential_for(principal, + CLOUDS[config.cloud][0]) + with credential: + credential.get_token(f'{audience}/.default') diff --git a/packages/kumo-adls/src/kumo_adls/sink.py b/packages/kumo-adls/src/kumo_adls/sink.py new file mode 100644 index 000000000..5525b30ff --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/sink.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from kumo_core.object_store_sink import ObjectStoreSinkBackend +from kumo_core.sources.cpu import COPY_WORKER_BYTES, memory_share +from kumoapi.data_source import DataSourceType +from pyarrow.fs import FSSpecHandler, PyFileSystem + +from .backend import ADLSSessions, canonical_uri +from .sdk_filesystem import ADLSFileSystem + + +class ADLSSinkBackend(ADLSSessions, ObjectStoreSinkBackend): + source_type = DataSourceType.ADLS + + def _uri(self, path: str) -> str: + return canonical_uri(path) + + def _upload_file(self, filesystem, source, destination): + if (isinstance(filesystem, PyFileSystem) + and isinstance(filesystem.handler, FSSpecHandler) + and isinstance(filesystem.handler.fs, ADLSFileSystem)): + filesystem.handler.fs.put_file(source, destination) + else: + super()._upload_file(filesystem, source, destination) + + def _download_file(self, filesystem, source, destination): + if (isinstance(filesystem, PyFileSystem) + and isinstance(filesystem.handler, FSSpecHandler) + and isinstance(filesystem.handler.fs, ADLSFileSystem)): + with memory_share(COPY_WORKER_BYTES, lambda ceiling: 1): + filesystem.handler.fs.fetch(source, destination) + else: + super()._download_file(filesystem, source, destination) diff --git a/packages/kumo-adls/src/kumo_adls/source_access.py b/packages/kumo-adls/src/kumo_adls/source_access.py new file mode 100644 index 000000000..344e4ce3a --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/source_access.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from kumo_core.sources.object_store_access import ObjectStoreAccess +from pyarrow.fs import FSSpecHandler, PyFileSystem + +from .sdk_filesystem import ADLSFileSystem +from .session import session_for + + +def source_access(principal, account: str, suffix: str) -> ObjectStoreAccess: + client = session_for(principal, account, suffix) + filesystem = ADLSFileSystem(client, account, suffix) + return ObjectStoreAccess(PyFileSystem(FSSpecHandler(filesystem)), + filesystem.fetch, filesystem.object_identity) diff --git a/packages/kumo-adls/src/kumo_adls/transfer.py b/packages/kumo-adls/src/kumo_adls/transfer.py new file mode 100644 index 000000000..9deb47a3f --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/transfer.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import threading +from contextlib import contextmanager + +from kumo_core.sources.cpu import copy_workers + +# The shared copy worker reserves 64 MiB. Serial downloads retain the initial +# response plus a current chunk and transport buffers; this is a payload +# estimate, not a cap on SDK allocations or process RSS. +INITIAL_GET_BYTES = 16 << 20 +CHUNK_BYTES = 4 << 20 +BULK_DOWNLOAD_CHUNK_BYTES = 16 << 20 +MAX_BULK_DOWNLOAD_CONCURRENCY = 1 +NETWORK_BLOCK_BYTES = 16 << 20 +MAX_FILE_CONCURRENCY = 4 +TRANSFER_WORKERS_ENV = 'KUMO_ADLS_TRANSFER_WORKERS' + + +def _transfer_workers(): + ceiling = copy_workers() + value = os.environ.get(TRANSFER_WORKERS_ENV) + if value is None: + return ceiling + try: + limit = int(value) + except ValueError: + raise ValueError( + f'{TRANSFER_WORKERS_ENV} must be a positive integer') from None + if limit < 1: + raise ValueError(f'{TRANSFER_WORKERS_ENV} must be a positive integer') + return min(ceiling, limit) + + +CONNECTIONS = _transfer_workers() +_SLOTS = threading.BoundedSemaphore(CONNECTIONS) + + +@contextmanager +def concurrency(size, *, upload=False, ranged=False): + """Share transfer slots across files/readers/accounts, excluding metadata.""" + initial = 0 if upload or ranged else INITIAL_GET_BYTES + chunk_bytes = CHUNK_BYTES if upload else BULK_DOWNLOAD_CHUNK_BYTES + maximum = MAX_BULK_DOWNLOAD_CONCURRENCY if ranged else MAX_FILE_CONCURRENCY + wanted = min( + maximum, + max(1, (max(0, size - initial) + chunk_bytes - 1) // chunk_bytes)) + _SLOTS.acquire() + acquired = 1 + try: + while acquired < wanted and _SLOTS.acquire(blocking=False): + acquired += 1 + yield acquired + finally: + for _ in range(acquired): + _SLOTS.release() diff --git a/packages/kumo-adls/tests/conftest.py b/packages/kumo-adls/tests/conftest.py new file mode 100644 index 000000000..18ea5e749 --- /dev/null +++ b/packages/kumo-adls/tests/conftest.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from fakes import FakeService +from kumoapi.data_source import ( + ADLSAnyURIConnectorTarget, + ADLSConnectorResourceConfig, + ADLSRootedConnectorTarget, + DataSourceType, +) + + +@pytest.fixture +def service(monkeypatch): + service = FakeService() + monkeypatch.setattr('kumo_adls.source_access.session_for', + lambda *args: service) + sink = pa.BufferOutputStream() + pq.write_table(pa.table({'id': [1, 2], 'value': [10.0, 20.0]}), sink) + service.files['curated/orders/events.parquet'] = sink.getvalue().to_pybytes( + ) + return service + + +@pytest.fixture +def connector(): + def build(*, account='companylake', prefix='orders', any_uri=False, + suffix='dfs.core.windows.net'): + target = ADLSAnyURIConnectorTarget( + ) if any_uri else ADLSRootedConnectorTarget( + storage_account=account, filesystem='curated', prefix=prefix, + endpoint_suffix=suffix) + return SimpleNamespace( + id='adls-test', config=ADLSConnectorResourceConfig(target=target), + identity=None, identity_id=None, source_type=DataSourceType.ADLS) + + return build diff --git a/packages/kumo-adls/tests/fakes.py b/packages/kumo-adls/tests/fakes.py new file mode 100644 index 000000000..04b4c5add --- /dev/null +++ b/packages/kumo-adls/tests/fakes.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +from datetime import datetime, timezone +from types import SimpleNamespace + +from azure.core import MatchConditions +from azure.core.exceptions import ResourceModifiedError, ResourceNotFoundError +from azure.storage.filedatalake import FileProperties, PathProperties + + +class FakeService: + """Azure SDK-shaped storage using real public property models.""" + def __init__(self): + self.files = {} + self.directories = {'curated', 'curated/orders'} + self.downloads = [] + self.created = [] + self.modified = datetime(2026, 1, 1, tzinfo=timezone.utc) + + def get_file_system_client(self, filesystem): + return FakeFilesystem(self, filesystem) + + +class FakeFilesystem: + def __init__(self, service, filesystem): + self.service = service + self.filesystem = filesystem + + def get_file_system_properties(self): + if self.filesystem not in self.service.directories: + raise ResourceNotFoundError('missing filesystem') + return {} + + def get_file_client(self, path): + return FakeFile(self.service, f'{self.filesystem}/{path}') + + get_directory_client = get_file_client + + def get_paths(self, path=None, recursive=True): + prefix = f'{self.filesystem}/' + (path.rstrip('/') + + '/' if path else '') + for name in sorted(set(self.service.files) | self.service.directories): + if not name.startswith(prefix): + continue + relative = name[len(prefix):] + if not relative or not recursive and '/' in relative: + continue + yield PathProperties( + name=name[len(self.filesystem) + 1:], is_directory=name + in self.service.directories, + content_length=len(self.service.files.get(name, b'')), + last_modified=self.service.modified) + + +class FakeFile: + def __init__(self, service, name): + self.service = service + self.name = name + + def get_file_properties(self): + if self.name not in self.service.files and self.name not in self.service.directories: + raise ResourceNotFoundError('missing object') + content = self.service.files.get(self.name, b'') + return FileProperties( + **{ + 'name': self.name, + 'ETag': hashlib.sha256(content).hexdigest(), + 'Content-Length': len(content), + 'Last-Modified': self.service.modified, + 'metadata': { + 'hdi_isfolder': 'true' + } if self.name in self.service.directories else {}, + }) + + def download_file(self, offset=None, length=None, **kwargs): + assert kwargs['match_condition'] == MatchConditions.IfNotModified + if kwargs['etag'] != self.get_file_properties().etag: + raise ResourceModifiedError('object was changed') + self.service.downloads.append((self.name, offset, length, kwargs)) + content = self.service.files[self.name] + content = content[offset or 0:] + if length is not None: + content = content[:length] + return SimpleNamespace(readall=lambda: content, + readinto=lambda stream: stream.write(content)) + + def upload_data(self, stream, length, overwrite, chunk_size, + max_concurrency, **kwargs): + if length == 0: + return {} + assert 1 <= max_concurrency <= 4 + if overwrite: + self.create_file() + self.service.files[self.name] = stream.read() + assert len(self.service.files[self.name]) == length + + def create_file(self): + self.service.created.append(self.name) + self.service.files[self.name] = b'' + return {'etag': self.get_file_properties().etag} + + def create_directory(self): + self.service.directories.add(self.name) + + def append_data(self, data, offset, length): + assert length == len(data) + assert offset == len(self.service.files[self.name]) + self.service.files[self.name] += data + + def flush_data(self, offset): + assert offset == len(self.service.files[self.name]) + + def delete_file(self, **kwargs): + if kwargs.get( + 'etag') and kwargs['etag'] != self.get_file_properties().etag: + raise ResourceModifiedError('object was changed') + self.service.files.pop(self.name) + + def delete_directory(self): + self.service.files = { + key: value + for key, value in self.service.files.items() + if not key.startswith(self.name + '/') + } + self.service.directories = { + key + for key in self.service.directories + if key != self.name and not key.startswith(self.name + '/') + } diff --git a/packages/kumo-adls/tests/live/test_live_filesystem.py b/packages/kumo-adls/tests/live/test_live_filesystem.py new file mode 100644 index 000000000..493c77e9d --- /dev/null +++ b/packages/kumo-adls/tests/live/test_live_filesystem.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Read-only Azure HNS checks; requires an explicitly selected existing directory.""" + +import os +from itertools import islice +from urllib.parse import unquote, urlsplit +from uuid import uuid4 + +import pytest +from kumo_adls.backend import ADLSSourceReader +from kumo_adls.sdk_filesystem import ADLSFileSystem +from kumo_adls.session import session_for +from kumoapi.adls import normalize_adls_root_uri +from pyarrow.fs import FileType, FSSpecHandler, PyFileSystem + +ROOT = os.environ.get('KUMO_ADLS_LIVE_ROOT') +if not ROOT: + pytest.skip( + 'Set KUMO_ADLS_LIVE_ROOT to an existing HNS directory; ' + 'live storage behavior was not validated.', allow_module_level=True) + + +@pytest.fixture +def live_directory(): + parsed = urlsplit(normalize_adls_root_uri(ROOT)) + _, host = parsed.netloc.split('@', 1) + account, suffix = host.split('.', 1) + filesystem = ADLSFileSystem(session_for(None, account, suffix), account, + suffix) + path = parsed.netloc + unquote(parsed.path).rstrip('/') + return filesystem, path + + +def test_existing_hns_directory_is_listable(live_directory): + filesystem, path = live_directory + assert filesystem.info(path)['type'] == 'directory' + assert isinstance(filesystem.ls(path), list) + + +def test_missing_hns_path_is_not_found_for_fsspec_and_arrow(live_directory): + filesystem, path = live_directory + missing = path + '/__kumo_missing_' + uuid4().hex + # Both names are absent; no Azure paths are created by this test. + for name in (missing, missing + '/child'): + with pytest.raises(FileNotFoundError): + filesystem.info(name) + with pytest.raises(FileNotFoundError): + filesystem.ls(name) + arrow = PyFileSystem(FSSpecHandler(filesystem)) + assert arrow.get_file_info(name).type == FileType.NotFound + + +def test_listing_identities_match_heads_and_conditional_reads(live_directory): + filesystem, path = live_directory + client, key, authority = filesystem._address(path) + if not key: + pytest.skip( + 'Listing identities require a directory below the filesystem root') + paths = [] + for item in islice(client.get_paths(path=key or None, recursive=False), + 1024): + if not item.is_directory: + paths.append(f'{authority}/{item.name}') + if len(paths) == 2: + break + if len(paths) < 2: + pytest.skip( + 'Live identity check needs two files among the first 1024 entries') + + # Require usable live listing metadata, rather than silently taking HEAD fallback. + listed = filesystem.object_identities(paths, root=path) + assert listed is not None + heads = [filesystem.object_identity(name) for name in paths] + assert listed == heads + arrow = PyFileSystem(FSSpecHandler(filesystem)) + reader = ADLSSourceReader(arrow, root=path, + object_identity=filesystem.object_identity) + found = arrow.get_file_info(paths) + assert reader._identities(found, 1) == heads + assert filesystem.object_identities(paths[:1], root=path) is None + assert reader._identities(found[:1], 1) == heads[:1] + for name, identity in zip(paths, heads): + with filesystem.open(name, 'rb', cache_type='none') as stream: + raw = stream.etag + token = raw[1:-1] if raw.startswith('"') and raw.endswith( + '"') else raw + assert token == identity[0] + assert len(stream.read(4)) == min(4, int(identity[-1])) diff --git a/packages/kumo-adls/tests/test_backend.py b/packages/kumo-adls/tests/test_backend.py new file mode 100644 index 000000000..dcfcfe7da --- /dev/null +++ b/packages/kumo-adls/tests/test_backend.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from kumo_adls.backend import ADLSSourceBackend +from kumo_adls.sink import ADLSSinkBackend +from kumo_core.sources.errors import UnsupportedReadModeError +from kumo_core.sources.types import ReadMode +from kumoapi.source_table import ADLSRelativeSource, ADLSURISource + +URI = 'abfss://curated@companylake.dfs.core.windows.net/orders/' + + +def test_discovery_schema_sample_and_ingestion(service, connector, tmp_path): + source = ADLSSourceBackend() + target = connector() + assert source.list_tables(target).table_names == ['events'] + config = source.get_table_config(target, 'events').source_table + assert isinstance(config.source, ADLSRelativeSource) + assert config.source.path == 'events' + assert [col.name for col in source.get_columns(target, 'events').cols + ] == ['id', 'value'] + assert '10.0' in source.get_sample(target, 'events', 1).sample_rows + version = source.get_data_version(target, 'events') + output = source.prepare_ingest(target, 'events', str(tmp_path / 'ingest'), + data_version=version) + assert output.version == version + assert pq.read_table(tmp_path / 'ingest').to_pydict() == { + 'id': [1, 2], + 'value': [10.0, 20.0] + } + assert source.get_download_urls(target, + 'events') == [URI + 'events.parquet'] + assert service.downloads + + +def test_rootless_exact_source_and_source_only_contract(service, connector): + source = ADLSSourceBackend() + target = connector(any_uri=True) + uri = URI + 'events.parquet' + config = source.get_table_config(target, uri).source_table + assert isinstance(config.source, ADLSURISource) + assert config.source.uri == uri + with pytest.raises(NotImplementedError, match='cannot list'): + source.list_tables(target) + with pytest.raises(NotImplementedError, match='source-only'): + ADLSSinkBackend().write_table(target, uri, pd.DataFrame({'id': [1]})) + + +def test_versions_separate_accounts_clouds_and_same_second_rewrites( + service, connector): + source = ADLSSourceBackend() + first = source.get_data_version(connector(), 'events') + assert first != source.get_data_version(connector(account='otherlake'), + 'events') + assert first != source.get_data_version( + connector(suffix='dfs.core.chinacloudapi.cn'), 'events') + content = service.files['curated/orders/events.parquet'] + service.files['curated/orders/events.parquet'] = content[:-1] + b'x' + assert first != source.get_data_version(connector(), 'events') + + +@pytest.mark.parametrize('path', [ + '../outside', 'foo/../outside', '%2e%2e/outside', + URI.replace('orders/', 'elsewhere/') + 'events.parquet', + URI.replace('companylake', 'otherlake') + 'events.parquet' +]) +def test_read_and_write_fences_before_storage_access(service, connector, path): + with pytest.raises(ValueError): + ADLSSourceBackend().get_columns(connector(), path) + with pytest.raises(ValueError): + ADLSSinkBackend().write_table(connector(), path, + pd.DataFrame({'id': [1]})) + assert not service.created + + +@pytest.mark.parametrize('filename', + ['scores.parquet', 'scores.csv', 'scores/']) +def test_output_roundtrip(service, connector, filename): + target = connector() + uri = ADLSSinkBackend().write_table(target, filename, + pd.DataFrame({'id': [3, 4]})) + assert uri.startswith(URI + filename.rstrip('/')) + sample = ADLSSourceBackend().get_sample(target, filename, 2).sample_rows + assert '3' in sample and '4' in sample + + +def test_encoded_uri_paths_and_decoded_relative_paths(service, connector): + service.files['curated/orders/space name.parquet'] = service.files[ + 'curated/orders/events.parquet'] + source = ADLSSourceBackend() + assert source.get_download_urls( + connector(), 'space name.parquet') == [URI + 'space%20name.parquet'] + assert source.get_columns(connector(), URI + 'space%20name.parquet').cols + assert source.get_columns(connector(any_uri=True), + URI + 'space%20name.parquet').cols + uri = ADLSSinkBackend().write_table(connector(), 'new name.parquet', + pd.DataFrame({'id': [1]})) + assert uri == URI + 'new%20name.parquet' + + +def test_model_bundle_roundtrip(service, connector, tmp_path): + bundle = tmp_path / 'bundle' + bundle.mkdir() + (bundle / 'model.bin').write_bytes(b'model-bytes') + sink = ADLSSinkBackend() + published = sink.publish_model(connector(), str(bundle), + path='model bundle') + assert published.uri == URI + 'model%20bundle' + result = sink.fetch_model(connector(), str(tmp_path / 'fetched'), + path='model bundle') + assert result.uri == published.uri + assert (tmp_path / 'fetched/model.bin').read_bytes() == b'model-bytes' + + +def test_rejects_sql_and_staged_reads(service, connector, tmp_path): + with pytest.raises(UnsupportedReadModeError): + ADLSSourceBackend().get_columns(connector(), 'events', sql='SELECT 1') + with pytest.raises(UnsupportedReadModeError): + ADLSSourceBackend(mode=ReadMode.STAGED).prepare_ingest( + connector(), 'events', str(tmp_path / 'ingest')) + + +def test_requested_version_is_stamped_without_refresh(service, connector, + tmp_path, monkeypatch): + from unittest.mock import Mock + + from kumo_adls.backend import ADLSSourceReader + + source = ADLSSourceBackend() + target = connector() + version = source.get_data_version(target, 'events') + original = service.files['curated/orders/events.parquet'] + # A new valid Parquet object with the same fixed last_modified timestamp. + sink = pa.BufferOutputStream() + pq.write_table(pa.table({'id': [3, 4], 'value': [30.0, 40.0]}), sink) + service.files['curated/orders/events.parquet'] = sink.getvalue().to_pybytes( + ) + assert original != service.files['curated/orders/events.parquet'] + refresh = Mock(side_effect=AssertionError('must not refresh the version')) + monkeypatch.setattr(ADLSSourceReader, 'data_version', refresh) + monkeypatch.setattr(ADLSSourceReader, '_data_version', refresh) + output = source.prepare_ingest(target, 'events', str(tmp_path / 'ingest'), + data_version=version) + assert output.version == version + assert pq.read_table(output.source['paths']).to_pydict() == { + 'id': [3, 4], + 'value': [30.0, 40.0] + } + refresh.assert_not_called() + + +def test_mutation_during_read_discards_staging(service, connector, tmp_path, + monkeypatch): + from kumo_adls.sdk_filesystem import ADLSFileSystem + from kumo_core.sources.errors import TransientSourceError + + sink = pa.BufferOutputStream() + pq.write_table(pa.table({'id': [3, 4], 'value': [30.0, 40.0]}), sink) + changed = sink.getvalue().to_pybytes() + original = ADLSFileSystem.fetch + + def fetch_then_mutate(self, path, destination): + result = original(self, path, destination) + service.files['curated/orders/events.parquet'] = changed + return result + + monkeypatch.setattr(ADLSFileSystem, 'fetch', fetch_then_mutate) + with pytest.raises(TransientSourceError, match='changed while'): + ADLSSourceBackend().prepare_ingest(connector(), 'events', + str(tmp_path / 'ingest')) + assert not list((tmp_path / 'ingest').rglob('*.parquet')) + + +def test_stream_file_export_and_retry_marker(service, connector, tmp_path): + from kumoapi.jobs import WriteMode + + source = tmp_path / 'predictions.parquet' + pq.write_table(pa.table({'score': [0.5, 0.75]}), source) + sink = ADLSSinkBackend() + target = connector() + first = sink.write_files(target, 'predictions/', [str(source)], + export_id='test-run') + second = sink.write_files(target, 'predictions/', [str(source)], + export_id='test-run') + assert first == second + assert ADLSSourceBackend().get_sample(target, 'predictions', 2).sample_rows + with pytest.raises(NotImplementedError): + sink.write_table(target, 'predictions', pd.DataFrame({'score': [1.0]}), + write_mode=WriteMode.APPEND) + + +def test_declared_ingest_dtype_preserves_physical_source_type( + service, connector, tmp_path): + from kumoapi.typing import Dtype + + source = ADLSSourceBackend() + target = connector() + assert source.source_types(target, 'events') == { + 'id': 'int64', + 'value': 'double', + } + output = source.prepare_ingest(target, 'events', str(tmp_path / 'ingest'), + dtypes={'id': Dtype.float64}) + staged = pq.read_table(output.source['paths']) + assert staged.schema.field('id').type == pa.int64() + assert staged.to_pydict() == {'id': [1, 2]} + + +@pytest.mark.parametrize('project', [False, True]) +def test_csv_ingestion_downloads_raw_files_before_local_conversion( + service, connector, tmp_path, project): + from kumoapi.typing import Dtype + + name = 'curated/orders/csv_events.csv' + service.files[name] = b'id,value\n1,10.5\n2,20.5\n' + source = ADLSSourceBackend() + target = connector() + version = source.get_data_version(target, 'csv_events.csv') + output = source.prepare_ingest( + target, 'csv_events.csv', str(tmp_path / 'ingest'), + dtypes={'id': Dtype.int64} if project else None, data_version=version) + + expected = {'id': [1, 2]} + if not project: + expected['value'] = [10.5, 20.5] + assert pq.read_table(output.source['paths']).to_pydict() == expected + assert output.version == version + assert any( + path == name and offset == 0 and length == len(service.files[name]) + for path, offset, length, _ in service.downloads) + assert not list(tmp_path.rglob('kumo-raw-source-*')) + + +@pytest.mark.parametrize('project', [False, True]) +def test_hive_parquet_ingestion_preserves_partition_values( + service, connector, tmp_path, project): + from kumoapi.typing import Dtype + + root = 'curated/orders/partitioned' + service.directories.add(root) + for region, row in [('north', 1), ('south', 2)]: + directory = f'{root}/region={region}' + service.directories.add(directory) + buffer = pa.BufferOutputStream() + pq.write_table(pa.table({'id': [row], 'value': [row * 10.0]}), buffer) + service.files[f'{directory}/part.parquet'] = buffer.getvalue( + ).to_pybytes() + + output = ADLSSourceBackend().prepare_ingest( + connector(), 'partitioned', str(tmp_path / 'ingest'), dtypes={ + 'id': Dtype.int64, + 'region': Dtype.string + } if project else None) + result = pq.read_table(output.source['paths']).to_pydict() + assert sorted(zip(result['id'], result['region'])) == [(1, 'north'), + (2, 'south')] + assert ('value' not in result) == project + if not project: + assert sorted(result['value']) == [10.0, 20.0] + + +@pytest.mark.parametrize('authority', [ + 'companylake.dfs.core.windows.net', + 'curated@companylake', + 'curated@', + '@companylake.dfs.core.windows.net', + 'curated@@companylake.dfs.core.windows.net', + 'curated@companylake.blob.core.windows.net', + 'curated@companylake.dfs.core.windows.net:443', + 'curated@companylake.dfs.core.windows.net:invalid', + 'curated:password@companylake.dfs.core.windows.net', + 'curated@.dfs.core.windows.net', + 'curated@bad.name.dfs.core.windows.net', + 'curated@[invalid', +]) +def test_rootless_invalid_authority_fails_before_source_access( + connector, monkeypatch, authority): + from unittest.mock import Mock + access = Mock(side_effect=AssertionError('must not access storage')) + monkeypatch.setattr('kumo_adls.backend.source_access', access) + with pytest.raises(ValueError): + ADLSSourceBackend().get_columns(connector(any_uri=True), + f'abfss://{authority}/events.parquet') + access.assert_not_called() + + +@pytest.mark.parametrize('path', [ + URI.replace('orders/', 'orders-sibling/') + 'events.parquet', + URI.replace('curated@', 'otherfs@') + 'events.parquet', +]) +def test_full_uri_sibling_and_filesystem_escape_is_rejected( + connector, monkeypatch, path): + from unittest.mock import Mock + access = Mock(side_effect=AssertionError('must not access storage')) + monkeypatch.setattr('kumo_adls.backend.source_access', access) + with pytest.raises(ValueError, match='source is outside'): + ADLSSourceBackend().get_columns(connector(), path) + with pytest.raises(ValueError, match='destination is outside'): + ADLSSinkBackend().write_table(connector(), path, + pd.DataFrame({'id': [1]})) + access.assert_not_called() diff --git a/packages/kumo-adls/tests/test_filesystem.py b/packages/kumo-adls/tests/test_filesystem.py new file mode 100644 index 000000000..34ce3d26c --- /dev/null +++ b/packages/kumo-adls/tests/test_filesystem.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import traceback +from unittest.mock import Mock + +import pytest +from azure.core.exceptions import ( + ClientAuthenticationError, + HttpResponseError, + ResourceModifiedError, + ResourceNotFoundError, +) +from kumo_adls.errors import translated +from kumo_adls.sdk_filesystem import ADLSFileSystem +from kumo_core.sources.errors import ( + PermanentSourceError, + ReconnectRequiredError, + SourceAccessDeniedError, + TransientSourceError, +) +from pyarrow.fs import FileType, FSSpecHandler, PyFileSystem + +AUTHORITY = 'curated@companylake.dfs.core.windows.net' + + +def test_real_sdk_property_models_and_conditional_ranges(service): + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + arrow = PyFileSystem(FSSpecHandler(filesystem)) + assert arrow.get_file_info(AUTHORITY + '/orders').type == FileType.Directory + assert arrow.get_file_info(AUTHORITY + + '/orders/events.parquet').type == FileType.File + with filesystem.open(AUTHORITY + '/orders/events.parquet', 'rb', + block_size=4, cache_type='none') as stream: + assert stream.read(4) == b'PAR1' + service.files['curated/orders/events.parquet'] = b'changed' + with pytest.raises(TransientSourceError, match='changed'): + stream.seek(5) + stream.read(4) + + +def test_chunked_output_and_empty_file(service): + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + for name, chunks in [('chunked', [b'abcd', b'efgh', b'ij']), ('empty', [])]: + with filesystem.open(AUTHORITY + '/' + name, 'wb', + block_size=4) as stream: + for chunk in chunks: + stream.write(chunk) + assert service.files['curated/' + name] == b''.join(chunks) + + +def test_arrow_wrapper_closes_partial_output_after_caller_error(service): + # Arrow does not forward the body exception to the fsspec __exit__. + # Keep this boundary explicit: direct-context abort tests do not cover it. + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + arrow = PyFileSystem(FSSpecHandler(filesystem)) + with pytest.raises(RuntimeError, match='serialization failed'): + with arrow.open_output_stream(AUTHORITY + '/partial') as writer: + writer.write(b'partial serialized content') + raise RuntimeError('serialization failed') + assert service.files['curated/partial'] == b'partial serialized content' + + +@pytest.mark.parametrize('error, expected', [ + (ClientAuthenticationError('token=secret'), ReconnectRequiredError), + (ResourceNotFoundError('token=secret'), FileNotFoundError), + (ResourceModifiedError('token=secret'), TransientSourceError), + (HttpResponseError( + 'token=secret', response=Mock(status_code=403, reason='denied', + headers={})), SourceAccessDeniedError), + (HttpResponseError( + 'token=secret', response=Mock(status_code=429, reason='throttled', + headers={})), TransientSourceError), + (HttpResponseError('token=secret', response=Mock( + status_code=400, reason='bad', headers={})), PermanentSourceError), +]) +def test_errors_are_classified_without_request_metadata(error, expected): + with pytest.raises(expected) as result: + with translated(): + raise error + assert 'token=secret' not in str(result.value) + assert result.value.__suppress_context__ + + +def test_memory_failure_preserves_allocation_traceback_and_adds_advice(): + error = MemoryError('allocation failed') + + def allocate_buffer(): + raise error + + with pytest.raises(MemoryError) as result: + with translated(): + allocate_buffer() + assert result.value is error + assert str(error) == 'allocation failed' + assert not error.__suppress_context__ + assert 'allocate_buffer' in [ + frame.name for frame in traceback.extract_tb(error.__traceback__) + ] + message = '\n'.join(error.__notes__) + assert 'If this occurred during a data transfer' in message + assert 'may help' in message + assert 'KUMO_ADLS_TRANSFER_WORKERS' in message + assert 'KUMO_SOURCE_COPY_WORKERS' in message + assert 'before importing' in message + assert 'restart' in message + + +@pytest.mark.parametrize('metadata_key', + ['hdi_isfolder', 'Hdi_isfolder', 'HDI_ISFOLDER']) +def test_directory_metadata_is_case_insensitive(service, connector, monkeypatch, + metadata_key): + from fakes import FakeFile + from kumo_adls.backend import ADLSSourceBackend + + original = FakeFile.get_file_properties + + def properties_with_header_spelling(self): + properties = original(self) + if properties.metadata.get('hdi_isfolder') == 'true': + properties.metadata = {metadata_key: 'true'} + return properties + + monkeypatch.setattr(FakeFile, 'get_file_properties', + properties_with_header_spelling) + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + arrow = PyFileSystem(FSSpecHandler(filesystem)) + directory = AUTHORITY + '/orders' + assert arrow.get_file_info(directory).type == FileType.Directory + assert ADLSSourceBackend().list_tables( + connector()).table_names == ['events'] + with pytest.raises(ValueError, match='versioned file'): + filesystem.object_identity(directory) + + +@pytest.mark.parametrize('failure', ['append', 'flush', 'body']) +@pytest.mark.parametrize('cleanup_fails', [False, True]) +def test_failed_write_aborts_without_retry_or_masking_original_error( + service, monkeypatch, failure, cleanup_fails): + from fakes import FakeFile + + attempts = [] + original_append = FakeFile.append_data + original_delete = FakeFile.delete_file + + def append(self, *args, **kwargs): + attempts.append('append') + if failure == 'append': + raise RuntimeError('original failure') + return original_append(self, *args, **kwargs) + + def flush(self, offset): + attempts.append('flush') + raise RuntimeError('original failure') + + def delete(self): + attempts.append('delete') + if cleanup_fails: + raise RuntimeError('cleanup failure') + original_delete(self) + + monkeypatch.setattr(FakeFile, 'append_data', append) + monkeypatch.setattr(FakeFile, 'delete_file', delete) + if failure == 'flush': + monkeypatch.setattr(FakeFile, 'flush_data', flush) + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + with pytest.raises(RuntimeError, match='original failure'): + with filesystem.open(AUTHORITY + '/failed', 'wb', + block_size=4) as stream: + stream.write(b'abcd') + if failure == 'body': + raise RuntimeError('original failure') + before = list(attempts) + stream.close() + assert attempts == before + assert attempts.count('delete') == 1 + assert stream.closed + if not cleanup_fails: + assert 'curated/failed' not in service.files + + +def test_exception_before_first_upload_does_not_create_file(service): + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + with pytest.raises(RuntimeError): + with filesystem.open(AUTHORITY + '/never-created', 'wb') as stream: + stream.write(b'small buffered write') + raise RuntimeError('caller failed') + assert not service.created + + +@pytest.mark.parametrize('started', [False, True]) +def test_discard_only_deletes_started_upload_and_is_idempotent( + service, started): + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + service.files['curated/existing'] = b'original' + stream = filesystem.open(AUTHORITY + '/existing', 'wb', block_size=4) + stream.write(b'abcd' if started else b'ab') + stream.discard() + stream.discard() + stream.close() + if started: + assert 'curated/existing' not in service.files + else: + assert service.files['curated/existing'] == b'original' diff --git a/packages/kumo-adls/tests/test_identity_listing.py b/packages/kumo-adls/tests/test_identity_listing.py new file mode 100644 index 000000000..86d5fa248 --- /dev/null +++ b/packages/kumo-adls/tests/test_identity_listing.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib +import json +import re +from types import SimpleNamespace +from unittest.mock import Mock +from urllib.parse import parse_qs, urlsplit + +import pyarrow.parquet as pq +import pytest +from azure.core.exceptions import HttpResponseError +from fakes import FakeFile, FakeFilesystem +from kumo_adls.backend import ADLSSourceReader +from kumo_adls.sdk_filesystem import ADLSFileSystem +from kumo_core.sources.errors import TransientSourceError +from kumo_core.sources.types import ReadMode +from pyarrow.fs import FSSpecHandler, PyFileSystem +from test_sdk_transfer import HEADERS, HTTPAdapter, client + +AUTHORITY = 'curated@companylake.dfs.core.windows.net' +PATHS = [AUTHORITY + '/orders/a.parquet', AUTHORITY + '/orders/b.parquet'] + + +def reader(filesystem): + return ADLSSourceReader(PyFileSystem(FSSpecHandler(filesystem)), + root=AUTHORITY + '/orders', + object_identity=filesystem.object_identity, + fetch=filesystem.fetch) + + +@pytest.fixture +def listing_service(service, monkeypatch): + service.files['curated/orders/a.parquet'] = b'aaaa' + service.files['curated/orders/b.parquet'] = b'bbbb' + original = FakeFilesystem.get_paths + calls = [] + + def listing(self, path=None, recursive=True): + calls.append((self.filesystem, path, recursive)) + for item in original(self, path=path, recursive=recursive): + if not item.is_directory: + item.etag = hashlib.sha256( + self.service.files[self.filesystem + '/' + + item.name]).hexdigest() + yield item + + monkeypatch.setattr(FakeFilesystem, 'get_paths', listing) + return service, calls + + +def test_fresh_listing_replaces_heads_and_keeps_order(listing_service, + monkeypatch): + service, calls = listing_service + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + identities = [filesystem.object_identity(path) for path in PATHS] + heads = Mock(side_effect=AssertionError('listing should avoid HEAD')) + monkeypatch.setattr(FakeFile, 'get_file_properties', heads) + source = reader(filesystem) + found = [SimpleNamespace(path=path) for path in reversed(PATHS)] + assert source._identities(found, 1) == list(reversed(identities)) + service.files['curated/orders/a.parquet'] = b'AAAA' + changed = source._identities(found, 1) + assert changed[0] == identities[1] + assert changed[1] != identities[0] + assert calls == [('curated', 'orders', True)] * 2 + heads.assert_not_called() + + +@pytest.mark.parametrize( + 'kind', + ['etag', 'size', 'missing', 'duplicate', 'directory', 'denied', 'bound']) +def test_incomplete_or_unusable_listing_falls_back(listing_service, monkeypatch, + kind): + service, _ = listing_service + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + expected = [filesystem.object_identity(path) for path in PATHS] + original = FakeFilesystem.get_paths + + def broken(self, **kwargs): + if kind == 'denied': + raise HttpResponseError( + 'denied', response=Mock(status_code=403, reason='denied', + headers={})) + items = list(original(self, **kwargs)) + selected = next(item for item in items + if item.name == 'orders/a.parquet') + if kind == 'etag': + selected.etag = None + elif kind == 'size': + selected.content_length = None + elif kind == 'directory': + selected.is_directory = True + elif kind == 'missing': + items.remove(selected) + elif kind == 'duplicate': + items.append(selected) + elif kind == 'bound': + items = [ + SimpleNamespace(name='orders/unrelated', is_directory=False, + etag='x', content_length=1) + ] * 1025 + items + yield from items + + monkeypatch.setattr(FakeFilesystem, 'get_paths', broken) + assert filesystem.object_identities(PATHS, + root=AUTHORITY + '/orders') is None + assert reader(filesystem)._identities( + [SimpleNamespace(path=path) for path in PATHS], 1) == expected + + +def test_scopes_nested_paths_and_never_scans_filesystem_root(listing_service): + service, calls = listing_service + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + nested = [ + AUTHORITY + '/orders/year=1/same.parquet', + AUTHORITY + '/orders/year=2/same.parquet' + ] + service.files['curated/orders/year=1/same.parquet'] = b'one' + service.files['curated/orders/year=2/same.parquet'] = b'two' + result = filesystem.object_identities(nested, root=AUTHORITY + '/orders') + assert result[0] != result[1] + assert calls[-1] == ('curated', 'orders', True) + count = len(calls) + assert filesystem.object_identities( + [AUTHORITY + '/one', AUTHORITY + '/two']) is None + assert filesystem.object_identities(PATHS, + root=AUTHORITY + '/private') is None + assert filesystem.object_identities(PATHS[:1]) is None + assert len(calls) == count + + +def test_custom_identity_callback_is_never_bypassed(listing_service): + service, calls = listing_service + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + source = reader(filesystem) + source.object_identity = Mock(return_value=('custom', '4')) + assert source._identities([SimpleNamespace(path=path) for path in PATHS], + 1) == [('custom', '4')] * 2 + assert not calls + assert source.object_identity.call_count == 2 + + +@pytest.mark.parametrize('listing_identities', [False, True]) +def test_prepare_stamps_and_reads_one_discovery(listing_service, tmp_path, + monkeypatch, + listing_identities): + service, _ = listing_service + # Keep source valid and use two parts so identities use listing. + original = service.files.pop('curated/orders/events.parquet') + service.files['curated/orders/a.parquet'] = original + service.files['curated/orders/b.parquet'] = original + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + if not listing_identities: + monkeypatch.setattr(filesystem, 'object_identities', + Mock(return_value=None)) + source = reader(filesystem) + version = source.data_version('') + objects = ADLSSourceReader._objects + versions = ADLSSourceReader._data_version + listings = [] + stamps = [] + + def listed(self, *args, **kwargs): + found = objects(self, *args, **kwargs) + listings.append(found) + return found + + def stamp(self, path, found): + stamps.append(found) + return versions(self, path, found) + + monkeypatch.setattr(ADLSSourceReader, '_objects', listed) + monkeypatch.setattr(ADLSSourceReader, '_data_version', stamp) + monkeypatch.setattr( + ADLSSourceReader, 'data_version', + Mock(side_effect=AssertionError( + 'must not relist the table for its version'))) + fetch = ADLSFileSystem.fetch + + def add_part(self, path, destination): + result = fetch(self, path, destination) + service.files['curated/orders/added.parquet'] = original + return result + + monkeypatch.setattr(ADLSFileSystem, 'fetch', add_part) + # The callback must be bound after patching the SDK fetch method. + source.fetch = filesystem.fetch + output = source.prepare('', str(tmp_path / 'out'), columns=None, + mode=ReadMode.STREAMED) + assert output.version == version + assert len(listings) == len(stamps) == 1 + assert stamps[0] is listings[0] + assert len(listings[0]) == 2 + assert pq.read_table(list(output.paths)).num_rows == 4 + assert 'curated/orders/added.parquet' in service.files + + +def test_real_sdk_json_pagination_and_quoted_head_etags(): + http = HTTPAdapter() + list_calls = [] + heads = [] + + def listing(request, context): + query = parse_qs(urlsplit(request.url).query) + assert query['directory'] == ['orders'] + assert query['recursive'] == ['true'] + continuation = query.get('continuation', [''])[0] + list_calls.append(continuation) + if not continuation: + context.headers['x-ms-continuation'] = 'next-page' + name = 'a' if not continuation else 'b' + return json.dumps({ + 'paths': [{ + 'name': f'orders/{name}.parquet', + 'etag': 'stable', + 'contentLength': '3', + 'lastModified': HEADERS['Last-Modified'], + }] + }).encode() + + def head(request, context): + heads.append(request.url) + return b'' + + def download(request, context): + assert request.headers['If-Match'] == '"stable"' + return b'abc' + + http.register_uri('GET', + re.compile(r'https://companylake\.dfs\.[^/]+/curated\?'), + headers={'Content-Type': + 'application/json'}, content=listing) + http.register_uri('HEAD', re.compile(r'https://companylake\.'), headers={ + **HEADERS, 'Content-Length': '3' + }, content=head) + http.register_uri( + 'GET', re.compile(r'https://companylake\.[^/]+/curated/orders/'), + headers={ + **HEADERS, 'Content-Length': '3', + 'Content-Range': 'bytes 0-2/3' + }, content=download) + with client(http) as service: + filesystem = ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net') + assert filesystem.object_identities(PATHS) == [('stable', '3')] * 2 + assert list_calls == ['', 'next-page'] + assert not heads + assert filesystem.object_identity(PATHS[0]) == ('stable', '3') + with filesystem.open(PATHS[0], 'rb') as stream: + assert stream.read() == b'abc' + + +def test_multiple_filesystems_keep_exact_authorities(listing_service): + service, calls = listing_service + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + other = [path.replace('curated@', 'another@') for path in PATHS] + service.files['another/orders/a.parquet'] = b'other-a' + service.files['another/orders/b.parquet'] = b'other-b' + selected = [PATHS[0], other[0], PATHS[1], other[1]] + assert filesystem.object_identities(selected) == [ + filesystem.object_identity(path) for path in selected + ] + assert calls == [('curated', 'orders', True), ('another', 'orders', True)] + + +def test_missing_selected_object_does_not_reuse_old_identity(listing_service): + service, _ = listing_service + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + source = reader(filesystem) + found = [SimpleNamespace(path=path) for path in PATHS] + assert source._identities(found, 1) + del service.files['curated/orders/b.parquet'] + with pytest.raises(FileNotFoundError): + source._identities(found, 1) + + +def test_non_sdk_filesystem_keeps_custom_identity_path(tmp_path): + from pyarrow.fs import LocalFileSystem + + identity = Mock(return_value=('version', '4')) + source = ADLSSourceReader(LocalFileSystem(), object_identity=identity) + assert source._identities([SimpleNamespace(path=str(tmp_path / 'one'))], + 1) == [('version', '4')] + identity.assert_called_once() + + +@pytest.mark.parametrize('status', [401, 429, 503]) +def test_listing_auth_and_transient_errors_propagate(listing_service, + monkeypatch, status): + from kumo_core.sources.errors import ReconnectRequiredError + + service, _ = listing_service + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + + def failed(*args, **kwargs): + raise HttpResponseError( + 'failed', response=Mock(status_code=status, reason='failed', + headers={})) + + monkeypatch.setattr(FakeFilesystem, 'get_paths', failed) + expected = ReconnectRequiredError if status == 401 else TransientSourceError + with pytest.raises(expected): + filesystem.object_identities(PATHS) diff --git a/packages/kumo-adls/tests/test_package_boundary.py b/packages/kumo-adls/tests/test_package_boundary.py new file mode 100644 index 000000000..9231f12ab --- /dev/null +++ b/packages/kumo-adls/tests/test_package_boundary.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +from sources.package_boundary import ( + assert_imports_nothing_foreign, + assert_top_level_import_is_cheap, +) + +ROOT = Path(__file__).parents[1] + + +def test_provider_import_boundary(): + assert_imports_nothing_foreign(ROOT / 'src/kumo_adls', vendor='azure') + + +def test_top_level_import_is_cheap(): + assert_top_level_import_is_cheap( + ROOT, 'kumo_adls', 'ADLS_NAMESPACE', + unloaded=('kumo_adls.backend', 'azure.identity', + 'azure.storage.filedatalake')) diff --git a/packages/kumo-adls/tests/test_range_cache.py b/packages/kumo-adls/tests/test_range_cache.py new file mode 100644 index 000000000..177d8f465 --- /dev/null +++ b/packages/kumo-adls/tests/test_range_cache.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import gc +import weakref +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import pytest +from azure.core.exceptions import ResourceModifiedError +from fakes import FakeFile +from kumo_adls.backend import ADLSSourceBackend +from kumo_adls.range_cache import MAX_BYTES, MAX_ENTRIES, RangeBytes +from kumo_adls.sdk_filesystem import ADLSFileSystem +from kumo_core.sources.errors import TransientSourceError + +AUTHORITY = 'curated@companylake.dfs.core.windows.net' +PATH = AUTHORITY + '/orders/events.parquet' + + +@pytest.fixture +def cache(monkeypatch): + value = RangeBytes() + monkeypatch.setattr('kumo_adls.sdk_filesystem.RANGE_BYTES', value) + return value + + +def filesystem(service): + return ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + + +def read_four(fs, path=PATH): + with fs.open(path, 'rb', cache_type='none') as stream: + return stream.read(4) + + +def test_reopened_file_checks_properties_and_changed_etag_misses( + service, cache, monkeypatch): + fs = filesystem(service) + heads = [] + original = FakeFile.get_file_properties + + def head(self): + heads.append(self.name) + return original(self) + + monkeypatch.setattr(FakeFile, 'get_file_properties', head) + first = read_four(fs) + initial_heads = len(heads) + assert read_four(fs) == first + assert len(heads) > initial_heads # Cached bytes never replace open's HEAD. + assert len(service.downloads) == 1 + content = service.files['curated/orders/events.parquet'] + service.files['curated/orders/events.parquet'] = b'NEWW' + content[4:] + assert read_four(fs) == b'NEWW' + assert len(service.downloads) == 2 + + +def test_open_snapshot_can_reuse_old_bytes_but_identity_guard_rejects_mutation( + service, connector, cache, tmp_path): + reader, name = ADLSSourceBackend()._reader(connector(), 'events') + fs = reader.filesystem.handler.fs + found = reader._objects(name, PATH) + identities = reader._identities(found, 1) + destination = tmp_path / 'staging' + destination.mkdir() + (destination / 'partial.parquet').write_bytes(b'partial') + with fs.open(PATH, 'rb', cache_type='none') as stream: + before = stream.read(4) + content = service.files['curated/orders/events.parquet'] + service.files['curated/orders/events.parquet'] = b'NEWW' + content[4:] + stream.seek(0) + assert stream.read(4) == before + assert len(service.downloads) == 1 + with pytest.raises(TransientSourceError, match='changed'): + reader._unchanged(name, found, str(destination), 1, identities) + assert not destination.exists() + + +@pytest.mark.parametrize('response', ['error', 'short', 'empty']) +def test_errors_and_incomplete_reads_never_enter_cache(service, cache, + monkeypatch, response): + def broken(self, **kwargs): + if response == 'error': + raise ResourceModifiedError('changed') + return SimpleNamespace( + readall=lambda: b'ab' if response == 'short' else b'') + + monkeypatch.setattr(FakeFile, 'download_file', broken) + fs = filesystem(service) + for _ in range(2): + with pytest.raises(TransientSourceError): + read_four(fs) + assert not cache._entries + assert cache._bytes == 0 + + +def test_zero_and_large_ranges_not_admitted(service, cache): + service.files['curated/orders/events.parquet'] = b'a' * (128 * 1024 + 1) + fs = filesystem(service) + with fs.open(PATH, 'rb', cache_type='none') as stream: + assert stream._fetch_range(0, 0) == b'' + assert not service.downloads + assert len(stream.read()) == 128 * 1024 + 1 + assert not cache._entries + + +def test_raw_etags_and_full_authorities_are_distinct(service, cache, + monkeypatch): + fs = filesystem(service) + original = FakeFile.get_file_properties + token = ['"opaque"'] + calls = [] + + def properties(self): + value = original(self) + value.etag = token[0] + return value + + def download(self, offset=None, length=None, **kwargs): + calls.append(kwargs['etag']) + return SimpleNamespace(readall=lambda: self.service.files[self.name][ + offset:offset + length]) + + monkeypatch.setattr(FakeFile, 'get_file_properties', properties) + monkeypatch.setattr(FakeFile, 'download_file', download) + assert read_four(fs) + token[0] = 'opaque' + assert read_four(fs) + service.files['another/orders/events.parquet'] = b'else' + assert read_four(fs, PATH.replace('curated@', 'another@')) == b'else' + assert calls == ['"opaque"', 'opaque', 'opaque'] + assert len(cache._entries) == 3 + + +def test_filesystems_do_not_share_bytes_and_gc_releases_namespace( + service, cache): + one = filesystem(service) + two = filesystem(service) + assert read_four(one) == read_four(two) + assert len(service.downloads) == 2 + namespace = one._range_namespace + reference = weakref.ref(one) + assert any(key[0] is namespace for key in cache._entries) + del one + gc.collect() + assert reference() is None + assert not any(key[0] is namespace for key in cache._entries) + assert any(key[0] is two._range_namespace for key in cache._entries) + + +def test_process_byte_and_entry_caps_are_global_and_thread_safe(): + cache = RangeBytes(max_bytes=32, max_entries=3, max_entry_bytes=16) + namespaces = [object() for _ in range(8)] + + def put(index): + key = (namespaces[index % 8], f'file-{index}', '"etag"', 0, 12) + cache.put(key, b'x' * 12) + cache.get(key) + + with ThreadPoolExecutor(max_workers=8) as pool: + list(pool.map(put, range(64))) + assert cache._bytes == sum(map(len, cache._entries.values())) + assert cache._bytes <= 32 + assert len(cache._entries) <= 3 + assert MAX_BYTES == 64 << 20 + assert MAX_ENTRIES == 1024 + + +def test_lru_updates_on_hits_replacements_and_namespace_cleanup(): + cache = RangeBytes(max_bytes=12, max_entries=2, max_entry_bytes=8) + namespace = object() + keys = [(namespace, str(index), '"etag"', 0, 4) for index in range(3)] + cache.put(keys[0], b'aaaa') + cache.put(keys[1], b'bbbb') + assert cache.get(keys[0]) == b'aaaa' + cache.put(keys[2], b'cccc') + assert cache.get(keys[1]) is None + cache.put(keys[0], b'AAAA') + assert cache._bytes == 8 + cache.put((namespace, 'oversize', 'etag', 0, 9), b'x' * 9) + cache.put((namespace, 'short', 'etag', 0, 4), b'x') + assert cache._bytes == 8 + cache.clear_namespace(namespace) + assert cache._bytes == 0 + assert not cache._entries + + +def test_reads_past_eof_use_pinned_size_and_do_not_cache_empty_ranges( + service, cache): + service.files['curated/orders/events.parquet'] = b'abc' + fs = filesystem(service) + with fs.open(PATH, 'rb', cache_type='none') as stream: + assert stream._fetch_range(0, 50000) == b'abc' + assert stream._fetch_range(0, 50000) == b'abc' + assert stream._fetch_range(3, 50000) == b'' + assert len(service.downloads) == 1 + assert next(iter(cache._entries))[-2:] == (0, 3) diff --git a/packages/kumo-adls/tests/test_sdk_transfer.py b/packages/kumo-adls/tests/test_sdk_transfer.py new file mode 100644 index 000000000..395f479b3 --- /dev/null +++ b/packages/kumo-adls/tests/test_sdk_transfer.py @@ -0,0 +1,432 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real Azure SDK transfer behavior against mocked HTTP; no cloud benchmarks.""" + +import base64 +import io +import re +import threading +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from urllib.parse import parse_qs, urlsplit + +import pytest +import requests +from azure.core.pipeline.transport import RequestsTransport +from azure.storage.filedatalake import DataLakeServiceClient +from kumo_adls import transfer +from kumo_adls.sdk_filesystem import ADLSFileSystem +from kumo_adls.transfer import ( + BULK_DOWNLOAD_CHUNK_BYTES, + CHUNK_BYTES, + INITIAL_GET_BYTES, + NETWORK_BLOCK_BYTES, +) +from kumo_core.sources.errors import SourceError +from pyarrow.fs import FSSpecHandler, PyFileSystem +from requests.adapters import BaseAdapter +from requests.structures import CaseInsensitiveDict +from urllib3.response import HTTPResponse + +AUTHORITY = 'curated@companylake.dfs.core.windows.net' +KEY = base64.b64encode(b'test-key').decode() +HEADERS = { + 'ETag': '"stable"', + 'Last-Modified': 'Wed, 21 Oct 2026 00:00:00 GMT', + 'x-ms-blob-type': 'BlockBlob', + 'x-ms-request-id': 'test' +} + + +class HTTPAdapter(BaseAdapter): + """In-memory HTTP responses without serializing concurrent SDK sends.""" + def __init__(self): + self.routes = [] + self.fragments = [] + + def register_uri(self, method, url, *, headers=None, status_code=200, + content=b''): + self.routes.append((method, url, headers or {}, status_code, content)) + + def send(self, request, **kwargs): + for method, url, headers, status, content in self.routes: + if request.method != method or not url.match(request.url): + continue + context = SimpleNamespace(headers=dict(headers), status_code=status) + body = content(request, context) if callable(content) else content + response = requests.Response() + response.status_code = context.status_code + response.headers = CaseInsensitiveDict(context.headers) + response.request = request + response.url = request.url + response.raw = HTTPResponse(body=io.BytesIO(body), + headers=context.headers, + request_method=request.method, + preload_content=False) + original_stream = response.raw.stream + + def stream(*args, **kwargs): + try: + for fragment in original_stream(*args, **kwargs): + self.fragments.append( + (request.method, request.headers.get('x-ms-range'), + len(fragment))) + yield fragment + finally: + closed = getattr(context, 'on_stream_closed', None) + if closed is not None: + closed() + + response.raw.stream = stream + return response + raise AssertionError( + f'Unexpected SDK request: {request.method} {request.url}') + + def close(self): + pass + + +@pytest.fixture +def mocked_http(): + return HTTPAdapter() + + +def client(adapter, *, network_block_bytes=NETWORK_BLOCK_BYTES, + read_chunk_bytes=BULK_DOWNLOAD_CHUNK_BYTES, + initial_get_bytes=INITIAL_GET_BYTES): + http = requests.Session() + http.mount('https://', adapter) + return DataLakeServiceClient( + 'https://companylake.dfs.core.windows.net', credential=KEY, + max_single_get_size=initial_get_bytes, + max_chunk_get_size=read_chunk_bytes, + connection_data_block_size=network_block_bytes, retry_total=0, + transport=RequestsTransport( + session=http, session_owner=True, + connection_data_block_size=network_block_bytes)) + + +@pytest.mark.parametrize( + 'initial_get_bytes,read_chunk_bytes,network_block_bytes', + [(8 << 20, read, block) for read in (4 << 20, 8 << 20) + for block in (1 << 20, 4 << 20, 8 << 20)] + + [(16 << 20, 16 << 20, 16 << 20), (16 << 20, 4 << 20, 16 << 20)]) +def test_sdk_large_arrow_range_preserves_geometry_and_transport_blocks( + mocked_http, initial_get_bytes, read_chunk_bytes, network_block_bytes): + offset = (2 << 20) + 17 + length = initial_get_bytes + read_chunk_bytes + 98_765 + size = offset + length + 123 + pattern = bytes(range(251)) + content = (pattern * ((size + len(pattern) - 1) // len(pattern)))[:size] + mocked_http.register_uri('HEAD', re.compile(r'https://companylake\.'), + headers={ + **HEADERS, 'Content-Length': str(size) + }) + ranges = [] + + def get(request, context): + assert request.headers['If-Match'] == '"stable"' + assert 'Range' not in request.headers + start, end = map(int, request.headers['x-ms-range'][6:].split('-')) + ranges.append((start, end)) + assert offset <= start <= end < offset + length + context.status_code = 206 + context.headers.update(HEADERS) + context.headers['Content-Range'] = f'bytes {start}-{end}/{size}' + context.headers['Content-Length'] = str(end - start + 1) + return content[start:end + 1] + + mocked_http.register_uri('GET', re.compile(r'https://companylake\.'), + content=get) + with client(mocked_http, network_block_bytes=network_block_bytes, + read_chunk_bytes=read_chunk_bytes, + initial_get_bytes=initial_get_bytes) as service: + filesystem = ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net') + arrow = PyFileSystem(FSSpecHandler(filesystem)) + with arrow.open_input_file(AUTHORITY + '/data') as stream: + assert stream.read_at(length, + offset) == content[offset:offset + length] + + expected = [(offset, offset + initial_get_bytes - 1) + ] + [(start, min(start + read_chunk_bytes, offset + length) - 1) + for start in range(offset + initial_get_bytes, offset + + length, read_chunk_bytes)] + assert sorted(ranges) == expected + for start, end in expected: + actual = [ + length for method, interval, length in mocked_http.fragments + if method == 'GET' and interval == f'bytes={start}-{end}' + ] + full, remainder = divmod(end - start + 1, network_block_bytes) + assert actual == ([network_block_bytes] * full + + ([remainder] if remainder else [])) + + +@pytest.mark.parametrize('failure', ['changed', 'truncated_body']) +@pytest.mark.parametrize( + 'initial_get_bytes,read_chunk_bytes,network_block_bytes', + [(8 << 20, 4 << 20, 1 << 20), (16 << 20, 16 << 20, 16 << 20), + (16 << 20, 4 << 20, 16 << 20)]) +def test_sdk_later_chunk_failure_preserves_destination_and_removes_temporary( + mocked_http, tmp_path, monkeypatch, failure, initial_get_bytes, + read_chunk_bytes, network_block_bytes): + monkeypatch.setattr(transfer, '_SLOTS', threading.BoundedSemaphore(1)) + size = initial_get_bytes + read_chunk_bytes + 113 + mocked_http.register_uri('HEAD', re.compile(r'https://companylake\.'), + headers={ + **HEADERS, 'Content-Length': str(size) + }) + ranges = [] + + def get(request, context): + assert request.headers['If-Match'] == '"stable"' + start, end = map(int, request.headers['x-ms-range'][6:].split('-')) + ranges.append((start, end)) + context.headers.update(HEADERS) + if start and failure == 'changed': + body = b'ConditionNotMetchanged' + context.status_code = 412 + context.headers.update({ + 'Content-Length': str(len(body)), + 'Content-Type': 'application/xml', + 'x-ms-error-code': 'ConditionNotMet' + }) + return body + context.status_code = 206 + context.headers['Content-Range'] = f'bytes {start}-{end}/{size}' + context.headers['Content-Length'] = str(end - start + 1) + # A later body ends early despite its declared full Content-Length. + return b'x' * (end - start + 1 - int(bool(start))) + + mocked_http.register_uri('GET', re.compile(r'https://companylake\.'), + content=get) + destination = tmp_path / 'existing' + destination.write_bytes(b'previous complete file') + with client(mocked_http, initial_get_bytes=initial_get_bytes, + read_chunk_bytes=read_chunk_bytes, + network_block_bytes=network_block_bytes) as service: + filesystem = ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net') + with pytest.raises(SourceError): + filesystem.fetch(AUTHORITY + '/data', str(destination)) + assert ranges[0] == (0, initial_get_bytes - 1) + assert any(start == initial_get_bytes for start, _ in ranges) + assert destination.read_bytes() == b'previous complete file' + assert not list(tmp_path.glob('.kumo-adls-*')) + assert transfer._SLOTS.acquire(blocking=False) + transfer._SLOTS.release() + + +@pytest.mark.parametrize('ceiling', [1, 2, 4]) +def test_sdk_download_obeys_transfer_ceiling_and_preserves_content( + mocked_http, tmp_path, monkeypatch, ceiling): + monkeypatch.setattr(transfer, '_SLOTS', threading.BoundedSemaphore(ceiling)) + size = INITIAL_GET_BYTES + BULK_DOWNLOAD_CHUNK_BYTES + 113 + names = [f'data-{i}' for i in range(4)] + content = { + name: + b''.join( + bytes([i + j]) * min(CHUNK_BYTES, size - offset) + for j, offset in enumerate(range(0, size, CHUNK_BYTES))) + for i, name in enumerate(names) + } + etags = {name: f'"{name}-stable"' for name in names} + + def filename(request): + return urlsplit(request.url).path.rsplit('/', 1)[-1] + + def head(request, context): + context.headers.update(HEADERS) + context.headers['ETag'] = etags[filename(request)] + context.headers['Content-Length'] = str(size) + return b'' + + mocked_http.register_uri('HEAD', re.compile(r'https://companylake\.'), + content=head) + barrier = threading.Barrier(ceiling, timeout=10) + lock = threading.Lock() + ranges = {name: [] for name in names} + active = {name: 0 for name in names} + request_threads = {name: set() for name in names} + fetch_threads = {} + maximum = 0 + + def get(request, context): + nonlocal maximum + name = filename(request) + assert request.headers['If-Match'] == etags[name] + start, end = map(int, request.headers['x-ms-range'][6:].split('-')) + with lock: + active[name] += 1 + assert active[name] == 1 + assert sum(active.values()) <= ceiling + maximum = max(maximum, sum(active.values())) + ranges[name].append((start, end)) + request_threads[name].add(threading.get_ident()) + + def closed(): + with lock: + active[name] -= 1 + + context.on_stream_closed = closed + if start == 0: + # Prove cross-file overlap while every file remains serial. + barrier.wait() + context.status_code = 206 + context.headers.update(HEADERS) + context.headers['ETag'] = etags[name] + context.headers['Content-Range'] = f'bytes {start}-{end}/{size}' + context.headers['Content-Length'] = str(end - start + 1) + return content[name][start:end + 1] + + mocked_http.register_uri('GET', re.compile(r'https://companylake\.'), + content=get) + with client(mocked_http) as service: + filesystem = ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net') + + def fetch(name): + fetch_threads[name] = threading.get_ident() + filesystem.fetch(AUTHORITY + '/' + name, str(tmp_path / name)) + + with ThreadPoolExecutor(max_workers=len(names)) as pool: + list(pool.map(fetch, names)) + expected = [(0, INITIAL_GET_BYTES - 1)] + [ + (start, min(start + BULK_DOWNLOAD_CHUNK_BYTES, size) - 1) + for start in range(INITIAL_GET_BYTES, size, BULK_DOWNLOAD_CHUNK_BYTES) + ] + for name in names: + assert (tmp_path / name).read_bytes() == content[name] + assert ranges[name] == expected + assert request_threads[name] == {fetch_threads[name]} + assert active[name] == 0 + assert maximum == ceiling + with transfer.concurrency(size, ranged=True) as workers: + assert workers == 1 + + +def test_sdk_huge_serial_download_fails_without_scheduling_chunk_futures( + mocked_http, tmp_path, monkeypatch): + import concurrent.futures + + size = 100 << 30 + mocked_http.register_uri('HEAD', re.compile(r'https://companylake\.'), + headers={ + **HEADERS, 'Content-Length': str(size) + }) + ranges = [] + + def get(request, context): + assert request.headers['If-Match'] == '"stable"' + start, end = map(int, request.headers['x-ms-range'][6:].split('-')) + ranges.append((start, end)) + context.headers.update(HEADERS) + if start: + body = b'ConditionNotMetchanged' + context.status_code = 412 + context.headers.update({ + 'Content-Length': str(len(body)), + 'Content-Type': 'application/xml', + 'x-ms-error-code': 'ConditionNotMet' + }) + return body + assert end == INITIAL_GET_BYTES - 1 + context.status_code = 206 + context.headers['Content-Range'] = f'bytes {start}-{end}/{size}' + context.headers['Content-Length'] = str(end - start + 1) + return b'x' * INITIAL_GET_BYTES + + def unexpected_pool(*args, **kwargs): + pytest.fail('Serial SDK download must not create a chunk executor') + + monkeypatch.setattr(concurrent.futures, 'ThreadPoolExecutor', + unexpected_pool) + mocked_http.register_uri('GET', re.compile(r'https://companylake\.'), + content=get) + destination = tmp_path / 'existing' + destination.write_bytes(b'previous complete file') + with client(mocked_http) as service: + filesystem = ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net') + with pytest.raises(SourceError): + filesystem.fetch(AUTHORITY + '/huge', str(destination)) + assert ranges == [(0, INITIAL_GET_BYTES - 1), + (INITIAL_GET_BYTES, + INITIAL_GET_BYTES + BULK_DOWNLOAD_CHUNK_BYTES - 1)] + assert destination.read_bytes() == b'previous complete file' + assert not list(tmp_path.glob('.kumo-adls-*')) + + +def test_sdk_seekable_upload_overlaps_chunks_and_flushes_after_all_appends( + mocked_http, tmp_path): + source = tmp_path / 'upload' + source.write_bytes(b''.join(bytes([i]) * CHUNK_BYTES for i in range(4))) + barrier = threading.Barrier(4, timeout=10) + lock = threading.Lock() + appended = {} + flushed = [] + mocked_http.register_uri('PUT', re.compile(r'https://companylake\.'), + status_code=201, headers=HEADERS) + + def patch(request, context): + query = parse_qs(urlsplit(request.url).query) + action = query['action'][0] + if action == 'append': + offset = int(query['position'][0]) + data = request.body + if hasattr(data, 'read'): + data = data.read() + barrier.wait() + with lock: + appended[offset] = data + elif action == 'flush': + assert len(appended) == 4 + flushed.append(int(query['position'][0])) + else: + raise AssertionError(action) + context.status_code = 202 if action == 'append' else 200 + context.headers.update(HEADERS) + return b'' + + mocked_http.register_uri('PATCH', re.compile(r'https://companylake\.'), + content=patch) + with client(mocked_http) as service: + filesystem = ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net') + filesystem.put_file(str(source), AUTHORITY + '/data') + assert b''.join(appended[offset] + for offset in sorted(appended)) == source.read_bytes() + assert flushed == [source.stat().st_size] + + +def test_sdk_empty_download_handles_azure_range_416(mocked_http, tmp_path): + mocked_http.register_uri('HEAD', re.compile(r'https://companylake\.'), + headers={ + **HEADERS, 'Content-Length': '0' + }) + attempts = [] + + def get(request, context): + attempts.append(request.headers.get('x-ms-range')) + context.headers.update(HEADERS) + context.headers['Content-Length'] = '0' + assert request.headers['If-Match'] == '"stable"' + if request.headers.get('x-ms-range'): + context.status_code = 416 + context.headers['x-ms-error-code'] = 'InvalidRange' + context.headers['Content-Range'] = 'bytes */0' + else: + context.status_code = 200 + return b'' + + mocked_http.register_uri('GET', re.compile(r'https://companylake\.'), + content=get) + with client(mocked_http) as service: + filesystem = ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net') + destination = tmp_path / 'empty' + filesystem.fetch(AUTHORITY + '/empty', str(destination)) + assert destination.read_bytes() == b'' + assert attempts[-1] is None diff --git a/packages/kumo-adls/tests/test_session.py b/packages/kumo-adls/tests/test_session.py new file mode 100644 index 000000000..36a91888f --- /dev/null +++ b/packages/kumo-adls/tests/test_session.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import ANY, Mock + +import pytest +from azure.identity import AzureAuthorityHosts +from kumo_adls import session +from kumo_core.sources.principals import Principal +from kumoapi.data_source import ADLSIdentityConfig, AzureCloud + + +@pytest.fixture(autouse=True) +def clean_sessions(): + session.close() + yield + session.close() + + +CLOUD_CASES = [ + (AzureCloud.PUBLIC, 'dfs.core.windows.net', + AzureAuthorityHosts.AZURE_PUBLIC_CLOUD), + (AzureCloud.US_GOVERNMENT, 'dfs.core.usgovcloudapi.net', + AzureAuthorityHosts.AZURE_GOVERNMENT), + (AzureCloud.CHINA, 'dfs.core.chinacloudapi.cn', + AzureAuthorityHosts.AZURE_CHINA), +] + + +@pytest.mark.parametrize('cloud, suffix, authority', CLOUD_CASES) +def test_explicit_workload_identity_and_cloud_match(monkeypatch, cloud, suffix, + authority): + principal = Principal( + 'test', + ADLSIdentityConfig(tenant_id='tenant', client_id='client', cloud=cloud), + None) + factory = Mock() + monkeypatch.setattr(session, 'WorkloadIdentityCredential', factory) + credential, scope = session.credential_for(principal, suffix) + assert credential is factory.return_value + assert scope == 'https://storage.azure.com' + factory.assert_called_once_with(tenant_id='tenant', client_id='client', + authority=authority) + other_suffix = next(value[0] for key, value in session.CLOUDS.items() + if key != cloud) + with pytest.raises(ValueError, match='cloud'): + session.credential_for(principal, other_suffix) + + +def test_ambient_identity_session_reuse_and_cleanup(monkeypatch): + credential = Mock() + credential_factory = Mock(return_value=credential) + factory = Mock() + monkeypatch.setattr(session, 'DefaultAzureCredential', credential_factory) + monkeypatch.setattr(session, 'DataLakeServiceClient', factory) + first = session.session_for(None, 'companylake', 'dfs.core.windows.net') + assert session.session_for(None, 'companylake', + 'dfs.core.windows.net') is first + credential_factory.assert_called_once() + factory.assert_called_once_with( + account_url='https://companylake.dfs.core.windows.net', + credential=credential, audience='https://storage.azure.com', + max_single_get_size=16 << 20, max_chunk_get_size=16 << 20, + connection_data_block_size=16 << 20, transport=ANY) + session.close() + factory.return_value.close.assert_called_once() + credential.close.assert_called_once() + + +def test_close_one_identity_preserves_other_sessions(monkeypatch): + factory = Mock(side_effect=lambda **kwargs: Mock()) + credentials = Mock(side_effect=lambda **kwargs: Mock()) + monkeypatch.setattr(session, 'WorkloadIdentityCredential', credentials) + monkeypatch.setattr(session, 'DataLakeServiceClient', factory) + config = ADLSIdentityConfig(tenant_id='tenant', client_id='client', + cloud=AzureCloud.PUBLIC) + first = Principal('first', config, None) + second = Principal('second', config, None) + client = session.session_for(first, 'companylake', 'dfs.core.windows.net') + other = session.session_for(second, 'companylake', 'dfs.core.windows.net') + session.close(identity_name=first.name) + client.close.assert_called_once() + other.close.assert_not_called() + + +@pytest.mark.parametrize('cloud, suffix, authority', CLOUD_CASES) +def test_probe_acquires_token_without_touching_storage(monkeypatch, cloud, + suffix, authority): + credential = Mock() + credential.__enter__ = Mock(return_value=credential) + credential.__exit__ = Mock() + monkeypatch.setattr(session, 'WorkloadIdentityCredential', + Mock(return_value=credential)) + principal = Principal( + 'test', + ADLSIdentityConfig(tenant_id='tenant', client_id='client', cloud=cloud), + None) + session.probe(principal) + credential.get_token.assert_called_once_with( + 'https://storage.azure.com/.default') + + +@pytest.mark.parametrize('cloud, suffix, authority', CLOUD_CASES) +def test_ambient_session_uses_common_storage_audience(monkeypatch, cloud, + suffix, authority): + credential = Mock() + credential_factory = Mock(return_value=credential) + client_factory = Mock() + monkeypatch.setattr(session, 'DefaultAzureCredential', credential_factory) + monkeypatch.setattr(session, 'DataLakeServiceClient', client_factory) + + session.session_for(None, 'companylake', suffix) + + credential_factory.assert_called_once_with(authority=authority) + client_factory.assert_called_once_with( + account_url=f'https://companylake.{suffix}', credential=credential, + audience='https://storage.azure.com', max_single_get_size=16 << 20, + max_chunk_get_size=16 << 20, connection_data_block_size=16 << 20, + transport=ANY) + + +@pytest.mark.parametrize('failing_resource', ['client', 'credential']) +def test_cleanup_failure_does_not_leave_other_sessions_open( + monkeypatch, failing_resource): + clients = [Mock(), Mock()] + credentials = [Mock(), Mock()] + factory = Mock(side_effect=clients) + credential_factory = Mock(side_effect=credentials) + monkeypatch.setattr(session, 'DefaultAzureCredential', credential_factory) + monkeypatch.setattr(session, 'DataLakeServiceClient', factory) + session.session_for(None, 'firstlake', 'dfs.core.windows.net') + session.session_for(None, 'otherlake', 'dfs.core.windows.net') + resource = clients[0] if failing_resource == 'client' else credentials[0] + resource.close.side_effect = RuntimeError('already closed') + + session.close() + + for item in clients + credentials: + item.close.assert_called_once() + assert not session._SESSIONS diff --git a/packages/kumo-adls/tests/test_transfer.py b/packages/kumo-adls/tests/test_transfer.py new file mode 100644 index 000000000..7291520b9 --- /dev/null +++ b/packages/kumo-adls/tests/test_transfer.py @@ -0,0 +1,534 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import threading +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from kumo_adls import transfer +from kumo_adls.sdk_filesystem import ADLSFileSystem +from pyarrow.fs import FileSelector, FSSpecHandler, PyFileSystem + +AUTHORITY = 'curated@companylake.dfs.core.windows.net' + + +@pytest.mark.parametrize('adls_limit, copy_limit, expected', [ + (None, None, 128), + (None, '4', 4), + ('4', None, 4), + ('16', '4', 4), + ('4', '16', 4), + ('1', '128', 1), + ('256', None, 128), +]) +def test_transfer_ceiling_is_provider_specific_and_respects_copy_limit( + monkeypatch, adls_limit, copy_limit, expected): + from kumo_core.sources import cpu + + monkeypatch.setattr(cpu, 'cpu_budget', lambda: 32) + for name, value in [(transfer.TRANSFER_WORKERS_ENV, adls_limit), + (cpu.COPY_WORKERS_ENV, copy_limit)]: + if value is None: + monkeypatch.delenv(name, raising=False) + else: + monkeypatch.setenv(name, value) + assert transfer._transfer_workers() == expected + # The ADLS ceiling does not change the shared pool used by other providers. + assert cpu.copy_workers() == (int(copy_limit) if copy_limit else 128) + + +@pytest.mark.parametrize('value', ['', '0', '-1', '1.5', 'four']) +def test_invalid_transfer_ceiling_is_rejected(monkeypatch, value): + monkeypatch.setenv(transfer.TRANSFER_WORKERS_ENV, value) + with pytest.raises(ValueError, match='must be a positive integer'): + transfer._transfer_workers() + + +def test_arrow_projection_does_not_read_unrequested_bytes(service): + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + path = AUTHORITY + '/orders/events.parquet' + with PyFileSystem( + FSSpecHandler(filesystem)).open_input_file(path) as stream: + stream.seek(4) + assert len(stream.read(8)) == 8 + assert [(offset, length) + for _, offset, length, _ in service.downloads] == [(4, 8)] + + +def test_recursive_arrow_listing_uses_one_sdk_traversal(service, monkeypatch): + from fakes import FakeFilesystem + + for i in range(20): + directory = f'curated/orders/year={i}' + service.directories.add(directory) + service.files[directory + '/part.parquet'] = b'PAR1' + original = FakeFilesystem.get_paths + calls = [] + + def paths(self, **kwargs): + calls.append(kwargs) + return original(self, **kwargs) + + monkeypatch.setattr(FakeFilesystem, 'get_paths', paths) + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + arrow = PyFileSystem(FSSpecHandler(filesystem)) + entries = arrow.get_file_info( + FileSelector(AUTHORITY + '/orders', recursive=True)) + assert len(entries) == 41 + assert calls == [{'path': 'orders', 'recursive': True}] + assert len(filesystem.find(AUTHORITY + '/orders', maxdepth=1)) == 1 + assert len(filesystem.find(AUTHORITY + '/orders', withdirs=True)) == 42 + + +def test_request_budget_is_shared_and_released_after_failure(monkeypatch): + monkeypatch.setattr(transfer, '_SLOTS', threading.BoundedSemaphore(4)) + entered = threading.Event() + release = threading.Event() + + def hold(): + with transfer.concurrency(100 << 20) as count: + assert count == 4 + entered.set() + assert release.wait(5) + + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(hold) + assert entered.wait(5) + assert not transfer._SLOTS.acquire(blocking=False) + release.set() + future.result() + with pytest.raises(RuntimeError): + with transfer.concurrency(100 << 20) as count: + assert count == 4 + raise RuntimeError('transfer failed') + with transfer.concurrency(1) as count: + assert count == 1 + with transfer.concurrency(100 << 20) as count: + assert count == 4 + + +@pytest.mark.parametrize('bulk', [True, False]) +def test_download_memory_failure_has_guidance_and_releases_resources( + service, tmp_path, monkeypatch, bulk): + from fakes import FakeFile + + def download(*args, **kwargs): + raise MemoryError + + monkeypatch.setattr(FakeFile, 'download_file', download) + monkeypatch.setattr(transfer, '_SLOTS', threading.BoundedSemaphore(4)) + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + destination = tmp_path / 'existing' + destination.write_bytes(b'previous') + path = AUTHORITY + '/orders/events.parquet' + with pytest.raises(MemoryError) as result: + if bulk: + filesystem.fetch(path, str(destination)) + else: + with filesystem.open(path, 'rb', cache_type='none') as stream: + stream.read(4) + assert 'KUMO_ADLS_TRANSFER_WORKERS' in '\n'.join(result.value.__notes__) + assert destination.read_bytes() == b'previous' + assert sorted(p.name for p in tmp_path.iterdir()) == ['existing'] + with transfer.concurrency(100 << 20, upload=True) as workers: + assert workers == 4 + + +def test_failed_fetch_keeps_previous_destination(service, tmp_path, + monkeypatch): + from fakes import FakeFile + from kumo_core.sources.errors import TransientSourceError + + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + target = tmp_path / 'existing' + target.write_bytes(b'previous') + monkeypatch.setattr( + FakeFile, 'download_file', lambda *a, **k: SimpleNamespace( + readinto=lambda stream: stream.write(b'bad'))) + with pytest.raises(TransientSourceError, match='truncated'): + filesystem.fetch(AUTHORITY + '/orders/events.parquet', str(target)) + assert target.read_bytes() == b'previous' + assert not list(tmp_path.glob('.kumo-adls-*')) + + +@pytest.mark.parametrize('failure', [None, 'download', 'promotion']) +def test_fetch_closes_temporary_file_before_promotion_and_cleanup( + service, tmp_path, monkeypatch, failure): + import os + import stat + + from fakes import FakeFile + from kumo_adls import sdk_filesystem + + opened = [] + make_temp = sdk_filesystem.NamedTemporaryFile + replace = os.replace + unlink = os.unlink + + def temporary(*args, **kwargs): + stream = make_temp(*args, **kwargs) + opened.append(stream) + return stream + + def promote(source, destination): + assert opened[0].closed + if failure == 'promotion': + raise OSError('promotion failed') + replace(source, destination) + + def remove(path, *args, **kwargs): + if opened and path == opened[0].name: + assert opened[0].closed + return unlink(path, *args, **kwargs) + + def download(*args, **kwargs): + raise OSError('download failed') + + monkeypatch.setattr(sdk_filesystem, 'NamedTemporaryFile', temporary) + monkeypatch.setattr(os, 'replace', promote) + monkeypatch.setattr(os, 'unlink', remove) + if failure == 'download': + monkeypatch.setattr(FakeFile, 'download_file', download) + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + destination = tmp_path / 'existing' + destination.write_bytes(b'previous') + if failure: + with pytest.raises(OSError, match=failure + ' failed'): + filesystem.fetch(AUTHORITY + '/orders/events.parquet', + str(destination)) + assert destination.read_bytes() == b'previous' + else: + filesystem.fetch(AUTHORITY + '/orders/events.parquet', str(destination)) + expected = service.files['curated/orders/events.parquet'] + assert destination.read_bytes() == expected + if os.name == 'posix': + assert stat.S_IMODE(destination.stat().st_mode) & 0o077 == 0 + assert opened[0].closed + assert not list(tmp_path.glob('.kumo-adls-*')) + + +def test_file_export_uses_seekable_sdk_upload_and_keeps_prepared_columns( + service, connector, tmp_path, monkeypatch): + import pyarrow as pa + import pyarrow.parquet as pq + from fakes import FakeFile + from kumo_adls.sink import ADLSSinkBackend + + source = tmp_path / 'source.parquet' + pq.write_table(pa.table({'id': [1, 2]}), source) + original = FakeFile.upload_data + uploads = [] + + def upload(self, stream, **kwargs): + assert stream.seekable() + uploads.append(self.name) + return original(self, stream, **kwargs) + + monkeypatch.setattr(FakeFile, 'upload_data', upload) + ADLSSinkBackend().write_files(connector(), 'predictions/', [str(source)], + extra_columns={'tag': + 'test'}, export_id='run') + assert len(uploads) == 1 + table = pq.read_table(pa.BufferReader(service.files[uploads[0]])) + assert table.to_pydict() == {'id': [1, 2], 'tag': ['test', 'test']} + + +def test_owned_http_pool_is_reused_and_closed(monkeypatch): + from kumo_adls import session + + session.close() + credential = Mock() + monkeypatch.setattr(session, 'DefaultAzureCredential', + Mock(return_value=credential)) + transport_session = Mock() + monkeypatch.setattr(session, 'Session', + Mock(return_value=transport_session)) + # Real SDK clients exercise ownership and child-client propagation. + original_transport = session.RequestsTransport + transports = [] + + def transport(**kwargs): + item = original_transport(**kwargs) + transports.append(item) + return item + + monkeypatch.setattr(session, 'RequestsTransport', transport) + client = session.session_for(None, 'companylake', 'dfs.core.windows.net') + assert session.session_for(None, 'companylake', + 'dfs.core.windows.net') is client + assert transports[ + 0].connection_config.data_block_size == transfer.NETWORK_BLOCK_BYTES + assert transport_session.mount.call_count == 2 + adapter = transport_session.mount.call_args.args[1] + assert adapter._pool_maxsize == transfer.CONNECTIONS + assert adapter._pool_block + session.close() + transport_session.close.assert_called_once() + credential.close.assert_called_once() + + +def test_empty_local_file_creates_and_truncates_remote_file(service, tmp_path): + source = tmp_path / 'empty' + source.touch() + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + for previous in (None, b'old data'): + service.files.pop('curated/empty', None) + if previous is not None: + service.files['curated/empty'] = previous + filesystem.put_file(str(source), AUTHORITY + '/empty') + assert service.files['curated/empty'] == b'' + + +def test_upload_parallelism_uses_all_chunks_without_download_initial_request( + monkeypatch): + monkeypatch.setattr(transfer, '_SLOTS', threading.BoundedSemaphore(4)) + with transfer.concurrency(12 << 20, upload=True) as workers: + assert workers == 3 + + +def test_failed_export_does_not_delete_another_writers_replacement( + service, monkeypatch, tmp_path): + from azure.core.exceptions import ResourceModifiedError + from fakes import FakeFile + from kumo_core.sources.errors import TransientSourceError + + def replaced(self, *args, **kwargs): + self.service.files[self.name] = b'other writer completed' + raise ResourceModifiedError('replaced during upload') + + monkeypatch.setattr(FakeFile, 'upload_data', replaced) + source = tmp_path / 'source' + source.write_bytes(b'new contents') + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + with pytest.raises(TransientSourceError): + filesystem.put_file(str(source), AUTHORITY + '/shared') + assert service.files['curated/shared'] == b'other writer completed' + + +def test_upload_reserves_and_releases_memory_on_failure(service, tmp_path, + monkeypatch): + from fakes import FakeFile + from kumo_core.sources import cpu + + before = cpu._reserved + + def fail(self, *args, **kwargs): + assert cpu._reserved >= before + cpu.COPY_WORKER_BYTES + raise RuntimeError('upload failed') + + monkeypatch.setattr(FakeFile, 'upload_data', fail) + source = tmp_path / 'source' + source.write_bytes(b'content') + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + with pytest.raises(RuntimeError, match='upload failed'): + filesystem.put_file(str(source), AUTHORITY + '/failed') + assert cpu._reserved == before + assert 'curated/failed' not in service.files + + +def test_version_checks_overlap_and_remain_independent_of_listing_order(): + from kumo_adls.backend import ADLSSourceReader + + barrier = threading.Barrier(2, timeout=5) + + def identity(path): + barrier.wait() + return ('etag:' + path, ) + + entries = [SimpleNamespace(path=f'object-{i}') for i in range(4)] + reader = ADLSSourceReader(None, object_identity=identity) + assert reader._data_version('table', entries) == reader._data_version( + 'table', entries[::-1]) + + +def test_model_download_uses_conditional_bulk_fetch(service, connector, + tmp_path): + from kumo_adls.sink import ADLSSinkBackend + + service.directories.add('curated/orders/bundle') + service.files['curated/orders/bundle/model.bin'] = b'model contents' + target = tmp_path / 'model' + ADLSSinkBackend().fetch_model(connector(), str(target), path='bundle') + assert (target / 'model.bin').read_bytes() == b'model contents' + _, offset, length, options = service.downloads[-1] + assert offset == 0 and length == len(b'model contents') + assert options['max_concurrency'] == 1 + + +def test_filesystem_client_cache_is_bounded_and_does_not_cache_metadata( + service, monkeypatch): + original = service.get_file_system_client + factory = Mock(side_effect=original) + monkeypatch.setattr(service, 'get_file_system_client', factory) + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + path = AUTHORITY + '/orders/events.parquet' + first = filesystem.object_identity(path) + service.files['curated/orders/events.parquet'] = b'rewritten' + assert filesystem.object_identity(path) != first + factory.assert_called_once_with('curated') + for i in range(40): + filesystem._address( + f'filesystem-{i}@companylake.dfs.core.windows.net/path') + assert len(filesystem._filesystems) == 32 + + +def test_large_fetch_uses_one_sdk_downloader_and_cleans_failure( + service, tmp_path, monkeypatch): + from azure.core import MatchConditions + from fakes import FakeFile + from kumo_core.sources.errors import TransientSourceError + + original = FakeFile.get_file_properties + calls = [] + + def properties(self): + result = original(self) + result.size = 100 << 30 + return result + + def download(self, **kwargs): + calls.append(kwargs) + assert kwargs == dict(offset=0, length=100 << 30, + etag=original(self).etag, + match_condition=MatchConditions.IfNotModified, + max_concurrency=1) + raise TransientSourceError('later download failed') + + monkeypatch.setattr(FakeFile, 'get_file_properties', properties) + monkeypatch.setattr(FakeFile, 'download_file', download) + monkeypatch.setattr(transfer, '_SLOTS', threading.BoundedSemaphore(4)) + destination = tmp_path / 'existing' + destination.write_bytes(b'previous') + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + with pytest.raises(TransientSourceError, match='later download failed'): + filesystem.fetch(AUTHORITY + '/orders/events.parquet', str(destination)) + assert len(calls) == 1 + assert destination.read_bytes() == b'previous' + assert not list(tmp_path.glob('.kumo-adls-*')) + with transfer.concurrency(100 << 20, upload=True) as workers: + assert workers == 4 + + +def test_empty_fetch_omits_range_arguments(service, tmp_path, monkeypatch): + from azure.core import MatchConditions + from fakes import FakeFile + + service.files['curated/empty'] = b'' + calls = [] + + def download(self, **kwargs): + calls.append(kwargs) + assert 'offset' not in kwargs and 'length' not in kwargs + assert kwargs['match_condition'] == MatchConditions.IfNotModified + assert kwargs['max_concurrency'] == 1 + return SimpleNamespace(readinto=lambda target: 0) + + monkeypatch.setattr(FakeFile, 'download_file', download) + destination = tmp_path / 'existing' + destination.write_bytes(b'previous') + ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net').fetch(AUTHORITY + '/empty', + str(destination)) + assert destination.read_bytes() == b'' + assert len(calls) == 1 + + +def test_serial_fetch_preserves_callers_context(service, tmp_path, monkeypatch): + import contextvars + + from azure.core import MatchConditions + from fakes import FakeFile + + tracing_context = contextvars.ContextVar('adls-test-span') + marker = object() + token = tracing_context.set(marker) + contents = service.files['curated/orders/events.parquet'] + calls = [] + original_properties = FakeFile.get_file_properties + + def download(self, **kwargs): + assert tracing_context.get() is marker + assert kwargs['etag'] == original_properties(self).etag + assert kwargs['match_condition'] == MatchConditions.IfNotModified + assert kwargs['max_concurrency'] == 1 + calls.append(kwargs) + + def readinto(target): + assert tracing_context.get() is marker + return target.write(contents) + + return SimpleNamespace(readinto=readinto) + + monkeypatch.setattr(FakeFile, 'download_file', download) + destination = tmp_path / 'download' + try: + ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net').fetch( + AUTHORITY + '/orders/events.parquet', str(destination)) + finally: + tracing_context.reset(token) + assert destination.read_bytes() == contents + assert len(calls) == 1 + + +@pytest.mark.parametrize('size', [0, 1, 16 << 20, (16 << 20) + 1, 100 << 30]) +def test_bulk_download_uses_one_slot_and_releases_on_failure(monkeypatch, size): + monkeypatch.setattr(transfer, '_SLOTS', threading.BoundedSemaphore(4)) + with pytest.raises(RuntimeError, match='failed'): + with transfer.concurrency(size, ranged=True) as workers: + assert workers == 1 + raise RuntimeError('failed') + with transfer.concurrency(100 << 20, upload=True) as workers: + assert workers == 4 + + +@pytest.mark.parametrize('size, expected', [(0, 1), (16 << 20, 1), + ((16 << 20) + 1, 1), (32 << 20, 1), + ((32 << 20) + 1, 2), (48 << 20, 2), + ((48 << 20) + 1, 3), (80 << 20, 4), + (64 << 20, 3), (100 << 30, 4)]) +def test_projection_download_counts_remaining_sdk_chunks( + monkeypatch, size, expected): + monkeypatch.setattr(transfer, '_SLOTS', threading.BoundedSemaphore(4)) + with transfer.concurrency(size) as workers: + assert workers == expected + with transfer.concurrency(100 << 20, upload=True) as workers: + assert workers == 4 + + +def test_bulk_tuning_preserves_upload_chunking_and_parallelism(monkeypatch): + monkeypatch.setattr(transfer, '_SLOTS', threading.BoundedSemaphore(4)) + assert transfer.INITIAL_GET_BYTES == transfer.BULK_DOWNLOAD_CHUNK_BYTES == 16 << 20 + assert transfer.NETWORK_BLOCK_BYTES == 16 << 20 + assert transfer.MAX_BULK_DOWNLOAD_CONCURRENCY == 1 + assert transfer.CHUNK_BYTES == 4 << 20 + with transfer.concurrency(12 << 20, upload=True) as workers: + assert workers == 3 + + +@pytest.mark.parametrize('short_write, wrong_report', [(True, False), + (False, True), + (True, True)]) +def test_fetch_checks_actual_position_and_reported_length( + service, tmp_path, monkeypatch, short_write, wrong_report): + from fakes import FakeFile + from kumo_core.sources.errors import TransientSourceError + + def download(self, **kwargs): + def readinto(target): + length = kwargs['length'] + target.write(b'x' * (length - 1 if short_write else length)) + return length - 1 if wrong_report else length + + return SimpleNamespace(readinto=readinto) + + monkeypatch.setattr(FakeFile, 'download_file', download) + destination = tmp_path / 'existing' + destination.write_bytes(b'previous') + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + with pytest.raises(TransientSourceError, match='truncated'): + filesystem.fetch(AUTHORITY + '/orders/events.parquet', str(destination)) + assert destination.read_bytes() == b'previous' + assert not list(tmp_path.glob('.kumo-adls-*')) diff --git a/packages/kumo-api/kumoapi/adls.py b/packages/kumo-api/kumoapi/adls.py index 5e85c5632..fd4f1366e 100644 --- a/packages/kumo-api/kumoapi/adls.py +++ b/packages/kumo-api/kumoapi/adls.py @@ -54,6 +54,15 @@ def normalize_adls_source_relative_path(value: str) -> str: def normalize_adls_source_uri(value: str) -> str: """Canonicalize an exact hierarchical-namespace ADLS source URI.""" + return _normalize_adls_uri(value, allow_root=False) + + +def normalize_adls_root_uri(value: str) -> str: + """Canonicalize a connector root, including a filesystem without a prefix.""" + return _normalize_adls_uri(value, allow_root=True) + + +def _normalize_adls_uri(value: str, *, allow_root: bool) -> str: if not isinstance(value, str) or not value or value != value.strip(): raise ValueError( 'ADLS URI must be non-empty without surrounding whitespace') @@ -74,6 +83,8 @@ def normalize_adls_source_uri(value: str) -> str: account = host[:-(len(matching_suffix) + 1)] filesystem = normalize_adls_filesystem(filesystem) account = normalize_adls_storage_account(account) + if allow_root and parsed.path in ('', '/'): + return f'abfss://{filesystem}@{account}.{matching_suffix}/' if not parsed.path.startswith('/') or parsed.path.startswith('//'): raise ValueError( 'ADLS URI must contain exactly one filesystem/path separator') diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index ec34f89f9..efab01dba 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -5,6 +5,7 @@ import posixpath import shutil from abc import ABC, abstractmethod +from functools import partial from pathlib import Path from typing import Mapping, Optional, Tuple @@ -31,6 +32,9 @@ def _destination(self, connector: SourceConnector, """Return the authenticated filesystem and validated bucket/key path.""" raise NotImplementedError + def _uri(self, path: str) -> str: + return f'{self.scheme}://{path}' + def write_table( self, connector: SourceConnector, @@ -41,7 +45,11 @@ def write_table( *, export_id: Optional[str] = None, ) -> str: - """Write a frame as one object, or parts when the path ends in a slash.""" + """Write a frame as one object, or parts when the path ends in a slash. + + Streaming output is not transactional: Arrow closes its output stream + after serialization errors, which can publish partial data. + """ del project_id if write_mode is not WriteMode.OVERWRITE: raise NotImplementedError( @@ -65,12 +73,14 @@ def write_files(self, connector, table_name, files, from kumo_core.sources import layout filesystem, path = self._destination(connector, table_name) kind = layout.file_type(path.rstrip('/')) or layout.PARQUET - storage = _ArrowExportFiles(filesystem, extra_columns) + storage = _ArrowExportFiles( + filesystem, extra_columns, upload=partial(self._upload_file, + filesystem)) _, output = export_selected_files(files, path, fs=storage, identity=export_id or str(files), file_type=kind, convert_csv=storage.convert_csv) - return f'{self.scheme}://{output}' + return self._uri(output) def publish_model( self, @@ -103,16 +113,26 @@ def publish_model( if item.is_file()): target = f'{destination}/{source.relative_to(root).as_posix()}' _prepared(filesystem, posixpath.dirname(target)) - with open(source, 'rb') as reader, \ - filesystem.open_output_stream(target) as writer: - shutil.copyfileobj(reader, writer, _CHUNK) + self._upload_file(filesystem, str(source), target) return PublishedModel( name=posixpath.basename(destination), version=destination, - uri=f'{self.scheme}://{destination}', + uri=self._uri(destination), metadata=dict(metadata), ) + def _upload_file(self, filesystem, source, destination): + """Providers may transfer an already prepared local file directly.""" + with open(source, 'rb') as reader, \ + filesystem.open_output_stream(destination) as writer: + shutil.copyfileobj(reader, writer, _CHUNK) + + def _download_file(self, filesystem, source, destination): + """Providers may transfer a model file without the streaming fallback.""" + with filesystem.open_input_stream(source) as reader, \ + open(destination, 'wb') as writer: + shutil.copyfileobj(reader, writer, _CHUNK) + def fetch_model( self, connector: SourceConnector, @@ -156,15 +176,13 @@ def fetch_model( # the local bundle, even when it contains parent components. target = destination / confined(info.path[len(source) + 1:]) target.parent.mkdir(parents=True, exist_ok=True) - with filesystem.open_input_stream(info.path) as reader, \ - open(target, 'wb') as writer: - shutil.copyfileobj(reader, writer, _CHUNK) + self._download_file(filesystem, info.path, str(target)) if not found: raise ValueError( f"Nothing is published at {self.scheme}://{source}. Check the " f"path the model was exported to.") return PublishedModel(name=posixpath.basename(source), version=source, - uri=f'{self.scheme}://{source}') + uri=self._uri(source)) def _write(self, filesystem: FileSystem, frame: pd.DataFrame, path: str, export_id: Optional[str] = None) -> str: @@ -186,12 +204,12 @@ def _write(self, filesystem: FileSystem, frame: pd.DataFrame, path: str, kind) with filesystem.open_output_stream(f'{directory}/_SUCCESS'): pass - return f'{self.scheme}://{directory}/' + return self._uri(f'{directory}/') if not path.casefold().endswith(f'.{kind}'): path += f'.{kind}' _prepared(filesystem, posixpath.dirname(path)) _write_one(filesystem, frame, path, kind) - return f'{self.scheme}://{path}' + return self._uri(path) def _occupied(self, filesystem: FileSystem, directory: str, marker: Optional[str]) -> bool: @@ -259,9 +277,10 @@ def _marker(export_id: Optional[str]) -> Optional[str]: class _ArrowExportFiles: """Adapt authenticated Arrow output to the shared hosted export operation.""" - def __init__(self, filesystem, extra_columns): + def __init__(self, filesystem, extra_columns, upload=None): self.filesystem = filesystem self.extra_columns = extra_columns + self.upload = upload def exists(self, path): from pyarrow.fs import FileType @@ -309,6 +328,9 @@ def copy_file(self, source, destination, overwrite=True): # CSV conversion already applied this destination's metadata. columns = self.extra_columns if source.endswith('.parquet') else None with prepared_artifact_file(source, columns) as prepared: + if self.upload is not None: + self.upload(prepared, destination) + return with open(prepared, 'rb') as reader, self.filesystem.open_output_stream( destination) as writer: diff --git a/packages/kumo-core/src/kumo_core/sources/object_store.py b/packages/kumo-core/src/kumo_core/sources/object_store.py index 4b7ac0a37..a88d24c82 100644 --- a/packages/kumo-core/src/kumo_core/sources/object_store.py +++ b/packages/kumo-core/src/kumo_core/sources/object_store.py @@ -159,7 +159,10 @@ def data_version(self, table: str, *, sql: str | None = None, exact=self.file_format is not None or not table) entries = ([info] if info.type is FileType.File else self._objects(table, info.path, include_hidden=True)) - return _source_version(info.path, entries) + return self._data_version(info.path, entries) + + def _data_version(self, path: str, entries: Sequence[Any]) -> str: + return _source_version(path, entries) def resolve(self, table: str) -> tuple[str, str]: """An explicit format addresses an exact path; discovery infers both.""" @@ -325,7 +328,7 @@ def prepare( include_hidden=not transform) # Stamp the same metadata listing the read consumes, before fetching # any files; later additions must not be stamped onto this output. - stamped = version or _source_version(path, found) + stamped = version or self._data_version(path, found) source_bytes = (sum(entry.size for entry in found) if all( entry.size is not None and entry.size >= 0 for entry in found) else None) diff --git a/packages/kumo-core/tests/sources/package_boundary.py b/packages/kumo-core/tests/sources/package_boundary.py index 530e3de0e..6f1a88cf2 100644 --- a/packages/kumo-core/tests/sources/package_boundary.py +++ b/packages/kumo-core/tests/sources/package_boundary.py @@ -38,11 +38,12 @@ "temporalio_commons", }) VENDOR_ROOTS = frozenset( - {"boto3", "botocore", "databricks", "google", "snowflake"}) + {"azure", "boto3", "botocore", "databricks", "google", "snowflake"}) PROVIDER_ROOTS = frozenset({ "kumo_bigquery", "kumo_databricks", "kumo_gcs", + "kumo_adls", "kumo_s3", "kumo_snowflake", }) diff --git a/packages/kumo-engine/README.md b/packages/kumo-engine/README.md index 38f193819..cd1257d53 100644 --- a/packages/kumo-engine/README.md +++ b/packages/kumo-engine/README.md @@ -29,6 +29,7 @@ The base install reads local CSV and Parquet files. For other data sources, inst | --- | --- | | Amazon S3 | `pip install "kumo-engine[s3]"` | | Google Cloud Storage | `pip install "kumo-engine[gcs]"` | +| Azure Data Lake Storage Gen2 | `pip install "kumo-engine[adls]"` | | BigQuery | `pip install "kumo-engine[bigquery]"` | | Snowflake | `pip install "kumo-engine[snowflake]"` | | Databricks | `pip install "kumo-engine[databricks]"` | @@ -91,6 +92,7 @@ To use your own data, point `FileConnector` at a directory of CSV or Parquet fil ## Documentation +- [ADLS connector guide](https://github.com/NVIDIA/kumo-engine/blob/main/packages/kumo-adls/README.md): setup, authentication, supported reads and exports. - [Kumo Engine documentation](https://docs.nvidia.com/sdgm/kumo-engine/get-started-with-kumo-engine/overview): prerequisites, GPU sizing, installation, and end-to-end workflows. Its User Code Guide covers connecting data, tables and graphs, training, and prediction. - [Predictive query examples](https://docs.nvidia.com/sdgm/kumo-engine/get-started-with-kumo-engine/workflows-and-solutions/predictive-query-examples): common prediction tasks written as predictive queries. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 4d89af512..4e5d33a2d 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -28,6 +28,8 @@ from kumo.engine._version import __version__ from kumo.engine.config import configure from kumo.engine.connector import ( + ADLSConnector, + ADLSIdentity, BigQueryConnector, BigQueryIdentity, Connector, @@ -57,6 +59,8 @@ ) __all__ = [ + "ADLSConnector", + "ADLSIdentity", "__version__", "configure", "EmbeddingModelStore", diff --git a/packages/kumo-engine/kumo/engine/_portable.py b/packages/kumo-engine/kumo/engine/_portable.py index a9c1f1238..f63d35ea8 100644 --- a/packages/kumo-engine/kumo/engine/_portable.py +++ b/packages/kumo-engine/kumo/engine/_portable.py @@ -49,6 +49,8 @@ from kumoapi.graph import Edge, graph_definition_from_edges from kumoapi.source_table import ( + ADLSRelativeSource, + ADLSSourceTable, BigQuerySourceTable, CustomSQLSourceTable, DatabricksSourceTable, @@ -110,6 +112,12 @@ def _source_table_name(source_table: LocalSourceTableType) -> str: source_table.source, S3RelativeSource) else source_table.source.uri) return '.'.join(('s3', _escape_dots(source_table.connector_id[:16]), _escape_dots(location))) + if isinstance(source_table, ADLSSourceTable): + location = (source_table.source.path if isinstance( + source_table.source, ADLSRelativeSource) else + source_table.source.uri) + return '.'.join(('adls', _escape_dots(source_table.connector_id[:16]), + _escape_dots(location))) if isinstance(source_table, GCSSourceTable): location = (source_table.source.path if isinstance( source_table.source, GCSRelativeSource) else diff --git a/packages/kumo-engine/kumo/engine/connector/__init__.py b/packages/kumo-engine/kumo/engine/connector/__init__.py index d328cf3ee..0f1413e45 100644 --- a/packages/kumo-engine/kumo/engine/connector/__init__.py +++ b/packages/kumo-engine/kumo/engine/connector/__init__.py @@ -17,12 +17,14 @@ `_response`) and the `_bind_connector` round-trip. See `kumo.engine.connector.base`. """ +from kumo.engine.connector.adls import ADLSConnector from kumo.engine.connector.base import Connector from kumo.engine.connector.bigquery import BigQueryConnector from kumo.engine.connector.databricks import DatabricksConnector from kumo.engine.connector.file import FileConnector from kumo.engine.connector.gcs import GCSConnector from kumo.engine.connector.identity import ( + ADLSIdentity, BigQueryIdentity, DatabricksIdentity, GCSIdentity, @@ -35,6 +37,8 @@ from kumo.engine.connector.source_table import SourceTable, SQLSourceTable __all__ = [ + "ADLSConnector", + "ADLSIdentity", "Connector", "FileConnector", "SnowflakeConnector", diff --git a/packages/kumo-engine/kumo/engine/connector/adls.py b/packages/kumo-engine/kumo/engine/connector/adls.py new file mode 100644 index 000000000..19c1a439f --- /dev/null +++ b/packages/kumo-engine/kumo/engine/connector/adls.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Local Azure Data Lake Storage Gen2 connector.""" + +from functools import cached_property +from typing import Optional +from urllib.parse import unquote, urlsplit + +from kumoapi.adls import normalize_adls_root_uri +from kumoapi.data_source import ( + ADLSAnyURIConnectorTarget, + ADLSConnectorResourceConfig, + ADLSRootedConnectorTarget, + DataSourceType, +) + +from kumo.engine.connector._extras import import_provider +from kumo.engine.connector.base import Connector +from kumo.engine.connector.identity import ADLSIdentity + + +class ADLSConnector(Connector): + """Read ADLS Gen2 files and write output beneath a configured root. + + Install ``kumo-engine[adls]``. Without an identity, Azure's ambient + credential chain signs requests. An explicit identity uses workload + federation. Rootless connectors accept full ``abfss://`` source URIs and + cannot discover tables or write output. + + Args: + root_dir: Optional filesystem/prefix URI, for example + ``abfss://curated@companylake.dfs.core.windows.net/orders/``. + identity: An ADLS workload identity; None selects ambient credentials. + """ + def __init__(self, *, root_dir: Optional[str] = None, + identity: Optional[ADLSIdentity] = None) -> None: + super().__init__(identity=identity) + if root_dir is None: + target = ADLSAnyURIConnectorTarget() + else: + parsed = urlsplit(normalize_adls_root_uri(root_dir)) + filesystem, host = parsed.netloc.split('@', 1) + storage_account, suffix = host.split('.', 1) + target = ADLSRootedConnectorTarget( + storage_account=storage_account, filesystem=filesystem, + prefix=unquote(parsed.path).strip('/'), endpoint_suffix=suffix) + self._config = ADLSConnectorResourceConfig(target=target) + self.target = target + self.any_uri = isinstance(target, ADLSAnyURIConnectorTarget) + self.storage_account = getattr(target, 'storage_account', None) + self.filesystem = getattr(target, 'filesystem', None) + self.prefix = getattr(target, 'prefix', None) + self.endpoint_suffix = getattr(target, 'endpoint_suffix', None) + + @property + def source_type(self) -> DataSourceType: + return DataSourceType.ADLS + + @property + def config(self) -> ADLSConnectorResourceConfig: + return self._config + + @cached_property + def source_backend(self): + backend = import_provider('kumo_adls.backend', 'adls').ADLSSourceBackend + return backend(identity=self.identity) + + @cached_property + def sink_backend(self): + backend = import_provider('kumo_adls.sink', 'adls').ADLSSinkBackend + return backend(identity=self.identity) + + def source_table_for(self, table_name: str): + return self.source_backend.get_table_config(self, + table_name).source_table + + def __repr__(self) -> str: + return ( + f'{type(self).__name__}(storage_account={self.storage_account!r}, ' + f'filesystem={self.filesystem!r}, prefix={self.prefix!r})') diff --git a/packages/kumo-engine/kumo/engine/connector/identity.py b/packages/kumo-engine/kumo/engine/connector/identity.py index 75047a3d4..816dbbc90 100644 --- a/packages/kumo-engine/kumo/engine/connector/identity.py +++ b/packages/kumo-engine/kumo/engine/connector/identity.py @@ -45,6 +45,8 @@ from kumo_core.sources.principals import store from kumoapi.data_source import ( + ADLSIdentityConfig, + AzureCloud, BigQueryIdentityConfig, CreateIdentityArgs, DatabricksAuthenticationKind, @@ -320,3 +322,18 @@ def __init__( authentication=authentication, ) super().__init__(config, credentials) + + +class ADLSIdentity(Identity): + """An Azure application authenticated through a projected workload token. + + ``AZURE_FEDERATED_TOKEN_FILE`` supplies the assertion. No secrets are + stored. Omit the connector identity to use Azure's ambient credentials. + Identity probing validates token acquisition, not storage permissions. + """ + + def __init__(self, *, tenant_id: str, client_id: str, + cloud: AzureCloud = AzureCloud.PUBLIC) -> None: + super().__init__( + ADLSIdentityConfig(tenant_id=tenant_id, client_id=client_id, + cloud=cloud)) diff --git a/packages/kumo-engine/pyproject.toml b/packages/kumo-engine/pyproject.toml index ce96f58f7..1a30967e9 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -65,6 +65,9 @@ s3=[ text=[ "kumo-glove-6b-50d>=1.0.0,<2", ] +adls=[ + "kumo-adls>=0.1.0.dev1", +] gcs=[ "kumo-gcs>=0.1.0.dev1", ] @@ -87,6 +90,7 @@ kumo-databricks = { workspace = true } kumo-bigquery = { workspace = true } kumo-s3 = { workspace = true } kumo-gcs = { workspace = true } +kumo-adls = { workspace = true } kumo-glove-6b-50d = { workspace = true } [build-system] diff --git a/packages/kumo-engine/test/test_adls_connector.py b/packages/kumo-engine/test/test_adls_connector.py new file mode 100644 index 000000000..7864e9dae --- /dev/null +++ b/packages/kumo-engine/test/test_adls_connector.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import hashlib + +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from kumo_core.sources.object_store_access import ObjectStoreAccess +from kumoapi.data_source import AzureCloud, DataSourceType +from kumoapi.source_table import ADLSRelativeSource, ADLSSourceTable +from pyarrow.fs import LocalFileSystem, SubTreeFileSystem + +from kumo.engine import ADLSConnector, ADLSIdentity, Table +from kumo.engine._portable import _source_table_name, portable_graph_config + +ROOT = 'abfss://curated@companylake.dfs.core.windows.net/orders/' + + +@pytest.fixture +def storage(monkeypatch, tmp_path): + directory = tmp_path / 'storage' + target = directory / 'curated@companylake.dfs.core.windows.net/orders' + target.mkdir(parents=True) + pq.write_table(pa.table({ + 'id': [1, 2], + 'value': [10., 20.] + }), target / 'events.parquet') + filesystem = SubTreeFileSystem(str(directory), LocalFileSystem()) + + def identity(path): + with filesystem.open_input_file(path) as stream: + return (hashlib.sha256(stream.read()).hexdigest(), ) + + monkeypatch.setattr( + 'kumo_adls.backend.source_access', + lambda *args: ObjectStoreAccess(filesystem, object_identity=identity)) + return directory + + +@pytest.mark.parametrize( + 'root, prefix', [(ROOT, 'orders'), (ROOT.rstrip('/'), 'orders'), + (ROOT.removesuffix('orders/'), ''), + (ROOT.removesuffix('/orders/'), ''), + (ROOT.replace('orders/', 'space%20name/'), 'space name')]) +def test_root_parsing_without_connecting(root, prefix): + connector = ADLSConnector(root_dir=root) + assert connector.prefix == prefix + assert connector.filesystem == 'curated' + assert connector.storage_account == 'companylake' + assert connector.source_type == DataSourceType.ADLS + + +@pytest.mark.parametrize('root', [ + 'abfs://curated@companylake.dfs.core.windows.net/orders/', + ROOT + '../outside', ROOT + '?token=secret', + ROOT.replace('orders/', 'orders//') +]) +def test_invalid_roots_fail_locally(root): + with pytest.raises(ValueError): + ADLSConnector(root_dir=root) + + +def test_workload_identity_is_registered_without_credentials(): + identity = ADLSIdentity(tenant_id='tenant', client_id='client', + cloud=AzureCloud.PUBLIC) + connector = ADLSConnector(root_dir=ROOT, identity=identity) + assert connector.identity is identity + assert connector.identity_id == identity.slug + assert identity._credentials is None + assert 'tenant' not in repr(identity) + assert connector.source_backend._identity is identity + assert connector.sink_backend._identity is identity + assert connector.source_backend is connector.source_backend + assert connector.sink_backend is connector.sink_backend + + +def test_table_inference_portable_graph_and_ingest(storage, tmp_path): + connector = ADLSConnector(root_dir=ROOT) + assert connector.table_names() == ['events'] + source = connector.source_table_for('events') + assert isinstance(source, ADLSSourceTable) + assert isinstance(source.source, ADLSRelativeSource) + assert _source_table_name(source).startswith('adls.') + assert len(_source_table_name(source).split('.')) == 3 + table = Table.from_source_table(connector['events'], + primary_key='id').infer_metadata() + graph = portable_graph_config({'events': table}, [], name='test') + assert graph['tables'][0]['source_table_name'].startswith('adls.') + ingest = connector.source_backend.prepare_ingest(connector, 'events', + str(tmp_path / 'ingest')) + assert pq.read_table(ingest.source['paths']).num_rows == 2 + connector.sink_backend.write_table(connector, 'scores', + pd.DataFrame({'score': [0.5]})) + assert pq.read_table( + storage / + 'curated@companylake.dfs.core.windows.net/orders/scores.parquet' + ).num_rows == 1 + + +def test_rootless_graph_coordinate_contains_exact_source(storage): + connector = ADLSConnector() + source = connector.source_table_for(ROOT + 'events.parquet') + coordinate = _source_table_name(source) + assert 'companylake' in coordinate + assert len(coordinate.split('.')) == 3 diff --git a/packages/kumo-engine/test/test_optional_imports.py b/packages/kumo-engine/test/test_optional_imports.py index c70368fed..22d40a95a 100644 --- a/packages/kumo-engine/test/test_optional_imports.py +++ b/packages/kumo-engine/test/test_optional_imports.py @@ -25,6 +25,7 @@ 'botocore', 's3fs', 'google', + 'azure', 'gcsfs', ) PROVIDER_PACKAGES = ( @@ -32,6 +33,7 @@ 'kumo_databricks', 'kumo_s3', 'kumo_gcs', + 'kumo_adls', 'kumo_bigquery', ) FORBIDDEN_AT_IMPORT = PROVIDER_PACKAGES + OPTIONAL_THIRD_PARTY @@ -92,6 +94,7 @@ def find_spec(self, name, path=None, target=None): # dependency, naming the extra, or the provider defers that import # further and fails later. for cls, extra in ( + (engine.ADLSConnector, 'adls'), (engine.SnowflakeConnector, 'snowflake'), (engine.DatabricksConnector, 'databricks'), (engine.S3Connector, 's3'), diff --git a/pyproject.toml b/pyproject.toml index 1a7180091..52cecc604 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ members = [ "packages/kumo-bigquery", "packages/kumo-databricks", "packages/kumo-gcs", + "packages/kumo-adls", "packages/kumo-s3", "packages/kumo-engine", "packages/kumo-sdk", @@ -85,6 +86,7 @@ kumo-pql = { workspace = true } kumo-core = { workspace = true } diskgraph = { workspace = true } kumo-gcs = { workspace = true } +kumo-adls = { workspace = true } kumo-s3 = { workspace = true } kumo-snowflake = { workspace = true } kumo-bigquery = { workspace = true } diff --git a/scripts/README.md b/scripts/README.md index 58e282a4b..c3a0f99d1 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -2,6 +2,14 @@ This directory contains scripts for developer setup, CI checks, operational utilities, and releases. See [release tooling](./release/README.md) for the release-specific commands. +## Object-store transfer diagnostics + +Shared S3/ADLS harnesses live in [`benchmarks/object_store/`](benchmarks/object_store/): +[`compare_live.py`](benchmarks/object_store/compare_live.py) measures production +raw bulk readers; [`compare_streams.py`](benchmarks/object_store/compare_streams.py) +is a [fixed-geometry loopback stream diagnostic](benchmarks/object_store/matched-transfer.md). +Their scopes and request geometry differ; keep their measurements separate. + ## Version stamping kumo-engine, kumo-sdk, the platform, the Studio images, and the Helm charts share one version, kept in the root [`VERSION`](../VERSION) file as a PEP 440 string such as `1.1.0.dev0` (also `1.1.0`, `1.1.0b1` and `1.1.0rc1`). After changing it, run: diff --git a/scripts/benchmarks/object_store/compare_live.py b/scripts/benchmarks/object_store/compare_live.py new file mode 100644 index 000000000..72e6065b0 --- /dev/null +++ b/scripts/benchmarks/object_store/compare_live.py @@ -0,0 +1,575 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compare existing, identical S3/ADLS tables through production bulk readers. + +This script only reads cloud objects; provisioning and fixture copying happen +separately. Supply trusted, complete-table manifests containing SHA256 and +footer row counts. Use --providers adls or --providers s3 for native-host +single-provider phases; the default runs alternating provider pairs. Each fresh +child prepares local Parquet through the actual provider fetch callback and +shared core raw-copy path, without the engine. + +Setup/authentication and manifest identity preflight are outside the timer. +The preparation timer includes production listing/version guards and transfers. +The nested transfer timer includes production fetch's own metadata, retries, +and checks; the remainder is NOT a pure metadata-latency measurement. SHA256, +footer verification, final manifest checks, and cleanup are outside both timers. +There are no request hooks/proxies: request/retry counts and token refresh +counts are explicitly uninstrumented. Cloud placement, TLS, auth, retries, and +provider request geometry differ; this measures end-to-end production paths. + +Per-table manifest schema (all objects, not a sampled prefix): +{schema_version:1, table, bucket, s3_prefix, adls_account, filesystem, + adls_prefix, endpoint_suffix?:"dfs.core.windows.net", + files:[{relative_path,s3_key,bytes,rows,sha256,source_etag,adls_etag}], + totals:{files,bytes,rows}} +""" + +import argparse +import hashlib +import importlib.metadata +import json +import os +import platform +import shutil +import socket +import subprocess +import sys +import time +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from datetime import datetime, timezone +from pathlib import Path, PurePosixPath +from uuid import uuid4 + + +def write_json(path, value): + path.write_text(json.dumps(value, separators=(',', ':')) + '\n') + + +def sha256(path): + digest = hashlib.sha256() + with path.open('rb') as source: + for block in iter(lambda: source.read(1 << 20), b''): + digest.update(block) + return digest.hexdigest() + + +OUTPUT_HASH_CHUNK_BYTES = 1 << 20 +MAX_VERIFICATION_WORKERS = 8 + + +def output_sha256(path): + """Hash only a downloaded output; advise away bytes after hashing them. + + Provenance/code/manifest hashes continue to use sha256 without cache advice. + One reusable buffer bounds each active hash worker's Python payload. + Advice is opportunistic and never changes file contents or correctness. + """ + digest = hashlib.sha256() + buffer = bytearray(OUTPUT_HASH_CHUNK_BYTES) + view = memoryview(buffer) + supported = hasattr(os, 'posix_fadvise') and hasattr( + os, 'POSIX_FADV_DONTNEED') + advice_calls = 0 + advice_errors = set() + offset = 0 + with path.open('rb', buffering=0) as source: + while True: + amount = source.readinto(buffer) + if not amount: + break + digest.update(view[:amount]) + if supported: + try: + os.posix_fadvise(source.fileno(), offset, amount, + os.POSIX_FADV_DONTNEED) + advice_calls += 1 + except OSError as error: + advice_errors.add(error.errno) + offset += amount + return digest.hexdigest(), dict(supported=supported, calls=advice_calls, + error_numbers=sorted(advice_errors)) + + +def verify_outputs(output, entries, requested_workers): + import pyarrow.parquet as pq + from kumo_core.sources.cpu import cpu_budget + + expected = {entry['relative_path'] for entry in entries} + if { + path.relative_to(output).as_posix() + for path in output.rglob('*') if path.is_file() + } != expected: + raise ValueError('Local staged membership differs from manifest') + # Prefer recently completed/warm files; retain deterministic tie breaking. + ordered = sorted( + entries, key=lambda entry: ((output / entry['relative_path']).stat( + ).st_mtime_ns, entry['relative_path']), reverse=True) + workers = min(requested_workers, MAX_VERIFICATION_WORKERS, cpu_budget(), + len(entries)) + if workers < 1: + raise ValueError( + 'Verification needs a positive worker count and nonempty table') + + def verify(entry): + path = output / entry['relative_path'] + if path.is_symlink(): + raise ValueError('Staged output must not be a symbolic link') + digest, advice = output_sha256(path) + actual = dict(relative_path=entry['relative_path'], + bytes=path.stat().st_size, + rows=pq.ParquetFile(path).metadata.num_rows, + sha256=digest) + if any(actual[key] != entry[key] + for key in ('bytes', 'rows', 'sha256')): + raise ValueError( + f'Local SHA256/size/footer mismatch: {entry["relative_path"]}') + return actual, advice + + verified = [] + advice_records = [] + entries_to_submit = iter(ordered) + with ThreadPoolExecutor(max_workers=workers) as pool: + pending = set() + try: + for entry in entries_to_submit: + pending.add(pool.submit(verify, entry)) + if len(pending) == workers: + break + while pending: + done, pending = wait(pending, return_when=FIRST_COMPLETED) + for future in done: + actual, advice = future.result() + verified.append(actual) + advice_records.append(advice) + for _ in done: + entry = next(entries_to_submit, None) + if entry is None: + break + pending.add(pool.submit(verify, entry)) + except BaseException: + for future in pending: + future.cancel() + # The executor waits for active hash readers before caller cleanup. + raise + verified.sort(key=lambda entry: entry['relative_path']) + config = dict( + requested_workers=requested_workers, effective_workers=workers, + maximum_workers=MAX_VERIFICATION_WORKERS, + hash_chunk_bytes=OUTPUT_HASH_CHUNK_BYTES, + order='newest output mtime_ns first, reverse relative_path for ties', + output_record_order='relative_path ascending', + cache_advice='POSIX_FADV_DONTNEED on output ranges after SHA256 update', + cache_advice_supported=all(record['supported'] + for record in advice_records), + cache_advice_calls=sum(record['calls'] for record in advice_records), + cache_advice_error_numbers=sorted({ + number + for record in advice_records + for number in record['error_numbers'] + })) + return verified, config + + +def relative_path(value): + path = PurePosixPath(value) + if (not value or value != path.as_posix() or path.is_absolute() + or any(part in ('.', '..') for part in path.parts) + or '\\' in value): + raise ValueError( + 'Manifest paths must be canonical relative POSIX paths') + return value + + +def read_manifest(path): + manifest = json.loads(path.read_text()) + if manifest['schema_version'] != 1: + raise ValueError('Unsupported live-benchmark manifest version') + for field in ('s3_prefix', 'adls_prefix'): + manifest[field] = relative_path(manifest[field].rstrip('/')) + if not manifest['files']: + raise ValueError('A full table must contain at least one file') + names = set() + for entry in manifest['files']: + name = relative_path(entry['relative_path']) + if name in names or not name.endswith('.parquet'): + raise ValueError('Manifest needs distinct Parquet paths') + names.add(name) + if entry['s3_key'] != manifest['s3_prefix'] + '/' + name: + raise ValueError('S3 key is outside the complete table prefix') + if (type(entry['bytes']) is not int or entry['bytes'] <= 0 + or type(entry['rows']) is not int or entry['rows'] < 0 + or len(entry['sha256']) != 64 + or any(char not in '0123456789abcdef' + for char in entry['sha256']) or not entry['source_etag'] + or not entry['adls_etag']): + raise ValueError( + 'Manifest needs sizes, rows, SHA256, and both pinned ETags') + expected = dict(files=len(names), + bytes=sum(item['bytes'] for item in manifest['files']), + rows=sum(item['rows'] for item in manifest['files'])) + if manifest['totals'] != expected: + raise ValueError('Manifest totals do not match its objects') + return manifest + + +def memory(): + result = {} + for filename, fields in [('status', {'VmHWM', 'VmRSS'}), + ('smaps_rollup', {'Pss'})]: + for line in Path('/proc/self', filename).read_text().splitlines(): + key, separator, value = line.partition(':') + if separator and key in fields: + result[key + '_kib'] = int(value.split()[0]) + return result + + +def metadata(): + result: dict[str, object] = dict(python=sys.version, + hostname=socket.gethostname(), + platform=platform.platform(), + cpu_count=os.cpu_count()) + result['dependencies'] = { + name: importlib.metadata.version(name) + for name in ('pyarrow', 'fsspec', 'boto3', 'botocore', 'azure-core', + 'azure-identity', 'azure-storage-file-datalake', + 'azure-storage-blob', 'requests', 'urllib3') + } + # Hash the implementation actually imported by the child, including dirty + # checkout changes that a commit identifier alone would not identify. + import kumo_adls.sdk_filesystem + import kumo_adls.transfer + import kumo_core.sources.object_store + import kumo_s3.access + result['provider_code'] = { + module.__name__: + dict(path=module.__file__, sha256=sha256(Path(module.__file__))) + for module in (kumo_adls.sdk_filesystem, kumo_adls.transfer, + kumo_core.sources.object_store, kumo_s3.access) + } + return result + + +def make_reader(provider, manifest, args): + import pyarrow as pa + from kumo_core.sources.cpu import copy_workers + from kumo_core.sources.object_store import ObjectStoreSourceReader + from kumo_core.types.enums import ConnectorType + + pa.set_io_thread_count(args.worker_count) + if copy_workers(args.worker_count) != args.worker_count: + raise ValueError('Host CPU/copy ceiling is below the requested workers') + if provider == 's3': + from kumo_s3.access import source_access + from kumo_s3.session import S3_MAX_RETRIES + if args.s3_anonymous: + from types import SimpleNamespace + session = SimpleNamespace(get_credentials=lambda: None, + region_name=args.s3_region) + else: + import boto3 + session = boto3.Session(region_name=args.s3_region) + access = source_access(manifest['bucket'], session=session, + default_region=args.s3_region, + allow_anonymous=args.s3_anonymous) + root = manifest['bucket'] + '/' + manifest['s3_prefix'] + reader = ObjectStoreSourceReader(access.filesystem, root=root, + file_format='parquet', + file_suffix='parquet', + fetch=access.fetch, + object_identity=access.object_identity, + source_type=ConnectorType.S3) + config = dict( + authentication='anonymous' + if args.s3_anonymous else 'ambient boto3 credentials', + retry_strategy='AwsDefaultS3RetryStrategy', + retry_max_attempts=S3_MAX_RETRIES, + arrow_io_threads=pa.io_thread_count(), + fetch='kumo_s3.access.source_access.fetch', + http_pool='native Arrow public API does not expose pool size') + return reader, config, lambda: None + + from kumo_adls import session, transfer + from kumo_adls.backend import ADLSSourceReader + from kumo_adls.source_access import source_access + if transfer.CONNECTIONS != args.worker_count: + raise ValueError( + 'ADLS transfer ceiling differs from the file-worker ceiling') + suffix = manifest.get('endpoint_suffix', 'dfs.core.windows.net') + access = source_access(None, manifest['adls_account'], suffix) + authority = f'{manifest["filesystem"]}@{manifest["adls_account"]}.{suffix}' + reader = ADLSSourceReader(access.filesystem, + root=authority + '/' + manifest['adls_prefix'], + file_format='parquet', file_suffix='parquet', + fetch=access.fetch, + object_identity=access.object_identity, + source_type=ConnectorType.ADLS) + config = dict( + authentication='ambient DefaultAzureCredential', + credential_selection=os.environ.get('AZURE_TOKEN_CREDENTIALS', + 'default chain'), + transfer_ceiling=transfer.CONNECTIONS, http_pool=transfer.CONNECTIONS, + arrow_io_threads=pa.io_thread_count(), + bulk_chunk_bytes=transfer.BULK_DOWNLOAD_CHUNK_BYTES, + bulk_max_workers=transfer.MAX_BULK_DOWNLOAD_CONCURRENCY, + fetch='kumo_adls.source_access.source_access.fetch', + retries='Azure SDK production defaults; actual retries uninstrumented') + return reader, config, session.close + + +def identity_token(value): + # SDK List Paths and HEAD can spell one opaque token with/without quotes. + return value[1:-1] if value.startswith('"') and value.endswith( + '"') else value + + +def inventory(reader, manifest, provider, workers): + found = reader._objects('', reader.root, + max_entries=4 * len(manifest['files']) + 1024, + include_hidden=True) + expected = {entry['relative_path']: entry for entry in manifest['files']} + actual = {entry.path[len(reader.root) + 1:]: entry for entry in found} + if set(actual) != set(expected): + raise ValueError( + 'Remote table membership differs from the complete manifest') + ordered = sorted(actual) + identities = reader._identities([actual[name] for name in ordered], workers) + key = 'source_etag' if provider == 's3' else 'adls_etag' + recorded = {} + for name, identity in zip(ordered, identities): + entry = expected[name] + if (actual[name].size != entry['bytes'] + or int(identity[-1]) != entry['bytes'] + or identity_token(identity[0]) != identity_token(entry[key])): + raise ValueError( + f'Remote object differs from its pinned manifest: {name}') + recorded[name] = identity + return recorded + + +def run_child(args): + from kumo_core.sources.types import ReadMode + + manifest = read_manifest(args.manifest[0]) + args.run_dir.mkdir(parents=True, exist_ok=False) + close = lambda: None + try: + reader, config, close = make_reader(args.provider, manifest, args) + preflight_start = time.perf_counter() + before_identity = inventory(reader, manifest, args.provider, + args.worker_count) + preflight_seconds = time.perf_counter() - preflight_start + provenance = metadata() + total_bytes = manifest['totals']['bytes'] + allowance = total_bytes + max(1 << 30, args.worker_count * (64 << 20)) + if shutil.disk_usage(args.run_dir).free < allowance: + raise ValueError( + 'Insufficient local free space for the complete table and headroom' + ) + transfer = {} + original = reader._download_objects + + def download(*positional, **keywords): + started = time.perf_counter() + cpu = time.process_time() + result = original(*positional, **keywords) + transfer['elapsed_seconds'] = time.perf_counter() - started + transfer['process_cpu_seconds'] = time.process_time() - cpu + return result + + reader._download_objects = download + before_memory = memory() + cpu = time.process_time() + started = time.perf_counter() + prepare_started_utc_ns = time.time_ns() + staged = reader.prepare('', str(args.run_dir / 'parquet'), columns=None, + mode=ReadMode.STREAMED, transform=False, + worker_ceiling=args.worker_count, + max_entries=4 * len(manifest['files']) + 1024, + max_bytes=total_bytes) + prepare_finished_utc_ns = time.time_ns() + elapsed = time.perf_counter() - started + cpu = time.process_time() - cpu + after_memory = memory() + if staged.source_size != total_bytes: + raise ValueError('Staged total size differs from manifest') + after_identity = inventory(reader, manifest, args.provider, + args.worker_count) + if before_identity != after_identity: + raise ValueError('Remote identities changed across the trial') + verification_started = time.perf_counter() + verified, verification_config = verify_outputs( + args.run_dir / 'parquet', manifest['files'], + args.verification_workers) + result = dict( + accepted=True, provider=args.provider, table=manifest['table'], + selected_providers=args.providers, trial_scope='paired' + if len(args.providers) == 2 else 'single-provider', + workers=args.worker_count, worker_limit_kind='ceiling', + effective_file_workers='uninstrumented', stored_bytes=total_bytes, + rows=manifest['totals']['rows'], files=len(verified), + prepare_seconds=elapsed, process_cpu_seconds=cpu, + prepare_started_utc_ns=prepare_started_utc_ns, + prepare_finished_utc_ns=prepare_finished_utc_ns, + transfer_phase=transfer, non_transfer_prepare_seconds=elapsed - + transfer['elapsed_seconds'], mib_per_second=total_bytes / elapsed / + (1 << 20), rows_per_second=manifest['totals']['rows'] / elapsed, + preflight_seconds=preflight_seconds, + verification_seconds=time.perf_counter() - verification_started, + verification_config=verification_config, + mode_used=staged.mode_used.value, data_version=staged.version, + memory_before=before_memory, memory_after=after_memory, + client_config=config, metadata=provenance, + request_counts='uninstrumented', retry_counts='uninstrumented', + token_refresh_counts='uninstrumented', + remote_identities_before=before_identity, + remote_identities_after=after_identity, verified_files=verified) + write_json(args.child_result, result) + finally: + close() + if args.run_dir.exists(): + shutil.rmtree(args.run_dir) + + +def trial_providers(providers, index): + """Alternate two-provider order; singleton trials never synthesize a pair.""" + if (not providers or len(set(providers)) != len(providers) + or any(provider not in ('adls', 's3') for provider in providers)): + raise ValueError('Select one or both distinct supported providers') + return tuple(providers) if index % 2 == 0 else tuple(reversed(providers)) + + +def run_parent(args): + manifests = [(path.resolve(strict=True), read_manifest(path)) + for path in args.manifest] + if len({manifest['table'] for _, manifest in manifests}) != len(manifests): + raise ValueError('Table names must be unique across manifests') + args.output.mkdir(parents=True, exist_ok=True) + args.results.mkdir(parents=True, exist_ok=False) + placement = json.loads(args.placement.read_text()) if args.placement else { + 'status': 'not supplied' + } + write_json( + args.results / 'arguments.json', + dict( + arguments={ + name: [str(item) for item in value] if isinstance(value, list) + else str(value) if isinstance(value, Path) else value + for name, value in vars(args).items() + }, created_utc=datetime.now(timezone.utc).isoformat(), + benchmark_sha256=sha256(Path(__file__)), placement=placement, + planned_download_bytes=len(args.providers) * args.trials * + len(args.workers) * sum(manifest['totals']['bytes'] + for _, manifest in manifests))) + pair = 0 + for index, (_, manifest) in enumerate(manifests): + pinned_manifest = args.results.resolve() / f'manifest-{index}.json' + write_json(pinned_manifest, manifest) + manifest_digest = sha256(pinned_manifest) + for workers in args.workers: + for trial in range(args.trials): + providers = trial_providers(args.providers, pair) + pair += 1 + for order, provider in enumerate(providers): + label = f'table{index}-w{workers}-t{trial}-{order}-{provider}' + run_dir = args.output.resolve() / (label + '-' + + uuid4().hex) + child_result = args.results / (label + '.json') + command = [ + sys.executable, + str(Path(__file__).resolve()), '--providers', + *args.providers, '--child', '--manifest', + str(pinned_manifest), '--provider', provider, + '--worker-count', + str(workers), '--run-dir', + str(run_dir), '--child-result', + str(child_result), '--verification-workers', + str(args.verification_workers) + ] + if args.s3_region: + command.extend(['--s3-region', args.s3_region]) + if args.s3_anonymous: + command.append('--s3-anonymous') + environment = os.environ.copy() + environment.update(KUMO_SOURCE_COPY_WORKERS=str(workers), + KUMO_ADLS_TRANSFER_WORKERS=str(workers)) + try: + completed = subprocess.run(command, check=False, + capture_output=True, + text=True, env=environment, + timeout=args.trial_timeout) + if completed.returncode: + raise RuntimeError(completed.stderr + or completed.stdout) + if sha256(pinned_manifest) != manifest_digest: + raise ValueError( + 'Pinned manifest changed during trial') + result = json.loads(child_result.read_text()) + result.update(trial=trial, provider_order=order, + placement=placement, + manifest_sha256=manifest_digest) + write_json(child_result, result) + print( + json.dumps( + dict( + trial=label, accepted=True, + prepare_seconds=result['prepare_seconds'])), + flush=True) + except Exception as error: + write_json( + args.results / (label + '.rejected.json'), + dict(accepted=False, table=manifest['table'], + provider=provider, reason=str(error))) + raise + finally: + if run_dir.exists(): + shutil.rmtree(run_dir) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--manifest', type=Path, nargs='+', required=True) + parser.add_argument('--output', type=Path) + parser.add_argument('--results', type=Path) + parser.add_argument('--placement', type=Path, + help='Measured client/AWS/Azure placement JSON') + parser.add_argument( + '--workers', type=int, nargs='+', default=[32], + help='File/transfer worker ceiling; shared memory budget may lower it') + parser.add_argument('--trials', type=int, default=3) + parser.add_argument( + '--providers', choices=['adls', + 's3'], nargs='+', default=['adls', 's3'], + help='Both alternating providers, or one native-host provider') + parser.add_argument( + '--verification-workers', type=int, default=8, + help='Untimed output hash workers, capped by CPU and eight') + parser.add_argument('--trial-timeout', type=float, default=14400) + parser.add_argument('--s3-region') + parser.add_argument('--s3-anonymous', action='store_true') + parser.add_argument('--child', action='store_true', help=argparse.SUPPRESS) + parser.add_argument('--provider', choices=['s3', 'adls'], + help=argparse.SUPPRESS) + parser.add_argument('--worker-count', type=int, help=argparse.SUPPRESS) + parser.add_argument('--run-dir', type=Path, help=argparse.SUPPRESS) + parser.add_argument('--child-result', type=Path, help=argparse.SUPPRESS) + args = parser.parse_args() + if args.child: + if args.provider not in args.providers: + parser.error('Child provider is outside the selected phase') + run_child(args) + else: + if len(set(args.providers)) != len(args.providers): + parser.error('--providers must not contain duplicates') + if args.output is None or args.results is None: + parser.error('--output and --results are required') + if (any(workers < 1 for workers in args.workers) or args.trials < 1 + or args.verification_workers < 1): + parser.error( + 'Workers, verification workers and trials must be positive') + run_parent(args) + + +if __name__ == '__main__': + main() diff --git a/scripts/benchmarks/object_store/compare_streams.py b/scripts/benchmarks/object_store/compare_streams.py new file mode 100644 index 000000000..36d9acd85 --- /dev/null +++ b/scripts/benchmarks/object_store/compare_streams.py @@ -0,0 +1,663 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Frozen-geometry ADLS/S3 stream diagnostic, not a production bulk benchmark. + +The fixed 8 MiB windows, 8 MiB SDK initial response, 4 MiB SDK chunk setting, +and 1 MiB transport blocks are diagnostic controls, not production defaults. +They intentionally do not track provider tuning. Both clients run the same +file worker loop: open (HEAD), sequential 8 MiB +reads, buffered atomic local output, then a fresh open (HEAD). ADLS explicitly +conditions reads on the opened ETag; native Arrow S3 has no public If-Match +option. The server records this difference rather than emulating it away. +Native S3 exposes no HTTP pool-size option; its Arrow I/O pool is configured +before construction. Every trial must prove exact request geometry and output +parity. Source hashing warms the OS cache before all timed trials. + +Run --check-protocol first: one real fixture file, one worker, both clients, +with no timing fields published. Full defaults are 24 fresh-process trials. +The caller must supply the source, local output, and results directories. +""" + +import argparse +import base64 +import hashlib +import importlib.metadata +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from types import SimpleNamespace +from urllib.parse import unquote, urlsplit +from urllib.request import urlopen +from uuid import uuid4 + +# Keep this diagnostic contract stable instead of following production defaults. +WINDOW = 8 << 20 +NETWORK_BLOCK = 1 << 20 +DIAGNOSTIC_GEOMETRY = dict( + scope='matched-stream diagnostic, not production bulk transfer', + derived_from_production=False, range_bytes=WINDOW, + adls_initial_get_bytes=WINDOW, adls_chunk_bytes=4 << 20, + network_block_bytes=NETWORK_BLOCK) +BUCKET = 'curated' +ACCOUNT = 'benchaccount' +AUTHORITY = f'{BUCKET}@{ACCOUNT}.dfs.core.windows.net' + + +def write_json(path, value): + path.write_text(json.dumps(value, separators=(',', ':')) + '\n') + + +def sha256(path): + digest = hashlib.sha256() + with path.open('rb') as source: + for block in iter(lambda: source.read(1 << 20), b''): + digest.update(block) + return digest.hexdigest() + + +def etag(stat): + return f'"{stat.st_size:x}-{stat.st_mtime_ns:x}"' + + +def make_manifest(source, maximum, expected): + import pyarrow.parquet as pq + + paths = sorted(source.rglob('*.parquet')) + if not paths or (maximum is None and len(paths) != expected): + raise ValueError( + f'Expected {expected} Parquet files; found {len(paths)}') + if maximum is not None: + # A one-file protocol check must exercise multiple range requests. + paths = ([max(paths, key=lambda path: path.stat().st_size)] + if maximum == 1 else paths[:maximum]) + entries = [] + for path in paths: + if path.is_symlink(): + raise ValueError(f'Source symlinks are unsupported: {path}') + before = path.stat() + digest = sha256(path) + rows = pq.ParquetFile(path).metadata.num_rows + after = path.stat() + if (before.st_size, before.st_mtime_ns) != (after.st_size, + after.st_mtime_ns): + raise ValueError(f'Source changed while preparing manifest: {path}') + entries.append( + dict(path=path.relative_to(source).as_posix(), size=after.st_size, + etag=etag(after), sha256=digest, rows=rows)) + return dict(source=str(source), files=entries, subset_diagnostic=maximum + is not None, + total_bytes=sum(entry['size'] for entry in entries)) + + +class RequestState: + def __init__(self, manifest): + self.source = Path(manifest['source']) + self.files = {entry['path']: entry for entry in manifest['files']} + self.condition = threading.Condition() + self.active = {'HEAD': 0, 'GET': 0} + self.reset(0) + + def idle(self): + if not self.condition.wait_for(lambda: not any(self.active.values()), + timeout=60): + raise TimeoutError('Object handlers did not become idle') + + def reset(self, latency): + with self.condition: + self.idle() + self.latency = latency + self.maximum = {'HEAD': 0, 'GET': 0} + self.successes = 0 + self.failures = 0 + self.body_bytes = 0 + self.unknown = [] + self.trace = {name: dict(heads=[], gets=[]) for name in self.files} + + def begin(self, method): + with self.condition: + self.active[method] += 1 + self.maximum[method] = max(self.maximum[method], + self.active[method]) + + def finish(self, method, name, detail, success, body_bytes): + with self.condition: + if name in self.trace and detail is not None: + self.trace[name]['heads' if method == + 'HEAD' else 'gets'].append(detail) + else: + self.unknown.append(dict(method=method, path=name)) + self.successes += int(success) + self.failures += int(not success) + self.body_bytes += body_bytes + self.active[method] -= 1 + self.condition.notify_all() + + def snapshot(self): + with self.condition: + self.idle() + return dict(max_active=dict(self.maximum), successes=self.successes, + failures=self.failures, + response_body_bytes=self.body_bytes, + unknown=list(self.unknown), files=self.trace) + + +class ObjectHandler(BaseHTTPRequestHandler): + protocol_version = 'HTTP/1.1' + disable_nagle_algorithm = True + + def setup(self): + super().setup() + self.connection.settimeout(60) + + def log_message(self, *args): + pass + + def control(self): + parsed = urlsplit(self.path) + if not parsed.path.startswith('/__control/'): + return False + try: + if parsed.path == '/__control/reset': + # Control requests come only from the parent loopback process. + from urllib.parse import parse_qs + self.server.state.reset( + float(parse_qs(parsed.query)['latency_ms'][0])) + value = {'reset': True} + elif parsed.path == '/__control/snapshot': + value = self.server.state.snapshot() + else: + raise ValueError('Unknown control endpoint') + body = json.dumps(value, separators=(',', ':')).encode() + self.send_response(200) + self.send_header('Content-Length', str(len(body))) + self.send_header('Content-Type', 'application/json') + self.send_header('Connection', 'close') + self.end_headers() + self.wfile.write(body) + self.close_connection = True + except (TimeoutError, ValueError, KeyError): + self.send_error(500, 'Control request failed') + return True + + def do_HEAD(self): + self.object_request('HEAD') + + def do_GET(self): + if not self.control(): + self.object_request('GET') + + def do_POST(self): + # Keep unexpected methods visible in the trial rejection evidence. + with self.server.state.condition: + self.server.state.unknown.append(dict(method='POST', + path=self.path)) + self.server.state.failures += 1 + self.send_error(405) + + def object_request(self, method): + state = self.server.state + state.begin(method) + name = self.path + detail = None + body_bytes = 0 + success = False + try: + path = unquote(urlsplit(self.path).path).lstrip('/') + if path.startswith(ACCOUNT + '/'): + path = path[len(ACCOUNT) + 1:] + prefix = BUCKET + '/' + if not path.startswith(prefix): + self.send_error(404) + return + name = path[len(prefix):] + if name not in state.files: + self.send_error(404) + return + source = state.source / name + stat = source.stat() + identity = etag(stat) + conditional = self.headers.get('If-Match') + detail = dict( + if_match=conditional, etag=identity, + auth_scheme=self.headers.get('Authorization', '').split(' ', + 1)[0]) + if conditional is not None and conditional not in ('*', identity): + self.send_error(412) + return + time.sleep(state.latency / 1000) + start, end = 0, stat.st_size - 1 + if method == 'GET': + standard = self.headers.get('Range') + azure = self.headers.get('x-ms-range') + if standard and azure and standard != azure: + raise ValueError('Conflicting range headers') + selected = standard or azure or '' + match = re.fullmatch(r'bytes=(\d+)-(\d+)', selected) + if match is None: + raise ValueError('Expected one explicit byte range') + start, end = map(int, match.groups()) + detail.update( + window=[start, end], + range_header='Range' if standard else 'x-ms-range') + if start > end or end >= stat.st_size: + self.send_error(416) + return + length = end - start + 1 + self.send_response(206 if method == 'GET' else 200) + self.send_header('Content-Length', str(length)) + self.send_header('ETag', identity) + self.send_header('Last-Modified', + self.date_time_string(stat.st_mtime)) + self.send_header('Accept-Ranges', 'bytes') + self.send_header('Content-Type', 'application/octet-stream') + self.send_header('x-ms-blob-type', 'BlockBlob') + self.send_header('x-ms-request-id', 'loopback') + self.send_header('x-amz-request-id', 'loopback') + if method == 'GET': + self.send_header('Content-Range', + f'bytes {start}-{end}/{stat.st_size}') + self.end_headers() + if method == 'GET': + with source.open('rb') as stream: + stream.seek(start) + remaining = length + while remaining: + data = stream.read(min(NETWORK_BLOCK, remaining)) + if not data: + raise OSError( + 'Source became shorter during response') + self.wfile.write(data) + body_bytes += len(data) + remaining -= len(data) + self.wfile.flush() + success = True + except (OSError, ValueError): + # A partial/failed response invalidates the entire trial. + self.close_connection = True + finally: + state.finish(method, name, detail, success, body_bytes) + + +class Server(ThreadingHTTPServer): + daemon_threads = True + block_on_close = False + request_queue_size = 128 + + +def make_client(provider, endpoint, workers): + import pyarrow as pa + import pyarrow.fs as fs + + pa.set_io_thread_count(workers) + if provider == 's3': + native = fs.S3FileSystem( + access_key=ACCOUNT, secret_key='benchmark', region='us-east-1', + scheme='http', endpoint_override=endpoint, + force_virtual_addressing=False, allow_delayed_open=False, + connect_timeout=30, request_timeout=60, + retry_strategy=fs.AwsStandardS3RetryStrategy(max_attempts=1)) + + def identity(stream): + metadata = stream.metadata() + metadata = { + (key.decode() if isinstance(key, bytes) else key).lower(): + value.decode() if isinstance(value, bytes) else value + for key, value in metadata.items() + } + return dict(size=stream.size(), etag=metadata['etag']) + + return SimpleNamespace( + open=lambda name: native.open_input_file(BUCKET + '/' + name), + identity=identity, + read=lambda stream, length, offset: stream.read_at(length, offset), + close=lambda: None, + config=dict(authentication='synthetic_sigv4', retry_max_attempts=1, + io_threads=workers, + http_pool='not configurable through native public API')) + + import requests + from azure.core.credentials import AzureNamedKeyCredential + from azure.core.pipeline.transport import RequestsTransport + from azure.storage.filedatalake import DataLakeServiceClient + from kumo_adls import transfer + from kumo_adls.sdk_filesystem import ADLSFileSystem + from requests.adapters import HTTPAdapter + + if transfer.CONNECTIONS != workers: + raise ValueError( + f'ADLS transfer ceiling {transfer.CONNECTIONS} differs from {workers} workers' + ) + session = requests.Session() + session.trust_env = False + adapter = HTTPAdapter(pool_connections=workers, pool_maxsize=workers, + pool_block=True, max_retries=0) + session.mount('http://', adapter) + transport = RequestsTransport(session=session, session_owner=True, + connection_data_block_size=NETWORK_BLOCK) + if transport.connection_config.data_block_size != NETWORK_BLOCK: + raise ValueError('Azure transport did not apply the network block size') + client = DataLakeServiceClient( + f'http://{endpoint}/{ACCOUNT}', credential=AzureNamedKeyCredential( + ACCOUNT, + base64.b64encode(b'benchmark').decode()), + max_single_get_size=WINDOW, max_chunk_get_size=4 << 20, + connection_data_block_size=NETWORK_BLOCK, retry_total=0, + connection_timeout=30, read_timeout=60, transport=transport) + filesystem = ADLSFileSystem(client, ACCOUNT, 'dfs.core.windows.net') + + def read(stream, length, offset): + stream.seek(offset) + return stream.read(length) + + return SimpleNamespace( + open=lambda name: filesystem.open(AUTHORITY + '/' + name, 'rb', + cache_type='none'), + identity=lambda stream: dict(size=stream.size, etag=stream.etag), + read=read, close=client.close, config=dict( + authentication='synthetic_shared_key', retry_total=0, + http_pool=workers, initial_get_bytes=WINDOW, chunk_bytes=4 << 20, + network_block_bytes=transport.connection_config.data_block_size, + transfer_ceiling=transfer.CONNECTIONS)) + + +def memory(): + # Capture before output hashing/footer verification, which can raise HWM. + result = {} + for filename, fields in [('status', {'VmHWM', 'VmRSS'}), + ('smaps_rollup', {'Pss'})]: + for line in Path('/proc/self', filename).read_text().splitlines(): + key, separator, value = line.partition(':') + if separator and key in fields: + result[key + '_kib'] = int(value.split()[0]) + return result + + +def versions(): + from kumo_adls import sdk_filesystem, transfer + + result = dict( + python=sys.version, platform=sys.platform, cpu_count=os.cpu_count(), + source_sha256={ + module.__name__: sha256(Path(module.__file__)) + for module in (sdk_filesystem, transfer) + }) + for name in ('pyarrow', 'fsspec', 'azure-core', + 'azure-storage-file-datalake', 'azure-storage-blob', + 'requests', 'urllib3'): + result[name] = importlib.metadata.version(name) + return result + + +def run_child(args): + import pyarrow.parquet as pq + + manifest = json.loads(args.manifest.read_text()) + files = manifest['files'] + client = make_client(args.provider, args.endpoint, args.worker_count) + args.run_dir.mkdir(parents=True, exist_ok=False) + before_memory = memory() + + def copy(entry): + name, size = entry['path'], entry['size'] + destination = args.run_dir / name + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = None + try: + with client.open(name) as source: + before = client.identity(source) + if before != dict(size=size, etag=entry['etag']): + raise ValueError( + f'Opened source differs from manifest: {name}') + with tempfile.NamedTemporaryFile(dir=destination.parent, + buffering=NETWORK_BLOCK, + delete=False) as target: + temporary = Path(target.name) + for offset in range(0, size, WINDOW): + length = min(WINDOW, size - offset) + data = client.read(source, length, offset) + if len(data) != length or target.write(data) != length: + raise OSError(f'Short read or local write: {name}') + del data + with client.open(name) as source: + after = client.identity(source) + if before != after: + raise ValueError(f'Source identity changed: {name}') + os.replace(temporary, destination) + temporary = None + return dict(path=name, before=before, after=after) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + try: + cpu = time.process_time() + started = time.perf_counter() + with ThreadPoolExecutor(max_workers=args.worker_count) as pool: + identities = list(pool.map(copy, files)) + elapsed = time.perf_counter() - started + cpu = time.process_time() - cpu + after_memory = memory() + verified = [] + for entry in files: + path = args.run_dir / entry['path'] + actual = dict(path=entry['path'], size=path.stat().st_size, + sha256=sha256(path), + rows=pq.ParquetFile(path).metadata.num_rows) + if any(actual[key] != entry[key] + for key in ('size', 'sha256', 'rows')): + raise ValueError(f'Output verification failed: {entry["path"]}') + verified.append(actual) + result = dict(provider=args.provider, workers=args.worker_count, + protocol_only=args.check_protocol, + geometry=DIAGNOSTIC_GEOMETRY, versions=versions(), + client_config=client.config, memory_before=before_memory, + memory_after=after_memory, identities=identities, + verified_files=verified, + stored_bytes=manifest['total_bytes']) + if not args.check_protocol: + result.update( + elapsed_seconds=elapsed, process_cpu_seconds=cpu, + mib_per_second=manifest['total_bytes'] / elapsed / (1 << 20)) + write_json(args.child_result, result) + finally: + client.close() + if args.run_dir.exists(): + shutil.rmtree(args.run_dir) + + +def rpc(endpoint, operation): + with urlopen(f'http://{endpoint}/__control/{operation}', + timeout=65) as response: + return json.load(response) + + +def validate_trace(trace, manifest, workers, provider): + if trace['unknown'] or trace['failures']: + raise ValueError('Unexpected or failed requests: ' + json.dumps(trace)) + expected_requests = 0 + if trace['max_active']['GET'] > workers or trace['max_active'][ + 'HEAD'] > workers: + raise ValueError( + 'Observed request concurrency exceeded file-worker limit') + for entry in manifest['files']: + observed = trace['files'][entry['path']] + expected = [[offset, min(offset + WINDOW, entry['size']) - 1] + for offset in range(0, entry['size'], WINDOW)] + actual = sorted(get['window'] for get in observed['gets']) + if len(observed['heads']) != 2 or actual != expected: + raise ValueError( + f'Unmatched request geometry: {entry["path"]}: {observed}') + if any(item['etag'] != entry['etag'] + for item in observed['heads'] + observed['gets']): + raise ValueError('Source ETag changed during request trace') + if provider == 'adls' and any(get['if_match'] != entry['etag'] + for get in observed['gets']): + raise ValueError('An ADLS GET lacked its raw ETag condition') + expected_auth = 'SharedKey' if provider == 'adls' else 'AWS4-HMAC-SHA256' + if any(item['auth_scheme'] != expected_auth + for item in observed['heads'] + observed['gets']): + raise ValueError('A storage request lacked its provider signature') + expected_requests += 2 + len(expected) + if (trace['successes'] != expected_requests + or trace['response_body_bytes'] != manifest['total_bytes']): + raise ValueError( + 'Request totals or transferred bytes differ from manifest') + + +def run_parent(args): + source = args.source.resolve(strict=True) + output = args.output.resolve() + results = args.results.resolve() + if output.is_relative_to(source) or source.is_relative_to(output): + raise ValueError( + 'Source and temporary output directories must be separate') + if results.is_relative_to(source): + raise ValueError('Results must be outside the read-only source fixture') + output.mkdir(parents=True, exist_ok=True) + results.mkdir(parents=True, exist_ok=False) + maximum = 1 if args.check_protocol else args.max_files + manifest = make_manifest(source, maximum, args.expected_files) + manifest_path = results / 'manifest.json' + write_json(manifest_path, manifest) + write_json( + results / 'arguments.json', + dict(arguments=vars_for_json(args), geometry=DIAGNOSTIC_GEOMETRY, + benchmark_sha256=sha256(Path(__file__)))) + state = RequestState(manifest) + server = Server(('127.0.0.1', 0), ObjectHandler) + server.state = state + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + endpoint = f'127.0.0.1:{server.server_port}' + workers = [1] if args.check_protocol else args.workers + latencies = [0] if args.check_protocol else args.latencies_ms + trials = 1 if args.check_protocol else args.trials + pair = 0 + try: + for worker_count in workers: + for latency in latencies: + for trial in range(trials): + providers = ('adls', 's3') if pair % 2 == 0 else ('s3', + 'adls') + pair += 1 + for order, provider in enumerate(providers): + label = f'w{worker_count}-l{latency}-t{trial}-{order}-{provider}' + run_dir = output / (label + '-' + uuid4().hex) + child_result = results / (label + '.child.json') + if run_dir.exists(): + raise FileExistsError( + f'Refusing an existing output: {run_dir}') + rpc(endpoint, f'reset?latency_ms={latency}') + command = [ + sys.executable, + str(Path(__file__).resolve()), '--child', + '--provider', provider, '--endpoint', endpoint, + '--worker-count', + str(worker_count), '--manifest', + str(manifest_path), '--run-dir', + str(run_dir), '--child-result', + str(child_result) + ] + if args.check_protocol: + command.append('--check-protocol') + try: + environment = os.environ.copy() + environment.update( + KUMO_ADLS_TRANSFER_WORKERS=str(worker_count), + KUMO_SOURCE_COPY_WORKERS=str(worker_count), + NO_PROXY='127.0.0.1,localhost', + no_proxy='127.0.0.1,localhost') + completed = subprocess.run( + command, check=False, capture_output=True, + text=True, timeout=args.trial_timeout, + env=environment) + trace = rpc(endpoint, 'snapshot') + write_json(results / (label + '.trace.json'), trace) + if completed.returncode: + raise RuntimeError(completed.stderr + or completed.stdout) + validate_trace(trace, manifest, worker_count, + provider) + result = json.loads(child_result.read_text()) + result.update( + latency_ms=latency, trial=trial, + provider_order=order, + subset_diagnostic=manifest['subset_diagnostic'], + trace_file=label + '.trace.json', accepted=True) + write_json(results / (label + '.json'), result) + child_result.unlink() + print( + json.dumps(dict(trial=label, accepted=True), + separators=(',', ':')), flush=True) + except Exception as error: + write_json(results / (label + '.rejected.json'), + dict(accepted=False, reason=str(error))) + raise + finally: + if run_dir.exists(): + shutil.rmtree(run_dir) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def vars_for_json(args): + return { + name: str(value) if isinstance(value, Path) else value + for name, value in vars(args).items() + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--source', type=Path, + help='Complete immutable local Parquet table root') + parser.add_argument('--output', type=Path, + help='Temporary local output root') + parser.add_argument('--results', type=Path, + help='New results directory (must not exist)') + parser.add_argument('--workers', type=int, nargs='+', default=[4, 32]) + parser.add_argument('--latencies-ms', type=float, nargs='+', + default=[0, 10]) + parser.add_argument('--trials', type=int, default=3) + parser.add_argument('--expected-files', type=int, default=580) + parser.add_argument( + '--max-files', type=int, + help='Subset diagnostic only; not a full-fixture result') + parser.add_argument('--check-protocol', action='store_true', + help='One-file pilot; publish no timings') + parser.add_argument('--trial-timeout', type=float, default=900) + parser.add_argument('--child', action='store_true', help=argparse.SUPPRESS) + for option in ('provider', 'endpoint'): + parser.add_argument('--' + option, help=argparse.SUPPRESS) + for option in ('manifest', 'run-dir', 'child-result'): + parser.add_argument('--' + option, type=Path, help=argparse.SUPPRESS) + parser.add_argument('--worker-count', type=int, help=argparse.SUPPRESS) + args = parser.parse_args() + if args.child: + run_child(args) + else: + if args.source is None or args.output is None or args.results is None: + parser.error('--source, --output, and --results are required') + if (any(workers < 1 for workers in args.workers) or args.trials < 1 + or any(latency < 0 for latency in args.latencies_ms) + or (args.max_files is not None and args.max_files < 1)): + parser.error( + 'Worker/trial/file counts must be positive and latencies nonnegative' + ) + run_parent(args) + + +if __name__ == '__main__': + main() diff --git a/scripts/benchmarks/object_store/matched-transfer.md b/scripts/benchmarks/object_store/matched-transfer.md new file mode 100644 index 000000000..97219a7a4 --- /dev/null +++ b/scripts/benchmarks/object_store/matched-transfer.md @@ -0,0 +1,108 @@ +# Frozen-geometry object-store stream diagnostic + +The [standalone harness](compare_streams.py) compares the actual ADLS +SDK/fsspec adapter and native Arrow S3 stream under a fixed diagnostic request +geometry and local write loop. The ADLS client is constructed with explicit +diagnostic settings; this does not measure production transfer defaults. It creates only a loopback server and local output; it does not +provision or write cloud resources. This page documents the method, not +historical timing results or cloud-service rankings. + +## Fixed diagnostic geometry + +These constants deliberately stay independent of production tuning: + +| Control | Diagnostic value | +| --- | ---: | +| Common sequential GET window | 8 MiB | +| ADLS SDK initial response | 8 MiB | +| ADLS SDK subsequent chunk setting | 4 MiB | +| HTTP transport / local write block | 1 MiB | + +The 8 MiB stream windows fit within the diagnostic SDK initial response. +These values are not the current production bulk-transfer geometry. Results +and arguments record `geometry.derived_from_production=false` with the exact +controls. Keep earlier frozen harness snapshots and their hashes with historical +evidence; relocating this maintained script does not change those snapshots. + +## Controls and acceptance + +Prepare one complete immutable Parquet table and a per-file manifest of SHA256, +byte lengths and footer rows. Both clients must read identical files and +encodings. Equal stored-byte totals from unrelated data do not establish equal +work. + +Each accepted trial must satisfy: + +- Exactly two HEADs per file: opening identity and a fresh identity after + reading. Discovery/schema sampling and provider listing shortcuts are excluded. +- Identical explicit 8 MiB GET windows, including partial final windows, and + exact response-body byte totals derived from the manifest. Retries, extra + requests, changed identities or failed responses reject the trial. +- One in-flight read per file and the same file-worker ceiling. Both clients + use the same buffered local writes, short-read/write checks, temporary files, + version recheck, close and atomic promotion. +- Every output's SHA256, size and footer rows match the source. Verification + runs after timing, CPU and RSS capture. + +ADLS disables fsspec read-ahead. S3 uses native Arrow +`open_input_file`/`read_at`, without replacing that path with boto3. ADLS sends +its opening ETag as a GET condition; the native S3 interface exposes no explicit +If-Match option here. Record observed headers and use a stable fixture with +fresh before/after identities. This does not establish identical mutation +semantics. + +ADLS slots and HTTP pooling match the requested worker count; Arrow I/O threads +use that count. Native S3 exposes no public HTTP-pool sizing option. Traces +record realized overlap and reject counts above the common ceiling. Equal +ceilings do not guarantee equal overlap. Both clients disable retries. + +A common HTTP/1.1 server handles both URL forms with identical file responses, +TCP_NODELAY and a common listen backlog. Include zero injected delay as well +as a latency condition. Server and clients share the host and local disk; +record page-cache state and server resource pressure rather than interpreting +server saturation as a client result. + +## Reproduction + +Use Linux with provider dependencies and an Arrow build supporting S3. Prepare +the complete local table separately; do not overlap hashing, tests, fixture +preparation or another transfer workload with timing. + +```bash +export PYTHONPATH=packages/kumo-adls/src:packages/kumo-core/src:packages/kumo-api +python scripts/benchmarks/object_store/compare_streams.py \ + --source /data/table --output /data/transfer-scratch \ + --results /data/protocol-check --check-protocol +python scripts/benchmarks/object_store/compare_streams.py \ + --source /data/table --output /data/transfer-scratch \ + --results /data/matched-transfers --workers 4 32 --latencies-ms 0 10 --trials 3 +``` + +Each result directory must be new. The protocol check selects the largest file +to exercise multiple ranges and emits no timing fields. The harness defaults +to 580 expected files; use `--expected-files` for another complete table. +`--max-files` is an explicitly marked subset diagnostic, not a complete-table +result. + +Trials alternate provider order in fresh child processes. Imports and client +construction are outside timing; source hashing warms the OS page cache. +Results retain source manifests, imported-code/dependency versions, harness +hash, identities, per-file traces, output fingerprints, requested and observed +concurrency, and individual metrics. Preserve rejections and explain exclusions. +RSS includes loaded dependencies and excludes server memory; it is not merely +range-buffer allocation. + +## Interpretation + +This common sequential range loop deliberately bypasses both providers' bulk +fetch callbacks and backend discovery/versioning +policies. It can diagnose stream-stack overhead under matched geometry; it +cannot establish complete connector ingestion, native deployment parity, +object-store durability, TLS/authentication cost, cloud networking or service +throttling. A smaller worker ceiling is not a fixed-RAM-quota benchmark. + +Use the [live bulk-reader harness](compare_live.py) and +[production transfer methodology](../../../packages/kumo-adls/docs/performance.md) +for production bulk-reader measurements on the appropriate native cloud hosts. Keep results from +these different scopes separate; historical raw traces/CSV remain provenance, +not current ADLS native-host measurements. diff --git a/scripts/check_python_license_metadata.py b/scripts/check_python_license_metadata.py index ffc7e8aa6..359ecdcbc 100644 --- a/scripts/check_python_license_metadata.py +++ b/scripts/check_python_license_metadata.py @@ -71,6 +71,8 @@ class DerivedSourceRequirement: wheel_source_root="packages/kumo-core/src"), DistributionSpec("packages/kumo-databricks/pyproject.toml", "kumo-databricks", ("LICENSE", )), + DistributionSpec("packages/kumo-adls/pyproject.toml", "kumo-adls", + ("LICENSE", )), DistributionSpec("packages/kumo-gcs/pyproject.toml", "kumo-gcs", ("LICENSE", )), DistributionSpec("packages/kumo-pql/pyproject.toml", "kumo-pql", diff --git a/scripts/ci/check_dependency_boundaries.sh b/scripts/ci/check_dependency_boundaries.sh index 23f403a23..14a1c44a6 100755 --- a/scripts/ci/check_dependency_boundaries.sh +++ b/scripts/ci/check_dependency_boundaries.sh @@ -42,12 +42,13 @@ PACKAGES=( "packages/kumo-bigquery:src:kumo_bigquery::kumoapi|google" "packages/kumo-databricks:src:kumo_databricks::databricks|kumoapi|pydantic" "packages/kumo-gcs:src:kumo_gcs::kumoapi|google" + "packages/kumo-adls:src:kumo_adls::" "packages/kumo-s3:src:kumo_s3::botocore|kumoapi|pydantic" ) # Module names deptry can't infer correctly from the distribution name. # pkg=a|b means the distribution provides multiple top-level modules. -MODULE_MAP="kumo-api=kumoapi,kumo-pql=kumopql,pytorch-frame=torch_frame,faiss-cpu=faiss,faiss-gpu=faiss,pytorch-lightning=pytorch_lightning|lightning_fabric,antlr4-python3-runtime=antlr4,pyyaml=yaml,python-dateutil=dateutil" +MODULE_MAP="azure-core=azure,azure-identity=azure,azure-storage-file-datalake=azure,kumo-api=kumoapi,kumo-pql=kumopql,pytorch-frame=torch_frame,faiss-cpu=faiss,faiss-gpu=faiss,pytorch-lightning=pytorch_lightning|lightning_fabric,antlr4-python3-runtime=antlr4,pyyaml=yaml,python-dateutil=dateutil" overall_status=0 diff --git a/scripts/gen_third_party.py b/scripts/gen_third_party.py index 88e2f5c1b..323a78281 100755 --- a/scripts/gen_third_party.py +++ b/scripts/gen_third_party.py @@ -49,6 +49,7 @@ "kumo-core", "kumo-databricks", "kumo-gcs", + "kumo-adls", "kumo-pql", "kumo-s3", "kumo-snowflake", diff --git a/scripts/test_public_connectors.py b/scripts/test_public_connectors.py index a87bdc831..24d5978e3 100644 --- a/scripts/test_public_connectors.py +++ b/scripts/test_public_connectors.py @@ -20,7 +20,7 @@ 'parquet': ('packages/kumo-core/tests/sources/parquet',), 's3': ('packages/kumo-s3/tests',), 'gcs': ('packages/kumo-gcs/tests',), - 'adls': ('test/io_utils/fsutil/adls',), + 'adls': ('test/io_utils/fsutil/adls', 'packages/kumo-adls/tests'), 'databricks': ('packages/kumo-databricks/tests',), 'snowflake': ('packages/kumo-snowflake/tests',), 'bigquery': ('packages/kumo-bigquery/tests',), diff --git a/test/scripts/test_engine_wheel.py b/test/scripts/test_engine_wheel.py index 5eb15dc84..e89935bf0 100644 --- a/test/scripts/test_engine_wheel.py +++ b/test/scripts/test_engine_wheel.py @@ -29,7 +29,7 @@ def test_wheel_builds_all_first_party_packages_from_monorepo( for name in [ 'kumo-api', 'kumo-core', 'kumo-pql', 'kumo-engine', 'kumo-sdk', 'kumo-databricks', 'kumo-snowflake', 'kumo-s3', 'kumo-gcs', - 'kumo-bigquery', 'diskgraph' + 'kumo-bigquery', 'kumo-adls', 'diskgraph' ]: package = repo / 'packages' / name package.mkdir(parents=True) @@ -74,7 +74,7 @@ def test_wheel_builds_all_first_party_packages_from_monorepo( sources = [call[-1] for call in recorded] expected = ['kumo-api', 'kumo-core', 'kumo-engine', 'kumo-pql'] expected += [ - 'kumo-bigquery', 'kumo-databricks', 'kumo-gcs', 'kumo-s3', + 'kumo-adls', 'kumo-bigquery', 'kumo-databricks', 'kumo-gcs', 'kumo-s3', 'kumo-snowflake' ] assert [Path(source) for source in sources] == [ @@ -138,11 +138,12 @@ def _wheel(directory, name, requires=(), extras=(), tag='py3-none-any', CORE_COMPONENTS = ['kumo-engine', 'kumo-api', 'kumo-pql', 'kumo-core', 'diskgraph'] -ENGINE_EXTRAS = ['snowflake', 'databricks', 's3', 'text', 'test'] +ENGINE_EXTRAS = ['adls', 'snowflake', 'databricks', 's3', 'text', 'test'] COMPONENTS = { 'kumo-engine': { 'requires': [ 'pandas', 'kumo-core[gpu]>=1.0', 'pytest; extra == "test"', + 'kumo-adls; extra == "adls"', 'kumo-snowflake[sdk]>=1.0; extra == "snowflake"', 'kumo-databricks; extra == "databricks"', 'kumo-s3; extra == "s3"', 'kumo-glove-6b-50d==1.0.0; extra == "text"' @@ -181,6 +182,10 @@ def _wheel(directory, name, requires=(), extras=(), tag='py3-none-any', 'kumo-databricks': { 'requires': ['kumo-core==1.0', 'databricks-sdk'], }, + 'kumo-adls': { + 'requires': ['kumo-core==1.0', 'azure-identity', + 'azure-storage-file-datalake'], + }, 'kumo-s3': { 'requires': ['kumo-core==1.0', 'boto3'], }, @@ -286,7 +291,7 @@ def test_engine_faiss_platform_selection_survives_merge( @pytest.mark.parametrize('providers', [ - ['kumo-databricks', 'kumo-s3', 'kumo-snowflake'], + ['kumo-adls', 'kumo-databricks', 'kumo-s3', 'kumo-snowflake'], ['kumo-databricks'], ]) def test_provider_dependencies_only_come_with_their_extra(tmp_path, providers): @@ -308,6 +313,10 @@ def test_provider_dependencies_only_come_with_their_extra(tmp_path, providers): assert provider.replace('-', '_') + '/__init__.py' in names # ...but its dependencies are conditional on the extra. expected_extras = { + 'kumo-adls': { + 'azure-identity; extra == "adls"', + 'azure-storage-file-datalake; extra == "adls"' + }, 'kumo-databricks': {'databricks-sdk; extra == "databricks"'}, 'kumo-s3': {'boto3; extra == "s3"'}, 'kumo-snowflake': { diff --git a/tests/public/Dockerfile.dockerignore b/tests/public/Dockerfile.dockerignore index 4bedfe32e..54d2bca52 100644 --- a/tests/public/Dockerfile.dockerignore +++ b/tests/public/Dockerfile.dockerignore @@ -19,6 +19,8 @@ !packages/kumo-core/src/kumo_core/io/__init__.py !packages/kumo-core/src/kumo_core/io/filesystem.py !packages/kumo-core/src/kumo_core/io/json.py +!packages/kumo-core/src/kumo_core/object_store_sink.py +!packages/kumo-core/src/kumo_core/artifact_export.py !packages/kumo-core/src/kumo_core/stype_inference.py !packages/kumo-core/src/kumo_core/types/__init__.py !packages/kumo-core/src/kumo_core/types/enums.py @@ -32,6 +34,8 @@ !packages/kumo-s3/tests/** !packages/kumo-gcs/src/** !packages/kumo-gcs/tests/** +!packages/kumo-adls/src/** +!packages/kumo-adls/tests/** !packages/kumo-databricks/src/** !packages/kumo-databricks/tests/** !packages/kumo-snowflake/src/** diff --git a/tests/public/pytest.ini b/tests/public/pytest.ini index 6b6c73a82..a4d9ecc77 100644 --- a/tests/public/pytest.ini +++ b/tests/public/pytest.ini @@ -9,6 +9,7 @@ pythonpath = ../../packages/kumo-core/tests ../../packages/kumo-s3/src ../../packages/kumo-gcs/src + ../../packages/kumo-adls/src ../../packages/kumo-bigquery/src ../../packages/kumo-databricks/src ../../packages/kumo-snowflake/src diff --git a/uv.lock b/uv.lock index 6c8b4649c..60686de0e 100644 --- a/uv.lock +++ b/uv.lock @@ -51,6 +51,7 @@ torch-geometric = "2026-10-13T00:00:00Z" members = [ "diskgraph", "kumo", + "kumo-adls", "kumo-api", "kumo-bigquery", "kumo-core", @@ -3110,6 +3111,40 @@ provides-extras = ["tracing", "drawing", "rest", "full", "test", "dev", "data-pl [package.metadata.requires-dev] public-tests = [{ name = "torch", index = "https://download.pytorch.org/whl/cpu", conflict = { package = "kumo", group = "public-tests" } }] +[[package]] +name = "kumo-adls" +version = "0.1.0.dev1" +source = { editable = "packages/kumo-adls" } +dependencies = [ + { name = "azure-core" }, + { name = "azure-identity" }, + { name = "azure-storage-file-datalake" }, + { name = "fsspec", extra = ["http"] }, + { name = "kumo-api" }, + { name = "kumo-core" }, + { name = "pyarrow" }, + { name = "requests" }, +] + +[package.optional-dependencies] +test = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "azure-core", specifier = ">=1.31,<2" }, + { name = "azure-identity", specifier = ">=1.25,<2" }, + { name = "azure-storage-file-datalake", specifier = ">=12.22,<13" }, + { name = "fsspec" }, + { name = "kumo-api", editable = "packages/kumo-api" }, + { name = "kumo-core", editable = "packages/kumo-core" }, + { name = "pyarrow", specifier = ">=14" }, + { name = "pytest", marker = "extra == 'test'", specifier = ">=7.4" }, + { name = "requests", specifier = ">=2.31,<3" }, +] +provides-extras = ["test"] + [[package]] name = "kumo-api" version = "3.0.0" @@ -3334,6 +3369,9 @@ dependencies = [ ] [package.optional-dependencies] +adls = [ + { name = "kumo-adls" }, +] bigquery = [ { name = "kumo-bigquery" }, ] @@ -3374,6 +3412,7 @@ requires-dist = [ { name = "faiss-cpu", marker = "sys_platform == 'darwin'", specifier = "==1.14.3" }, { name = "faiss-gpu", marker = "sys_platform == 'linux'", specifier = "==1.14.3" }, { name = "filelock", specifier = ">=3.0.0" }, + { name = "kumo-adls", marker = "extra == 'adls'", editable = "packages/kumo-adls" }, { name = "kumo-api", editable = "packages/kumo-api" }, { name = "kumo-bigquery", marker = "extra == 'bigquery'", editable = "packages/kumo-bigquery" }, { name = "kumo-core", editable = "packages/kumo-core" }, @@ -3391,7 +3430,7 @@ requires-dist = [ { name = "torch", marker = "extra == 'test'", specifier = ">=2.9,<2.13" }, { name = "typing-extensions", specifier = ">=4.5.0" }, ] -provides-extras = ["test", "snowflake", "databricks", "bigquery", "s3", "text", "gcs"] +provides-extras = ["test", "snowflake", "databricks", "bigquery", "s3", "text", "adls", "gcs"] [[package]] name = "kumo-gcs" From b279ec2af6d8d69864fc41d61ccb67a112e91c44 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Thu, 22 Oct 2026 19:35:09 +0000 Subject: [PATCH 02/15] Align ADLS README with connector guides --- packages/kumo-adls/README.md | 191 +++++++--------------- packages/kumo-adls/docs/implementation.md | 111 +++++++++++++ 2 files changed, 171 insertions(+), 131 deletions(-) create mode 100644 packages/kumo-adls/docs/implementation.md diff --git a/packages/kumo-adls/README.md b/packages/kumo-adls/README.md index 391fa4ae8..5795d6282 100644 --- a/packages/kumo-adls/README.md +++ b/packages/kumo-adls/README.md @@ -1,155 +1,84 @@ # kumo-adls Azure Data Lake Storage Gen2 source and sink provider for `kumo-engine`. -Install through `kumo-engine[adls]`. The package registers source, sink, and -identity plugins using the same entry-point groups as other providers. +The provider supplies an authenticated Arrow filesystem and conditional SDK +file transfers to the shared object-store machinery in `kumo-core`. Core owns +CSV/Parquet layout, conversion, local staging, and output. + +`backend.py` and `sink.py` implement the local `SourceBackend` and +`ArtifactSink` contracts. Installing this distribution registers source, sink, +and identity plugins. Engine imports the provider lazily through its extra. + +## Local workflow + +```sh +pip install "kumo-engine[adls]" +``` + +Use a wheel or source checkout containing this provider; this does not establish +availability in an older published release. ```python -from kumo.engine import ADLSConnector, Table +from kumo.engine import ADLSConnector, Graph, Table lake = ADLSConnector( root_dir="abfss://curated@companylake.dfs.core.windows.net/orders/", ) orders = Table.from_source_table(lake["orders.parquet"], primary_key="id") +snapshot = Graph(tables=[orders]).snapshot() ``` -A rooted connector discovers and reads CSV/Parquet tables beneath its -filesystem/prefix and writes predictions, embeddings, and model bundles there. -A rootless connector reads exact canonical `abfss://` URIs; discovery and output -are unavailable. SQL and append output are unsupported. +Reads use shared `ObjectStoreSourceReader.prepare` in `ReadMode.STREAMED`, +including local file staging. `Graph.snapshot()` passes the staged Parquet to +DiskGraph through the existing ingestion interface; DiskGraph needs no Azure +credentials for those local files. + +A rooted connector discovers CSV/Parquet tables and writes predictions, +embeddings, and model bundles beneath its filesystem/prefix. A rootless +connector reads exact canonical `abfss://` URIs; discovery and output are +unavailable. SQL, warehouse-style `ReadMode.STAGED`, and append output are +unsupported. Prepared exports use parallel SDK uploads; publication is not +atomic, so consumers must wait for export completion. + +## Authentication and resources Without an identity, requests use Azure's ambient credential chain. An explicit `ADLSIdentity(tenant_id=..., client_id=..., cloud=...)` uses workload federation -with `AZURE_FEDERATED_TOKEN_FILE`. Identity secrets, storage keys, and SAS tokens -are not connector arguments. The selected identity cloud must match the URI's -public, US Government, or China endpoint. See the shared -[ADLS storage contract](../../docs/cloud-storage/adls.md). - -The provider adapts Azure SDK operations to shared Arrow source and sink -machinery in `kumo-core`. ETags guard range/full-file reads and source versions; -account-qualified paths distinguish accounts and clouds. Azure Identity refreshes -tokens, and identity cleanup closes cached clients. Sessions are cached for the -process lifetime, without automatic eviction; identity cleanup or -`kumo_adls.session.close()` releases them. Construction runs under a global lock -to prevent duplicate clients, so concurrent first use of different accounts can -serialize. Arrow/fsspec wrappers are lightweight and created per access. - -Identity probing checks token acquisition only. It does not verify storage -existence, RBAC/ACL permissions, or connectivity; list or read the configured -source to verify those. - -Bulk downloads use one conditional SDK stream per file, with 16 MiB initial -and subsequent responses and 16 MiB HTTP transport blocks. Files download -concurrently; each file holds one transfer slot and uses no per-file executor -or queued chunk futures. The SDK writes each response to the local temporary -file before requesting the next chunk. The initial response remains buffered -for the stream's lifetime; this is not zero-copy network-to-disk streaming. -Prepared uploads retain up to four 4 MiB chunks. Large projected Arrow reads -can use up to four SDK workers, bounded by the shared transfer slots. A 64 MiB -range uses up to three workers; ranges above 64 MiB can use four. The initial -response plus four 16 MiB chunks can account for about 80 MiB before transport -buffers, SDK copies, and the complete requested range returned by `readall()`. -Projected reads are not covered by the 64 MiB bulk-copy reservation. - -Transfers across files and accounts share a CPU-derived concurrency ceiling -(at most 128, lowered by `KUMO_SOURCE_COPY_WORKERS`). The HTTP connection pool -matches that ceiling. Metadata bypasses transfer slots and shares the HTTP -transport. Uploads reserve 64 MiB per active file through the shared memory -ledger; source bulk copies use the core copy reservation. These reservations -are accounting estimates, not caps on total process RSS. Larger responses -reduce requests and fragment assembly while increasing memory per active file. -Plan roughly 48–64 MiB of payload buffers per active bulk download, with -additional headroom for SDK copies and other allocations. At 128 active files, -that estimate alone is 6–8 GiB; the reservation does not guarantee that this -fits a host's memory limit. Actual memory depends on allocation lifetimes. - -For memory-constrained hosts, set `KUMO_ADLS_TRANSFER_WORKERS=4` before -importing the provider to limit ADLS transfer requests and its HTTP pool -without lowering other providers' copy limits. The value must be a positive -integer and cannot raise the shared copy/CPU ceiling. The limit is process-wide -across ADLS accounts and readers; it is read once at import. Even with one slot, -files use the same serial SDK stream. Arrow buffers, cached bytes, SDK copies, -loaded libraries, and Python overhead still consume memory. Reduced ceilings -are covered by correctness tests; a fixed-RAM host requires separate validation. -An invalid value fails at the provider's first import. Set these variables in -the launching environment and restart an existing process after changing them. -Catchable allocation failures inside ADLS SDK operations retain their original -exception and traceback, with a note suggesting these controls if the failure -occurred during a transfer. They may not help with listing or session setup. -An operating-system OOM kill cannot provide a Python error message. -Reducing the shared copy ceiling -also limits waiting file workers and their buffers; the ADLS-only limit does not. - -Projected Arrow reads disable fsspec read-ahead, avoiding unrelated column data. -Successful small ranges can reuse immutable bytes under the opened file's raw -ETag. The process-wide LRU holds at most 64 MiB of payload/1,024 entries, admits -ranges up to 128 KiB, and isolates filesystem/principal namespaces. Opening a -file still performs fresh properties checks; metadata is never cached here. -Existing opened streams may reuse their pinned version, while ingest's final -identity checks still reject source changes. Filesystem cleanup releases its -cached ranges; Python container/key overhead is additional. Each source access -creates a new filesystem namespace. This cache avoids repeated footer reads -within one reader, not across separate reads or connector operations. The -64 MiB limit bounds aggregate retained payload across concurrent readers; it -is not a persistent connector cache or a preallocated buffer. -Recursive discovery uses Azure's paginated recursive listing. Version ETag -checks use fresh bounded listings when complete ETags are available, falling -back to concurrent HEADs for missing metadata, narrow permissions, and exact -single files. ADLS uses the shared object-store `prepare`: discover files once, -stamp the version from that same selected list, and read those files. A supplied -version is stamped without refreshing it, matching S3/GCS. Shared per-file -identity guards check selected objects before and after reading; files added -after discovery belong to a subsequent read. Identity checks may still perform -listings or HEADs, so one discovery does not mean one metadata request. -ETag quote and directory-path normalization can invalidate existing version -hashes once, requiring restaging. -Prepared export files and model files use seekable parallel uploads; Arrow -writers keep the buffered streaming path. Failed prepared uploads delete only -the object matching their creation ETag, preserving another writer's output. -Prepared uploads create the destination before uploading; concurrent readers -can observe an empty or incomplete object. Publishing an individual file is -not atomic, so consumers must wait for export completion. -Streaming writes are not transactional: Arrow closes its wrapper without -forwarding serialization exceptions to the fsspec context manager, so failed -serialization can leave partial output. Direct fsspec contexts can discard -on caller errors; upload failures also trigger best-effort cleanup. - -Downloaded files remain owner-private, matching the S3 staging path. Local -engine ingestion uses the same OS user; sharing scratch files across users -requires explicitly managed permissions. Temporary files are closed before -promotion with `os.replace`, including on failure cleanup paths. This avoids -requiring POSIX open-file replacement; it does not establish full Windows -engine support. - -These settings remove known serial paths and request amplification; they are -not evidence of optimal live Azure throughput. Validate full logical tables, -client placement, resource limits, retries, and service throttling before making -production performance claims. -See the [performance methodology](docs/performance.md) for complete-table -validation, native-host comparisons, per-reader cache scope, and resource -tradeoffs. The [matched-request stream method](../../scripts/benchmarks/object_store/matched-transfer.md) -separately controls HEADs, range GETs, bytes, and worker limits for ADLS and -native S3; it uses fixed historical diagnostic geometry rather than the -production 16 MiB download configuration or full connector ingestion. -For live comparisons on Linux using your own storage locations, the -[bulk-reader transfer harness](../../scripts/benchmarks/object_store/compare_live.py) accepts complete-table -manifests and verifies staged files after each production bulk-reader trial. - -Run offline provider and existing filesystem tests from the repository root: - -```bash +with `AZURE_FEDERATED_TOKEN_FILE`. Secrets, storage keys, and SAS tokens are not +connector arguments. The identity cloud must match the public, US Government, +or China URI endpoint. Probing checks token acquisition only; list or read the +configured source to verify storage permissions and connectivity. + +Set a positive `KUMO_ADLS_TRANSFER_WORKERS` before the provider's first import +to lower ADLS transfer concurrency and HTTP pooling. Restart existing processes +after changing it. `KUMO_SOURCE_COPY_WORKERS` lowers the shared file-copy +ceiling. Both controls limit concurrency rather than total memory; SDK and +projected-read buffers need additional headroom. + +See the [storage contract](../../docs/cloud-storage/adls.md), +[implementation and resource controls](docs/implementation.md), and +[performance methodology](docs/performance.md) for authentication scope, +consistency, memory estimates, export behavior, and benchmark limits. + +## Tests + +From the repository root, run public tests without cloud credentials: + +```sh python3 scripts/test_public.py --group adls ``` -Live access, token renewal, RBAC/ACL propagation, and sovereign endpoints require -validation against an explicitly configured Azure environment. Offline tests do -not establish live cloud support. +See the [public test guide](../../docs/development/public-tests.md) for setup +and the [test matrix](../../docs/connector-test-matrix.md#reader-groups) for coverage. -Run the read-only live filesystem checks with an existing HNS directory: +Read-only live filesystem checks require an existing HNS directory and ambient +credentials; they do not create or delete cloud paths: -```bash +```sh KUMO_ADLS_LIVE_ROOT='abfss://filesystem@account.dfs.core.windows.net/existing-directory/' \ python -m pytest packages/kumo-adls/tests/live --confcutdir=packages/kumo-adls/tests ``` -The live checks use ambient credentials and do not create or delete cloud paths. +Live token renewal, RBAC/ACL propagation, sovereign endpoints, and full graph +workloads require separate validation. Offline tests and raw-copy measurements +do not establish end-to-end graph performance or live cloud support. diff --git a/packages/kumo-adls/docs/implementation.md b/packages/kumo-adls/docs/implementation.md new file mode 100644 index 000000000..40f92f82b --- /dev/null +++ b/packages/kumo-adls/docs/implementation.md @@ -0,0 +1,111 @@ +# ADLS implementation and resource controls + +The provider adapts Azure SDK operations to shared Arrow source and sink +machinery in `kumo-core`. ETags guard range/full-file reads and source versions; +account-qualified paths distinguish accounts and clouds. Azure Identity refreshes +tokens, and identity cleanup closes cached clients. Sessions are cached for the +process lifetime, without automatic eviction; identity cleanup or +`kumo_adls.session.close()` releases them. Construction runs under a global lock +to prevent duplicate clients, so concurrent first use of different accounts can +serialize. Arrow/fsspec wrappers are lightweight and created per access. + +Identity probing checks token acquisition only. It does not verify storage +existence, RBAC/ACL permissions, or connectivity; list or read the configured +source to verify those. + +Bulk downloads use one conditional SDK stream per file, with 16 MiB initial +and subsequent responses and 16 MiB HTTP transport blocks. Files download +concurrently; each file holds one transfer slot and uses no per-file executor +or queued chunk futures. The SDK writes each response to the local temporary +file before requesting the next chunk. The initial response remains buffered +for the stream's lifetime; this is not zero-copy network-to-disk streaming. +Prepared uploads retain up to four 4 MiB chunks. Large projected Arrow reads +can use up to four SDK workers, bounded by the shared transfer slots. A 64 MiB +range uses up to three workers; ranges above 64 MiB can use four. The initial +response plus four 16 MiB chunks can account for about 80 MiB before transport +buffers, SDK copies, and the complete requested range returned by `readall()`. +Projected reads are not covered by the 64 MiB bulk-copy reservation. + +Transfers across files and accounts share a CPU-derived concurrency ceiling +(at most 128, lowered by `KUMO_SOURCE_COPY_WORKERS`). The HTTP connection pool +matches that ceiling. Metadata bypasses transfer slots and shares the HTTP +transport. Uploads reserve 64 MiB per active file through the shared memory +ledger; source bulk copies use the core copy reservation. These reservations +are accounting estimates, not caps on total process RSS. Larger responses +reduce requests and fragment assembly while increasing memory per active file. +Plan roughly 48–64 MiB of payload buffers per active bulk download, with +additional headroom for SDK copies and other allocations. At 128 active files, +that estimate alone is 6–8 GiB; the reservation does not guarantee that this +fits a host's memory limit. Actual memory depends on allocation lifetimes. + +For memory-constrained hosts, set `KUMO_ADLS_TRANSFER_WORKERS=4` before +importing the provider to limit ADLS transfer requests and its HTTP pool +without lowering other providers' copy limits. The value must be a positive +integer and cannot raise the shared copy/CPU ceiling. The limit is process-wide +across ADLS accounts and readers; it is read once at import. Even with one slot, +files use the same serial SDK stream. Arrow buffers, cached bytes, SDK copies, +loaded libraries, and Python overhead still consume memory. Reduced ceilings +are covered by correctness tests; a fixed-RAM host requires separate validation. +An invalid value fails at the provider's first import. Set these variables in +the launching environment and restart an existing process after changing them. +Catchable allocation failures inside ADLS SDK operations retain their original +exception and traceback, with a note suggesting these controls if the failure +occurred during a transfer. They may not help with listing or session setup. +An operating-system OOM kill cannot provide a Python error message. +Reducing the shared copy ceiling +also limits waiting file workers and their buffers; the ADLS-only limit does not. + +Projected Arrow reads disable fsspec read-ahead, avoiding unrelated column data. +Successful small ranges can reuse immutable bytes under the opened file's raw +ETag. The process-wide LRU holds at most 64 MiB of payload/1,024 entries, admits +ranges up to 128 KiB, and isolates filesystem/principal namespaces. Opening a +file still performs fresh properties checks; metadata is never cached here. +Existing opened streams may reuse their pinned version, while ingest's final +identity checks still reject source changes. Filesystem cleanup releases its +cached ranges; Python container/key overhead is additional. Each source access +creates a new filesystem namespace. This cache avoids repeated footer reads +within one reader, not across separate reads or connector operations. The +64 MiB limit bounds aggregate retained payload across concurrent readers; it +is not a persistent connector cache or a preallocated buffer. +Recursive discovery uses Azure's paginated recursive listing. Version ETag +checks use fresh bounded listings when complete ETags are available, falling +back to concurrent HEADs for missing metadata, narrow permissions, and exact +single files. ADLS uses the shared object-store `prepare`: discover files once, +stamp the version from that same selected list, and read those files. A supplied +version is stamped without refreshing it, matching S3/GCS. Shared per-file +identity guards check selected objects before and after reading; files added +after discovery belong to a subsequent read. Identity checks may still perform +listings or HEADs, so one discovery does not mean one metadata request. +ETag quote and directory-path normalization can invalidate existing version +hashes once, requiring restaging. +Prepared export files and model files use seekable parallel uploads; Arrow +writers keep the buffered streaming path. Failed prepared uploads delete only +the object matching their creation ETag, preserving another writer's output. +Prepared uploads create the destination before uploading; concurrent readers +can observe an empty or incomplete object. Publishing an individual file is +not atomic, so consumers must wait for export completion. +Streaming writes are not transactional: Arrow closes its wrapper without +forwarding serialization exceptions to the fsspec context manager, so failed +serialization can leave partial output. Direct fsspec contexts can discard +on caller errors; upload failures also trigger best-effort cleanup. + +Downloaded files remain owner-private, matching the S3 staging path. Local +engine ingestion uses the same OS user; sharing scratch files across users +requires explicitly managed permissions. Temporary files are closed before +promotion with `os.replace`, including on failure cleanup paths. This avoids +requiring POSIX open-file replacement; it does not establish full Windows +engine support. + +These settings remove known serial paths and request amplification; they are +not evidence of optimal live Azure throughput. Validate full logical tables, +client placement, resource limits, retries, and service throttling before making +production performance claims. +See the [performance methodology](performance.md) for complete-table +validation, native-host comparisons, per-reader cache scope, and resource +tradeoffs. The [matched-request stream method](../../../scripts/benchmarks/object_store/matched-transfer.md) +separately controls HEADs, range GETs, bytes, and worker limits for ADLS and +native S3; it uses fixed historical diagnostic geometry rather than the +production 16 MiB download configuration or full connector ingestion. +For live comparisons on Linux using your own storage locations, the +[bulk-reader transfer harness](../../../scripts/benchmarks/object_store/compare_live.py) accepts complete-table +manifests and verifies staged files after each production bulk-reader trial. From 8093bbf2e3b4711d1bf8b7d26f2c9b734d52547c Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Thu, 22 Oct 2026 19:40:23 +0000 Subject: [PATCH 03/15] Keep ADLS performance methodology in local documentation --- packages/kumo-adls/README.md | 7 +- packages/kumo-adls/docs/implementation.md | 14 -- packages/kumo-adls/docs/performance.md | 137 ------------------ .../object_store/matched-transfer.md | 3 +- 4 files changed, 4 insertions(+), 157 deletions(-) delete mode 100644 packages/kumo-adls/docs/performance.md diff --git a/packages/kumo-adls/README.md b/packages/kumo-adls/README.md index 5795d6282..799060258 100644 --- a/packages/kumo-adls/README.md +++ b/packages/kumo-adls/README.md @@ -55,10 +55,9 @@ after changing it. `KUMO_SOURCE_COPY_WORKERS` lowers the shared file-copy ceiling. Both controls limit concurrency rather than total memory; SDK and projected-read buffers need additional headroom. -See the [storage contract](../../docs/cloud-storage/adls.md), -[implementation and resource controls](docs/implementation.md), and -[performance methodology](docs/performance.md) for authentication scope, -consistency, memory estimates, export behavior, and benchmark limits. +See the [storage contract](../../docs/cloud-storage/adls.md) +and [implementation and resource controls](docs/implementation.md) for +authentication scope, consistency, memory estimates, and export behavior. ## Tests diff --git a/packages/kumo-adls/docs/implementation.md b/packages/kumo-adls/docs/implementation.md index 40f92f82b..76b07f86d 100644 --- a/packages/kumo-adls/docs/implementation.md +++ b/packages/kumo-adls/docs/implementation.md @@ -95,17 +95,3 @@ requires explicitly managed permissions. Temporary files are closed before promotion with `os.replace`, including on failure cleanup paths. This avoids requiring POSIX open-file replacement; it does not establish full Windows engine support. - -These settings remove known serial paths and request amplification; they are -not evidence of optimal live Azure throughput. Validate full logical tables, -client placement, resource limits, retries, and service throttling before making -production performance claims. -See the [performance methodology](performance.md) for complete-table -validation, native-host comparisons, per-reader cache scope, and resource -tradeoffs. The [matched-request stream method](../../../scripts/benchmarks/object_store/matched-transfer.md) -separately controls HEADs, range GETs, bytes, and worker limits for ADLS and -native S3; it uses fixed historical diagnostic geometry rather than the -production 16 MiB download configuration or full connector ingestion. -For live comparisons on Linux using your own storage locations, the -[bulk-reader transfer harness](../../../scripts/benchmarks/object_store/compare_live.py) accepts complete-table -manifests and verifies staged files after each production bulk-reader trial. diff --git a/packages/kumo-adls/docs/performance.md b/packages/kumo-adls/docs/performance.md deleted file mode 100644 index 4f7b5cd2c..000000000 --- a/packages/kumo-adls/docs/performance.md +++ /dev/null @@ -1,137 +0,0 @@ -# Object-store transfer performance methodology - -Measure complete, verified workloads before attributing performance to a -connector. This guide describes engineering tradeoffs and reproducible controls; -it contains no deployment-specific results or resource identifiers. - -## Implementation and tradeoffs - -- Bound queued tasks as well as active threads. A worker limit alone can still - retain one future per chunk of a large object. Lazy submission bounds that - overhead independently of object size. -- Budget concurrency across files. Larger chunks or fewer workers per file - do not necessarily reduce aggregate memory when the request ceiling stays - unchanged. Capture CPU and RSS alongside elapsed time. -- Avoid redundant buffering when the SDK already holds a response. Bulk ADLS - transfers use one serial SDK downloader per file, with files copied in - parallel. Matching response and transport-block sizes reduces fragment - assembly. This is not zero-copy or direct network-to-disk streaming. -- Cache immutable bytes under a version identity with an explicit lifetime. - The small-range cache avoids repeated footer reads within one reader; it - does not cache across separate reader/source-access instances or replace fresh - metadata checks. -- Separate experiments. Do not pool measurements from different revisions, - layouts, client placements, or request policies into a cumulative speedup. - -The [filesystem adapter](../src/kumo_adls/sdk_filesystem.py), -[session setup](../src/kumo_adls/session.py), and -[transfer limits](../src/kumo_adls/transfer.py) own these behaviors. - -## Resource controls - -Bulk ADLS downloads hold one transfer slot per file and use serial conditional -16 MiB responses, with no per-file executor or chunk-future backlog. The SDK -retains its initial response while processing later chunks. Transport blocks -are also 16 MiB; prepared uploads retain up to four 4 MiB chunks. The SDK still -buffers responses; a 64 MiB shared copy reservation is not a cap on total RSS. -Arrow buffers, cached bytes, metadata, loaded libraries, and Python overhead -consume additional memory. -Budget roughly 48–64 MiB of bulk payload buffers per active file, plus SDK-copy -headroom; 128 active files can therefore imply 6–8 GiB of payload alone. These -are planning estimates, not measured peaks or strict allocation bounds. - -Projected reads use `readall()` and materialize the complete requested range. -A 64 MiB range uses up to three SDK workers after its initial 16 MiB response; -ranges above 64 MiB can use four. The initial response plus four 16 MiB chunks -can account for about 80 MiB before transport blocks, SDK copies and the output -range itself. Parallel SDK range reads can also queue futures for the requested -range. Neither the 80 MiB payload estimate nor the 64 MiB bulk-copy reservation -is a memory bound for projected reads. - -`KUMO_SOURCE_COPY_WORKERS` lowers the shared copy ceiling. Set the positive -startup value `KUMO_ADLS_TRANSFER_WORKERS` before importing the provider to -lower only ADLS transfer slots and HTTP pooling, capped by the shared CPU/copy -ceiling. Bulk downloads use the same serial SDK path at every ceiling. -Metadata requests bypass transfer slots but share the HTTP pool. These controls bound requests, -not total RSS; a reduced ceiling alone is not a fixed-RAM-quota benchmark. -Both settings belong in the launching environment; restart a process after -changing the ADLS limit, which is fixed at import. Invalid ADLS values fail at -first provider import. Catchable allocation failures inside ADLS SDK operations -retain their original exception/traceback and add a note suggesting both controls -if the failure occurred during transfer; they may not help with listing or -session setup. An OS OOM kill cannot show this guidance. For lower memory across whole reads, -reduce shared file concurrency as well as ADLS transfer slots. - -The process-wide small-range LRU caps payload at 64 MiB and 1,024 entries, with -Python key/container overhead additional. It admits complete successful ranges -up to 128 KiB, keyed by filesystem namespace, full path, raw ETag, and exact -range. Each source access creates a new namespace and cleanup releases it. -Opens and final source checks remain fresh. Prepared exports do not use this -reader cache. - -## Live comparison recipe - -Use identical complete tables and physical files on both stores. Prepare -trusted manifests containing every relative path, size, SHA256, footer row -count, and pinned provider ETag. Fixture copying and equivalence validation -must finish before timing. Do not substitute unrelated objects of equal -compressed size for equivalent work. - -The [live bulk-reader harness](../../../scripts/benchmarks/object_store/compare_live.py) only reads cloud -objects and writes local results/scratch. Use authenticated, existing sources -and a Linux environment with provider dependencies and Arrow S3 support. -For native deployment comparisons, run S3 on the AWS client and ADLS on the -Azure client in the respective storage regions. Record actual CPU/RAM, OS, -network, disk and storage-class differences; matching vCPU counts does not -make the hosts identical. - -Example single-provider phases, executed on their respective hosts: - -```bash -# AWS host -python scripts/benchmarks/object_store/compare_live.py --providers s3 \ - --manifest /data/manifests/table.json --placement /data/aws-placement.json \ - --output /data/transfer-scratch --results /data/s3-results \ - --workers 32 --trials 3 --verification-workers 8 --s3-region us-west-2 - -# Azure host, using its configured ambient identity -python scripts/benchmarks/object_store/compare_live.py --providers adls \ - --manifest /data/manifests/table.json --placement /data/azure-placement.json \ - --output /data/transfer-scratch --results /data/adls-results \ - --workers 32 --trials 3 --verification-workers 8 -``` - -Results directories must be new and scratch must fit the complete table plus -headroom. Run one timed workload per host. Preserve every individual trial, -rejection, manifest/harness/imported-source hash, dependency version and -placement record. Cloud-side cache state remains unknown in a fresh process. - -## Timing, acceptance, and limits - -Setup, authentication preflight, provenance and initial source identity checks -occur outside timing. The preparation timer covers the production reader's -`prepare(transform=False)` metadata guards and raw copying. A nested download -phase includes fetch's own metadata/checks/retries; it is not pure wire time, -and the remainder is not pure metadata latency. This workload does not measure -Graph construction, typed decoding, projection, conversion, training, or export. - -Capture process CPU and RSS/HWM before untimed output verification. Loaded -provider stacks and setup allocations can affect the high-water mark; kernel -page cache and remote service memory are excluded. Local buffered writes, -close and atomic promotion do not establish durable-disk throughput without -an explicit durability requirement. - -Require unchanged exact remote membership/identities and every output file's -SHA256, size and footer rows before accepting a trial. Verification is bounded -and untimed; its scheduling/cache advice must remain in provenance. Request, -retry, token-refresh counts and effective worker overlap are uninstrumented -in the live harness, so do not infer them from configured ceilings. - -The [matched-request harness](../../../scripts/benchmarks/object_store/matched-transfer.md) separately isolates stream -client stacks under loopback HTTP. It bypasses production bulk paths and does -not replace native cloud measurements. Its fixed 8 MiB windows, 4 MiB SDK -chunks and 1 MiB transport blocks model different geometry from production -ADLS downloads (16 MiB for all three). Historical raw artifacts are audit -provenance, not current deployment results. Projections/conversions, CSV, -fixed-RAM limits, sustained refresh/throttling, DataFrame writes, model bundles, -and exports require their own representative evidence. diff --git a/scripts/benchmarks/object_store/matched-transfer.md b/scripts/benchmarks/object_store/matched-transfer.md index 97219a7a4..2a6d6452d 100644 --- a/scripts/benchmarks/object_store/matched-transfer.md +++ b/scripts/benchmarks/object_store/matched-transfer.md @@ -101,8 +101,7 @@ cannot establish complete connector ingestion, native deployment parity, object-store durability, TLS/authentication cost, cloud networking or service throttling. A smaller worker ceiling is not a fixed-RAM-quota benchmark. -Use the [live bulk-reader harness](compare_live.py) and -[production transfer methodology](../../../packages/kumo-adls/docs/performance.md) +Use the [live bulk-reader harness](compare_live.py) for production bulk-reader measurements on the appropriate native cloud hosts. Keep results from these different scopes separate; historical raw traces/CSV remain provenance, not current ADLS native-host measurements. From 61ec7cfc271db932cedfaad222eb3ac608fcd166 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Thu, 22 Oct 2026 19:50:35 +0000 Subject: [PATCH 04/15] Preserve concurrent replacements during ADLS write cleanup --- packages/kumo-adls/docs/implementation.md | 3 + .../kumo-adls/src/kumo_adls/sdk_filesystem.py | 14 +++- packages/kumo-adls/tests/fakes.py | 11 ++- packages/kumo-adls/tests/test_filesystem.py | 68 ++++++++++++++++++- packages/kumo-adls/tests/test_sdk_transfer.py | 33 +++++++++ 5 files changed, 121 insertions(+), 8 deletions(-) diff --git a/packages/kumo-adls/docs/implementation.md b/packages/kumo-adls/docs/implementation.md index 76b07f86d..8db208c93 100644 --- a/packages/kumo-adls/docs/implementation.md +++ b/packages/kumo-adls/docs/implementation.md @@ -88,6 +88,9 @@ Streaming writes are not transactional: Arrow closes its wrapper without forwarding serialization exceptions to the fsspec context manager, so failed serialization can leave partial output. Direct fsspec contexts can discard on caller errors; upload failures also trigger best-effort cleanup. +Streaming cleanup deletes only an object matching the original creation ETag. +It leaves changed or uncertain objects in place. Appends are not fenced, so +this cleanup guard does not provide isolation between concurrent writers. Downloaded files remain owner-private, matching the S3 staging path. Local engine ingestion uses the same OS user; sharing scratch files across users diff --git a/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py b/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py index a939b1680..418b6e40c 100644 --- a/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py +++ b/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py @@ -84,11 +84,15 @@ class _WriteFile(AbstractBufferedFile): """ def __init__(self, filesystem, path, client, **kwargs): self.client = client + self._creation_etag = None super().__init__(filesystem, path, mode='wb', **kwargs) def _initiate_upload(self): with translated(): - self.client.create_file() + self._creation_etag = self.client.create_file().get('etag') + if not self._creation_etag: + raise TransientSourceError( + 'ADLS upload returned no object identity') def _upload_chunk(self, final=False): data = self.buffer.getvalue() @@ -108,10 +112,14 @@ def discard(self): if self.closed: return self.closed = True - if self.offset is None: + if not self._creation_etag: return try: - self.client.delete_file() + # Append is unfenced: its response cannot prove we still own the + # path. Delete only the original creation, never a replacement. + self.client.delete_file( + etag=self._creation_etag, + match_condition=MatchConditions.IfNotModified) except Exception: # Cleanup must preserve the original upload failure. pass diff --git a/packages/kumo-adls/tests/fakes.py b/packages/kumo-adls/tests/fakes.py index 04b4c5add..0b3a3b74c 100644 --- a/packages/kumo-adls/tests/fakes.py +++ b/packages/kumo-adls/tests/fakes.py @@ -15,6 +15,7 @@ class FakeService: """Azure SDK-shaped storage using real public property models.""" def __init__(self): self.files = {} + self.uncommitted = {} self.directories = {'curated', 'curated/orders'} self.downloads = [] self.created = [] @@ -95,11 +96,13 @@ def upload_data(self, stream, length, overwrite, chunk_size, if overwrite: self.create_file() self.service.files[self.name] = stream.read() + self.service.uncommitted.pop(self.name, None) assert len(self.service.files[self.name]) == length def create_file(self): self.service.created.append(self.name) self.service.files[self.name] = b'' + self.service.uncommitted[self.name] = b'' return {'etag': self.get_file_properties().etag} def create_directory(self): @@ -107,17 +110,19 @@ def create_directory(self): def append_data(self, data, offset, length): assert length == len(data) - assert offset == len(self.service.files[self.name]) - self.service.files[self.name] += data + assert offset == len(self.service.uncommitted[self.name]) + self.service.uncommitted[self.name] += data def flush_data(self, offset): - assert offset == len(self.service.files[self.name]) + assert offset == len(self.service.uncommitted[self.name]) + self.service.files[self.name] = self.service.uncommitted.pop(self.name) def delete_file(self, **kwargs): if kwargs.get( 'etag') and kwargs['etag'] != self.get_file_properties().etag: raise ResourceModifiedError('object was changed') self.service.files.pop(self.name) + self.service.uncommitted.pop(self.name, None) def delete_directory(self): self.service.files = { diff --git a/packages/kumo-adls/tests/test_filesystem.py b/packages/kumo-adls/tests/test_filesystem.py index 34ce3d26c..0fa605584 100644 --- a/packages/kumo-adls/tests/test_filesystem.py +++ b/packages/kumo-adls/tests/test_filesystem.py @@ -154,11 +154,11 @@ def flush(self, offset): attempts.append('flush') raise RuntimeError('original failure') - def delete(self): + def delete(self, **kwargs): attempts.append('delete') if cleanup_fails: raise RuntimeError('cleanup failure') - original_delete(self) + original_delete(self, **kwargs) monkeypatch.setattr(FakeFile, 'append_data', append) monkeypatch.setattr(FakeFile, 'delete_file', delete) @@ -180,6 +180,70 @@ def delete(self): assert 'curated/failed' not in service.files +@pytest.mark.parametrize('failure', ['append', 'flush', 'body']) +def test_failed_streaming_write_preserves_replacement(service, monkeypatch, + failure): + from azure.core import MatchConditions + from fakes import FakeFile + + replacement = b'other writer completed' + append = FakeFile.append_data + deleted = [] + delete = FakeFile.delete_file + + def append_or_fail(self, *args, **kwargs): + if failure == 'append': + service.files[self.name] = replacement + raise RuntimeError('original failure') + return append(self, *args, **kwargs) + + def flush_or_fail(self, offset): + service.files[self.name] = replacement + raise RuntimeError('original failure') + + def conditional_delete(self, **kwargs): + assert kwargs['match_condition'] == MatchConditions.IfNotModified + deleted.append(kwargs['etag']) + return delete(self, **kwargs) + + monkeypatch.setattr(FakeFile, 'append_data', append_or_fail) + monkeypatch.setattr(FakeFile, 'delete_file', conditional_delete) + if failure == 'flush': + monkeypatch.setattr(FakeFile, 'flush_data', flush_or_fail) + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + with pytest.raises(RuntimeError, match='original failure'): + with filesystem.open(AUTHORITY + '/shared', 'wb', + block_size=4) as stream: + stream.write(b'abcd') + if failure == 'body': + service.files['curated/shared'] = replacement + raise RuntimeError('original failure') + assert len(deleted) == 1 + assert service.files['curated/shared'] == replacement + assert stream.closed + + +def test_streaming_creation_without_etag_never_deletes_unknown_object( + service, monkeypatch): + from fakes import FakeFile + + create = FakeFile.create_file + + def unidentified(self): + create(self) + return {} + + delete = Mock(side_effect=AssertionError('must not delete unknown object')) + monkeypatch.setattr(FakeFile, 'create_file', unidentified) + monkeypatch.setattr(FakeFile, 'delete_file', delete) + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + with pytest.raises(TransientSourceError, match='object identity'): + with filesystem.open(AUTHORITY + '/unknown', 'wb', + block_size=4) as stream: + stream.write(b'abcd') + delete.assert_not_called() + + def test_exception_before_first_upload_does_not_create_file(service): filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') with pytest.raises(RuntimeError): diff --git a/packages/kumo-adls/tests/test_sdk_transfer.py b/packages/kumo-adls/tests/test_sdk_transfer.py index 395f479b3..ce92d0f88 100644 --- a/packages/kumo-adls/tests/test_sdk_transfer.py +++ b/packages/kumo-adls/tests/test_sdk_transfer.py @@ -359,6 +359,39 @@ def unexpected_pool(*args, **kwargs): assert not list(tmp_path.glob('.kumo-adls-*')) +def test_sdk_streaming_cleanup_uses_creation_etag_not_unfenced_append_etag( + mocked_http): + path = re.compile( + r'https://companylake\.dfs\.core\.windows\.net/curated/streamed(?:\?|$)' + ) + deleted = [] + + def reject_delete(request, context): + deleted.append(request.headers['If-Match']) + context.status_code = 412 + return b'' + + mocked_http.register_uri('PUT', path, headers=HEADERS, status_code=201) + mocked_http.register_uri('PATCH', path, headers={ + **HEADERS, 'ETag': '"replacement"' + }, status_code=202) + mocked_http.register_uri('DELETE', path, headers=HEADERS, + content=reject_delete) + service = client(mocked_http) + try: + filesystem = ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net') + with pytest.raises(RuntimeError, match='caller failed'): + with filesystem.open(AUTHORITY + '/streamed', 'wb', + block_size=4) as stream: + stream.write(b'abcd') + raise RuntimeError('caller failed') + assert stream.closed + assert deleted == ['"stable"'] + finally: + service.close() + + def test_sdk_seekable_upload_overlaps_chunks_and_flushes_after_all_appends( mocked_http, tmp_path): source = tmp_path / 'upload' From ee9036bfa0383a31c7842dc719c30949820fad3d Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Thu, 22 Oct 2026 21:58:21 +0000 Subject: [PATCH 05/15] Use native Arrow streams for ADLS output --- packages/kumo-adls/README.md | 23 +- packages/kumo-adls/docs/implementation.md | 61 ++-- packages/kumo-adls/pyproject.toml | 2 +- .../kumo-adls/src/kumo_adls/native_output.py | 99 +++++++ packages/kumo-adls/src/kumo_adls/sink.py | 19 +- packages/kumo-adls/tests/conftest.py | 21 ++ packages/kumo-adls/tests/test_native_auth.py | 229 +++++++++++++++ .../kumo-adls/tests/test_native_output.py | 275 ++++++++++++++++++ packages/kumo-adls/tests/test_transfer.py | 24 +- uv.lock | 4 +- 10 files changed, 713 insertions(+), 44 deletions(-) create mode 100644 packages/kumo-adls/src/kumo_adls/native_output.py create mode 100644 packages/kumo-adls/tests/test_native_auth.py create mode 100644 packages/kumo-adls/tests/test_native_output.py diff --git a/packages/kumo-adls/README.md b/packages/kumo-adls/README.md index 799060258..450b72be5 100644 --- a/packages/kumo-adls/README.md +++ b/packages/kumo-adls/README.md @@ -37,8 +37,8 @@ A rooted connector discovers CSV/Parquet tables and writes predictions, embeddings, and model bundles beneath its filesystem/prefix. A rootless connector reads exact canonical `abfss://` URIs; discovery and output are unavailable. SQL, warehouse-style `ReadMode.STAGED`, and append output are -unsupported. Prepared exports use parallel SDK uploads; publication is not -atomic, so consumers must wait for export completion. +unsupported. Prepared exports and generated output use native Arrow streams, +matching S3/GCS; consumers must wait for export completion. ## Authentication and resources @@ -46,14 +46,23 @@ Without an identity, requests use Azure's ambient credential chain. An explicit `ADLSIdentity(tenant_id=..., client_id=..., cloud=...)` uses workload federation with `AZURE_FEDERATED_TOKEN_FILE`. Secrets, storage keys, and SAS tokens are not connector arguments. The identity cloud must match the public, US Government, -or China URI endpoint. Probing checks token acquisition only; list or read the -configured source to verify storage permissions and connectivity. +or China URI endpoint. Native writes require the process's `AZURE_TENANT_ID` +and `AZURE_CLIENT_ID` to match an explicit identity and +`AZURE_FEDERATED_TOKEN_FILE` to name its projected token. Other per-connector +identities remain read-only. Native writes also require `AZURE_AUTHORITY_HOST` +to match the cloud (public is the default). Unsupported configurations fail +explicitly; writes do not switch identity or fall back to SDK uploads. +Probing checks Python token acquisition only; it does not attest native-write +compatibility or storage permissions. Set a positive `KUMO_ADLS_TRANSFER_WORKERS` before the provider's first import -to lower ADLS transfer concurrency and HTTP pooling. Restart existing processes +to lower SDK read concurrency and HTTP pooling. Restart existing processes after changing it. `KUMO_SOURCE_COPY_WORKERS` lowers the shared file-copy -ceiling. Both controls limit concurrency rather than total memory; SDK and -projected-read buffers need additional headroom. +ceiling. These controls do not bound native Arrow uploads or their background +buffers. Native writes use Arrow's C++ credentials and upload scheduling; +ambient credential selection can differ from the Python SDK read chain. +Directory exports use the shared eight-file scheduler, while model bundles +upload files sequentially. Reserve memory headroom for native uploads. See the [storage contract](../../docs/cloud-storage/adls.md) and [implementation and resource controls](docs/implementation.md) for diff --git a/packages/kumo-adls/docs/implementation.md b/packages/kumo-adls/docs/implementation.md index 8db208c93..c1d649588 100644 --- a/packages/kumo-adls/docs/implementation.md +++ b/packages/kumo-adls/docs/implementation.md @@ -9,7 +9,7 @@ process lifetime, without automatic eviction; identity cleanup or to prevent duplicate clients, so concurrent first use of different accounts can serialize. Arrow/fsspec wrappers are lightweight and created per access. -Identity probing checks token acquisition only. It does not verify storage +Identity probing checks Python token acquisition only, not native output compatibility. It does not verify storage existence, RBAC/ACL permissions, or connectivity; list or read the configured source to verify those. @@ -19,7 +19,7 @@ concurrently; each file holds one transfer slot and uses no per-file executor or queued chunk futures. The SDK writes each response to the local temporary file before requesting the next chunk. The initial response remains buffered for the stream's lifetime; this is not zero-copy network-to-disk streaming. -Prepared uploads retain up to four 4 MiB chunks. Large projected Arrow reads +Large projected Arrow reads can use up to four SDK workers, bounded by the shared transfer slots. A 64 MiB range uses up to three workers; ranges above 64 MiB can use four. The initial response plus four 16 MiB chunks can account for about 80 MiB before transport @@ -29,8 +29,8 @@ Projected reads are not covered by the 64 MiB bulk-copy reservation. Transfers across files and accounts share a CPU-derived concurrency ceiling (at most 128, lowered by `KUMO_SOURCE_COPY_WORKERS`). The HTTP connection pool matches that ceiling. Metadata bypasses transfer slots and shares the HTTP -transport. Uploads reserve 64 MiB per active file through the shared memory -ledger; source bulk copies use the core copy reservation. These reservations +transport. Source bulk copies use the core copy reservation. These SDK +limits and reservations do not apply to native Arrow uploads. These reservations are accounting estimates, not caps on total process RSS. Larger responses reduce requests and fragment assembly while increasing memory per active file. Plan roughly 48–64 MiB of payload buffers per active bulk download, with @@ -40,7 +40,7 @@ fits a host's memory limit. Actual memory depends on allocation lifetimes. For memory-constrained hosts, set `KUMO_ADLS_TRANSFER_WORKERS=4` before importing the provider to limit ADLS transfer requests and its HTTP pool -without lowering other providers' copy limits. The value must be a positive +without lowering other providers' copy limits. This does not limit native writes. The value must be a positive integer and cannot raise the shared copy/CPU ceiling. The limit is process-wide across ADLS accounts and readers; it is read once at import. Even with one slot, files use the same serial SDK stream. Arrow buffers, cached bytes, SDK copies, @@ -78,19 +78,44 @@ after discovery belong to a subsequent read. Identity checks may still perform listings or HEADs, so one discovery does not mean one metadata request. ETag quote and directory-path normalization can invalidate existing version hashes once, requiring restaging. -Prepared export files and model files use seekable parallel uploads; Arrow -writers keep the buffered streaming path. Failed prepared uploads delete only -the object matching their creation ETag, preserving another writer's output. -Prepared uploads create the destination before uploading; concurrent readers -can observe an empty or incomplete object. Publishing an individual file is -not atomic, so consumers must wait for export completion. -Streaming writes are not transactional: Arrow closes its wrapper without -forwarding serialization exceptions to the fsspec context manager, so failed -serialization can leave partial output. Direct fsspec contexts can discard -on caller errors; upload failures also trigger best-effort cleanup. -Streaming cleanup deletes only an object matching the original creation ETag. -It leaves changed or uncertain objects in place. Appends are not fenced, so -this cleanup guard does not provide isolation between concurrent writers. +Prepared exports, generated Parquet/CSV output, and model bundles use native +Arrow output streams. The shared sink copies prepared bytes in 8 MiB blocks, +with up to eight files per directory export; model bundles upload files +sequentially. Arrow owns per-file buffering and background upload scheduling. +The ADLS SDK transfer semaphore and 64 MiB copy reservation do not constrain +these native buffers. Copy block size is not an upload-memory limit. Consumers +must wait for export completion; streaming writes and export directories are +not transactional. Arrow close after a caller failure may publish partial +output, as with the shared S3/GCS sink. The SDK writer's conditional cleanup +and final-flush guarantees do not apply to native output. + +Native output is bound lazily on first write, leaving SDK metadata, source +reads, model downloads, and their account-qualified paths unchanged. Only +output streams drop the virtual account qualifier to use Arrow's container/key +path. Existing root containment checks run before filesystem access. +A native filesystem is reused within one sink access, with thread-safe first +construction; it is not cached across accesses or identities. The provider +requires Arrow 24, the version whose native credential/URI APIs are validated. + +Without an explicit identity, writes use Arrow's C++ ambient credential chain; +reads use Python Azure Identity. These chains can select different ambient +credentials. Explicit workload writes require process `AZURE_TENANT_ID` and +`AZURE_CLIENT_ID` to equal the selected identity, plus +`AZURE_FEDERATED_TOKEN_FILE`. They select only Arrow's workload credential, +without ambient fallback. Arrow captures these settings at construction and +rereads projected token contents when obtaining a fresh assertion. A different +per-connector identity remains supported for reads, but its native writes fail +with a configuration error. The provider never changes process environment or +mints delegation SAS to bridge credentials. + +Native writes validate `AZURE_AUTHORITY_HOST` against the selected cloud. Public +Azure defaults to `https://login.microsoftonline.com`; US Government needs +`https://login.microsoftonline.us`, and China needs +`https://login.chinacloudapi.cn`. Both native Blob and DFS endpoints use the +selected storage suffix. Offline tests construct real native filesystems for +all three clouds and check identity/authority rejection and endpoint mapping. +They do not establish live federation, token renewal, sovereign connectivity, +or RBAC/ACL propagation; those require suitable live environments. Downloaded files remain owner-private, matching the S3 staging path. Local engine ingestion uses the same OS user; sharing scratch files across users diff --git a/packages/kumo-adls/pyproject.toml b/packages/kumo-adls/pyproject.toml index 726d67e17..0361bbf3d 100644 --- a/packages/kumo-adls/pyproject.toml +++ b/packages/kumo-adls/pyproject.toml @@ -17,7 +17,7 @@ license-files = ["LICENSE"] dependencies = [ "kumo-core==0.1.0.dev1", "kumo-api>=3.0.0", - "pyarrow>=14", + "pyarrow>=24,<25", "fsspec", "requests>=2.31,<3", "azure-core>=1.31,<2", diff --git a/packages/kumo-adls/src/kumo_adls/native_output.py b/packages/kumo-adls/src/kumo_adls/native_output.py new file mode 100644 index 000000000..05cc6a6b7 --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/native_output.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import threading +from urllib.parse import urlencode, urlsplit + +from kumoapi.adls import ( + normalize_adls_endpoint_suffix, + normalize_adls_filesystem, + normalize_adls_storage_account, +) +from kumoapi.data_source import ADLSIdentityConfig +from pyarrow.fs import AzureFileSystem, FileSystem, FSSpecHandler + +from .session import CLOUDS + + +def filesystem_for(principal, account, suffix, container): + """Bind native output without changing or substituting an identity.""" + account = normalize_adls_storage_account(account) + suffix = normalize_adls_endpoint_suffix(suffix) + container = normalize_adls_filesystem(container) + cloud, (_, _, authority) = next((cloud, values) + for cloud, values in CLOUDS.items() + if values[0] == suffix) + if principal is not None: + config = principal.config + if not isinstance(config, ADLSIdentityConfig): + raise ValueError('The selected identity is not an ADLS identity') + if config.cloud != cloud: + raise ValueError( + 'ADLS endpoint does not match the selected identity cloud') + if (os.environ.get('AZURE_TENANT_ID') != config.tenant_id + or os.environ.get('AZURE_CLIENT_ID') != config.client_id or + not os.environ.get('AZURE_FEDERATED_TOKEN_FILE', '').strip()): + raise NotImplementedError( + 'Native Arrow ADLS writes require AZURE_TENANT_ID and ' + 'AZURE_CLIENT_ID to match the selected workload identity, ' + 'and AZURE_FEDERATED_TOKEN_FILE to name its projected token. ' + 'Configure these in the process environment before writing. ' + 'Per-connector workload identities that differ from the ' + 'process identity remain supported for reads only.') + + host = os.environ.get('AZURE_AUTHORITY_HOST', + 'https://login.microsoftonline.com') + parsed = urlsplit(host) + if (parsed.scheme != 'https' or parsed.netloc.casefold() != authority + or parsed.path not in ('', '/') or parsed.query or parsed.fragment): + raise ValueError( + 'Native Arrow ADLS writes require AZURE_AUTHORITY_HOST=' + + f'https://{authority} for the selected cloud') + options = dict( + blob_storage_authority='.' + suffix.replace('dfs.', 'blob.', 1), + dfs_storage_authority='.' + suffix) + if principal is None: + return AzureFileSystem(account_name=account, **options) + + # Arrow's workload credential reads environment configuration at creation; + # unlike DefaultAzureCredential it cannot fall back to another principal. + # The short URI form keeps account/container correct in sovereign clouds. + query = urlencode(dict(credential_kind='workload_identity', **options)) + filesystem, _ = FileSystem.from_uri( + f'abfss://{account}@{container}/?{query}') + return filesystem + + +class NativeOutputHandler(FSSpecHandler): + """Keep SDK metadata/reads and return native Arrow output streams.""" + def __init__(self, filesystem, output_filesystem=None, *, + output_factory=None): + super().__init__(filesystem) + self.output_filesystem = output_filesystem + self.output_factory = output_factory + self._output_lock = threading.Lock() + + def _output(self): + with self._output_lock: + if self.output_filesystem is None: + if self.output_factory is None: + raise TypeError('Native output needs a filesystem factory') + self.output_filesystem = self.output_factory() + return self.output_filesystem + + def create_dir(self, path, recursive): + self.fs._address(path) + # Shared sinks prepare parent directories before opening output. + # Reject unsupported credentials before that SDK mutation as well. + self._output() + super().create_dir(path, recursive) + + def open_output_stream(self, path, metadata): + # Validate the virtual authority before dropping its account qualifier. + self.fs._address(path) + authority, _, key = path.strip('/').partition('/') + container, _, _ = authority.partition('@') + return self._output().open_output_stream(f'{container}/{key}', + metadata=metadata) diff --git a/packages/kumo-adls/src/kumo_adls/sink.py b/packages/kumo-adls/src/kumo_adls/sink.py index 5525b30ff..e92885330 100644 --- a/packages/kumo-adls/src/kumo_adls/sink.py +++ b/packages/kumo-adls/src/kumo_adls/sink.py @@ -2,12 +2,15 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from functools import partial + from kumo_core.object_store_sink import ObjectStoreSinkBackend from kumo_core.sources.cpu import COPY_WORKER_BYTES, memory_share from kumoapi.data_source import DataSourceType from pyarrow.fs import FSSpecHandler, PyFileSystem from .backend import ADLSSessions, canonical_uri +from .native_output import NativeOutputHandler, filesystem_for from .sdk_filesystem import ADLSFileSystem @@ -17,13 +20,21 @@ class ADLSSinkBackend(ADLSSessions, ObjectStoreSinkBackend): def _uri(self, path: str) -> str: return canonical_uri(path) - def _upload_file(self, filesystem, source, destination): + def _filesystem(self, connector, bucket, region): + filesystem = super()._filesystem(connector, bucket, region) if (isinstance(filesystem, PyFileSystem) and isinstance(filesystem.handler, FSSpecHandler) and isinstance(filesystem.handler.fs, ADLSFileSystem)): - filesystem.handler.fs.put_file(source, destination) - else: - super()._upload_file(filesystem, source, destination) + container, host = bucket.split('@', 1) + account, suffix = host.split('.', 1) + # Model reads use this filesystem too; bind output credentials + # only when a native output stream is actually requested. + native = partial(filesystem_for, self._principal(connector), + account, suffix, container) + return PyFileSystem( + NativeOutputHandler(filesystem.handler.fs, + output_factory=native)) + return filesystem def _download_file(self, filesystem, source, destination): if (isinstance(filesystem, PyFileSystem) diff --git a/packages/kumo-adls/tests/conftest.py b/packages/kumo-adls/tests/conftest.py index 18ea5e749..8f1acfa13 100644 --- a/packages/kumo-adls/tests/conftest.py +++ b/packages/kumo-adls/tests/conftest.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import io from types import SimpleNamespace import pyarrow as pa @@ -21,6 +22,8 @@ def service(monkeypatch): service = FakeService() monkeypatch.setattr('kumo_adls.source_access.session_for', lambda *args: service) + monkeypatch.setattr('kumo_adls.sink.filesystem_for', + lambda *args: FakeNativeFilesystem(service)) sink = pa.BufferOutputStream() pq.write_table(pa.table({'id': [1, 2], 'value': [10.0, 20.0]}), sink) service.files['curated/orders/events.parquet'] = sink.getvalue().to_pybytes( @@ -41,3 +44,21 @@ def build(*, account='companylake', prefix='orders', any_uri=False, identity=None, identity_id=None, source_type=DataSourceType.ADLS) return build + + +class FakeNativeFilesystem: + """Native-output-shaped fake; commit only when Arrow closes the stream.""" + def __init__(self, service): + self.service = service + + def open_output_stream(self, path, metadata=None): + service = self.service + + class Output(io.BytesIO): + def close(self): + if not self.closed: + service.files[path] = self.getvalue() + super().close() + + service.created.append(path) + return pa.PythonFile(Output(), mode='w') diff --git a/packages/kumo-adls/tests/test_native_auth.py b/packages/kumo-adls/tests/test_native_auth.py new file mode 100644 index 000000000..3f5d2a646 --- /dev/null +++ b/packages/kumo-adls/tests/test_native_auth.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from types import SimpleNamespace +from unittest.mock import Mock +from urllib.parse import parse_qs, urlsplit + +import pytest +from kumo_adls import native_output +from kumo_adls.session import CLOUDS +from kumoapi.data_source import ADLSIdentityConfig, AzureCloud +from pyarrow.fs import AzureFileSystem, FileSystem + +TENANT = 'private-tenant-value' +CLIENT = 'private-client-value' +ACCOUNT = 'companylake' +CONTAINER = 'curated' + + +@pytest.fixture +def auth_environment(monkeypatch, tmp_path): + for name in list(os.environ): + if name.startswith('AZURE_'): + monkeypatch.delenv(name) + token = tmp_path / 'private-token-file' + token.write_text('private-token-contents') + monkeypatch.setenv('AZURE_TENANT_ID', TENANT) + monkeypatch.setenv('AZURE_CLIENT_ID', CLIENT) + monkeypatch.setenv('AZURE_FEDERATED_TOKEN_FILE', str(token)) + return token + + +def principal(cloud=AzureCloud.PUBLIC): + return SimpleNamespace(config=ADLSIdentityConfig( + tenant_id=TENANT, client_id=CLIENT, cloud=cloud)) + + +def assert_native_mapping(filesystem, suffix): + assert isinstance(filesystem, AzureFileSystem) + _, (options, ) = filesystem.__reduce__() + assert options['account_name'] == ACCOUNT + assert options['dfs_storage_authority'] == '.' + suffix + assert options['blob_storage_authority'] == '.' + suffix.replace( + 'dfs.', 'blob.', 1) + assert options['dfs_storage_scheme'] == 'https' + assert options['blob_storage_scheme'] == 'https' + + +@pytest.mark.parametrize('cloud', list(CLOUDS)) +def test_workload_native_construction_pins_identity_and_cloud( + auth_environment, monkeypatch, cloud): + suffix, _, authority = CLOUDS[cloud] + monkeypatch.setenv('AZURE_AUTHORITY_HOST', 'https://' + authority) + calls = [] + + def from_uri(uri): + calls.append(uri) + parsed = urlsplit(uri) + assert parsed.scheme == 'abfss' + assert parsed.netloc == f'{ACCOUNT}@{CONTAINER}' + assert parsed.path == '/' + assert parse_qs(parsed.query) == { + 'credential_kind': ['workload_identity'], + 'dfs_storage_authority': ['.' + suffix], + 'blob_storage_authority': + ['.' + suffix.replace('dfs.', 'blob.', 1)] + } + assert TENANT not in uri + assert CLIENT not in uri + assert str(auth_environment) not in uri + return FileSystem.from_uri(uri) + + monkeypatch.setattr(native_output, 'FileSystem', + SimpleNamespace(from_uri=from_uri)) + before = dict(os.environ) + filesystem = native_output.filesystem_for(principal(cloud), ACCOUNT, suffix, + CONTAINER) + assert_native_mapping(filesystem, suffix) + assert len(calls) == 1 + assert dict(os.environ) == before + # Arrow does not serialize credential_kind in __reduce__; the URI spy + # above checks the selected credential rather than claiming round-trip proof. + + +@pytest.mark.parametrize('cloud', list(CLOUDS)) +def test_ambient_native_construction_uses_selected_cloud_without_env_mutation( + auth_environment, monkeypatch, cloud): + suffix, _, authority = CLOUDS[cloud] + monkeypatch.setenv('AZURE_AUTHORITY_HOST', 'https://' + authority + '/') + before = dict(os.environ) + filesystem = native_output.filesystem_for(None, ACCOUNT, suffix, CONTAINER) + assert_native_mapping(filesystem, suffix) + assert dict(os.environ) == before + + +def forbid_construction(monkeypatch): + factory = Mock( + side_effect=AssertionError('must reject before construction')) + monkeypatch.setattr(native_output, 'AzureFileSystem', factory) + monkeypatch.setattr(native_output, 'FileSystem', + SimpleNamespace(from_uri=factory)) + return factory + + +@pytest.mark.parametrize('name,value', [ + ('AZURE_TENANT_ID', None), + ('AZURE_TENANT_ID', 'another-private-tenant'), + ('AZURE_CLIENT_ID', None), + ('AZURE_CLIENT_ID', 'another-private-client'), + ('AZURE_FEDERATED_TOKEN_FILE', None), + ('AZURE_FEDERATED_TOKEN_FILE', ''), + ('AZURE_FEDERATED_TOKEN_FILE', ' '), +]) +def test_workload_mismatch_never_falls_back_or_discloses_identity( + auth_environment, monkeypatch, name, value): + if value is None: + monkeypatch.delenv(name) + else: + monkeypatch.setenv(name, value) + factory = forbid_construction(monkeypatch) + before = dict(os.environ) + with pytest.raises(NotImplementedError) as error: + native_output.filesystem_for(principal(), ACCOUNT, + CLOUDS[AzureCloud.PUBLIC][0], CONTAINER) + factory.assert_not_called() + assert dict(os.environ) == before + for secret in [ + TENANT, CLIENT, + str(auth_environment), 'private-token-contents', + 'another-private-tenant', 'another-private-client' + ]: + assert secret not in str(error.value) + + +@pytest.mark.parametrize('selected', list(CLOUDS)) +def test_explicit_identity_cloud_mismatch_fails_before_construction( + auth_environment, monkeypatch, selected): + other = next(cloud for cloud in CLOUDS if cloud != selected) + factory = forbid_construction(monkeypatch) + before = dict(os.environ) + with pytest.raises(ValueError, match='identity cloud'): + native_output.filesystem_for(principal(selected), ACCOUNT, + CLOUDS[other][0], CONTAINER) + factory.assert_not_called() + assert dict(os.environ) == before + + +@pytest.mark.parametrize('selected', [None, AzureCloud.PUBLIC]) +@pytest.mark.parametrize('authority', [ + 'http://login.microsoftonline.com', + 'login.microsoftonline.com', + 'https://private-user:private-password@login.microsoftonline.com', + 'https://login.microsoftonline.com:443', + 'https://login.microsoftonline.com/private-path', + 'https://login.microsoftonline.com?token=private-query', + 'https://login.microsoftonline.com#private-fragment', + 'https://login.microsoftonline.us', + 'https://login.microsoftonline.com.evil.example', + '', +]) +def test_authority_rejected_before_any_factory_without_secret_echo( + auth_environment, monkeypatch, selected, authority): + monkeypatch.setenv('AZURE_AUTHORITY_HOST', authority) + factory = forbid_construction(monkeypatch) + before = dict(os.environ) + with pytest.raises(ValueError) as error: + native_output.filesystem_for(None if selected is None else principal(), + ACCOUNT, CLOUDS[AzureCloud.PUBLIC][0], + CONTAINER) + factory.assert_not_called() + assert dict(os.environ) == before + for secret in [ + 'private-user', 'private-password', 'private-path', 'private-query', + 'private-fragment', TENANT, CLIENT, + str(auth_environment) + ]: + assert secret not in str(error.value) + + +@pytest.mark.parametrize('cloud', [AzureCloud.US_GOVERNMENT, AzureCloud.CHINA]) +def test_sovereign_cloud_rejects_unset_public_default_authority( + auth_environment, monkeypatch, cloud): + factory = forbid_construction(monkeypatch) + with pytest.raises(ValueError, match='AZURE_AUTHORITY_HOST'): + native_output.filesystem_for(None, ACCOUNT, CLOUDS[cloud][0], CONTAINER) + factory.assert_not_called() + + +def test_default_public_authority_is_accepted(auth_environment): + filesystem = native_output.filesystem_for(principal(), ACCOUNT, + CLOUDS[AzureCloud.PUBLIC][0], + CONTAINER) + assert_native_mapping(filesystem, CLOUDS[AzureCloud.PUBLIC][0]) + + +def test_changed_environment_is_revalidated_without_cached_filesystem( + auth_environment, monkeypatch, tmp_path): + factory = Mock(wraps=FileSystem.from_uri) + monkeypatch.setattr(native_output, 'FileSystem', + SimpleNamespace(from_uri=factory)) + selected = principal() + suffix = CLOUDS[AzureCloud.PUBLIC][0] + first = native_output.filesystem_for(selected, ACCOUNT, suffix, CONTAINER) + monkeypatch.setenv('AZURE_CLIENT_ID', 'changed-private-client') + before = dict(os.environ) + with pytest.raises(NotImplementedError): + native_output.filesystem_for(selected, ACCOUNT, suffix, CONTAINER) + assert dict(os.environ) == before + assert factory.call_count == 1 + monkeypatch.setenv('AZURE_CLIENT_ID', CLIENT) + next_token = tmp_path / 'another-private-token' + next_token.write_text('another-private-token-contents') + monkeypatch.setenv('AZURE_FEDERATED_TOKEN_FILE', str(next_token)) + second = native_output.filesystem_for(selected, ACCOUNT, suffix, CONTAINER) + assert second is not first + assert factory.call_count == 2 + + +def test_other_identity_type_fails_before_native_construction( + auth_environment, monkeypatch): + factory = forbid_construction(monkeypatch) + before = dict(os.environ) + with pytest.raises(ValueError, match='not an ADLS identity'): + native_output.filesystem_for(SimpleNamespace(config=object()), ACCOUNT, + CLOUDS[AzureCloud.PUBLIC][0], CONTAINER) + factory.assert_not_called() + assert dict(os.environ) == before diff --git a/packages/kumo-adls/tests/test_native_output.py b/packages/kumo-adls/tests/test_native_output.py new file mode 100644 index 000000000..50de27939 --- /dev/null +++ b/packages/kumo-adls/tests/test_native_output.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from kumo_adls.native_output import NativeOutputHandler +from kumo_adls.sdk_filesystem import ADLSFileSystem +from kumo_core.object_store_sink import ObjectStoreSinkBackend +from pyarrow.fs import ( + FileSelector, + FileType, + LocalFileSystem, + PyFileSystem, + SubTreeFileSystem, +) + +AUTHORITY = 'curated@companylake.dfs.core.windows.net' + + +@pytest.fixture +def storage(service, tmp_path): + native = SubTreeFileSystem(str(tmp_path), LocalFileSystem()) + native.create_dir('curated/orders') + sdk = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + handler = NativeOutputHandler(sdk, native) + return PyFileSystem(handler), native, sdk + + +def test_metadata_and_reads_keep_sdk_authority(storage, service, monkeypatch): + arrow, native, sdk = storage + paths = [] + original = sdk.info + + def info(path, **kwargs): + paths.append(path) + return original(path, **kwargs) + + monkeypatch.setattr(sdk, 'info', info) + path = AUTHORITY + '/orders/events.parquet' + result = arrow.get_file_info(path) + assert result.path == path + assert result.type == FileType.File + assert result.size == len(service.files['curated/orders/events.parquet']) + assert path in paths + entries = arrow.get_file_info(FileSelector(AUTHORITY + '/orders')) + assert [entry.path for entry in entries] == [path] + with arrow.open_input_file(path) as reader: + assert reader.read() == service.files['curated/orders/events.parquet'] + assert native.get_file_info( + 'curated/orders/events.parquet').type == FileType.NotFound + + +@pytest.mark.parametrize('key', + ['orders/plain.bin', 'orders/a space %2F # ü.bin']) +def test_output_is_native_stream_and_preserves_literal_key(storage, key): + arrow, native, _ = storage + with arrow.open_output_stream(AUTHORITY + '/' + key) as writer: + assert isinstance(writer, pa.NativeFile) + writer.write(b'first') + writer.write(b'\x00second') + writer.flush() + with native.open_input_file('curated/' + key) as reader: + assert reader.read() == b'first\x00second' + + +def test_generated_parquet_writer_uses_native_output(storage): + arrow, native, _ = storage + table = pa.table({'id': [1, 2, 3], 'value': ['alpha', None, 'ü']}) + path = AUTHORITY + '/orders/generated.parquet' + with arrow.open_output_stream(path) as writer: + with pq.ParquetWriter(writer, table.schema) as parquet: + parquet.write_table(table.slice(0, 1)) + parquet.write_table(table.slice(1)) + with native.open_input_file('curated/orders/generated.parquet') as reader: + assert pq.read_table(reader).equals(table) + + +def test_shared_sink_copies_prepared_file_through_native_stream( + storage, tmp_path): + arrow, native, _ = storage + # Cross the shared sink's 8 MiB copy boundary, including a final short write. + payload = bytes(range(256)) * ((8 << 20) // 256) + b'last chunk' + source = tmp_path / 'prepared.bin' + source.write_bytes(payload) + ObjectStoreSinkBackend._upload_file(None, arrow, str(source), + AUTHORITY + '/orders/copied.bin') + with native.open_input_file('curated/orders/copied.bin') as reader: + assert reader.read() == payload + + +def test_wrong_account_is_rejected_before_native_output_opens( + storage, monkeypatch): + arrow, _, _ = storage + opened = [] + + class UnopenedOutput: + def open_output_stream(self, *args, **kwargs): + opened.append((args, kwargs)) + pytest.fail('Invalid authority reached native output') + + arrow.handler.output_filesystem = UnopenedOutput() + with pytest.raises(ValueError, match='bound storage account'): + arrow.open_output_stream( + 'curated@otherlake.dfs.core.windows.net/orders/x') + assert opened == [] + + +@pytest.mark.parametrize('extension', ['csv', 'parquet']) +def test_sink_generated_table_routes_native_and_reads_sdk( + service, connector, monkeypatch, extension): + import pandas as pd + from fakes import FakeFile + from kumo_adls.sink import ADLSSinkBackend + from kumoapi.jobs import WriteMode + + def unexpected_upload(*args, **kwargs): + pytest.fail('Native output must not call an SDK upload') + + monkeypatch.setattr(ADLSFileSystem, 'put_file', unexpected_upload) + monkeypatch.setattr(FakeFile, 'upload_data', unexpected_upload) + backend = ADLSSinkBackend() + backend.write_table(connector(), 'result.' + extension, + pd.DataFrame({'id': [1, 2]}), + write_mode=WriteMode.OVERWRITE) + path = 'curated/orders/result.' + extension + payload = service.files[path] + if extension == 'parquet': + assert pq.read_table(pa.BufferReader(payload)).to_pydict() == { + 'id': [1, 2] + } + else: + assert payload.decode().splitlines() == ['"id"', '1', '2'] + filesystem, destination = backend._destination(connector(), + 'result.' + extension) + assert filesystem.get_file_info(destination).size == len(payload) + with filesystem.open_input_file(destination) as stream: + assert stream.read() == payload + + +def test_sink_factory_is_lazy_for_reads_and_reused_for_concurrent_outputs( + service, connector, monkeypatch): + import threading + from concurrent.futures import ThreadPoolExecutor + + from conftest import FakeNativeFilesystem + from kumo_adls.sink import ADLSSinkBackend + + calls = [] + barrier = threading.Barrier(4) + + def factory(*args): + calls.append(args) + return FakeNativeFilesystem(service) + + monkeypatch.setattr('kumo_adls.sink.filesystem_for', factory) + filesystem, _ = ADLSSinkBackend()._destination(connector(), 'unused') + with filesystem.open_input_file(AUTHORITY + + '/orders/events.parquet') as stream: + assert stream.read(4) == b'PAR1' + assert calls == [] + + def write(index): + barrier.wait(timeout=5) + with filesystem.open_output_stream(AUTHORITY + + f'/orders/out-{index}') as stream: + stream.write(str(index).encode()) + + with ThreadPoolExecutor(max_workers=4) as pool: + list(pool.map(write, range(4))) + assert len(calls) == 1 + for index in range(4): + assert service.files[f'curated/orders/out-{index}'] == str( + index).encode() + + +def test_sink_unsupported_output_factory_does_not_block_model_reads( + service, connector, monkeypatch, tmp_path): + from kumo_adls.sink import ADLSSinkBackend + + def unsupported(*args): + raise NotImplementedError('unsupported native output identity') + + monkeypatch.setattr('kumo_adls.sink.filesystem_for', unsupported) + backend = ADLSSinkBackend() + filesystem, path = backend._destination(connector(), 'events.parquet') + output = tmp_path / 'model.parquet' + backend._download_file(filesystem, path, str(output)) + assert output.read_bytes() == service.files['curated/orders/events.parquet'] + with pytest.raises(NotImplementedError, match='unsupported native output'): + filesystem.open_output_stream(AUTHORITY + '/orders/new') + + +def test_sink_integrates_real_local_native_stream_and_shared_copy( + service, connector, monkeypatch, tmp_path): + from kumo_adls.sink import ADLSSinkBackend + + native = SubTreeFileSystem(str(tmp_path), LocalFileSystem()) + native.create_dir('curated/orders') + monkeypatch.setattr('kumo_adls.sink.filesystem_for', lambda *args: native) + backend = ADLSSinkBackend() + filesystem, destination = backend._destination(connector(), 'native.bin') + source = tmp_path / 'prepared.bin' + source.write_bytes(b'native output') + assert ADLSSinkBackend._upload_file is ObjectStoreSinkBackend._upload_file + backend._upload_file(filesystem, str(source), destination) + with native.open_input_file('curated/orders/native.bin') as stream: + assert stream.read() == b'native output' + assert 'curated/orders/native.bin' not in service.files + + +def test_export_metadata_fence_rejects_other_identity_before_native_open( + service, connector, monkeypatch, tmp_path): + from kumo_adls.sink import ADLSSinkBackend + from kumo_core.artifact_export import ( + ArtifactDestinationConflict, + export_marker, + ) + + source = tmp_path / 'source.parquet' + pq.write_table(pa.table({'id': [1]}), source) + backend = ADLSSinkBackend() + backend.write_files(connector(), 'predictions/', [str(source)], + export_id='first') + marker = 'curated/orders/predictions/_KUMO_EXPORT' + assert service.files[marker].decode() == export_marker('first') + committed = dict(service.files) + + def unexpected_factory(*args): + pytest.fail('Conflicting export must fail before opening native output') + + monkeypatch.setattr('kumo_adls.sink.filesystem_for', unexpected_factory) + with pytest.raises(ArtifactDestinationConflict, match='another export'): + backend.write_files(connector(), 'predictions/', [str(source)], + export_id='second') + assert service.files == committed + + +@pytest.mark.parametrize('operation', + ['write_table', 'write_files', 'publish_model']) +def test_unsupported_native_config_fails_before_sdk_mutation( + service, connector, monkeypatch, tmp_path, operation): + import pandas as pd + from kumo_adls.sink import ADLSSinkBackend + + calls = [] + + def unsupported(*args): + calls.append(args) + raise NotImplementedError('unsupported native output identity') + + monkeypatch.setattr('kumo_adls.sink.filesystem_for', unsupported) + initial_directories = set(service.directories) + initial_files = dict(service.files) + initial_created = list(service.created) + backend = ADLSSinkBackend() + source = tmp_path / 'source.parquet' + pq.write_table(pa.table({'id': [1]}), source) + bundle = tmp_path / 'bundle' + bundle.mkdir() + (bundle / 'model.bin').write_bytes(b'model') + with pytest.raises(NotImplementedError, match='unsupported native output'): + if operation == 'write_table': + backend.write_table(connector(), 'new/output.parquet', + pd.DataFrame({'id': [1]})) + elif operation == 'write_files': + backend.write_files(connector(), 'new/predictions/', [str(source)], + export_id='test') + else: + backend.publish_model(connector(), str(bundle), path='new/model') + assert len(calls) == 1 + assert service.directories == initial_directories + assert service.files == initial_files + assert service.created == initial_created diff --git a/packages/kumo-adls/tests/test_transfer.py b/packages/kumo-adls/tests/test_transfer.py index 7291520b9..d6f429efe 100644 --- a/packages/kumo-adls/tests/test_transfer.py +++ b/packages/kumo-adls/tests/test_transfer.py @@ -210,29 +210,29 @@ def download(*args, **kwargs): assert not list(tmp_path.glob('.kumo-adls-*')) -def test_file_export_uses_seekable_sdk_upload_and_keeps_prepared_columns( +def test_file_export_uses_native_copy_and_keeps_prepared_columns( service, connector, tmp_path, monkeypatch): import pyarrow as pa import pyarrow.parquet as pq from fakes import FakeFile from kumo_adls.sink import ADLSSinkBackend + def unexpected_upload(*args, **kwargs): + pytest.fail('Prepared export must use native Arrow output') + + monkeypatch.setattr(FakeFile, 'upload_data', unexpected_upload) + monkeypatch.setattr(ADLSFileSystem, 'put_file', unexpected_upload) source = tmp_path / 'source.parquet' pq.write_table(pa.table({'id': [1, 2]}), source) - original = FakeFile.upload_data - uploads = [] - - def upload(self, stream, **kwargs): - assert stream.seekable() - uploads.append(self.name) - return original(self, stream, **kwargs) - - monkeypatch.setattr(FakeFile, 'upload_data', upload) ADLSSinkBackend().write_files(connector(), 'predictions/', [str(source)], extra_columns={'tag': 'test'}, export_id='run') - assert len(uploads) == 1 - table = pq.read_table(pa.BufferReader(service.files[uploads[0]])) + outputs = [ + name for name in service.files + if '/predictions/' in name and name.endswith('.parquet') + ] + assert len(outputs) == 1 + table = pq.read_table(pa.BufferReader(service.files[outputs[0]])) assert table.to_pydict() == {'id': [1, 2], 'tag': ['test', 'test']} diff --git a/uv.lock b/uv.lock index 60686de0e..c5bf6defa 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 5 +revision = 3 requires-python = ">=3.11" resolution-markers = [ "python_full_version >= '3.14' and sys_platform != 'darwin' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests'", @@ -3139,7 +3139,7 @@ requires-dist = [ { name = "fsspec" }, { name = "kumo-api", editable = "packages/kumo-api" }, { name = "kumo-core", editable = "packages/kumo-core" }, - { name = "pyarrow", specifier = ">=14" }, + { name = "pyarrow", specifier = ">=24,<25" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=7.4" }, { name = "requests", specifier = ">=2.31,<3" }, ] From e5337205499f5536e5f5611b1fb44639bef1c329 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Thu, 22 Oct 2026 22:39:16 +0000 Subject: [PATCH 06/15] Pin ADLS SDK upload conditions and clarify transfer settings --- packages/kumo-adls/README.md | 8 +-- packages/kumo-adls/tests/test_sdk_transfer.py | 53 ++++++++++++++++++- packages/kumo-adls/tests/test_transfer.py | 43 +++++++++++++++ packages/kumo-api/kumoapi/adls.py | 5 +- 4 files changed, 102 insertions(+), 7 deletions(-) diff --git a/packages/kumo-adls/README.md b/packages/kumo-adls/README.md index 450b72be5..dfa5657fe 100644 --- a/packages/kumo-adls/README.md +++ b/packages/kumo-adls/README.md @@ -56,9 +56,11 @@ Probing checks Python token acquisition only; it does not attest native-write compatibility or storage permissions. Set a positive `KUMO_ADLS_TRANSFER_WORKERS` before the provider's first import -to lower SDK read concurrency and HTTP pooling. Restart existing processes -after changing it. `KUMO_SOURCE_COPY_WORKERS` lowers the shared file-copy -ceiling. These controls do not bound native Arrow uploads or their background +to lower SDK read concurrency and HTTP pooling. The connection pool and transfer +slots are fixed at import; changing the variable afterward has no effect until +the process restarts. `KUMO_SOURCE_COPY_WORKERS` lowers the shared file-copy +ceiling and its contribution to the ADLS limit is also captured at import. +These controls do not bound native Arrow uploads or their background buffers. Native writes use Arrow's C++ credentials and upload scheduling; ambient credential selection can differ from the Python SDK read chain. Directory exports use the shared eight-file scheduler, while model bundles diff --git a/packages/kumo-adls/tests/test_sdk_transfer.py b/packages/kumo-adls/tests/test_sdk_transfer.py index ce92d0f88..d52988b0e 100644 --- a/packages/kumo-adls/tests/test_sdk_transfer.py +++ b/packages/kumo-adls/tests/test_sdk_transfer.py @@ -400,12 +400,19 @@ def test_sdk_seekable_upload_overlaps_chunks_and_flushes_after_all_appends( lock = threading.Lock() appended = {} flushed = [] + created = [] + + def create(request, context): + created.append(request) + return b'' + mocked_http.register_uri('PUT', re.compile(r'https://companylake\.'), - status_code=201, headers=HEADERS) + status_code=201, headers=HEADERS, content=create) def patch(request, context): query = parse_qs(urlsplit(request.url).query) action = query['action'][0] + assert 'If-None-Match' not in request.headers if action == 'append': offset = int(query['position'][0]) data = request.body @@ -416,6 +423,7 @@ def patch(request, context): appended[offset] = data elif action == 'flush': assert len(appended) == 4 + assert request.headers['If-Match'] == HEADERS['ETag'] flushed.append(int(query['position'][0])) else: raise AssertionError(action) @@ -432,6 +440,49 @@ def patch(request, context): assert b''.join(appended[offset] for offset in sorted(appended)) == source.read_bytes() assert flushed == [source.stat().st_size] + # overwrite=False must not recreate the object made by put_file(). + assert len(created) == 1 + + +def test_sdk_prepared_upload_conflict_keeps_cleanup_conditional( + mocked_http, tmp_path): + from kumo_core.sources.errors import TransientSourceError + + source = tmp_path / 'upload' + source.write_bytes(b'prepared bytes') + actions = [] + path = re.compile(r'https://companylake\.') + mocked_http.register_uri('PUT', path, status_code=201, headers=HEADERS) + + def patch(request, context): + action = parse_qs(urlsplit(request.url).query)['action'][0] + actions.append(action) + assert 'If-None-Match' not in request.headers + context.headers.update(HEADERS) + if action == 'append': + context.status_code = 202 + context.headers['ETag'] = '"append-response"' + else: + assert action == 'flush' + assert request.headers['If-Match'] == HEADERS['ETag'] + context.status_code = 412 + return b'' + + def delete(request, context): + actions.append('delete') + assert request.headers['If-Match'] == HEADERS['ETag'] + # Another writer's replacement must not be removed by cleanup. + context.status_code = 412 + return b'' + + mocked_http.register_uri('PATCH', path, content=patch) + mocked_http.register_uri('DELETE', path, content=delete, headers=HEADERS) + with client(mocked_http) as service: + filesystem = ADLSFileSystem(service, 'companylake', + 'dfs.core.windows.net') + with pytest.raises(TransientSourceError): + filesystem.put_file(str(source), AUTHORITY + '/data') + assert actions == ['append', 'flush', 'delete'] def test_sdk_empty_download_handles_azure_range_416(mocked_http, tmp_path): diff --git a/packages/kumo-adls/tests/test_transfer.py b/packages/kumo-adls/tests/test_transfer.py index d6f429efe..6e3843c69 100644 --- a/packages/kumo-adls/tests/test_transfer.py +++ b/packages/kumo-adls/tests/test_transfer.py @@ -329,6 +329,49 @@ def fail(self, *args, **kwargs): assert 'curated/failed' not in service.files +@pytest.mark.parametrize('upload_fails', [False, True]) +def test_put_file_conditions_bind_upload_and_cleanup_to_creation( + tmp_path, upload_fails): + from azure.core import MatchConditions + from azure.core.exceptions import ResourceModifiedError + from azure.storage.filedatalake import DataLakeFileClient + from kumo_core.sources.errors import TransientSourceError + + source = tmp_path / 'source' + source.write_bytes(b'prepared bytes') + client = Mock(spec=DataLakeFileClient) + client.create_file.return_value = {'etag': '"created"'} + service = Mock() + service.get_file_system_client.return_value.get_file_client.return_value = client + + def upload(stream, **kwargs): + assert not stream.closed + assert stream.read() == b'prepared bytes' + assert kwargs == { + 'length': source.stat().st_size, + 'overwrite': False, + 'etag': '"created"', + 'match_condition': MatchConditions.IfNotModified, + 'chunk_size': transfer.CHUNK_BYTES, + 'max_concurrency': 1, + } + if upload_fails: + raise ResourceModifiedError('object replaced') + + client.upload_data.side_effect = upload + filesystem = ADLSFileSystem(service, 'companylake', 'dfs.core.windows.net') + if upload_fails: + with pytest.raises(TransientSourceError): + filesystem.put_file(str(source), AUTHORITY + '/prepared') + client.delete_file.assert_called_once_with( + etag='"created"', match_condition=MatchConditions.IfNotModified) + else: + filesystem.put_file(str(source), AUTHORITY + '/prepared') + client.delete_file.assert_not_called() + client.create_file.assert_called_once_with() + client.upload_data.assert_called_once() + + def test_version_checks_overlap_and_remain_independent_of_listing_order(): from kumo_adls.backend import ADLSSourceReader diff --git a/packages/kumo-api/kumoapi/adls.py b/packages/kumo-api/kumoapi/adls.py index fd4f1366e..9f3ab7035 100644 --- a/packages/kumo-api/kumoapi/adls.py +++ b/packages/kumo-api/kumoapi/adls.py @@ -24,7 +24,7 @@ def normalize_adls_storage_account(value: str) -> str: if not isinstance(value, str) or not value.strip(): raise ValueError('Required provider field must not be blank') - if not isinstance(value, str) or _ADLS_ACCOUNT_RE.fullmatch(value) is None: + if _ADLS_ACCOUNT_RE.fullmatch(value) is None: raise ValueError('Invalid canonical Azure Data Lake storage account') return value @@ -32,8 +32,7 @@ def normalize_adls_storage_account(value: str) -> str: def normalize_adls_filesystem(value: str) -> str: if not isinstance(value, str) or not value.strip(): raise ValueError('Required provider field must not be blank') - if (not isinstance(value, str) - or _ADLS_FILESYSTEM_RE.fullmatch(value) is None or '--' in value): + if _ADLS_FILESYSTEM_RE.fullmatch(value) is None or '--' in value: raise ValueError('Invalid canonical Azure Data Lake filesystem') return value From 9bf5d5961a31290c99a079ab263129bfd2faa682 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Thu, 22 Oct 2026 23:58:16 +0000 Subject: [PATCH 07/15] Fix ADLS provider integration in public CI --- .gitleaks.toml | 8 ++++++++ test/scripts/public_tests/test_runner.py | 16 ++++++++++++++++ test/scripts/test_python_license_metadata.py | 2 +- tests/public_setup/Dockerfile | 2 +- tests/public_setup/Dockerfile.dockerignore | 1 + 5 files changed, 27 insertions(+), 2 deletions(-) diff --git a/.gitleaks.toml b/.gitleaks.toml index 4facfbd1f..0f6d413b5 100644 --- a/.gitleaks.toml +++ b/.gitleaks.toml @@ -34,3 +34,11 @@ condition = "AND" regexTarget = "secret" regexes = ['''^YOUR_API_KEY$'''] paths = ['''^packages/kumo-ui-catalog/src/stories/components/ui/Markdown\.stories\.tsx$'''] + +[[allowlists]] +description = "This loopback benchmark authentication-scheme label is not a credential." +targetRules = ["generic-api-key"] +condition = "AND" +regexTarget = "secret" +regexes = ['''^synthetic_sigv4$'''] +paths = ['''^scripts/benchmarks/object_store/compare_streams\.py$'''] diff --git a/test/scripts/public_tests/test_runner.py b/test/scripts/public_tests/test_runner.py index 94949231a..1a90b3e2b 100644 --- a/test/scripts/public_tests/test_runner.py +++ b/test/scripts/public_tests/test_runner.py @@ -10,6 +10,7 @@ import subprocess import sys import tempfile +import tomllib import unittest from pathlib import Path from unittest.mock import patch @@ -115,6 +116,18 @@ def test_docker_includes_new_tests_but_not_developer_credentials(self): target = context / core_path / module target.parent.mkdir(parents=True, exist_ok=True) target.write_text('# artifact dependency fixture\n') + manifests = [] + if profile == 'public_setup': + workspace = tomllib.loads( + (runner.ROOT / 'pyproject.toml').read_text()) + manifests = [ + Path(member) / 'pyproject.toml' for member in + workspace['tool']['uv']['workspace']['members'] + ] + for manifest in manifests: + target = context / manifest + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text('# workspace metadata fixture\n') (context / 'Dockerfile').write_text('FROM scratch\nCOPY . /\n') shutil.copyfile( runner.ROOT / 'tests' / profile / 'Dockerfile.dockerignore', @@ -131,6 +144,9 @@ def test_docker_includes_new_tests_but_not_developer_credentials(self): for module in artifact_modules: target = root / 'output' / core_path / module self.assertTrue(target.is_file(), str(module)) + for manifest in manifests: + self.assertTrue((root / 'output' / manifest).is_file(), + str(manifest)) @unittest.skipUnless(shutil.which('docker'), 'Docker is required') def test_wheel_rebuild_replaces_previous_output(self): diff --git a/test/scripts/test_python_license_metadata.py b/test/scripts/test_python_license_metadata.py index 76f1e2a3e..752b8558c 100644 --- a/test/scripts/test_python_license_metadata.py +++ b/test/scripts/test_python_license_metadata.py @@ -52,7 +52,7 @@ def synthetic_spec( def test_repository_manifest_inventory_passes() -> None: - assert len(check_python_license_metadata.DISTRIBUTIONS) == 12 + assert len(check_python_license_metadata.DISTRIBUTIONS) == 13 assert check_python_license_metadata.validate_repository( REPOSITORY_ROOT) == [] diff --git a/tests/public_setup/Dockerfile b/tests/public_setup/Dockerfile index b25a4e3b3..35dbf9fa6 100644 --- a/tests/public_setup/Dockerfile +++ b/tests/public_setup/Dockerfile @@ -34,7 +34,7 @@ RUN uv export --frozen --no-dev --no-emit-workspace --no-hashes \ COPY --chown=65534:65534 . . RUN uv pip install --system --no-deps ./packages/kumo-api ./packages/kumo-sdk \ ./packages/kumo-pql ./packages/kumo-core ./packages/kumo-engine \ - ./packages/kumo-s3 ./packages/kumo-gcs ./packages/kumo-snowflake \ + ./packages/kumo-s3 ./packages/kumo-gcs ./packages/kumo-adls ./packages/kumo-snowflake \ ./packages/kumo-databricks ./packages/kumo-bigquery '.[full,test]' \ && uv pip check ENV PYTHONPATH=/src DEPLOY_STAGE=test KUMO_TESTING=True \ diff --git a/tests/public_setup/Dockerfile.dockerignore b/tests/public_setup/Dockerfile.dockerignore index c8b5a2905..a2e21ac20 100644 --- a/tests/public_setup/Dockerfile.dockerignore +++ b/tests/public_setup/Dockerfile.dockerignore @@ -25,6 +25,7 @@ !packages/kumo-pql/** !packages/kumo-s3/** !packages/kumo-gcs/** +!packages/kumo-adls/** !packages/kumo-snowflake/** !packages/kumo-databricks/** !packages/kumo-bigquery/** From 0d8e669149e596f6296558f0604934b143ca9545 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Thu, 22 Oct 2026 23:58:18 +0000 Subject: [PATCH 08/15] Preserve dynamic API type descriptors with current Pydantic --- packages/kumo-api/kumoapi/json_serde.py | 4 ++-- packages/kumo-api/test/test_json_serde.py | 19 ++++++++++++++++++- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/packages/kumo-api/kumoapi/json_serde.py b/packages/kumo-api/kumoapi/json_serde.py index 8f3356019..b7900afd5 100644 --- a/packages/kumo-api/kumoapi/json_serde.py +++ b/packages/kumo-api/kumoapi/json_serde.py @@ -9,7 +9,7 @@ from decimal import Decimal from enum import Enum from pathlib import PurePath -from typing import Any, Dict, Type, TypeVar, Union +from typing import Any, Dict, Type, TypeVar, Union, cast from uuid import UUID from pydantic import BaseModel, SecretStr @@ -139,5 +139,5 @@ def from_json(obj: Any, cls: Union[Type[T], object]) -> T: if isinstance(obj, str): obj = json.loads(obj) from pydantic import TypeAdapter - adapter = TypeAdapter(cls) + adapter: TypeAdapter[T] = TypeAdapter(cast(Any, cls)) return adapter.validate_python(obj) diff --git a/packages/kumo-api/test/test_json_serde.py b/packages/kumo-api/test/test_json_serde.py index 0592e2415..9276d9bee 100644 --- a/packages/kumo-api/test/test_json_serde.py +++ b/packages/kumo-api/test/test_json_serde.py @@ -4,8 +4,10 @@ import dataclasses from dataclasses import field +from typing import Annotated, Union -from pydantic import BaseModel +import pytest +from pydantic import BaseModel, Field from pydantic.dataclasses import dataclass from kumoapi.json_serde import dataclass_to_dict, from_json @@ -17,6 +19,21 @@ class MetadataEnvelope: metadata: dict +@pytest.mark.parametrize('descriptor,obj,expected', [ + (int, '"7"', 7), + (list[int], '["1", 2]', [1, 2]), + (Union[int, None], '"3"', 3), + (Annotated[int, Field(gt=0)], '"4"', 4), + (dict[str, int], { + 'value': '5' + }, { + 'value': 5 + }), +]) +def test_from_json_accepts_type_descriptors(descriptor, obj, expected): + assert from_json(obj, descriptor) == expected + + def test_deserialization_does_not_migrate_unrelated_provider_metadata(): metadata = { 'type': 'gcs', From 875ffd9c0a7ce84ef026f7d82ad5aeb3783ec4eb Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Fri, 23 Oct 2026 22:32:14 +0000 Subject: [PATCH 09/15] Reject unsupported ADLS output before prediction scoring --- .../kumo-adls/src/kumo_adls/native_output.py | 11 ++- packages/kumo-adls/src/kumo_adls/sink.py | 18 ++++- .../kumo-adls/tests/test_sink_preflight.py | 45 +++++++++++ .../kumo-engine/test/test_adls_connector.py | 77 +++++++++++++++++++ 4 files changed, 148 insertions(+), 3 deletions(-) create mode 100644 packages/kumo-adls/tests/test_sink_preflight.py diff --git a/packages/kumo-adls/src/kumo_adls/native_output.py b/packages/kumo-adls/src/kumo_adls/native_output.py index 05cc6a6b7..c198ffa47 100644 --- a/packages/kumo-adls/src/kumo_adls/native_output.py +++ b/packages/kumo-adls/src/kumo_adls/native_output.py @@ -17,8 +17,8 @@ from .session import CLOUDS -def filesystem_for(principal, account, suffix, container): - """Bind native output without changing or substituting an identity.""" +def validate_output_configuration(principal, account, suffix, container): + """Validate native output locally, without constructing storage clients.""" account = normalize_adls_storage_account(account) suffix = normalize_adls_endpoint_suffix(suffix) container = normalize_adls_filesystem(container) @@ -51,6 +51,13 @@ def filesystem_for(principal, account, suffix, container): raise ValueError( 'Native Arrow ADLS writes require AZURE_AUTHORITY_HOST=' + f'https://{authority} for the selected cloud') + return account, suffix, container + + +def filesystem_for(principal, account, suffix, container): + """Bind native output without changing or substituting an identity.""" + account, suffix, container = validate_output_configuration( + principal, account, suffix, container) options = dict( blob_storage_authority='.' + suffix.replace('dfs.', 'blob.', 1), dfs_storage_authority='.' + suffix) diff --git a/packages/kumo-adls/src/kumo_adls/sink.py b/packages/kumo-adls/src/kumo_adls/sink.py index e92885330..06e748fe6 100644 --- a/packages/kumo-adls/src/kumo_adls/sink.py +++ b/packages/kumo-adls/src/kumo_adls/sink.py @@ -7,16 +7,32 @@ from kumo_core.object_store_sink import ObjectStoreSinkBackend from kumo_core.sources.cpu import COPY_WORKER_BYTES, memory_share from kumoapi.data_source import DataSourceType +from kumoapi.jobs import WriteMode from pyarrow.fs import FSSpecHandler, PyFileSystem from .backend import ADLSSessions, canonical_uri -from .native_output import NativeOutputHandler, filesystem_for +from .native_output import ( + NativeOutputHandler, + filesystem_for, + validate_output_configuration, +) from .sdk_filesystem import ADLSFileSystem class ADLSSinkBackend(ADLSSessions, ObjectStoreSinkBackend): source_type = DataSourceType.ADLS + def preflight_write_files(self, connector, table_name, + write_mode=WriteMode.OVERWRITE): + """Reject unsupported output before prediction work or Azure access.""" + if write_mode is not WriteMode.OVERWRITE: + raise NotImplementedError('ADLS output only supports OVERWRITE') + bucket, _, _ = self._resolved(connector, table_name) + container, host = bucket.split('@', 1) + account, suffix = host.split('.', 1) + validate_output_configuration(self._principal(connector), account, + suffix, container) + def _uri(self, path: str) -> str: return canonical_uri(path) diff --git a/packages/kumo-adls/tests/test_sink_preflight.py b/packages/kumo-adls/tests/test_sink_preflight.py new file mode 100644 index 000000000..9403bffad --- /dev/null +++ b/packages/kumo-adls/tests/test_sink_preflight.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from kumo_adls import native_output +from kumo_adls.session import CLOUDS +from kumo_adls.sink import ADLSSinkBackend +from kumoapi.data_source import ADLSIdentityConfig + + +@pytest.mark.parametrize('cloud', list(CLOUDS)) +@pytest.mark.parametrize('explicit', [False, True]) +@pytest.mark.parametrize('full_uri', [False, True]) +def test_supported_preflight_constructs_no_clients(monkeypatch, connector, + cloud, explicit, full_uri): + suffix, _, authority = CLOUDS[cloud] + monkeypatch.setenv('AZURE_AUTHORITY_HOST', 'https://' + authority) + monkeypatch.setenv('AZURE_TENANT_ID', 'tenant') + monkeypatch.setenv('AZURE_CLIENT_ID', 'client') + monkeypatch.setenv('AZURE_FEDERATED_TOKEN_FILE', '/projected/token') + principal = SimpleNamespace( + config=ADLSIdentityConfig(tenant_id='tenant', client_id='client', + cloud=cloud)) if explicit else None + backend = ADLSSinkBackend() + monkeypatch.setattr(backend, '_principal', lambda _: principal) + forbidden = Mock( + side_effect=AssertionError('preflight constructed a client')) + monkeypatch.setattr('kumo_adls.backend.source_access', forbidden) + monkeypatch.setattr(native_output, 'AzureFileSystem', forbidden) + monkeypatch.setattr(native_output, 'FileSystem', + SimpleNamespace(from_uri=forbidden)) + target = 'predictions/' + if full_uri: + target = f'abfss://curated@companylake.{suffix}/orders/{target}' + before = dict(os.environ) + + backend.preflight_write_files(connector(suffix=suffix), target) + + forbidden.assert_not_called() + assert dict(os.environ) == before diff --git a/packages/kumo-engine/test/test_adls_connector.py b/packages/kumo-engine/test/test_adls_connector.py index 7864e9dae..9bfcf3683 100644 --- a/packages/kumo-engine/test/test_adls_connector.py +++ b/packages/kumo-engine/test/test_adls_connector.py @@ -3,6 +3,9 @@ # SPDX-License-Identifier: Apache-2.0 import hashlib +from types import SimpleNamespace +from typing import cast +from unittest.mock import Mock import pandas as pd import pyarrow as pa @@ -10,6 +13,8 @@ import pytest from kumo_core.sources.object_store_access import ObjectStoreAccess from kumoapi.data_source import AzureCloud, DataSourceType +from kumoapi.jobs import PredictionArtifactType, WriteMode +from kumoapi.model_plan import ModelPlan from kumoapi.source_table import ADLSRelativeSource, ADLSSourceTable from pyarrow.fs import LocalFileSystem, SubTreeFileSystem @@ -106,3 +111,75 @@ def test_rootless_graph_coordinate_contains_exact_source(storage): coordinate = _source_table_name(source) assert 'companylake' in coordinate assert len(coordinate.split('.')) == 3 + + +@pytest.mark.parametrize('case,error,match', [ + ('rootless', NotImplementedError, 'source-only'), + ('client', NotImplementedError, 'match the selected workload identity'), + ('tenant', NotImplementedError, 'match the selected workload identity'), + ('token', NotImplementedError, 'projected token'), + ('authority', ValueError, 'AZURE_AUTHORITY_HOST'), + ('cloud', ValueError, 'selected identity cloud'), + ('append', NotImplementedError, 'OVERWRITE'), + ('outside', ValueError, 'outside the connector root'), +]) +def test_invalid_prediction_output_fails_before_scoring(monkeypatch, tmp_path, + snapshot, case, error, + match): + from kumo_core.types import TaskType + + from kumo.engine import PredictionOutputConfig, PredictionTable, Trainer + + monkeypatch.setenv('AZURE_TENANT_ID', 'tenant') + monkeypatch.setenv('AZURE_CLIENT_ID', 'client') + monkeypatch.setenv('AZURE_FEDERATED_TOKEN_FILE', '/projected/token') + monkeypatch.setenv('AZURE_AUTHORITY_HOST', + 'https://login.microsoftonline.com') + if case in ('client', 'tenant'): + monkeypatch.setenv('AZURE_' + case.upper() + '_ID', 'different') + elif case == 'token': + monkeypatch.delenv('AZURE_FEDERATED_TOKEN_FILE') + elif case == 'authority': + monkeypatch.setenv('AZURE_AUTHORITY_HOST', + 'https://login.microsoftonline.us') + identity = ADLSIdentity( + tenant_id='tenant', client_id='client', cloud=AzureCloud.US_GOVERNMENT + if case == 'cloud' else AzureCloud.PUBLIC) + connector = ADLSConnector(root_dir=None if case == 'rootless' else ROOT, + identity=identity) + output = PredictionOutputConfig( + connector, artifact_type=PredictionArtifactType.PREDICTIONS, + table_name=ROOT.replace('orders/', 'outside/') + if case == 'outside' else 'predictions/', write_mode=WriteMode.APPEND + if case == 'append' else WriteMode.OVERWRITE) + trainer = Trainer(ModelPlan(), accelerator='cpu') + monkeypatch.setattr(trainer, '_training_result', + SimpleNamespace(artifact_paths={})) + trainer._training_query = 'trained query' + monkeypatch.setattr(Trainer, 'scratch_dir', property(lambda _: tmp_path)) + # Isolate output validation from query parsing and model execution. + monkeypatch.setattr( + 'kumo_core.prediction_options.prediction_query_contract', lambda *args: + (None, TaskType.REGRESSION.value)) + monkeypatch.setattr('kumo_core.pquery.normalize_graph', lambda _: None) + monkeypatch.setattr( + 'kumo_core.pquery.parse_and_validate_query', + lambda *args: SimpleNamespace(query=SimpleNamespace(top_k=None))) + work = Mock(side_effect=AssertionError('prediction work started')) + monkeypatch.setattr( + 'kumo.engine.trainer.trainer.build_generated_training_task_dataset', + work) + monkeypatch.setattr('kumo_core.prediction.materialize_prediction_task', + work) + monkeypatch.setattr('kumo_core.prediction.score_prediction', work) + azure = Mock(side_effect=AssertionError('Azure access started')) + monkeypatch.setattr('kumo_adls.backend.source_access', azure) + monkeypatch.setattr('kumo_adls.sink.filesystem_for', azure) + table = cast(PredictionTable, SimpleNamespace(query='prediction query')) + + with pytest.raises(error, match=match): + trainer.predict(table, snapshot=snapshot, output_config=output) + + work.assert_not_called() + azure.assert_not_called() + assert not (tmp_path / 'runs').exists() From daa0b26c65928665b69a36728c457c905ee8c2e3 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Fri, 23 Oct 2026 22:41:40 +0000 Subject: [PATCH 10/15] Skip SDK directory preparation for native ADLS exports --- .../kumo-adls/src/kumo_adls/native_output.py | 5 +- packages/kumo-adls/tests/conftest.py | 4 + .../kumo-adls/tests/test_native_output.py | 82 ++++++++++++++++++- 3 files changed, 84 insertions(+), 7 deletions(-) diff --git a/packages/kumo-adls/src/kumo_adls/native_output.py b/packages/kumo-adls/src/kumo_adls/native_output.py index c198ffa47..84b2e576f 100644 --- a/packages/kumo-adls/src/kumo_adls/native_output.py +++ b/packages/kumo-adls/src/kumo_adls/native_output.py @@ -92,10 +92,9 @@ def _output(self): def create_dir(self, path, recursive): self.fs._address(path) - # Shared sinks prepare parent directories before opening output. - # Reject unsupported credentials before that SDK mutation as well. + # Native Blob uploads handle parents; shared preparation only needs + # to validate the path and bind the selected output credentials. self._output() - super().create_dir(path, recursive) def open_output_stream(self, path, metadata): # Validate the virtual authority before dropping its account qualifier. diff --git a/packages/kumo-adls/tests/conftest.py b/packages/kumo-adls/tests/conftest.py index 8f1acfa13..3f6add4f1 100644 --- a/packages/kumo-adls/tests/conftest.py +++ b/packages/kumo-adls/tests/conftest.py @@ -53,6 +53,10 @@ def __init__(self, service): def open_output_stream(self, path, metadata=None): service = self.service + # Blob uploads materialize missing HNS parents without DFS mkdir calls. + parts = path.split('/') + service.directories.update('/'.join(parts[:i]) + for i in range(1, len(parts))) class Output(io.BytesIO): def close(self): diff --git a/packages/kumo-adls/tests/test_native_output.py b/packages/kumo-adls/tests/test_native_output.py index 50de27939..a20060c28 100644 --- a/packages/kumo-adls/tests/test_native_output.py +++ b/packages/kumo-adls/tests/test_native_output.py @@ -90,8 +90,9 @@ def test_shared_sink_copies_prepared_file_through_native_stream( assert reader.read() == payload -def test_wrong_account_is_rejected_before_native_output_opens( - storage, monkeypatch): +@pytest.mark.parametrize('operation', ['open_output_stream', 'create_dir']) +def test_wrong_account_is_rejected_before_native_output_binding( + storage, monkeypatch, operation): arrow, _, _ = storage opened = [] @@ -102,8 +103,8 @@ def open_output_stream(self, *args, **kwargs): arrow.handler.output_filesystem = UnopenedOutput() with pytest.raises(ValueError, match='bound storage account'): - arrow.open_output_stream( - 'curated@otherlake.dfs.core.windows.net/orders/x') + getattr(arrow, + operation)('curated@otherlake.dfs.core.windows.net/orders/x') assert opened == [] @@ -273,3 +274,76 @@ def unsupported(*args): assert service.directories == initial_directories assert service.files == initial_files assert service.created == initial_created + + +@pytest.mark.parametrize('operation', + ['write_table', 'write_files', 'publish_model']) +def test_native_exports_skip_sdk_directory_preparation(service, connector, + monkeypatch, tmp_path, + operation): + import pandas as pd + from fakes import FakeFile + from kumo_adls.sink import ADLSSinkBackend + + def forbidden(*args, **kwargs): + pytest.fail( + 'Native exports must not create directories through the SDK') + + monkeypatch.setattr(FakeFile, 'create_directory', forbidden) + # Preparing this root must not attempt to create its ancestor "orders". + target = connector(prefix='orders/team') + prefix = 'curated/orders/team/' + backend = ADLSSinkBackend() + if operation == 'write_table': + backend.write_table(target, 'new/results.parquet', + pd.DataFrame({'id': [1]})) + contents = service.files[prefix + 'new/results.parquet'] + assert pq.read_table(pa.BufferReader(contents)).to_pydict() == { + 'id': [1] + } + elif operation == 'write_files': + source = tmp_path / 'source.parquet' + pq.write_table(pa.table({'id': [1]}), source) + backend.write_files(target, 'new/predictions/', [str(source)], + export_id='test') + outputs = [ + data for name, data in service.files.items() + if name.startswith(prefix + 'new/predictions/') + ] + assert source.read_bytes() in outputs + assert prefix + 'new/predictions/_SUCCESS' in service.files + else: + bundle = tmp_path / 'bundle' + (bundle / 'nested').mkdir(parents=True) + (bundle / 'model.bin').write_bytes(b'model') + (bundle / 'nested/weights.bin').write_bytes(b'weights') + backend.publish_model(target, str(bundle), path='new/model') + assert service.files[prefix + 'new/model/model.bin'] == b'model' + assert service.files[prefix + + 'new/model/nested/weights.bin'] == b'weights' + + +def test_directory_preparation_validates_without_creating_output( + service, connector, monkeypatch): + from conftest import FakeNativeFilesystem + from kumo_adls.sink import ADLSSinkBackend + + calls = [] + + def factory(*args): + calls.append(args) + return FakeNativeFilesystem(service) + + monkeypatch.setattr('kumo_adls.sink.filesystem_for', factory) + filesystem, path = ADLSSinkBackend()._destination(connector(), + 'new/nested/') + initial_directories = set(service.directories) + initial_files = dict(service.files) + + filesystem.create_dir(path) + filesystem.create_dir(path, recursive=False) + + assert len(calls) == 1 + assert service.directories == initial_directories + assert service.files == initial_files + assert service.created == [] From bf699065a7b61c191dc8dadcd104dc442b7f1ee3 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Fri, 23 Oct 2026 22:52:49 +0000 Subject: [PATCH 11/15] Validate rooted ADLS URI destinations consistently --- packages/kumo-adls/src/kumo_adls/backend.py | 22 +++-- packages/kumo-adls/tests/test_coordinates.py | 96 +++++++++++++++++++ .../kumo-engine/test/test_adls_connector.py | 9 +- 3 files changed, 115 insertions(+), 12 deletions(-) create mode 100644 packages/kumo-adls/tests/test_coordinates.py diff --git a/packages/kumo-adls/src/kumo_adls/backend.py b/packages/kumo-adls/src/kumo_adls/backend.py index 6cd4a6235..60a8b01db 100644 --- a/packages/kumo-adls/src/kumo_adls/backend.py +++ b/packages/kumo-adls/src/kumo_adls/backend.py @@ -105,15 +105,19 @@ def _contained_key(root, uri, *, purpose): raise ValueError(f'ADLS {purpose} is outside the connector root') return key + def _relative_coordinate(self, root, coordinate, *, purpose): + if '://' in coordinate: + key = self._contained_key(root, coordinate, purpose=purpose) + prefix = root.prefix.rstrip('/') + coordinate = key[len(prefix):].lstrip('/') if prefix else key + return normalize_adls_source_relative_path(coordinate) + def _placement(self, connector, table_name): root = self._root_for(connector) if root is None: return super()._placement(connector, table_name) - if '://' in table_name: - key = self._contained_key(root, table_name, purpose='source') - prefix = root.prefix.rstrip('/') - table_name = key[len(prefix):].lstrip('/') if prefix else key - table_name = normalize_adls_source_relative_path(table_name) + table_name = self._relative_coordinate(root, table_name, + purpose='source') return root.bucket, root.path, table_name, None def _source_access(self, connector, bucket, region): @@ -137,11 +141,9 @@ def _resolved(self, connector, coordinate): if root is None: raise NotImplementedError( 'Rootless ADLS connectors are source-only') - if '://' in coordinate: - key = self._contained_key(root, coordinate, purpose='destination') - else: - relative = normalize_adls_source_relative_path(coordinate) - key = '/'.join(part for part in (root.prefix, relative) if part) + relative = self._relative_coordinate(root, coordinate, + purpose='destination') + key = '/'.join(part for part in (root.prefix, relative) if part) return root.bucket, f'{root.bucket}/{key}', None diff --git a/packages/kumo-adls/tests/test_coordinates.py b/packages/kumo-adls/tests/test_coordinates.py new file mode 100644 index 000000000..a5ae30589 --- /dev/null +++ b/packages/kumo-adls/tests/test_coordinates.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pandas as pd +import pytest +from kumo_adls.backend import ADLSSourceBackend +from kumo_adls.sink import ADLSSinkBackend +from kumoapi.source_table import ADLSRelativeSource, ADLSURISource + +URI = 'abfss://curated@companylake.dfs.core.windows.net/orders/' + + +@pytest.mark.parametrize('prefix', ['', 'orders', 'orders/team']) +@pytest.mark.parametrize('filename', [ + 'scores%23review.parquet', + 'scores%3Freview.parquet', + 'scores%25review.parquet', + 'scores%3Areview.parquet', +]) +def test_full_uri_output_obeys_rooted_relative_policy_before_storage( + connector, monkeypatch, prefix, filename): + from unittest.mock import Mock + + target = connector(prefix=prefix) + root = URI.removesuffix('orders/') + (prefix + '/' if prefix else '') + uri = root + filename + access = Mock(side_effect=AssertionError('must not access storage')) + native = Mock(side_effect=AssertionError('must not bind native output')) + monkeypatch.setattr('kumo_adls.backend.source_access', access) + monkeypatch.setattr('kumo_adls.sink.filesystem_for', native) + source = ADLSSourceBackend() + sink = ADLSSinkBackend() + + with pytest.raises(ValueError): + source.get_table_config(target, uri) + with pytest.raises(ValueError): + source.get_columns(target, uri) + with pytest.raises(ValueError): + sink.preflight_write_files(target, uri) + with pytest.raises(ValueError): + sink.write_table(target, uri, pd.DataFrame({'id': [1]})) + + access.assert_not_called() + native.assert_not_called() + + +@pytest.mark.parametrize('prefix', ['', 'orders/team']) +@pytest.mark.parametrize('filename', + ['scores review.parquet', 'résultats ü.parquet']) +def test_valid_full_uri_export_roundtrips_through_source_schema( + service, connector, prefix, filename): + from urllib.parse import quote + + target = connector(prefix=prefix) + root = URI.removesuffix('orders/') + (prefix + '/' if prefix else '') + uri = root + quote(filename) + returned = ADLSSinkBackend().write_table(target, uri, + pd.DataFrame({'id': [3, 4]})) + assert returned == uri + source = ADLSSourceBackend() + config = source.get_table_config(target, returned).source_table + assert isinstance(config.source, ADLSRelativeSource) + assert config.source.path == filename + assert config.source.path == source.get_table_config( + target, filename).source_table.source.path + assert '3' in source.get_sample(target, returned, 2).sample_rows + assert '4' in source.get_sample(target, filename, 2).sample_rows + + +@pytest.mark.parametrize('coordinate', ['', 'orders', 'orders/']) +def test_empty_root_relative_output_fails_before_storage( + connector, monkeypatch, coordinate): + from unittest.mock import Mock + + target = connector() + uri = URI.removesuffix('orders/') + coordinate if coordinate else '' + access = Mock(side_effect=AssertionError('must not access storage')) + monkeypatch.setattr('kumo_adls.backend.source_access', access) + with pytest.raises(ValueError): + ADLSSourceBackend().get_table_config(target, uri) + with pytest.raises(ValueError): + ADLSSinkBackend().write_table(target, uri, pd.DataFrame({'id': [1]})) + access.assert_not_called() + + +def test_rootless_full_uri_keeps_exact_source_policy(service, connector): + uri = URI + 'scores%23review.parquet' + service.files['curated/orders/scores#review.parquet'] = service.files[ + 'curated/orders/events.parquet'] + source = ADLSSourceBackend() + target = connector(any_uri=True) + config = source.get_table_config(target, uri).source_table + assert isinstance(config.source, ADLSURISource) + assert config.source.uri == uri + assert '10.0' in source.get_sample(target, uri, 1).sample_rows diff --git a/packages/kumo-engine/test/test_adls_connector.py b/packages/kumo-engine/test/test_adls_connector.py index 9bfcf3683..529bf9c1d 100644 --- a/packages/kumo-engine/test/test_adls_connector.py +++ b/packages/kumo-engine/test/test_adls_connector.py @@ -122,6 +122,7 @@ def test_rootless_graph_coordinate_contains_exact_source(storage): ('cloud', ValueError, 'selected identity cloud'), ('append', NotImplementedError, 'OVERWRITE'), ('outside', ValueError, 'outside the connector root'), + ('encoded_uri', ValueError, 'must not be a URI, query, or fragment'), ]) def test_invalid_prediction_output_fails_before_scoring(monkeypatch, tmp_path, snapshot, case, error, @@ -147,10 +148,14 @@ def test_invalid_prediction_output_fails_before_scoring(monkeypatch, tmp_path, if case == 'cloud' else AzureCloud.PUBLIC) connector = ADLSConnector(root_dir=None if case == 'rootless' else ROOT, identity=identity) + destination = 'predictions/' + if case == 'outside': + destination = ROOT.replace('orders/', 'outside/') + elif case == 'encoded_uri': + destination = ROOT + 'scores%23review.parquet' output = PredictionOutputConfig( connector, artifact_type=PredictionArtifactType.PREDICTIONS, - table_name=ROOT.replace('orders/', 'outside/') - if case == 'outside' else 'predictions/', write_mode=WriteMode.APPEND + table_name=destination, write_mode=WriteMode.APPEND if case == 'append' else WriteMode.OVERWRITE) trainer = Trainer(ModelPlan(), accelerator='cpu') monkeypatch.setattr(trainer, '_training_result', From 2c1c321a70d36f3d86a6162d9fba3b961b290333 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Fri, 23 Oct 2026 23:18:08 +0000 Subject: [PATCH 12/15] Expose shared ADLS readers for caller-owned Azure clients --- packages/kumo-adls/README.md | 26 +++ packages/kumo-adls/docs/implementation.md | 59 +++++ packages/kumo-adls/src/kumo_adls/__init__.py | 3 + packages/kumo-adls/src/kumo_adls/backend.py | 40 +--- packages/kumo-adls/src/kumo_adls/client.py | 83 +++++++ packages/kumo-adls/src/kumo_adls/clouds.py | 17 ++ .../kumo-adls/src/kumo_adls/native_output.py | 2 +- packages/kumo-adls/src/kumo_adls/reader.py | 77 +++++++ .../kumo-adls/src/kumo_adls/sdk_filesystem.py | 2 + packages/kumo-adls/src/kumo_adls/session.py | 53 +---- .../kumo-adls/src/kumo_adls/source_access.py | 6 +- packages/kumo-adls/tests/conftest.py | 9 +- packages/kumo-adls/tests/fakes.py | 8 +- packages/kumo-adls/tests/test_client.py | 115 ++++++++++ packages/kumo-adls/tests/test_sdk_transfer.py | 3 +- packages/kumo-adls/tests/test_session.py | 22 +- .../kumo-adls/tests/test_source_access.py | 217 ++++++++++++++++++ packages/kumo-adls/tests/test_transfer.py | 10 +- .../kumo-engine/test/test_adls_connector.py | 1 + .../benchmarks/object_store/compare_live.py | 3 +- 20 files changed, 657 insertions(+), 99 deletions(-) create mode 100644 packages/kumo-adls/src/kumo_adls/client.py create mode 100644 packages/kumo-adls/src/kumo_adls/clouds.py create mode 100644 packages/kumo-adls/src/kumo_adls/reader.py create mode 100644 packages/kumo-adls/tests/test_client.py create mode 100644 packages/kumo-adls/tests/test_source_access.py diff --git a/packages/kumo-adls/README.md b/packages/kumo-adls/README.md index dfa5657fe..20cba35b1 100644 --- a/packages/kumo-adls/README.md +++ b/packages/kumo-adls/README.md @@ -40,6 +40,32 @@ unavailable. SQL, warehouse-style `ReadMode.STAGED`, and append output are unsupported. Prepared exports and generated output use native Arrow streams, matching S3/GCS; consumers must wait for export completion. +## Shared SDK reads + +Callers with broker-selected credentials or an existing Azure service client +can use the same reader without engine identities or Azure authentication +environment variables: + +```python +from kumo_adls import create_client, source_reader +from kumo_core.sources.types import ReadMode + +# credential is supplied by the caller's authentication system. +with create_client(credential, "companylake") as client: + reader = source_reader(client, "companylake", filesystem="curated", + prefix="orders") + staged = reader.prepare("events.parquet", "/local/staging", + columns=None, mode=ReadMode.STREAMED) +``` + +`source_reader` and `source_access` borrow clients and never close them. +`create_client` creates a caller-owned client/transport and borrows its +credential; the caller closes that credential separately. Existing clients +must use the matching account's canonical HTTPS DFS service endpoint. Native +Arrow output restrictions belong to the engine adapter, not this read API. +See [implementation details](docs/implementation.md#shared-sdk-read-boundary) +for validation, ownership, and Studio adoption follow-ups. + ## Authentication and resources Without an identity, requests use Azure's ambient credential chain. An explicit diff --git a/packages/kumo-adls/docs/implementation.md b/packages/kumo-adls/docs/implementation.md index c1d649588..648b650e1 100644 --- a/packages/kumo-adls/docs/implementation.md +++ b/packages/kumo-adls/docs/implementation.md @@ -123,3 +123,62 @@ requires explicitly managed permissions. Temporary files are closed before promotion with `os.replace`, including on failure cleanup paths. This avoids requiring POSIX open-file replacement; it does not establish full Windows engine support. + + +## Shared SDK read boundary + +`kumo_adls.client.create_client(credential, account, suffix)` constructs the +same tuned Azure SDK transport used by engine sessions. It never selects an +identity or acquires a token during construction. It owns the created transport +and cleans it up on construction failure. The returned client is caller-owned; +closing it does not close the supplied credential. Credential selection, +refresh capability, and credential lifetime belong to the caller. + +`kumo_adls.source_access.source_access(client, account, suffix)` borrows a +`DataLakeServiceClient` and supplies `ObjectStoreAccess` with conditional SDK +fetches and identity callbacks. `kumo_adls.reader.source_reader` builds the +same `ADLSSourceReader` used by the engine, with a canonical decoded filesystem +prefix. These entry points import no engine identity registry, Azure Identity +credential chain, or native-output adapter. They neither inspect process +Azure authentication settings nor cache or close borrowed clients. + +Before I/O, the filesystem checks the SDK client's account name and service +URL against the requested account and cloud. The URL must be HTTPS, use the +canonical DFS authority without userinfo or an explicit port, and have no +filesystem/path or fragment. SAS query credentials are accepted as opaque SDK +configuration and are never echoed in validation errors. Custom authorities, +emulator URLs, filesystem clients, and mismatched accounts/clouds are rejected. +A caller providing its own client also owns its transport/retry configuration; +use `create_client` for the provider's tuned defaults. Read transfer slots still +apply, but caller-selected SDK buffer sizes can change actual memory use. + +The engine adapter resolves its principal, owns the resulting credential and +cached client, then uses this same source-access/reader implementation. Engine +session cleanup closes only those owned resources. Borrowed clients stay live +until their own caller closes them. Account-qualified paths, ETag version +hashes, conditional downloads, and listing-to-HEAD fallback remain shared. +Moving the reader does not change versions or invalidate staged data by itself. +Native Arrow writes continue to bind credentials separately in the engine sink; +process-identity requirements do not apply to caller-supplied SDK readers. + +Offline contract tests cover real SDK requests over an in-memory HTTP transport, +equivalent engine/shared rows and versions, unchanged conditional ETags, +concurrent account/identity isolation, endpoint rejection before requests, +borrowed-client lifetime, and owned-resource cleanup. A fresh-process import +guard rejects identity-adapter imports and constructs the shared reader with +Azure authentication variables removed. This is not live Azure federation or +sovereign-cloud validation. + +Follow-ups: + +- Studio read adoption: use broker-selected credentials with this reader, retire + the old Arrow adapter and duplicated download machinery, and define path + compatibility and migration from modification-time versions to ETag versions. +- Shared cloud configuration: consolidate hosted and package endpoint, + authority, and audience mappings. This package uses the generic storage + audience for all three clouds; existing hosted mappings differ and require + deliberate integration and live validation. +- Write primitives: identify which hosted SDK operations should become shared + credential-bound functions and validate them individually. Existing SDK + writer helpers are not the Studio default. SAS generation, multipart upload, + and atomic promotion are separate work; native engine output is unchanged. diff --git a/packages/kumo-adls/src/kumo_adls/__init__.py b/packages/kumo-adls/src/kumo_adls/__init__.py index 2a7c0132f..6a28b238d 100644 --- a/packages/kumo-adls/src/kumo_adls/__init__.py +++ b/packages/kumo-adls/src/kumo_adls/__init__.py @@ -7,6 +7,9 @@ ADLS_NAMESPACE = 'azure-adls' _LAZY_EXPORTS = { + 'ADLSSourceReader': 'reader', + 'create_client': 'client', + 'source_reader': 'reader', 'ADLSSourceBackend': 'backend', 'ADLSSinkBackend': 'sink', 'probe': 'session', diff --git a/packages/kumo-adls/src/kumo_adls/backend.py b/packages/kumo-adls/src/kumo_adls/backend.py index 60a8b01db..d8b423b7f 100644 --- a/packages/kumo-adls/src/kumo_adls/backend.py +++ b/packages/kumo-adls/src/kumo_adls/backend.py @@ -2,12 +2,9 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import hashlib -import json from typing import Optional from urllib.parse import quote, unquote, urlsplit -from kumo_core.sources.object_store import ObjectStoreSourceReader from kumo_core.sources.object_store_backend import ObjectStoreSourceBackend from kumo_core.sources.object_store_sessions import ( ObjectStorePaths, @@ -31,10 +28,10 @@ ADLSSourceTable, ADLSURISource, ) -from pyarrow.fs import FSSpecHandler, PyFileSystem from . import ADLS_NAMESPACE -from .sdk_filesystem import ADLSFileSystem +from .reader import ADLSSourceReader +from .session import session_for from .source_access import source_access @@ -44,36 +41,6 @@ def canonical_uri(path: str) -> str: f'abfss://{authority}/{quote(key, safe="/-._~!$&()*+,;=:@")}') -class ADLSSourceReader(ObjectStoreSourceReader): - source_type = ConnectorType.ADLS - - def _identities(self, found, worker_ceiling): - if isinstance(self.filesystem, PyFileSystem): - handler = self.filesystem.handler - adapter = handler.fs if isinstance(handler, FSSpecHandler) else None - if (isinstance(adapter, ADLSFileSystem) and getattr( - self.object_identity, '__self__', None) is adapter - and getattr(self.object_identity, '__func__', - None) is ADLSFileSystem.object_identity): - identities = adapter.object_identities( - [entry.path for entry in found], root=self.root) - if identities is not None: - return identities - return super()._identities(found, worker_ceiling) - - def _data_version(self, path, entries): - # The shared single-file version is only an mtime. Include authority - # and ETags even for one object, so accounts and same-second rewrites - # cannot alias the same cached source. - ordered = sorted(entries, key=lambda entry: entry.path) - metadata = list( - zip((entry.path for entry in ordered), - self._identities(ordered, None))) - # Discovery and preparation spell the same directory with/without '/'. - return hashlib.sha256( - json.dumps((path.rstrip('/'), metadata)).encode()).hexdigest() - - class ADLSSessions(ObjectStoreSessions): namespace = ADLS_NAMESPACE rooted = ADLSRootedConnectorTarget @@ -124,7 +91,8 @@ def _source_access(self, connector, bucket, region): del region _, host = bucket.split('@', 1) account, suffix = host.split('.', 1) - return source_access(self._principal(connector), account, suffix) + client = session_for(self._principal(connector), account, suffix) + return source_access(client, account, suffix) def _over(self, connector, bucket, root, region): access = self._source_access(connector, bucket, region) diff --git a/packages/kumo-adls/src/kumo_adls/client.py b/packages/kumo-adls/src/kumo_adls/client.py new file mode 100644 index 000000000..809db4a37 --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/client.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Credential-bound SDK clients; authentication selection belongs to callers.""" + +from urllib.parse import urlsplit + +from azure.core.pipeline.transport import RequestsTransport +from azure.storage.filedatalake import DataLakeServiceClient +from kumoapi.adls import ( + normalize_adls_endpoint_suffix, + normalize_adls_storage_account, +) +from requests import Session +from requests.adapters import HTTPAdapter + +from .clouds import CLOUDS +from .transfer import ( + BULK_DOWNLOAD_CHUNK_BYTES, + CONNECTIONS, + INITIAL_GET_BYTES, + NETWORK_BLOCK_BYTES, +) + + +def create_client( + credential, account: str, + suffix: str = 'dfs.core.windows.net') -> DataLakeServiceClient: + """Create a caller-owned client/transport, borrowing the supplied credential. + + Closing this client closes its transport, not the credential. No ambient + identity is selected; callers must supply a credential, including for SAS. + """ + account = normalize_adls_storage_account(account) + suffix = normalize_adls_endpoint_suffix(suffix) + if credential is None: + raise TypeError( + 'ADLS client construction requires a resolved credential') + audience = next(values[1] for values in CLOUDS.values() + if values[0] == suffix) + transport_session = Session() + try: + adapter = HTTPAdapter(pool_connections=CONNECTIONS, + pool_maxsize=CONNECTIONS, pool_block=True) + transport_session.mount('https://', adapter) + transport_session.mount('http://', adapter) + transport = RequestsTransport( + session=transport_session, session_owner=True, + connection_data_block_size=NETWORK_BLOCK_BYTES) + except BaseException: + transport_session.close() + raise + try: + return DataLakeServiceClient( + account_url=f'https://{account}.{suffix}', credential=credential, + audience=audience, max_single_get_size=INITIAL_GET_BYTES, + max_chunk_get_size=BULK_DOWNLOAD_CHUNK_BYTES, + connection_data_block_size=NETWORK_BLOCK_BYTES, transport=transport) + except BaseException: + transport.close() + raise + + +def validate_client(client, account: str, suffix: str) -> tuple[str, str]: + """Validate a borrowed service client without I/O or credential disclosure.""" + account = normalize_adls_storage_account(account) + suffix = normalize_adls_endpoint_suffix(suffix) + message = 'ADLS client does not match the requested account and HTTPS endpoint' + url = getattr(client, 'url', None) + if not isinstance(url, str) or any( + ord(c) <= 32 or ord(c) == 127 for c in url): + raise ValueError(message) + try: + parsed = urlsplit(url) + except ValueError: + raise ValueError(message) from None + # Query parameters can hold SAS credentials; neither inspect nor echo them. + if (parsed.scheme != 'https' + or parsed.netloc.casefold() != f'{account}.{suffix}' + or parsed.path not in ('', '/') or parsed.fragment + or getattr(client, 'account_name', None) != account): + raise ValueError(message) + return account, suffix diff --git a/packages/kumo-adls/src/kumo_adls/clouds.py b/packages/kumo-adls/src/kumo_adls/clouds.py new file mode 100644 index 000000000..3a1ed2675 --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/clouds.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Cloud endpoints shared by SDK clients and the engine's credential adapters.""" + +from kumoapi.data_source import AzureCloud + +# Storage uses the generic audience in all clouds; login authorities differ. +CLOUDS = { + AzureCloud.PUBLIC: ('dfs.core.windows.net', 'https://storage.azure.com', + 'login.microsoftonline.com'), + AzureCloud.US_GOVERNMENT: + ('dfs.core.usgovcloudapi.net', 'https://storage.azure.com', + 'login.microsoftonline.us'), + AzureCloud.CHINA: ('dfs.core.chinacloudapi.cn', 'https://storage.azure.com', + 'login.chinacloudapi.cn'), +} diff --git a/packages/kumo-adls/src/kumo_adls/native_output.py b/packages/kumo-adls/src/kumo_adls/native_output.py index 84b2e576f..1b09f0df7 100644 --- a/packages/kumo-adls/src/kumo_adls/native_output.py +++ b/packages/kumo-adls/src/kumo_adls/native_output.py @@ -14,7 +14,7 @@ from kumoapi.data_source import ADLSIdentityConfig from pyarrow.fs import AzureFileSystem, FileSystem, FSSpecHandler -from .session import CLOUDS +from .clouds import CLOUDS def validate_output_configuration(principal, account, suffix, container): diff --git a/packages/kumo-adls/src/kumo_adls/reader.py b/packages/kumo-adls/src/kumo_adls/reader.py new file mode 100644 index 000000000..88cf549dd --- /dev/null +++ b/packages/kumo-adls/src/kumo_adls/reader.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ADLS readers over borrowed clients, independent of engine identities.""" + +import hashlib +import json +from typing import Literal + +from kumo_core.sources.csv_format import CsvOptions +from kumo_core.sources.object_store import ObjectStoreSourceReader +from kumo_core.types.enums import ConnectorType +from kumoapi.adls import ( + normalize_adls_endpoint_suffix, + normalize_adls_filesystem, + normalize_adls_storage_account, +) +from kumoapi.object_store import normalize_object_store_root_prefix +from pyarrow.fs import FSSpecHandler, PyFileSystem + +from .sdk_filesystem import ADLSFileSystem +from .source_access import source_access + + +class ADLSSourceReader(ObjectStoreSourceReader): + source_type = ConnectorType.ADLS + + def _identities(self, found, worker_ceiling): + if isinstance(self.filesystem, PyFileSystem): + handler = self.filesystem.handler + adapter = handler.fs if isinstance(handler, FSSpecHandler) else None + if (isinstance(adapter, ADLSFileSystem) and getattr( + self.object_identity, '__self__', None) is adapter + and getattr(self.object_identity, '__func__', + None) is ADLSFileSystem.object_identity): + identities = adapter.object_identities( + [entry.path for entry in found], root=self.root) + if identities is not None: + return identities + return super()._identities(found, worker_ceiling) + + def _data_version(self, path, entries): + # The shared single-file version is only an mtime. Include authority + # and ETags even for one object, so accounts and same-second rewrites + # cannot alias the same cached source. + ordered = sorted(entries, key=lambda entry: entry.path) + metadata = list( + zip((entry.path for entry in ordered), + self._identities(ordered, None))) + # Discovery and preparation spell the same directory with/without '/'. + return hashlib.sha256( + json.dumps((path.rstrip('/'), metadata)).encode()).hexdigest() + + +def source_reader(client, account: str, suffix: str = 'dfs.core.windows.net', *, + filesystem: str, prefix: str = '', + file_format: Literal['csv', 'parquet'] | None = None, + file_suffix: str | None = None, + csv_options: CsvOptions | None = None) -> ADLSSourceReader: + """Build the engine's reader over a borrowed, already-authenticated client. + + The caller owns the client and credential for the reader's entire lifetime. + Prefix is decoded, relative to the filesystem; empty selects its root. + """ + account = normalize_adls_storage_account(account) + suffix = normalize_adls_endpoint_suffix(suffix) + filesystem = normalize_adls_filesystem(filesystem) + prefix = normalize_object_store_root_prefix(prefix, + provider_name='Azure Data Lake') + access = source_access(client, account, suffix) + authority = f'{filesystem}@{account}.{suffix}' + root = authority + ('/' + prefix if prefix else '') + return ADLSSourceReader(access.filesystem, root=root, + file_format=file_format, file_suffix=file_suffix, + csv_options=csv_options, + source_type=ConnectorType.ADLS, fetch=access.fetch, + object_identity=access.object_identity) diff --git a/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py b/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py index 418b6e40c..820a3e695 100644 --- a/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py +++ b/packages/kumo-adls/src/kumo_adls/sdk_filesystem.py @@ -21,6 +21,7 @@ from kumo_core.sources.layout import MAX_SCANNED_ENTRIES from kumoapi.adls import normalize_adls_filesystem +from .client import validate_client from .errors import translated from .range_cache import MAX_ENTRY_BYTES, RANGE_BYTES from .transfer import CHUNK_BYTES, concurrency @@ -141,6 +142,7 @@ class ADLSFileSystem(AbstractFileSystem): cachable = False def __init__(self, client, account: str, suffix: str): + account, suffix = validate_client(client, account, suffix) super().__init__() self._range_namespace = object() weakref.finalize(self, RANGE_BYTES.clear_namespace, diff --git a/packages/kumo-adls/src/kumo_adls/session.py b/packages/kumo-adls/src/kumo_adls/session.py index f348b0839..98c56c72d 100644 --- a/packages/kumo-adls/src/kumo_adls/session.py +++ b/packages/kumo-adls/src/kumo_adls/session.py @@ -5,37 +5,13 @@ import atexit import threading -from azure.core.pipeline.transport import RequestsTransport -from azure.identity import ( - AzureAuthorityHosts, - DefaultAzureCredential, - WorkloadIdentityCredential, -) -from azure.storage.filedatalake import DataLakeServiceClient +from azure.identity import DefaultAzureCredential, WorkloadIdentityCredential from kumo_core.sources.principals import Principal -from kumoapi.data_source import ADLSIdentityConfig, AzureCloud -from requests import Session -from requests.adapters import HTTPAdapter +from kumoapi.data_source import ADLSIdentityConfig +from .client import create_client +from .clouds import CLOUDS from .errors import translated -from .transfer import ( - BULK_DOWNLOAD_CHUNK_BYTES, - CONNECTIONS, - INITIAL_GET_BYTES, - NETWORK_BLOCK_BYTES, -) - -# Storage uses the same generic token audience in every cloud. -# DFS endpoints and Entra authorities remain cloud-specific. -CLOUDS = { - AzureCloud.PUBLIC: ('dfs.core.windows.net', 'https://storage.azure.com', - AzureAuthorityHosts.AZURE_PUBLIC_CLOUD), - AzureCloud.US_GOVERNMENT: - ('dfs.core.usgovcloudapi.net', 'https://storage.azure.com', - AzureAuthorityHosts.AZURE_GOVERNMENT), - AzureCloud.CHINA: ('dfs.core.chinacloudapi.cn', 'https://storage.azure.com', - AzureAuthorityHosts.AZURE_CHINA), -} def credential_for(principal: Principal | None, endpoint_suffix: str): @@ -72,25 +48,10 @@ def session_for(principal: Principal | None, account: str, suffix: str): with _LOCK: if key not in _SESSIONS: with translated(): - credential, audience = credential_for(principal, suffix) - transport_session = Session() - adapter = HTTPAdapter(pool_connections=CONNECTIONS, - pool_maxsize=CONNECTIONS, pool_block=True) - transport_session.mount('https://', adapter) - transport_session.mount('http://', adapter) - transport = RequestsTransport( - session=transport_session, session_owner=True, - connection_data_block_size=NETWORK_BLOCK_BYTES) + credential, _ = credential_for(principal, suffix) try: - client = DataLakeServiceClient( - account_url=f'https://{account}.{suffix}', - credential=credential, audience=audience, - max_single_get_size=INITIAL_GET_BYTES, - max_chunk_get_size=BULK_DOWNLOAD_CHUNK_BYTES, - connection_data_block_size=NETWORK_BLOCK_BYTES, - transport=transport) - except Exception: - transport.close() + client = create_client(credential, account, suffix) + except BaseException: credential.close() raise _SESSIONS[key] = client, credential diff --git a/packages/kumo-adls/src/kumo_adls/source_access.py b/packages/kumo-adls/src/kumo_adls/source_access.py index 344e4ce3a..d1b052749 100644 --- a/packages/kumo-adls/src/kumo_adls/source_access.py +++ b/packages/kumo-adls/src/kumo_adls/source_access.py @@ -6,11 +6,11 @@ from pyarrow.fs import FSSpecHandler, PyFileSystem from .sdk_filesystem import ADLSFileSystem -from .session import session_for -def source_access(principal, account: str, suffix: str) -> ObjectStoreAccess: - client = session_for(principal, account, suffix) +def source_access(client, account: str, + suffix: str = 'dfs.core.windows.net') -> ObjectStoreAccess: + """Borrow a service client; wrappers never close it or its credential.""" filesystem = ADLSFileSystem(client, account, suffix) return ObjectStoreAccess(PyFileSystem(FSSpecHandler(filesystem)), filesystem.fetch, filesystem.object_identity) diff --git a/packages/kumo-adls/tests/conftest.py b/packages/kumo-adls/tests/conftest.py index 3f6add4f1..790ece34c 100644 --- a/packages/kumo-adls/tests/conftest.py +++ b/packages/kumo-adls/tests/conftest.py @@ -20,8 +20,13 @@ @pytest.fixture def service(monkeypatch): service = FakeService() - monkeypatch.setattr('kumo_adls.source_access.session_for', - lambda *args: service) + + def client_for(principal, account, suffix): + return SimpleNamespace( + account_name=account, url=f'https://{account}.{suffix}/', + get_file_system_client=service.get_file_system_client) + + monkeypatch.setattr('kumo_adls.backend.session_for', client_for) monkeypatch.setattr('kumo_adls.sink.filesystem_for', lambda *args: FakeNativeFilesystem(service)) sink = pa.BufferOutputStream() diff --git a/packages/kumo-adls/tests/fakes.py b/packages/kumo-adls/tests/fakes.py index 0b3a3b74c..9f2d7b1bd 100644 --- a/packages/kumo-adls/tests/fakes.py +++ b/packages/kumo-adls/tests/fakes.py @@ -13,7 +13,10 @@ class FakeService: """Azure SDK-shaped storage using real public property models.""" - def __init__(self): + def __init__(self, account='companylake', suffix='dfs.core.windows.net'): + self.account_name = account + self.url = f'https://{account}.{suffix}/' + self.closed = False self.files = {} self.uncommitted = {} self.directories = {'curated', 'curated/orders'} @@ -21,6 +24,9 @@ def __init__(self): self.created = [] self.modified = datetime(2026, 1, 1, tzinfo=timezone.utc) + def close(self): + self.closed = True + def get_file_system_client(self, filesystem): return FakeFilesystem(self, filesystem) diff --git a/packages/kumo-adls/tests/test_client.py b/packages/kumo-adls/tests/test_client.py new file mode 100644 index 000000000..e6c4c93e8 --- /dev/null +++ b/packages/kumo-adls/tests/test_client.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import Mock + +import pytest +from kumo_adls import client as clients +from kumo_adls.clouds import CLOUDS +from kumo_adls.source_access import source_access +from kumo_adls.transfer import NETWORK_BLOCK_BYTES + + +class Credential: + def __init__(self): + self.closed = False + + def get_token(self, *scopes, **kwargs): + raise AssertionError('Construction must not acquire a token') + + def close(self): + self.closed = True + + +@pytest.mark.parametrize('cloud', list(CLOUDS)) +def test_factory_borrows_credential_and_owns_transport(monkeypatch, cloud): + suffix, _, _ = CLOUDS[cloud] + credential = Credential() + owned_session = clients.Session() + closed = Mock(wraps=owned_session.close) + monkeypatch.setattr(owned_session, 'close', closed) + monkeypatch.setattr(clients, 'Session', lambda: owned_session) + client = clients.create_client(credential, 'companylake', suffix) + assert clients.validate_client(client, 'companylake', + suffix) == ('companylake', suffix) + assert client._config.max_single_get_size == 16 << 20 + assert client._config.max_chunk_get_size == 16 << 20 + assert client._pipeline._transport.connection_config.data_block_size == NETWORK_BLOCK_BYTES + client.close() + closed.assert_called_once() + assert not credential.closed + credential.close() + + +@pytest.mark.parametrize('stage', ['adapter', 'mount', 'transport', 'client']) +def test_factory_failure_closes_owned_resources_not_credential( + monkeypatch, stage): + credential = Credential() + owned_session = Mock() + monkeypatch.setattr(clients, 'Session', lambda: owned_session) + failure = Mock(side_effect=RuntimeError('construction failed')) + if stage == 'adapter': + monkeypatch.setattr(clients, 'HTTPAdapter', failure) + elif stage == 'mount': + owned_session.mount.side_effect = RuntimeError('construction failed') + elif stage == 'transport': + monkeypatch.setattr(clients, 'RequestsTransport', failure) + else: + monkeypatch.setattr(clients, 'DataLakeServiceClient', failure) + with pytest.raises(RuntimeError, match='construction failed'): + clients.create_client(credential, 'companylake') + owned_session.close.assert_called_once() + assert not credential.closed + + +def test_factory_requires_credential_before_creating_transport(monkeypatch): + factory = Mock(side_effect=AssertionError('must not create transport')) + monkeypatch.setattr(clients, 'Session', factory) + with pytest.raises(TypeError, match='resolved credential'): + clients.create_client(None, 'companylake') + factory.assert_not_called() + + +@pytest.mark.parametrize('url', [ + 'https://otherlake.dfs.core.windows.net/', + 'https://companylake.dfs.core.usgovcloudapi.net/', + 'https://companylake.dfs.core.windows.net.evil.example/', + 'https://companylake.blob.core.windows.net/', + 'http://companylake.dfs.core.windows.net/', + 'https://private-user:private-password@companylake.dfs.core.windows.net/', + 'https://companylake.dfs.core.windows.net:443/', + 'https://companylake.dfs.core.windows.net/filesystem/', + 'https://companylake.dfs.core.windows.net/#private-fragment', + 'https://[invalid/?sig=private-sas', + 'https://companylake.dfs.core.windows.net/\n', + None, +]) +def test_borrowed_endpoint_rejected_without_io_or_secret_echo(url): + borrowed = Mock(account_name='companylake', url=url) + with pytest.raises(ValueError, match='requested account') as error: + source_access(borrowed, 'companylake') + borrowed.get_file_system_client.assert_not_called() + borrowed.close.assert_not_called() + for value in ('private-user', 'private-password', 'private-fragment', + 'private-sas'): + assert value not in str(error.value) + + +def test_borrowed_client_account_name_must_match_endpoint(): + borrowed = Mock(account_name='otherlake', + url='https://companylake.dfs.core.windows.net/') + with pytest.raises(ValueError, match='requested account'): + source_access(borrowed, 'companylake') + borrowed.get_file_system_client.assert_not_called() + borrowed.close.assert_not_called() + + +def test_sas_client_url_is_supported_without_identity_resolution(): + from azure.storage.filedatalake import DataLakeServiceClient + + with DataLakeServiceClient( + 'https://companylake.dfs.core.windows.net/?sv=test&sig=synthetic', + credential=None) as borrowed: + access = source_access(borrowed, 'companylake') + assert access.filesystem.handler.fs.client is borrowed diff --git a/packages/kumo-adls/tests/test_sdk_transfer.py b/packages/kumo-adls/tests/test_sdk_transfer.py index d52988b0e..31e8e8132 100644 --- a/packages/kumo-adls/tests/test_sdk_transfer.py +++ b/packages/kumo-adls/tests/test_sdk_transfer.py @@ -9,6 +9,7 @@ import threading from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace +from typing import Any, Callable from urllib.parse import parse_qs, urlsplit import pytest @@ -46,7 +47,7 @@ def __init__(self): self.fragments = [] def register_uri(self, method, url, *, headers=None, status_code=200, - content=b''): + content: bytes | Callable[[Any, Any], bytes] = b''): self.routes.append((method, url, headers or {}, status_code, content)) def send(self, request, **kwargs): diff --git a/packages/kumo-adls/tests/test_session.py b/packages/kumo-adls/tests/test_session.py index 36a91888f..2961dbd61 100644 --- a/packages/kumo-adls/tests/test_session.py +++ b/packages/kumo-adls/tests/test_session.py @@ -6,6 +6,7 @@ import pytest from azure.identity import AzureAuthorityHosts +from kumo_adls import client as client_module from kumo_adls import session from kumo_core.sources.principals import Principal from kumoapi.data_source import ADLSIdentityConfig, AzureCloud @@ -53,7 +54,7 @@ def test_ambient_identity_session_reuse_and_cleanup(monkeypatch): credential_factory = Mock(return_value=credential) factory = Mock() monkeypatch.setattr(session, 'DefaultAzureCredential', credential_factory) - monkeypatch.setattr(session, 'DataLakeServiceClient', factory) + monkeypatch.setattr(client_module, 'DataLakeServiceClient', factory) first = session.session_for(None, 'companylake', 'dfs.core.windows.net') assert session.session_for(None, 'companylake', 'dfs.core.windows.net') is first @@ -72,7 +73,7 @@ def test_close_one_identity_preserves_other_sessions(monkeypatch): factory = Mock(side_effect=lambda **kwargs: Mock()) credentials = Mock(side_effect=lambda **kwargs: Mock()) monkeypatch.setattr(session, 'WorkloadIdentityCredential', credentials) - monkeypatch.setattr(session, 'DataLakeServiceClient', factory) + monkeypatch.setattr(client_module, 'DataLakeServiceClient', factory) config = ADLSIdentityConfig(tenant_id='tenant', client_id='client', cloud=AzureCloud.PUBLIC) first = Principal('first', config, None) @@ -108,7 +109,7 @@ def test_ambient_session_uses_common_storage_audience(monkeypatch, cloud, credential_factory = Mock(return_value=credential) client_factory = Mock() monkeypatch.setattr(session, 'DefaultAzureCredential', credential_factory) - monkeypatch.setattr(session, 'DataLakeServiceClient', client_factory) + monkeypatch.setattr(client_module, 'DataLakeServiceClient', client_factory) session.session_for(None, 'companylake', suffix) @@ -128,7 +129,7 @@ def test_cleanup_failure_does_not_leave_other_sessions_open( factory = Mock(side_effect=clients) credential_factory = Mock(side_effect=credentials) monkeypatch.setattr(session, 'DefaultAzureCredential', credential_factory) - monkeypatch.setattr(session, 'DataLakeServiceClient', factory) + monkeypatch.setattr(client_module, 'DataLakeServiceClient', factory) session.session_for(None, 'firstlake', 'dfs.core.windows.net') session.session_for(None, 'otherlake', 'dfs.core.windows.net') resource = clients[0] if failing_resource == 'client' else credentials[0] @@ -139,3 +140,16 @@ def test_cleanup_failure_does_not_leave_other_sessions_open( for item in clients + credentials: item.close.assert_called_once() assert not session._SESSIONS + + +def test_client_construction_failure_releases_engine_owned_credential( + monkeypatch): + credential = Mock() + monkeypatch.setattr(session, 'credential_for', lambda *args: + (credential, 'https://storage.azure.com')) + monkeypatch.setattr(session, 'create_client', + Mock(side_effect=RuntimeError('construction failed'))) + with pytest.raises(RuntimeError, match='construction failed'): + session.session_for(None, 'companylake', 'dfs.core.windows.net') + credential.close.assert_called_once() + assert not session._SESSIONS diff --git a/packages/kumo-adls/tests/test_source_access.py b/packages/kumo-adls/tests/test_source_access.py new file mode 100644 index 000000000..ce4e73b22 --- /dev/null +++ b/packages/kumo-adls/tests/test_source_access.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import gc +import os +import re +import subprocess +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from unittest.mock import Mock + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from azure.core.credentials import AccessToken +from fakes import FakeService +from kumo_adls import client as clients +from kumo_adls.backend import ADLSSourceBackend +from kumo_adls.reader import ADLSSourceReader, source_reader +from kumo_core.sources.types import ReadMode +from test_sdk_transfer import HEADERS, HTTPAdapter + + +@pytest.fixture +def no_auth_environment(monkeypatch): + for key in list(os.environ): + if key.startswith('AZURE_'): + monkeypatch.delenv(key) + + +def parquet(value): + stream = pa.BufferOutputStream() + pq.write_table(pa.table({'id': [value]}), stream) + return stream.getvalue().to_pybytes() + + +def reader_for(client): + return source_reader(client, client.account_name, filesystem='curated', + prefix='orders') + + +@pytest.mark.usefixtures('no_auth_environment') +def test_shared_reader_matches_engine_rows_versions_and_borrowed_lifetime( + monkeypatch, connector, tmp_path): + credential = SimpleNamespace(closed=False, scopes=[]) + + def get_token(*scopes, **kwargs): + credential.scopes.append(scopes) + return AccessToken('resolved-token', int(time.time()) + 3600) + + credential.get_token = get_token + payload = parquet(7) + adapter = HTTPAdapter() + properties = dict(HEADERS, **{'Content-Length': str(len(payload))}) + adapter.register_uri('HEAD', re.compile(r'https://companylake\.'), + headers=properties) + conditions = [] + + def download(request, context): + conditions.append(request.headers.get('If-Match')) + requested = request.headers.get('x-ms-range', 'bytes=0-') + start, end = requested.removeprefix('bytes=').split('-', 1) + first, last = int(start), int(end) if end else len(payload) - 1 + context.headers.update({ + 'Content-Length': + str(last - first + 1), + 'Content-Range': + f'bytes {first}-{last}/{len(payload)}' + }) + context.status_code = 206 + return payload[first:last + 1] + + adapter.register_uri( + 'GET', + re.compile(r'https://companylake\.[^/]+/curated/orders/events.parquet'), + headers=properties, content=download) + monkeypatch.setattr(clients, 'HTTPAdapter', lambda **kwargs: adapter) + borrowed = clients.create_client(credential, 'companylake') + close = Mock(wraps=borrowed.close) + monkeypatch.setattr(borrowed, 'close', close) + selected = Mock(return_value=borrowed) + monkeypatch.setattr('kumo_adls.backend.session_for', selected) + shared = reader_for(borrowed) + engine, name = ADLSSourceBackend()._reader(connector(), 'events.parquet') + selected.assert_called_once_with(None, 'companylake', + 'dfs.core.windows.net') + assert type(shared) is type(engine) is ADLSSourceReader + version = shared.data_version(name) + assert version == engine.data_version(name) + outputs = [ + reader.prepare(name, str(tmp_path / str(index)), columns=None, + mode=ReadMode.STREAMED, transform=False) + for index, reader in enumerate((shared, engine)) + ] + for output in outputs: + assert output.version == version + assert pq.read_table(list(output.paths)).to_pydict() == {'id': [7]} + assert conditions and all(condition == '"stable"' + for condition in conditions) + assert credential.scopes and all( + scopes == ('https://storage.azure.com/.default', ) + for scopes in credential.scopes) + del shared, engine + gc.collect() + close.assert_not_called() + assert not credential.closed + # The caller can still use and close its client after both readers disappear. + assert reader_for(borrowed).data_version(name) == version + borrowed.close() + close.assert_called_once() + assert not credential.closed + + +@pytest.mark.usefixtures('no_auth_environment') +def test_injected_clients_isolate_concurrent_identities_and_accounts(tmp_path): + cases = [] + for account in ('companylake', 'otherlake'): + for identity in (1, 2): + value = len(cases) + 1 + client = FakeService(account) + client.files['curated/orders/events.parquet'] = parquet(value) + cases.append((account, identity, client, value)) + + def read(case): + account, identity, client, value = case + reader = reader_for(client) + version = reader.data_version('events.parquet') + output = reader.prepare('events.parquet', + str(tmp_path / f'{account}-{identity}'), + columns=None, mode=ReadMode.STREAMED, + transform=False) + assert pq.read_table(list(output.paths)).to_pydict() == {'id': [value]} + assert output.version == version + assert not client.closed + return version + + with ThreadPoolExecutor(max_workers=4) as executor: + versions = list(executor.map(read, cases)) + assert len(set(versions)) == 4 + for _, _, client, _ in cases: + client.close() + + +@pytest.mark.usefixtures('no_auth_environment') +def test_engine_cleanup_does_not_close_injected_client(monkeypatch): + from kumo_adls import session + + borrowed = FakeService() + borrowed.files['curated/orders/events.parquet'] = parquet(1) + reader = reader_for(borrowed) + version = reader.data_version('events.parquet') + owned = FakeService('ownedlake') + credential = Mock() + monkeypatch.setattr(session, 'credential_for', lambda *args: + (credential, 'https://storage.azure.com')) + monkeypatch.setattr(session, 'create_client', lambda *args: owned) + session.close() + session.session_for(None, 'ownedlake', 'dfs.core.windows.net') + session.close() + assert owned.closed + credential.close.assert_called_once() + assert not borrowed.closed + assert reader.data_version('events.parquet') == version + + +@pytest.mark.parametrize('prefix', ['', 'orders/team']) +def test_shared_reader_preserves_root_spelling(connector, monkeypatch, prefix): + client = FakeService() + monkeypatch.setattr('kumo_adls.backend.session_for', lambda *args: client) + shared = source_reader(client, 'companylake', filesystem='curated', + prefix=prefix) + engine = ADLSSourceBackend()._discovery(connector(prefix=prefix)) + assert shared.root == engine.root + + +@pytest.mark.parametrize('prefix', ['../outside', 'orders/', 'orders#fragment']) +def test_reader_rejects_noncanonical_prefix_before_client_use(prefix): + client = Mock() + with pytest.raises(ValueError): + source_reader(client, 'companylake', filesystem='curated', + prefix=prefix) + client.get_file_system_client.assert_not_called() + + +def test_shared_reader_imports_and_construction_do_not_load_identity_adapters(): + script = ''' +import importlib.abc, os, sys +for key in list(os.environ): + if key.startswith('AZURE_'): del os.environ[key] +forbidden=('azure.identity', 'kumo_adls.session', 'kumo_adls.backend', + 'kumo_adls.native_output', 'kumo_core.sources.principals', 'kumo.engine') +class Guard(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if any(fullname==name or fullname.startswith(name+'.') for name in forbidden): + raise AssertionError('Unexpected identity import: '+fullname) +sys.meta_path.insert(0, Guard()) +from kumo_adls import create_client, source_reader +from kumo_adls.source_access import source_access +from kumo_core.sources.types import ReadMode +class Credential: + def get_token(self, *args, **kwargs): + raise AssertionError('Unexpected authentication during construction') +client=create_client(Credential(), 'companylake') +reader=source_reader(client, 'companylake', filesystem='curated', prefix='orders') +assert reader.root=='curated@companylake.dfs.core.windows.net/orders' +assert source_access(client, 'companylake').filesystem is not None +assert not any(name in sys.modules for name in forbidden) +client.close() +''' + completed = subprocess.run([sys.executable, '-c', script], + capture_output=True, text=True, + env=dict(os.environ, + PYTHONPATH=os.pathsep.join(sys.path))) + assert completed.returncode == 0, completed.stderr diff --git a/packages/kumo-adls/tests/test_transfer.py b/packages/kumo-adls/tests/test_transfer.py index 6e3843c69..5e7f324de 100644 --- a/packages/kumo-adls/tests/test_transfer.py +++ b/packages/kumo-adls/tests/test_transfer.py @@ -237,6 +237,7 @@ def unexpected_upload(*args, **kwargs): def test_owned_http_pool_is_reused_and_closed(monkeypatch): + from kumo_adls import client as client_module from kumo_adls import session session.close() @@ -244,10 +245,10 @@ def test_owned_http_pool_is_reused_and_closed(monkeypatch): monkeypatch.setattr(session, 'DefaultAzureCredential', Mock(return_value=credential)) transport_session = Mock() - monkeypatch.setattr(session, 'Session', + monkeypatch.setattr(client_module, 'Session', Mock(return_value=transport_session)) # Real SDK clients exercise ownership and child-client propagation. - original_transport = session.RequestsTransport + original_transport = client_module.RequestsTransport transports = [] def transport(**kwargs): @@ -255,7 +256,7 @@ def transport(**kwargs): transports.append(item) return item - monkeypatch.setattr(session, 'RequestsTransport', transport) + monkeypatch.setattr(client_module, 'RequestsTransport', transport) client = session.session_for(None, 'companylake', 'dfs.core.windows.net') assert session.session_for(None, 'companylake', 'dfs.core.windows.net') is client @@ -341,7 +342,8 @@ def test_put_file_conditions_bind_upload_and_cleanup_to_creation( source.write_bytes(b'prepared bytes') client = Mock(spec=DataLakeFileClient) client.create_file.return_value = {'etag': '"created"'} - service = Mock() + service = Mock(account_name='companylake', + url='https://companylake.dfs.core.windows.net/') service.get_file_system_client.return_value.get_file_client.return_value = client def upload(stream, **kwargs): diff --git a/packages/kumo-engine/test/test_adls_connector.py b/packages/kumo-engine/test/test_adls_connector.py index 529bf9c1d..c3d0c6603 100644 --- a/packages/kumo-engine/test/test_adls_connector.py +++ b/packages/kumo-engine/test/test_adls_connector.py @@ -39,6 +39,7 @@ def identity(path): with filesystem.open_input_file(path) as stream: return (hashlib.sha256(stream.read()).hexdigest(), ) + monkeypatch.setattr('kumo_adls.backend.session_for', lambda *args: object()) monkeypatch.setattr( 'kumo_adls.backend.source_access', lambda *args: ObjectStoreAccess(filesystem, object_identity=identity)) diff --git a/scripts/benchmarks/object_store/compare_live.py b/scripts/benchmarks/object_store/compare_live.py index 72e6065b0..831f1ffed 100644 --- a/scripts/benchmarks/object_store/compare_live.py +++ b/scripts/benchmarks/object_store/compare_live.py @@ -297,7 +297,8 @@ def make_reader(provider, manifest, args): raise ValueError( 'ADLS transfer ceiling differs from the file-worker ceiling') suffix = manifest.get('endpoint_suffix', 'dfs.core.windows.net') - access = source_access(None, manifest['adls_account'], suffix) + client = session.session_for(None, manifest['adls_account'], suffix) + access = source_access(client, manifest['adls_account'], suffix) authority = f'{manifest["filesystem"]}@{manifest["adls_account"]}.{suffix}' reader = ADLSSourceReader(access.filesystem, root=authority + '/' + manifest['adls_prefix'], From d3ffe58fef638267503be2656c8f1d0ef12d32af Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Mon, 26 Oct 2026 16:51:19 +0000 Subject: [PATCH 13/15] Restore exact-root ADLS model bundle destinations --- packages/kumo-adls/README.md | 7 +- packages/kumo-adls/docs/implementation.md | 7 ++ packages/kumo-adls/src/kumo_adls/backend.py | 4 +- packages/kumo-adls/src/kumo_adls/sink.py | 13 ++++ packages/kumo-adls/tests/test_coordinates.py | 73 +++++++++++++++++++ .../src/kumo_core/object_store_sink.py | 9 ++- 6 files changed, 109 insertions(+), 4 deletions(-) diff --git a/packages/kumo-adls/README.md b/packages/kumo-adls/README.md index 20cba35b1..ed8ccac10 100644 --- a/packages/kumo-adls/README.md +++ b/packages/kumo-adls/README.md @@ -37,7 +37,12 @@ A rooted connector discovers CSV/Parquet tables and writes predictions, embeddings, and model bundles beneath its filesystem/prefix. A rootless connector reads exact canonical `abfss://` URIs; discovery and output are unavailable. SQL, warehouse-style `ReadMode.STAGED`, and append output are -unsupported. Prepared exports and generated output use native Arrow streams, +unsupported. Model bundles can be published and reloaded at the connector's +full root URI, with or without a trailing slash. Table sources and prediction +exports require a nonempty child path. Use a dedicated bundle directory: +model reload reads all files below its selected directory. + +Prepared exports and generated output use native Arrow streams, matching S3/GCS; consumers must wait for export completion. ## Shared SDK reads diff --git a/packages/kumo-adls/docs/implementation.md b/packages/kumo-adls/docs/implementation.md index 648b650e1..ade852cef 100644 --- a/packages/kumo-adls/docs/implementation.md +++ b/packages/kumo-adls/docs/implementation.md @@ -93,6 +93,13 @@ Native output is bound lazily on first write, leaving SDK metadata, source reads, model downloads, and their account-qualified paths unchanged. Only output streams drop the virtual account qualifier to use Arrow's container/key path. Existing root containment checks run before filesystem access. +Model publish/reload use a directory-specific destination hook. A full URI +matching the rooted connector's exact filesystem/prefix is accepted, including +a filesystem root and either trailing-slash spelling. Table and prediction +coordinates still require a nonempty child path; the model exception does not +enable filename completion outside a prefix. Other model coordinates retain +the same account/root fencing and relative-path validation. Model reload reads +all files under that directory, so use a dedicated bundle location. A native filesystem is reused within one sink access, with thread-safe first construction; it is not cached across accesses or identities. The provider requires Arrow 24, the version whose native credential/URI APIs are validated. diff --git a/packages/kumo-adls/src/kumo_adls/backend.py b/packages/kumo-adls/src/kumo_adls/backend.py index d8b423b7f..0793eb56a 100644 --- a/packages/kumo-adls/src/kumo_adls/backend.py +++ b/packages/kumo-adls/src/kumo_adls/backend.py @@ -14,6 +14,7 @@ from kumo_core.sources.principals import IdentityRef from kumo_core.sources.types import ReadMode from kumoapi.adls import ( + normalize_adls_root_uri, normalize_adls_source_relative_path, normalize_adls_source_uri, ) @@ -36,8 +37,9 @@ def canonical_uri(path: str) -> str: + """Serialize an account-qualified file or directory path.""" authority, _, key = path.partition('/') - return normalize_adls_source_uri( + return normalize_adls_root_uri( f'abfss://{authority}/{quote(key, safe="/-._~!$&()*+,;=:@")}') diff --git a/packages/kumo-adls/src/kumo_adls/sink.py b/packages/kumo-adls/src/kumo_adls/sink.py index 06e748fe6..dc8977ca3 100644 --- a/packages/kumo-adls/src/kumo_adls/sink.py +++ b/packages/kumo-adls/src/kumo_adls/sink.py @@ -3,9 +3,11 @@ # SPDX-License-Identifier: Apache-2.0 from functools import partial +from urllib.parse import unquote, urlsplit from kumo_core.object_store_sink import ObjectStoreSinkBackend from kumo_core.sources.cpu import COPY_WORKER_BYTES, memory_share +from kumoapi.adls import normalize_adls_root_uri from kumoapi.data_source import DataSourceType from kumoapi.jobs import WriteMode from pyarrow.fs import FSSpecHandler, PyFileSystem @@ -22,6 +24,17 @@ class ADLSSinkBackend(ADLSSessions, ObjectStoreSinkBackend): source_type = DataSourceType.ADLS + def _model_destination(self, connector, coordinate): + """Model bundles may name the root directory; tables must name a child.""" + root = self._root_for(connector) + if root is not None and '://' in coordinate: + parsed = urlsplit(normalize_adls_root_uri(coordinate)) + key = unquote(parsed.path).strip('/') + if parsed.netloc == root.bucket and key == root.prefix: + return self._filesystem(connector, root.bucket, + root.region), root.path + return self._destination(connector, coordinate) + def preflight_write_files(self, connector, table_name, write_mode=WriteMode.OVERWRITE): """Reject unsupported output before prediction work or Azure access.""" diff --git a/packages/kumo-adls/tests/test_coordinates.py b/packages/kumo-adls/tests/test_coordinates.py index a5ae30589..1245f3d5f 100644 --- a/packages/kumo-adls/tests/test_coordinates.py +++ b/packages/kumo-adls/tests/test_coordinates.py @@ -94,3 +94,76 @@ def test_rootless_full_uri_keeps_exact_source_policy(service, connector): assert isinstance(config.source, ADLSURISource) assert config.source.uri == uri assert '10.0' in source.get_sample(target, uri, 1).sample_rows + + +@pytest.mark.parametrize('prefix', ['', 'orders', 'orders/team']) +@pytest.mark.parametrize('trailing_slash', [False, True]) +def test_full_root_model_directory_roundtrip(service, connector, tmp_path, + prefix, trailing_slash): + target = connector(prefix=prefix) + root_uri = URI.removesuffix('orders/').rstrip('/') + if prefix: + root_uri += '/' + prefix + if trailing_slash: + root_uri += '/' + bundle = tmp_path / 'bundle' + (bundle / 'nested').mkdir(parents=True) + (bundle / 'model.bin').write_bytes(b'model') + (bundle / 'nested/weights.bin').write_bytes(b'weights') + sink = ADLSSinkBackend() + original_files = dict(service.files) + + published = sink.publish_model(target, str(bundle), path=root_uri) + assert published.uri.rstrip('/') == root_uri.rstrip('/') + destination = tmp_path / 'fetched' + restored = sink.fetch_model(target, str(destination), path=published.uri) + + assert restored.uri == published.uri + assert all(service.files[name] == content + for name, content in original_files.items()) + assert (destination / 'model.bin').read_bytes() == b'model' + assert (destination / 'nested/weights.bin').read_bytes() == b'weights' + # Directory-only model resolution must not enable filename completion at root. + with pytest.raises(ValueError): + sink.preflight_write_files(target, root_uri) + with pytest.raises(ValueError): + sink.write_table(target, root_uri, pd.DataFrame({'id': [1]})) + with pytest.raises(ValueError): + ADLSSourceBackend().get_table_config(target, root_uri) + + +@pytest.mark.parametrize('coordinate', [ + URI.replace('companylake', 'otherlake'), + URI.replace('curated@', 'otherfs@'), + URI.replace('orders/', 'orders-sibling/'), + URI + 'scores%23review', + URI + '?sig=private-value', +]) +def test_model_directory_rejection_precedes_storage_and_local_replacement( + connector, monkeypatch, tmp_path, coordinate): + from unittest.mock import Mock + + access = Mock(side_effect=AssertionError('must not access storage')) + monkeypatch.setattr('kumo_adls.backend.session_for', access) + target = connector() + destination = tmp_path / 'existing' + destination.mkdir() + retained = destination / 'retained.bin' + retained.write_bytes(b'keep') + sink = ADLSSinkBackend() + + with pytest.raises(ValueError): + sink.publish_model(target, str(destination), path=coordinate) + with pytest.raises(ValueError): + sink.fetch_model(target, str(destination), path=coordinate) + + access.assert_not_called() + assert retained.read_bytes() == b'keep' + + +def test_rootless_model_directory_remains_unsupported(connector, tmp_path): + target = connector(any_uri=True) + sink = ADLSSinkBackend() + for operation in (sink.publish_model, sink.fetch_model): + with pytest.raises(NotImplementedError, match='source-only'): + operation(target, str(tmp_path / 'bundle'), path=URI) diff --git a/packages/kumo-core/src/kumo_core/object_store_sink.py b/packages/kumo-core/src/kumo_core/object_store_sink.py index efab01dba..3a510adab 100644 --- a/packages/kumo-core/src/kumo_core/object_store_sink.py +++ b/packages/kumo-core/src/kumo_core/object_store_sink.py @@ -32,6 +32,11 @@ def _destination(self, connector: SourceConnector, """Return the authenticated filesystem and validated bucket/key path.""" raise NotImplementedError + def _model_destination(self, connector: SourceConnector, + coordinate: str) -> Tuple[FileSystem, str]: + """Resolve a model directory, which may have provider-specific roots.""" + return self._destination(connector, coordinate) + def _uri(self, path: str) -> str: return f'{self.scheme}://{path}' @@ -106,7 +111,7 @@ def publish_model( f"Publishing a model to {self.provider} needs a destination " f"path. Set output_path on the ModelExportConfig.") - filesystem, destination = self._destination(connector, path) + filesystem, destination = self._model_destination(connector, path) destination = destination.rstrip('/') root = Path(bundle_dir) for source in sorted(item for item in root.rglob('*') @@ -159,7 +164,7 @@ def fetch_model( f"was published to. Reuse the output_path from the export " f"config.") - filesystem, source = self._destination(connector, path) + filesystem, source = self._model_destination(connector, path) source = source.rstrip('/') destination = Path(local_dir) if destination.exists(): From 5850fd60e479430a22464feba7fd796d13262115 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Mon, 26 Oct 2026 21:03:19 +0000 Subject: [PATCH 14/15] Raise ADLS dependency minimums past known vulnerabilities Require azure-core >=1.38 and requests >=2.33 to exclude their known vulnerable releases from ADLS installs. Regenerate the workspace lockfile without changing resolved package versions. Validation: 320 offline ADLS tests at the proposed minimums; pip-audit passes for both the minimums with locked transitives and the locked ADLS dependency tree; uv lock --check. --- packages/kumo-adls/pyproject.toml | 4 ++-- uv.lock | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/packages/kumo-adls/pyproject.toml b/packages/kumo-adls/pyproject.toml index 0361bbf3d..5976721dc 100644 --- a/packages/kumo-adls/pyproject.toml +++ b/packages/kumo-adls/pyproject.toml @@ -19,8 +19,8 @@ dependencies = [ "kumo-api>=3.0.0", "pyarrow>=24,<25", "fsspec", - "requests>=2.31,<3", - "azure-core>=1.31,<2", + "requests>=2.33,<3", + "azure-core>=1.38,<2", "azure-identity>=1.25,<2", "azure-storage-file-datalake>=12.22,<13", ] diff --git a/uv.lock b/uv.lock index c5bf6defa..43ad4dace 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.11" resolution-markers = [ "python_full_version >= '3.14' and sys_platform != 'darwin' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests'", @@ -3133,7 +3133,7 @@ test = [ [package.metadata] requires-dist = [ - { name = "azure-core", specifier = ">=1.31,<2" }, + { name = "azure-core", specifier = ">=1.38,<2" }, { name = "azure-identity", specifier = ">=1.25,<2" }, { name = "azure-storage-file-datalake", specifier = ">=12.22,<13" }, { name = "fsspec" }, @@ -3141,7 +3141,7 @@ requires-dist = [ { name = "kumo-core", editable = "packages/kumo-core" }, { name = "pyarrow", specifier = ">=24,<25" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=7.4" }, - { name = "requests", specifier = ">=2.31,<3" }, + { name = "requests", specifier = ">=2.33,<3" }, ] provides-extras = ["test"] @@ -3357,8 +3357,8 @@ provides-extras = ["test"] name = "kumo-engine" source = { editable = "packages/kumo-engine" } dependencies = [ - { name = "faiss-cpu", marker = "sys_platform == 'darwin' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, - { name = "faiss-gpu", marker = "sys_platform == 'linux' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "faiss-cpu", marker = "sys_platform == 'darwin' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "faiss-gpu", marker = "sys_platform == 'linux' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, { name = "filelock" }, { name = "kumo-api" }, { name = "kumo-core" }, From ca8fc880457099c52e1da7494069049f81d0b239 Mon Sep 17 00:00:00 2001 From: Elvin Gu Date: Mon, 26 Oct 2026 21:03:58 +0000 Subject: [PATCH 15/15] Revert "Raise ADLS dependency minimums past known vulnerabilities" This reverts commit b92d9b65f3018c1fc042c5cbbcc42c819958d19e. --- packages/kumo-adls/pyproject.toml | 4 ++-- uv.lock | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/packages/kumo-adls/pyproject.toml b/packages/kumo-adls/pyproject.toml index 5976721dc..0361bbf3d 100644 --- a/packages/kumo-adls/pyproject.toml +++ b/packages/kumo-adls/pyproject.toml @@ -19,8 +19,8 @@ dependencies = [ "kumo-api>=3.0.0", "pyarrow>=24,<25", "fsspec", - "requests>=2.33,<3", - "azure-core>=1.38,<2", + "requests>=2.31,<3", + "azure-core>=1.31,<2", "azure-identity>=1.25,<2", "azure-storage-file-datalake>=12.22,<13", ] diff --git a/uv.lock b/uv.lock index 43ad4dace..c5bf6defa 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 5 +revision = 3 requires-python = ">=3.11" resolution-markers = [ "python_full_version >= '3.14' and sys_platform != 'darwin' and extra != 'extra-4-kumo-cpu' and extra != 'extra-4-kumo-cu130' and extra != 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests'", @@ -3133,7 +3133,7 @@ test = [ [package.metadata] requires-dist = [ - { name = "azure-core", specifier = ">=1.38,<2" }, + { name = "azure-core", specifier = ">=1.31,<2" }, { name = "azure-identity", specifier = ">=1.25,<2" }, { name = "azure-storage-file-datalake", specifier = ">=12.22,<13" }, { name = "fsspec" }, @@ -3141,7 +3141,7 @@ requires-dist = [ { name = "kumo-core", editable = "packages/kumo-core" }, { name = "pyarrow", specifier = ">=24,<25" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=7.4" }, - { name = "requests", specifier = ">=2.33,<3" }, + { name = "requests", specifier = ">=2.31,<3" }, ] provides-extras = ["test"] @@ -3357,8 +3357,8 @@ provides-extras = ["test"] name = "kumo-engine" source = { editable = "packages/kumo-engine" } dependencies = [ - { name = "faiss-cpu", marker = "sys_platform == 'darwin' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, - { name = "faiss-gpu", marker = "sys_platform == 'linux' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'extra-4-kumo-cu132') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu132' and extra == 'group-4-kumo-public-tests')" }, + { name = "faiss-cpu", marker = "sys_platform == 'darwin' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, + { name = "faiss-gpu", marker = "sys_platform == 'linux' or (extra == 'extra-4-kumo-cpu' and extra == 'extra-4-kumo-cu130') or (extra == 'extra-4-kumo-cpu' and extra == 'group-4-kumo-public-tests') or (extra == 'extra-4-kumo-cu130' and extra == 'group-4-kumo-public-tests')" }, { name = "filelock" }, { name = "kumo-api" }, { name = "kumo-core" },