Skip to content

Allow ignored-field dataloaders to use spawned workers - #2760

Open
林SO (Linxiushen) wants to merge 1 commit into
microsoft:mainfrom
Linxiushen:codex/pickle-ignored-field-collator
Open

林SO (Linxiushen) wants to merge 1 commit into
microsoft:mainfrom
Linxiushen:codex/pickle-ignored-field-collator

Conversation

@Linxiushen

Copy link
Copy Markdown

Describe your changes

dataloader_with_ignored_batch_fields creates a function-local collator that cannot be pickled by spawned DataLoader workers. With num_workers=1 and multiprocessing_context="spawn", iteration fails before the first batch, including through the configured DataConfig/DataContainer path on Windows.

Move the collator to module scope and bind the ignored fields with functools.partial, preserving scalar-field conversion, tensors, labels and forwarded options.

Validation on Windows CPU:

  • Identical focused suite on pristine base: 5 failed, 7 passed; all five failures are the unpicklable collator. Patched suite: 12 passed.
  • Nine new cases include pickle round trips, string/list fields, serial loaders and three real spawned-worker cases checking six batches and child PIDs; three existing DataConfig controls also pass.
  • Configured lintrunner -a, final lintrunner including Pylint, and git diff --check passed.

Regression tests use a 120-second worker startup allowance after the initial 20-second allowance timed out on cold imports; those failures are retained. Product timeout defaults are unchanged. Focused pytest uses --confcutdir=test/data_container to exclude unrelated model fixtures. The full model/provider suite and performance/GPU validation were not run.

Checklist before requesting a review

  • Add unit tests for this change.
  • Make sure all tests can pass. Focused suite passes; full suite not run locally.
  • Update documents if necessary. Existing API and configuration are preserved.
  • Lint and apply fixes to your code by running lintrunner -a
  • Is this a user-facing change? If yes, give a description of this change to be included in the release notes.

Release note: Fix ignored-field DataLoader collation failing with spawned worker processes, including the configured DataContainer path on Windows.

(Optional) Issue link

Independently reproduced; no linked issue.

AI assistance: implementation, tests and this description were prepared with Codex and reviewed against the exact diff and native test results.

Copilot AI balanced review requested due to automatic review settings October 27, 2026 15:26
@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Approval recommended

The narrow refactor preserves collation behavior and includes targeted regression coverage, with no blocking issues identified.

No open issues 🎉

What changed in this PR

Fixes ignored-field DataLoader collation with spawned workers, including the configured DataContainer path on Windows, while preserving existing batching behavior.

Changes:

  • Moves the collator to module scope and binds ignored fields with functools.partial so it can be pickled.
  • Adds regression coverage for serialization, serial and spawned workers, configured loaders, and existing batching behavior.
File Description
test/​data_container/​test_ignored_batch_fields.py Adds pickle, spawned-worker, configuration, and batching regression tests.
olive/​data/​component/​dataloader.py Replaces the local collator with a picklable module-level helper.

Review effort: Balanced

This branch has not been deployed

No deployments
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.

2 participants