From 84467d29960f1616f0e99fc47dd48456d8bc7084 Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Fri, 16 Oct 2026 21:37:12 +0000 Subject: [PATCH 1/3] launcher: re-export partial or stale DFlash checkpoints after training dflash_online_training.sh skipped any exported-checkpoint-* directory that existed, so an export interrupted mid-write was never redone, and a resumed run kept the previous exported-checkpoint-final. An export now counts as done only once its config.json exists (exporters write weights first); otherwise the directory is removed and exported again. exported-checkpoint-final is also redone when modelopt_state.pth is newer than its config.json. Signed-off-by: Zhehao Hu --- .../common/specdec/dflash_online_training.sh | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/tools/launcher/common/specdec/dflash_online_training.sh b/tools/launcher/common/specdec/dflash_online_training.sh index e70800e069a..5b8ed3fb236 100644 --- a/tools/launcher/common/specdec/dflash_online_training.sh +++ b/tools/launcher/common/specdec/dflash_online_training.sh @@ -156,13 +156,15 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then if [ -n "$OUTPUT_DIR" ]; then # Export all checkpoints that don't already have an exported version EXPORTED=0 - for CKPT in $(ls -d ${OUTPUT_DIR}/checkpoint-* 2>/dev/null | sort -t- -k2 -n); do + for CKPT in "${OUTPUT_DIR}"/checkpoint-*; do + [ -d "$CKPT" ] || continue STEP=$(basename "$CKPT" | sed 's/checkpoint-//') EXPORT_DIR="${OUTPUT_DIR}/exported-checkpoint-${STEP}" - if [ -d "$EXPORT_DIR" ]; then + if [ -f "${EXPORT_DIR}/config.json" ]; then echo "Already exported: ${EXPORT_DIR}" continue fi + rm -rf "${EXPORT_DIR}" echo "=== Exporting: ${CKPT} → ${EXPORT_DIR} ===" CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ --model_path "${CKPT}" \ @@ -170,12 +172,14 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then --trust_remote_code EXPORTED=$((EXPORTED + 1)) done - # Also export final model if saved directly to output_dir - if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && [ ! -d "${OUTPUT_DIR}/exported-checkpoint-final" ]; then - echo "=== Exporting: ${OUTPUT_DIR} → ${OUTPUT_DIR}/exported-checkpoint-final ===" + FINAL_EXPORT="${OUTPUT_DIR}/exported-checkpoint-final" + if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && { [ ! -f "${FINAL_EXPORT}/config.json" ] || \ + [ "${OUTPUT_DIR}/modelopt_state.pth" -nt "${FINAL_EXPORT}/config.json" ]; }; then + rm -rf "${FINAL_EXPORT}" + echo "=== Exporting: ${OUTPUT_DIR} → ${FINAL_EXPORT} ===" CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ --model_path "${OUTPUT_DIR}" \ - --export_path "${OUTPUT_DIR}/exported-checkpoint-final" \ + --export_path "${FINAL_EXPORT}" \ --trust_remote_code EXPORTED=$((EXPORTED + 1)) fi From abe2e9f2ec308a923fa6119794f1fb0ade977179 Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Thu, 29 Oct 2026 04:52:04 +0000 Subject: [PATCH 2/3] launcher: count a DFlash export complete only once its config parses, and never serve a stale final An interrupted export can leave config.json empty or truncated, which the existence check counted as done. An export now counts as complete only when its config.json parses to a non-empty object, and each export is checked again after it is written; a failed or invalid export fails the task. The full-HF exporter cannot load FSDP2 shards, so a sharded checkpoint whose export is incomplete now fails the task instead of being re-exported. latest_checkpoint always serves exported-checkpoint-final, so a final left in a reused output_dir would win over the newest step. The final model is exported only when the root model was saved no earlier than the newest checkpoint; otherwise an existing final export is removed once the newest step export is complete. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zhehao Hu --- .../common/specdec/dflash_online_training.sh | 84 ++++++-- .../tests/test_dflash_online_training.py | 183 +++++++++++++++++- 2 files changed, 242 insertions(+), 25 deletions(-) diff --git a/tools/launcher/common/specdec/dflash_online_training.sh b/tools/launcher/common/specdec/dflash_online_training.sh index 5b8ed3fb236..77bebc78c31 100644 --- a/tools/launcher/common/specdec/dflash_online_training.sh +++ b/tools/launcher/common/specdec/dflash_online_training.sh @@ -134,6 +134,38 @@ if [ -n "${ACCELERATE_CONFIG:-}" ] && [ -f "${ACCELERATE_CONFIG}" ]; then echo "Using accelerate config: ${ACCELERATE_CONFIG}" fi +# Config creation precedes serialization, so existence alone does not prove completion. +export_complete() { + python3 - "$1/config.json" <<'PY' +import json +import sys + +try: + with open(sys.argv[1]) as f: + config = json.load(f) +except (OSError, ValueError): + sys.exit(1) + +sys.exit(0 if isinstance(config, dict) and config else 1) +PY +} + +export_draft() { + local source=$1 destination=$2 + + # Callers use ||, which suppresses ERR handling inside this function. + rm -rf "$destination" || return 1 + echo "=== Exporting: ${source} → ${destination} ===" + CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ + --model_path "$source" \ + --export_path "$destination" \ + --trust_remote_code || return 1 + if ! export_complete "$destination"; then + echo "ERROR: export produced an invalid config: ${destination}" >&2 + return 1 + fi +} + set -x start_time=$(date +%s) # The ERR trap only records a failure and lets the script go on, so test the status here: @@ -160,28 +192,48 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then [ -d "$CKPT" ] || continue STEP=$(basename "$CKPT" | sed 's/checkpoint-//') EXPORT_DIR="${OUTPUT_DIR}/exported-checkpoint-${STEP}" - if [ -f "${EXPORT_DIR}/config.json" ]; then + if export_complete "$EXPORT_DIR"; then echo "Already exported: ${EXPORT_DIR}" continue fi - rm -rf "${EXPORT_DIR}" - echo "=== Exporting: ${CKPT} → ${EXPORT_DIR} ===" - CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ - --model_path "${CKPT}" \ - --export_path "${EXPORT_DIR}" \ - --trust_remote_code + if compgen -G "${CKPT}/pytorch_model_fsdp_*" >/dev/null; then + echo "ERROR: ${CKPT} needs re-exporting, but the HF exporter cannot load FSDP shards." >&2 + exit 1 + fi + export_draft "$CKPT" "$EXPORT_DIR" || exit 1 EXPORTED=$((EXPORTED + 1)) done + + LATEST_CKPT=$(latest_checkpoint "$OUTPUT_DIR" checkpoint-) + LATEST_SAVE="${LATEST_CKPT}/trainer_state.json" + if [ ! -f "$LATEST_SAVE" ]; then + LATEST_SAVE="$LATEST_CKPT" + fi + + FINAL_IS_CURRENT=0 + if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && { + [ -z "$LATEST_CKPT" ] || + ! [ "$LATEST_SAVE" -nt "${OUTPUT_DIR}/modelopt_state.pth" ] + }; then + FINAL_IS_CURRENT=1 + fi + FINAL_EXPORT="${OUTPUT_DIR}/exported-checkpoint-final" - if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && { [ ! -f "${FINAL_EXPORT}/config.json" ] || \ - [ "${OUTPUT_DIR}/modelopt_state.pth" -nt "${FINAL_EXPORT}/config.json" ]; }; then - rm -rf "${FINAL_EXPORT}" - echo "=== Exporting: ${OUTPUT_DIR} → ${FINAL_EXPORT} ===" - CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ - --model_path "${OUTPUT_DIR}" \ - --export_path "${FINAL_EXPORT}" \ - --trust_remote_code - EXPORTED=$((EXPORTED + 1)) + if [ "$FINAL_IS_CURRENT" = 1 ]; then + if ! export_complete "$FINAL_EXPORT" || + [ "${OUTPUT_DIR}/modelopt_state.pth" -nt "${FINAL_EXPORT}/config.json" ]; then + export_draft "$OUTPUT_DIR" "$FINAL_EXPORT" || exit 1 + EXPORTED=$((EXPORTED + 1)) + fi + elif [ -e "$FINAL_EXPORT" ]; then + # Retire an obsolete final only once its replacement is available. + if [ -z "$LATEST_CKPT" ] || + ! export_complete "${OUTPUT_DIR}/exported-${LATEST_CKPT##*/}"; then + echo "ERROR: obsolete final export has no complete newest-step replacement." >&2 + exit 1 + fi + echo "Removing obsolete final export: ${FINAL_EXPORT}" + rm -rf "$FINAL_EXPORT" || exit 1 fi if [ "$EXPORTED" -eq 0 ]; then echo "No new checkpoints to export in ${OUTPUT_DIR}" diff --git a/tools/launcher/tests/test_dflash_online_training.py b/tools/launcher/tests/test_dflash_online_training.py index 1a065f48abb..205bf87e1cd 100644 --- a/tools/launcher/tests/test_dflash_online_training.py +++ b/tools/launcher/tests/test_dflash_online_training.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""dflash_online_training.sh: the regression verdict, its strict mode and its validation.""" +"""dflash_online_training.sh: export validaion and regression checks.""" import json import os @@ -25,37 +25,42 @@ _COMMON = Path(__file__).resolve().parents[1] / "common" _EXPORT = "modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py" +_CONFIG = '{"model_type": "qwen3"}' -def _run(tmp_path, train_exit=0, loss=4.0, **env_overrides): +def _run(tmp_path, train_exit=0, loss=4.0, checkpoint_time=None, **env_overrides): """Run the script in a sandbox: stub pip/nvidia-smi/accelerate/export/torch, real gate. Returns (process, verdict or None, whether training was launched). """ - (tmp_path / "common" / "specdec").mkdir(parents=True) + (tmp_path / "common" / "specdec").mkdir(parents=True, exist_ok=True) for name in ("service_utils.sh", "check_regression.py", "specdec/dflash_online_training.sh"): shutil.copy(_COMMON / name, tmp_path / "common" / name) (tmp_path / "modules/Model-Optimizer/modelopt").mkdir( - parents=True + parents=True, exist_ok=True ) # service_utils appends here export = tmp_path / _EXPORT - export.parent.mkdir(parents=True) + export.parent.mkdir(parents=True, exist_ok=True) export.write_text( "import os, sys\n" "d = sys.argv[sys.argv.index('--export_path') + 1]\n" "os.makedirs(d, exist_ok=True)\n" - "open(os.path.join(d, 'config.json'), 'w').write('{}')\n" + "open(os.path.join(d, 'model.safetensors'), 'w').write('new')\n" + "if os.environ.get('EXPORT_FAIL'):\n" + " sys.exit(1)\n" + f"config = os.environ.get('EXPORT_CONFIG', {_CONFIG!r})\n" + "open(os.path.join(d, 'config.json'), 'w').write(config)\n" ) # The script counts GPUs with torch; without it that line trips the ERR trap and every # run exits 1, which would hide the exit code under test. pystub = tmp_path / "pystub" - pystub.mkdir() + pystub.mkdir(exist_ok=True) (pystub / "torch.py").write_text( "class cuda:\n @staticmethod\n def device_count():\n return 1\n" ) launched = tmp_path / "launched" bin_dir = tmp_path / "bin" - bin_dir.mkdir() + bin_dir.mkdir(exist_ok=True) for tool, body in ( ("pip", "exit 0"), ("nvidia-smi", "echo 1"), @@ -66,10 +71,15 @@ def _run(tmp_path, train_exit=0, loss=4.0, **env_overrides): # A checkpoint with the given loss, and a passing verdict from an earlier run. out = tmp_path / "out" - (out / "checkpoint-2000").mkdir(parents=True) + (out / "checkpoint-2000").mkdir(parents=True, exist_ok=True) (out / "checkpoint-2000" / "trainer_state.json").write_text( json.dumps({"log_history": [{"loss": loss}]}) ) + if checkpoint_time is not None: + os.utime( + out / "checkpoint-2000" / "trainer_state.json", + (checkpoint_time, checkpoint_time), + ) (out / "regression_check.json").write_text('{"pass": true}') env = { @@ -91,6 +101,161 @@ def _run(tmp_path, train_exit=0, loss=4.0, **env_overrides): return proc, json.loads(verdict.read_text()) if verdict.exists() else None, launched.exists() +def _seed_export(path, config): + path.mkdir(parents=True) + (path / "model.safetensors").write_text("old") + if config is not None: + (path / "config.json").write_text(config) + + +@pytest.mark.parametrize("final", [False, True]) +@pytest.mark.parametrize("config", [None, "", '{"model_ty', "{}", "[]", _CONFIG]) +def test_incomplete_export_is_repaired(tmp_path, final, config): + """Repair invalid exports and leave complete exports untouched.""" + out = tmp_path / "out" + if final: + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + (out / "modelopt_state.pth").write_text("") + export = out / f"exported-checkpoint-{'final' if final else '2000'}" + _seed_export(export, config) + proc, verdict, _ = _run(tmp_path, checkpoint_time=1) + assert proc.returncode == 0, proc.stderr + assert verdict["pass"] is True + assert (export / "config.json").read_text() == _CONFIG + assert (export / "model.safetensors").read_text() == ("old" if config == _CONFIG else "new") + + +@pytest.mark.parametrize("stale", [False, True]) +def test_final_export_tracks_root_model(tmp_path, stale): + """Refresh a valid final only when the root model is newer.""" + out = tmp_path / "out" + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + export = out / "exported-checkpoint-final" + _seed_export(export, _CONFIG) + state = out / "modelopt_state.pth" + state.write_text("") + t = (export / "config.json").stat().st_mtime + (10 if stale else -10) + os.utime(state, (t, t)) + proc, _, _ = _run(tmp_path, checkpoint_time=1) + assert proc.returncode == 0, proc.stderr + assert (export / "model.safetensors").read_text() == ("new" if stale else "old") + + +@pytest.mark.parametrize("final", [False, True]) +@pytest.mark.parametrize("failure", [{"EXPORT_FAIL": "1"}, {"EXPORT_CONFIG": ""}]) +def test_failed_export_writes_no_verdict(tmp_path, final, failure): + """Exporter errors and invalid output stop the script before grading training.""" + out = tmp_path / "out" + out.mkdir() + (out / "modelopt_state.pth").write_text("") + if final: + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + proc, verdict, _ = _run(tmp_path, checkpoint_time=1, **failure) + assert proc.returncode == 1, proc.stderr + assert verdict is None + if not final: + assert not (out / "exported-checkpoint-final").exists() + + +def test_repeat_run_keeps_complete_exports(tmp_path): + """A repeated run does not rewrite completed step or final exports.""" + out = tmp_path / "out" + out.mkdir() + (out / "modelopt_state.pth").write_text("") + proc, _, _ = _run(tmp_path, checkpoint_time=1) + assert proc.returncode == 0, proc.stderr + weights = [ + out / f"exported-checkpoint-{step}" / "model.safetensors" for step in ("2000", "final") + ] + for path in weights: + os.utime(path, (1, 1)) + proc, _, _ = _run(tmp_path, checkpoint_time=1) + assert proc.returncode == 0, proc.stderr + assert all(path.stat().st_mtime == 1 for path in weights) + + +@pytest.mark.parametrize("config", [None, ""]) +def test_sharded_export_failure_preserves_leftover_final(tmp_path, config): + """A leftover final cannot bypass an unexportable newest checkpoint.""" + out = tmp_path / "out" + checkpoint = out / "checkpoint-2000" + (checkpoint / "pytorch_model_fsdp_0").mkdir(parents=True) + + state = out / "modelopt_state.pth" + state.write_text("") + os.utime(state, (100, 100)) + + final = out / "exported-checkpoint-final" + _seed_export(final, _CONFIG) + os.utime(final / "config.json", (110, 110)) + + step_export = out / "exported-checkpoint-2000" + if config is not None: + _seed_export(step_export, config) + + # The stub exporter would succeed, so failure must come from the shard guard. + proc, verdict, _ = _run(tmp_path, checkpoint_time=200) + + assert proc.returncode == 1, proc.stderr + assert verdict is None + assert (final / "model.safetensors").read_text() == "old" + assert (final / "config.json").read_text() == _CONFIG + if config is None: + assert not step_export.exists() + else: + assert (step_export / "model.safetensors").read_text() == "old" + assert (step_export / "config.json").read_text() == config + + +@pytest.mark.parametrize("existing_final", [False, True]) +@pytest.mark.parametrize("root_time", [100, 200, 300]) +def test_final_selection_tracks_latest_checkpoint(tmp_path, existing_final, root_time): + """Only a current root model may supply the preferred final export.""" + out = tmp_path / "out" + step_export = out / "exported-checkpoint-2000" + _seed_export(step_export, _CONFIG) + + state = out / "modelopt_state.pth" + state.write_text("") + current_root = root_time >= 200 + os.utime(state, (root_time, root_time)) + + final = out / "exported-checkpoint-final" + if existing_final: + _seed_export(final, _CONFIG) + os.utime(final / "config.json", (root_time + 10, root_time + 10)) + + proc, verdict, _ = _run(tmp_path, checkpoint_time=200) + + assert proc.returncode == 0, proc.stderr + assert verdict["pass"] is True + assert verdict["trainer_state"].endswith("checkpoint-2000/trainer_state.json") + assert final.exists() == current_root + assert (step_export / "model.safetensors").read_text() == "old" + if current_root: + assert (final / "model.safetensors").read_text() == ("old" if existing_final else "new") + + selected = subprocess.run( + [ + "bash", + "-c", + "source common/service_utils.sh >/dev/null; " + 'latest_checkpoint "$1" exported-checkpoint-', + "bash", + str(out), + ], + cwd=tmp_path, + env={ + **os.environ, + "PATH": f"{tmp_path / 'bin'}:{os.environ['PATH']}", + }, + capture_output=True, + text=True, + check=True, + ) + assert selected.stdout.strip() == str(final if current_root else step_export) + + @pytest.mark.parametrize("train_exit", [0, 1]) def test_failed_training_writes_no_verdict(tmp_path, train_exit): """A crash must not let a surviving checkpoint, or an old verdict, pass the gate.""" From a2127efd68978024771709d2fa47ef57011e13fb Mon Sep 17 00:00:00 2001 From: Zhehao Hu Date: Thu, 29 Oct 2026 21:40:00 +0000 Subject: [PATCH 3/3] launcher: keep the previous DFlash export until its replacement is published, and skip only unneeded FSDP exports export_draft deleted the previous export before exporting, so a failed export left nothing to serve. Each export is now written to .tmp, checked, and swapped in while the previous export waits as .old; a failed export or publish keeps the previous one, and the next run restores a .old left by a job stopped mid-swap. An incomplete export of an FSDP-sharded checkpoint no longer fails the task by itself. It is skipped with a warning, and the task fails only when the newest checkpoint lacks a complete export and no current final export will be served. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zhehao Hu --- .../common/specdec/dflash_online_training.sh | 48 +++++- .../tests/test_dflash_online_training.py | 142 +++++++++++++++--- 2 files changed, 164 insertions(+), 26 deletions(-) diff --git a/tools/launcher/common/specdec/dflash_online_training.sh b/tools/launcher/common/specdec/dflash_online_training.sh index 77bebc78c31..83ae850bfa0 100644 --- a/tools/launcher/common/specdec/dflash_online_training.sh +++ b/tools/launcher/common/specdec/dflash_online_training.sh @@ -150,20 +150,43 @@ sys.exit(0 if isinstance(config, dict) and config else 1) PY } +recover_export() { + local destination=$1 backup="$1.old" + + [ -e "$backup" ] || return 0 + if [ -e "$destination" ]; then + # Only a complete publication supersedes the backup. + export_complete "$destination" || return 1 + rm -rf "$backup" || return 1 + else + mv -T "$backup" "$destination" || return 1 + fi +} + export_draft() { local source=$1 destination=$2 + local scratch="$2.tmp" backup="$2.old" # Callers use ||, which suppresses ERR handling inside this function. - rm -rf "$destination" || return 1 + recover_export "$destination" || return 1 + rm -rf "$scratch" || return 1 echo "=== Exporting: ${source} → ${destination} ===" CUDA_VISIBLE_DEVICES=0 python3 modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py \ --model_path "$source" \ - --export_path "$destination" \ + --export_path "$scratch" \ --trust_remote_code || return 1 - if ! export_complete "$destination"; then - echo "ERROR: export produced an invalid config: ${destination}" >&2 + if ! export_complete "$scratch"; then + echo "ERROR: export produced an invalid config: ${scratch}" >&2 + return 1 + fi + if [ -e "$destination" ]; then + mv -T "$destination" "$backup" || return 1 + fi + if ! mv -T "$scratch" "$destination"; then + recover_export "$destination" || return 1 return 1 fi + rm -rf "$backup" || return 1 } set -x @@ -188,6 +211,10 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then if [ -n "$OUTPUT_DIR" ]; then # Export all checkpoints that don't already have an exported version EXPORTED=0 + for BACKUP in "${OUTPUT_DIR}"/exported-checkpoint-*.old; do + [ -d "$BACKUP" ] || continue + recover_export "${BACKUP%.old}" || exit 1 + done for CKPT in "${OUTPUT_DIR}"/checkpoint-*; do [ -d "$CKPT" ] || continue STEP=$(basename "$CKPT" | sed 's/checkpoint-//') @@ -197,19 +224,19 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then continue fi if compgen -G "${CKPT}/pytorch_model_fsdp_*" >/dev/null; then - echo "ERROR: ${CKPT} needs re-exporting, but the HF exporter cannot load FSDP shards." >&2 - exit 1 + echo "WARNING: cannot re-export FSDP shards in ${CKPT}; checking the required export below." >&2 + continue fi export_draft "$CKPT" "$EXPORT_DIR" || exit 1 EXPORTED=$((EXPORTED + 1)) done - LATEST_CKPT=$(latest_checkpoint "$OUTPUT_DIR" checkpoint-) LATEST_SAVE="${LATEST_CKPT}/trainer_state.json" if [ ! -f "$LATEST_SAVE" ]; then LATEST_SAVE="$LATEST_CKPT" fi + # Equal timestamps cannot establish ordering; only a newer checkpoint makes the root stale. FINAL_IS_CURRENT=0 if [ -f "${OUTPUT_DIR}/modelopt_state.pth" ] && { [ -z "$LATEST_CKPT" ] || @@ -218,6 +245,13 @@ elif [ "${SLURM_PROCID:-0}" = "0" ]; then FINAL_IS_CURRENT=1 fi + # Skipped sharded exports are safe only if the current serving artifact is available. + if [ "$FINAL_IS_CURRENT" != 1 ] && [ -n "$LATEST_CKPT" ] && + ! export_complete "${OUTPUT_DIR}/exported-${LATEST_CKPT##*/}"; then + echo "ERROR: newest checkpoint has no complete export: ${LATEST_CKPT}" >&2 + exit 1 + fi + FINAL_EXPORT="${OUTPUT_DIR}/exported-checkpoint-final" if [ "$FINAL_IS_CURRENT" = 1 ]; then if ! export_complete "$FINAL_EXPORT" || diff --git a/tools/launcher/tests/test_dflash_online_training.py b/tools/launcher/tests/test_dflash_online_training.py index 205bf87e1cd..4ded06482da 100644 --- a/tools/launcher/tests/test_dflash_online_training.py +++ b/tools/launcher/tests/test_dflash_online_training.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""dflash_online_training.sh: export validaion and regression checks.""" +"""dflash_online_training.sh: export validation and regression checks.""" import json import os @@ -26,6 +26,7 @@ _COMMON = Path(__file__).resolve().parents[1] / "common" _EXPORT = "modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py" _CONFIG = '{"model_type": "qwen3"}' +_MV = shutil.which("mv") def _run(tmp_path, train_exit=0, loss=4.0, checkpoint_time=None, **env_overrides): @@ -43,7 +44,9 @@ def _run(tmp_path, train_exit=0, loss=4.0, checkpoint_time=None, **env_overrides export.parent.mkdir(parents=True, exist_ok=True) export.write_text( "import os, sys\n" - "d = sys.argv[sys.argv.index('--export_path') + 1]\n" + "src = sys.argv[sys.argv.index('--model_path') + 1]\n" + f"open({str(tmp_path / 'exports.log')!r}, 'a').write(src + '\\n')\n" + "d =sys.argv[sys.argv.index('--export_path') + 1]\n" "os.makedirs(d, exist_ok=True)\n" "open(os.path.join(d, 'model.safetensors'), 'w').write('new')\n" "if os.environ.get('EXPORT_FAIL'):\n" @@ -65,6 +68,14 @@ def _run(tmp_path, train_exit=0, loss=4.0, checkpoint_time=None, **env_overrides ("pip", "exit 0"), ("nvidia-smi", "echo 1"), ("accelerate", f"touch {launched}; exit {train_exit}"), + ( + "mv", + ( + '[ -n "$PUBLISH_FAIL" ] && [[ "$2" == *.tmp ]] && exit 1\n' + '[ -n "$RECOVER_FAIL" ] && [[ "$2" == *.old ]] && exit 1\n' + f'exec {_MV} "$@"' + ), + ), ): (bin_dir / tool).write_text(f"#!/bin/bash\n{body}\n") (bin_dir / tool).chmod(0o755) @@ -180,26 +191,23 @@ def test_sharded_export_failure_preserves_leftover_final(tmp_path, config): out = tmp_path / "out" checkpoint = out / "checkpoint-2000" (checkpoint / "pytorch_model_fsdp_0").mkdir(parents=True) - state = out / "modelopt_state.pth" state.write_text("") os.utime(state, (100, 100)) - final = out / "exported-checkpoint-final" _seed_export(final, _CONFIG) os.utime(final / "config.json", (110, 110)) - step_export = out / "exported-checkpoint-2000" if config is not None: _seed_export(step_export, config) - # The stub exporter would succeed, so failure must come from the shard guard. + # The stub exporter would succeed, so the script must reject the missing required export. proc, verdict, _ = _run(tmp_path, checkpoint_time=200) - assert proc.returncode == 1, proc.stderr assert verdict is None assert (final / "model.safetensors").read_text() == "old" assert (final / "config.json").read_text() == _CONFIG + assert not (tmp_path / "exports.log").exists() if config is None: assert not step_export.exists() else: @@ -207,26 +215,119 @@ def test_sharded_export_failure_preserves_leftover_final(tmp_path, config): assert (step_export / "config.json").read_text() == config +@pytest.mark.parametrize( + ("incomplete_step", "current_final"), + [("1600", None), ("2000", "existing"), ("2000", "export")], +) +def test_incomplete_shards_allow_current_export(tmp_path, incomplete_step, current_final): + """An obsolete shard export must not block a complete newest step or current final.""" + out = tmp_path / "out" + for step in ("1600", "2000"): + (out / f"checkpoint-{step}" / "pytorch_model_fsdp_0").mkdir(parents=True) + _seed_export( + out / f"exported-checkpoint-{step}", "" if step == incomplete_step else _CONFIG + ) + final = out / "exported-checkpoint-final" + if current_final: + state = out / "modelopt_state.pth" + state.write_text("") + os.utime(state, (300, 300)) + if current_final == "existing": + _seed_export(final, _CONFIG) + os.utime(final / "config.json", (310, 310)) + + proc, verdict, _ = _run(tmp_path, checkpoint_time=200) + assert proc.returncode == 0, proc.stderr + assert verdict["pass"] is True + required = final if current_final else out / "exported-checkpoint-2000" + assert json.loads((required / "config.json").read_text()) + log = tmp_path / "exports.log" + assert (log.read_text().splitlines() if log.exists() else []) == ( + [str(out)] if current_final == "export" else [] + ) + assert _latest_export(tmp_path) == str(required) + + +@pytest.mark.parametrize( + "failure", + [ + {"EXPORT_FAIL": "1"}, + {"EXPORT_CONFIG": ""}, + {"PUBLISH_FAIL": "1"}, + {"PUBLISH_FAIL": "1", "RECOVER_FAIL": "1"}, + ], +) +def test_failed_replacement_preserves_final(tmp_path, failure): + """Export and publication failures retain the previous usable final.""" + out = tmp_path / "out" + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + final = out / "exported-checkpoint-final" + _seed_export(final, _CONFIG) + os.utime(final / "config.json", (250, 250)) + state = out / "modelopt_state.pth" + state.write_text("") + os.utime(state, (300, 300)) + proc, verdict, _ = _run(tmp_path, checkpoint_time=200, **failure) + assert proc.returncode == 1, proc.stderr + assert verdict is None + backup = final.with_name(final.name + ".old") + preserved = backup if failure.get("RECOVER_FAIL") else final + assert (preserved / "model.safetensors").read_text() == "old" + assert (preserved / "config.json").read_text() == _CONFIG + assert backup.exists() == bool(failure.get("RECOVER_FAIL")) + + +@pytest.mark.parametrize( + ("published", "recovery_fails"), [(False, False), (True, False), (False, True)] +) +def test_interrupted_publication_recovery(tmp_path, published, recovery_fails): + """Recover a swap, stopping without deleting the backup if restoration fails.""" + out = tmp_path / "out" + _seed_export(out / "exported-checkpoint-2000", _CONFIG) + state = out / "modelopt_state.pth" + state.write_text("") + os.utime(state, (300, 300)) + final = out / "exported-checkpoint-final" + backup = final.with_name(final.name + ".old") + _seed_export(backup, _CONFIG) + os.utime(backup / "config.json", (310, 310)) + if published: + _seed_export(final, _CONFIG) + (final / "model.safetensors").write_text("published") + os.utime(final / "config.json", (310, 310)) + proc, verdict, _ = _run( + tmp_path, checkpoint_time=200, **({"RECOVER_FAIL": "1"} if recovery_fails else {}) + ) + assert proc.returncode == int(recovery_fails), proc.stderr + assert not (tmp_path / "exports.log").exists() + if recovery_fails: + assert verdict is None + assert not final.exists() + assert (backup / "model.safetensors").read_text() == "old" + assert (backup / "config.json").read_text() == _CONFIG + else: + assert verdict["pass"] is True + assert not backup.exists() + assert (final / "model.safetensors").read_text() == ("published" if published else "old") + + @pytest.mark.parametrize("existing_final", [False, True]) @pytest.mark.parametrize("root_time", [100, 200, 300]) def test_final_selection_tracks_latest_checkpoint(tmp_path, existing_final, root_time): - """Only a current root model may supply the preferred final export.""" + """Retire finals only when the checkpoint is strictly newer than the root.""" out = tmp_path / "out" step_export = out / "exported-checkpoint-2000" _seed_export(step_export, _CONFIG) - state = out / "modelopt_state.pth" state.write_text("") current_root = root_time >= 200 os.utime(state, (root_time, root_time)) - final = out / "exported-checkpoint-final" if existing_final: _seed_export(final, _CONFIG) os.utime(final / "config.json", (root_time + 10, root_time + 10)) proc, verdict, _ = _run(tmp_path, checkpoint_time=200) - assert proc.returncode == 0, proc.stderr assert verdict["pass"] is True assert verdict["trainer_state"].endswith("checkpoint-2000/trainer_state.json") @@ -235,25 +336,28 @@ def test_final_selection_tracks_latest_checkpoint(tmp_path, existing_final, root if current_root: assert (final / "model.safetensors").read_text() == ("old" if existing_final else "new") + assert _latest_export(tmp_path) == str(final if current_root else step_export) + + +def _latest_export(tmp_path): selected = subprocess.run( [ "bash", "-c", - "source common/service_utils.sh >/dev/null; " - 'latest_checkpoint "$1" exported-checkpoint-', + ( + "source common/service_utils.sh >/dev/null; " + 'latest_checkpoint "$1" exported-checkpoint-' + ), "bash", - str(out), + str(tmp_path / "out"), ], cwd=tmp_path, - env={ - **os.environ, - "PATH": f"{tmp_path / 'bin'}:{os.environ['PATH']}", - }, + env={**os.environ, "PATH": f"{tmp_path / 'bin'}:{os.environ['PATH']}"}, capture_output=True, text=True, check=True, ) - assert selected.stdout.strip() == str(final if current_root else step_export) + return selected.stdout.strip() @pytest.mark.parametrize("train_exit", [0, 1])