Skip to content

Fix merge_index dropping nested directories from shard paths#983

Open
discobot wants to merge 1 commit into
mosaicml:mainfrom
discobot:fix/959-merge-index-nested-paths
Open

Fix merge_index dropping nested directories from shard paths#983
discobot wants to merge 1 commit into
mosaicml:mainfrom
discobot:fix/959-merge-index-nested-paths

Conversation

@discobot

Copy link
Copy Markdown

Description of changes:

merge_index drops intermediate directories from shard paths in the merged index.json when
partitions are nested more than one level under the root. Root cause: _merge_index_from_list
(streaming/base/util.py) rebuilds each merged basename as
os.path.join(os.path.basename(p.parent), basename) — exactly one directory level of the
partition index's path — so anything between the merge root and the partition is silently
dropped.

This PR resolves each partition's directory relative to the merge root instead, via a new
_get_partition_dirname helper that compares the partition index url against out (local part
against local root, remote against remote, same bucket only) and takes os.path.relpath on the
url paths. Partitions that don't reside under out (arbitrary index_file_urls) fall back to
the previous single-level behavior, so existing one-level layouts are unaffected.

Testing:

  • New test_merge_index_from_root_local_nested writes 4 MDS partitions at
    out/group{1,2}/subdir{1,2}, merges, and asserts the merged basenames keep the full relative
    paths and that a StreamingDataset on the merged root iterates all samples. It fails on main
    ("expected 4 shard files but got 2") and passes with this change.
  • New test_get_partition_dirname covers local, remote, tuple, root-level, cross-bucket, and
    outside-root fallback cases.
  • All pre-existing tests/test_util.py tests (incl. the spark-based
    test_merge_index_from_list_local / test_merge_index_from_root_local matrices) and
    tests/base/converters/test_dataframe_to_mds.py pass locally.

Issue #, if available:

Merge Checklist:

Put an x without space in the boxes that apply. If you are unsure about any checklist, please don't hesitate to ask. We are here to help! This is simply a reminder of what we are going to look for before merging your pull request.

General

  • I have read the contributor guidelines
  • This is a documentation change or typo fix. If so, skip the rest of this checklist.
  • I certify that the changes I am introducing will be backward compatible, and I have discussed concerns about this, if any, with the MosaicML team.
  • I have updated any necessary documentation, including README and API docs (if appropriate).

Tests

  • I ran pre-commit on my change. (check out the pre-commit section of prerequisites)
  • I have added tests that prove my fix is effective or that my feature works (if appropriate).
  • I ran the tests locally to make sure it pass. (check out testing)
  • I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes.

merge_index prefixed each shard basename with only the immediate parent
directory of its partition index, so partitions nested more than one level
under the merge root ended up with shard paths that do not exist, breaking
downstream streaming with 404s (issue mosaicml#959). Resolve each partition's
directory relative to the merge root instead, falling back to the previous
single-level behavior when a partition does not reside under the root.
Add a spark-free nested-layout regression test and unit tests for the new
path resolution helper.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

merge_index drops full directory path in merged index.json, causing broken shard references and 404s

1 participant