[bugfix] restore ODPS sessions by partition index on resume

[tiankongdeguiji]

Summary

OdpsReader._restore_sessions rebuilt the per-input session list on resume as restored_sess_reqs + current_sessions[n_restored:], which assumes the sessions found in the checkpoint are exactly the leading partitions in order. That does not hold:

• A partition that produced no rows before the checkpoint (e.g. an empty one under is_orderby_partition) never gets a checkpoint key, so the consumed set skips it.
• Session order in dataloader_state.json follows the rank-major merge in save_dataloader_state, not partition order.

With partitions [P0, P1(empty), P2] consumed into P2, restore produced [s0, s2, s2_new]: P2's remaining rows were read via the old session and then P2 was read again in full via the fresh session, with no error.

The get_read_session response carries no partition spec, so the reader cannot map a session back to its partition on its own. The checkpoint source key now carries the session position, {input_path}#{sess_idx}#{session_id}:{start}, and restore puts each validated session back at that index. Unconsumed partitions keep their freshly created session at their own position.

Error handling on resume:

• A key of one of the job's own input paths without a session index comes from an older version and raises a clear RuntimeError asking to restart from scratch (those sessions expire within a day of the writer dying anyway).
• Keys of unknown input paths keep the existing warn-and-skip, so a fine_tune_checkpoint state referencing other tables, or parquet paths containing #, still do not abort startup.
• An index beyond the current session list for the same input path (e.g. is_orderby_partition toggled off) raises.

Alternative considered: putting the partition spec in the key instead of the index. The index is simpler to parse and is what positional restore needs.

Test Plan

• New test_odps_dataset_checkpoint_resume_empty_partition (ODPS-credentialed): table with dt=20240319 (10000 rows), an empty dt=20240318, and dt=20240320 (10000 rows) under is_orderby_partition=True. Consumes into the third partition, restores into a fresh dataset, asserts sessions 0 and 2 are the checkpoint sessions and session 1 is the dataset's own fresh one, then reads to exhaustion and asserts consumed plus remaining rows equal 20000. This fails on master with the duplicate read.
• New offline OdpsRestoreSessionsTest with a stub client: positional placement with reserved keys mixed in, unknown-path warn-and-skip, old-format key raise, out-of-range index raise. Runs in every CI lane.
• test_odps_dataset_checkpoint_resume_orderby_partition adjusted for the new key layout; test_odps_dataset_checkpoint_metadata now pins the three-segment key.
• pre-commit run --files ... and pyrefly check clean. ODPS tests run on the CPU CI for this PR.

🤖 Generated with Claude Code

https://claude.ai/code/session_01MgZZz8Baa365GMyVuRFvZb

[github-actions]

Low-med: this raise fires for any key shaped ...#non-digit:..., including keys this reader never produced, which pre-diff code warned-and-skipped. Two concrete paths:

• main.py:839-840 also loads dataloader_state.json from an external fine_tune_checkpoint source. An older-tzrec source checkpoint whose ODPS keys reference unrelated tables (fine-tuning onto different inputs) previously fell into the unknown-input_path warning below and was skipped; now it aborts startup on every rank with a misleading "restart from scratch".
• Parquet source ids are {file_path}:{start} (parquet_dataset.py:276) and # is legal in POSIX/OSS paths, so e.g. /data/train#v2/part-0.parquet:1024 in a mixed-history state file hard-fails an ODPS job.

Suggestion: scope the raise to keys that claim one of this job's own inputs — for an old-format key, input_and_idx is the old input path, so raise only when it (or the parsed input_path) is in self._input_to_sess, otherwise fall through to the existing unknown-path warning. That keeps the intended hard failure for genuinely old checkpoints of the resumed table while restoring the tolerant skip for foreign keys.

[github-actions]

Nit: the paragraph was rewritten but the Raises sentence still only covers the expired/invalid case. The function now also raises for a key without a session index (line 571) and for an index beyond the current session list (line 591), and warns-and-skips unknown input paths (line 580). Since this docstring was touched anyway, worth completing, e.g.: "Raises RuntimeError if a key carries no session index (written by an older version), if an index exceeds the current number of partition sessions, or if any session is expired/invalid."

[github-actions]

The two new RuntimeError branches — old-format key (odps_dataset.py:571) and out-of-range index (:591) — are the user-facing guardrails of this fix but have no committed coverage; every test that exercises _restore_sessions is gated behind ODPS credentials, and even those never hit the error paths.

