Skip to content

Fix Arrow dataset loading with Windows workers - #529

Open
kiwoongyoon wants to merge 2 commits into
amazon-science:mainfrom
kiwoongyoon:gitcontribute/20260823-040816-028611-amazon-science-chronos-forecasting-issue-285
Open

Fix Arrow dataset loading with Windows workers#529
kiwoongyoon wants to merge 2 commits into
amazon-science:mainfrom
kiwoongyoon:gitcontribute/20260823-040816-028611-amazon-science-chronos-forecasting-issue-285

Conversation

@kiwoongyoon

Copy link
Copy Markdown

Title: Fix Arrow dataset loading with Windows workers

Summary

  • Defer GluonTS FileDataset construction until a data-loader worker starts iterating.
  • Keep the existing filtering, sampling probabilities, and worker-sharding behavior unchanged while avoiding pickling PyArrow's open reader on Windows.

The supplied issue snapshot is unassigned and contains no comments or related-PR indication. Please recheck open and recently merged PRs before opening this draft.

Test evidence

  • python3 -c 'import ast,pickle,sys,types; from pathlib import Path; tree=ast.parse(Path("scripts/training/train.py").read_text()); node=next(n for n in tree.body if isinstance(n,ast.ClassDef) and n.name=="LazyFileDataset"); mod=types.ModuleType("lazy_dataset_test"); mod.Path=Path; opened=[]; mod.FileDataset=lambda **kwargs: opened.append(kwargs) or ["entry"]; sys.modules[mod.__name__]=mod; exec(compile(ast.Module(body=[node],type_ignores=[]),"scripts/training/train.py","exec"),mod.__dict__); dataset=mod.LazyFileDataset(Path("sample.arrow")); assert opened==[]; restored=pickle.loads(pickle.dumps(dataset)); assert opened==[]; assert list(restored)==["entry"]; assert opened==[{"path":Path("sample.arrow"),"freq":"h"}]'
  • python3 -c 'from pathlib import Path; compile(Path("scripts/training/train.py").read_text(), "scripts/training/train.py", "exec")'
  • git diff --check

All commands exited 0.

Risks or notes for maintainers

  • The focused assertion verifies that the production wrapper stays unopened through pickling and opens the expected dataset only on iteration.
  • A native Windows training run and the full pytest suite were not available in the Linux Builder environment, which lacks pytest, torch, and gluonts.
  • The change does not disable workers or reject the reported configuration.

Fixes #285

@kiwoongyoon
kiwoongyoon marked this pull request as ready for review August 23, 2026 10:07
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.

trying to finetune I get EOFError: Ran out of input

1 participant