Both are cheap to test offline in the style of OdpsStorageErrorLogTest below. The old-format raise fires during key parsing before any attribute access, so it works with a bare stub:

with self.assertRaisesRegex(RuntimeError, "older TorchEasyRec"):
    odps_dataset.OdpsReader._restore_sessions(
        SimpleNamespace(), {"odps://p/tables/t/dt=x#sessid:10": 0}
    )

The bounds check only needs _input_to_sess / _table_to_cli stubs (a SimpleNamespace with a 1-element session list and a key claiming index 5). Worth adding so the guards don't silently regress in non-credentialed CI lanes.

[github-actions]

The comment now claims the three-field format, but none of the assertions below actually distinguish it — they all pass identically under the old {input_path}#{session_id}:{start} layout, so a regression that drops sess_idx from to_batches would only be caught by the credential-gated resume tests. Cheap pin, since input paths contain no #: segments = prefix.split("#") → assertEqual(len(segments), 3) and assertTrue(segments[1].isdigit()).

[github-actions]

Multi-area review summary (code quality, performance, test coverage, documentation, security — all five completed)

The core fix is correct. Independently traced and confirmed:

• Positional restore handles empty leading/middle/trailing partitions; multiple slice keys sharing one session merge harmlessly (ids are rank-broadcast, so always identical); repeated load_state_dict is idempotent.
• Rank-uniform failure: every rank parses the same merged dataloader_state.json, and both new RuntimeErrors are deterministic functions of (state, config) — they fire on all ranks or none, and _restore_sessions introduces no collectives, so no hung-barrier risk.
• Producer/consumer agreement: the to_batches prefixes match exactly what the rpartition(":") grouping in calc_slice_intervals and the exact-prefix match in calc_remaining_intervals expect; reserved keys (__epochs_completed__, __data_ts_watermark__) carry no : and are skipped before parsing.
• The new empty-partition test is deterministic and discriminating: traced the worker split (w0 gets P0[0,5120)+P2[0,5120), w1 gets P0[5120,10000)+P2[5120,10000)), the #2# break lands at iteration ~10 well inside the 20-iteration cap, and num_consumed + num_remaining == 20000 holds exactly for any delivery interleaving (prefetched-but-undelivered batches are re-read on resume). It fails on master with the duplicate read.

Inline comments posted (none blocking):

1. odps_dataset.py:570 — the old-format raise also catches foreign keys the pre-diff code warned-and-skipped (fine-tune checkpoints referencing unrelated tables; parquet ids with # in the path). Suggest scoping the raise to keys claiming this job own input paths.
2. odps_dataset.py:554 — docstring Raises clause not updated for the two new error cases.
3. odps_dataset_test.py:449 — both new RuntimeError branches lack coverage; they are offline-testable in the existing OdpsStorageErrorLogTest stub style.
4. odps_dataset_test.py:258 — metadata test assertions pass under the old key format too; suggest pinning the sess_idx field.

Non-blocking observations (pre-existing, not introduced by this PR):

• Restored sessions are never registered with _refresh_sessions_daemon — it keeps refreshing the replaced sessions instead. A resumed run reading restored sessions longer than their residual lifetime (~24h from the last refresh by the dead job) will hit session expiry mid-training with no self-healing path. Same failure domain as this PR; likely worth a follow-up.
• PR-body nit: "a changed partition list already raises" — actually a changed partition list changes input_path, so those keys hit the unknown-path warning and are skipped (silent full re-read); the raise covers only same-path-with-fewer-sessions (e.g. is_orderby_partition toggled off).

[tiankongdeguiji]

Addressed the four inline comments in 13fa120:

• Old-format keys now raise only when they belong to one of this job's input paths; keys of unknown paths (fine-tune checkpoints on other tables, parquet paths containing #) keep the warn-and-skip.
• Docstring Raises clause covers all three error cases and the unknown-path skip.
• Added offline OdpsRestoreSessionsTest (stub client): positional placement, unknown-path skip, old-format raise, out-of-range raise.
• test_odps_dataset_checkpoint_metadata pins the {input_path}#{sess_idx}#{session_id} layout.

The restored sessions not being registered with _refresh_sessions_daemon is pre-existing and left for a follow-up.