From eac1b775511f96a73ff7bbae571ff297d57fd578 Mon Sep 17 00:00:00 2001 From: Gsunshine Date: Mon, 23 Sep 2024 18:46:35 -0400 Subject: [PATCH 1/4] Add AMP gradscalar. --- ct_train.py | 5 ++-- training/ct_training_loop.py | 52 +++++++++++++++++++----------------- 2 files changed, 30 insertions(+), 27 deletions(-) diff --git a/ct_train.py b/ct_train.py index d9976a5..f725187 100644 --- a/ct_train.py +++ b/ct_train.py @@ -78,7 +78,8 @@ def convert(self, value, param, ctx): @click.option('--fp16', help='Enable mixed-precision training', metavar='BOOL', type=bool, default=False, show_default=True) @click.option('--tf32', help='Enable tf32 for A100/H100 training speed', metavar='BOOL', type=bool, default=False, show_default=True) @click.option('--ls', help='Loss scaling', metavar='FLOAT', type=click.FloatRange(min=0, min_open=True), default=1, show_default=True) -@click.option('--enable_gradscaler', help='Enable torch.cuda.amp.GradScaler, NOTE overwritting loss_scale set by --ls', metavar='BOOL', type=bool, default=False, show_default=True) +@click.option('--amp', help='Enable torch.cuda.amp.GradScaler for mixed precision training, \ + NOTE overwritting loss_scale set by --ls', metavar='BOOL', type=bool, default=False, show_default=True) @click.option('--bench', help='Enable cuDNN benchmarking', metavar='BOOL', type=bool, default=True, show_default=True) @click.option('--cache', help='Cache dataset in CPU memory', metavar='BOOL', type=bool, default=True, show_default=True) @click.option('--workers', help='DataLoader worker processes', metavar='INT', type=click.IntRange(min=1), default=1, show_default=True) @@ -165,7 +166,7 @@ def main(**kwargs): c.ema_halflife_kimg = int(opts.ema * 1000) if opts.ema is not None else opts.ema c.ema_beta = opts.ema_beta c.update(batch_size=opts.batch, batch_gpu=opts.batch_gpu) - c.update(loss_scaling=opts.ls, cudnn_benchmark=opts.bench, enable_tf32=opts.tf32, enable_gradscaler=opts.enable_gradscaler) + c.update(loss_scaling=opts.ls, cudnn_benchmark=opts.bench, enable_tf32=opts.tf32, enable_gradscaler=opts.enable_amp) c.update(kimg_per_tick=opts.tick, snapshot_ticks=opts.snap, state_dump_ticks=opts.dump, ckpt_ticks=opts.ckpt, double_ticks=opts.double) c.update(mid_t=opts.mid_t, metrics=opts.metrics, sample_ticks=opts.sample_every, eval_ticks=opts.eval_every) diff --git a/training/ct_training_loop.py b/training/ct_training_loop.py index 4f5f0b8..10c4d4b 100644 --- a/training/ct_training_loop.py +++ b/training/ct_training_loop.py @@ -128,8 +128,8 @@ def training_loop( mid_t = None, # Intermediate t for few-step generation. metrics = None, # Metrics for evaluation. cudnn_benchmark = True, # Enable torch.backends.cudnn.benchmark? - enable_tf32 = False, # Enable tf32 for A100/H100 GPUs? - enable_gradscaler = False, # Enable torch.cuda.amp.GradScaler + enable_tf32 = False, # Enable tf32 for A100/H100 GPUs? + enable_amp = False, # Enable torch.cuda.amp.GradScaler device = torch.device('cuda'), ): # Initialize. @@ -169,8 +169,9 @@ def training_loop( optimizer = dnnlib.util.construct_class_by_name(params=net.parameters(), **optimizer_kwargs) # subclass of torch.optim.Optimizer augment_pipe = dnnlib.util.construct_class_by_name(**augment_kwargs) if augment_kwargs is not None else None # training.augment.AugmentPipe - dist.print0(f'GradScaler enabled: {enable_gradscaler}') - if enable_gradscaler: + # Automatic Mixed Precision + dist.print0(f'GradScaler enabled: {enable_amp} for mixed preicision training') + if enable_amp: # https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html#adding-gradscaler # https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation dist.print0(f'Setting up GradScaler...') @@ -206,10 +207,11 @@ def training_loop( data = torch.load(resume_state_dump, map_location=torch.device('cpu')) misc.copy_params_and_buffers(src_module=data['net'], dst_module=net, require_all=True) optimizer.load_state_dict(data['optimizer_state']) - if enable_gradscaler: + if enable_amp: if 'gradscaler_state' in data: + # NOTE(aiihn): Although not loading the state_dict of the GradScaler works well, + # loading it can improve reproducibility. dist.print0(f'Loading GradScaler state from "{resume_state_dump}"...') - # Although not loading the state_dict of the GradScaler works well, loading it can improve reproducibility. scaler.load_state_dict(data['gradscaler_state']) else: dist.print0(f'GradScaler state is not found in "{resume_state_dump}", using the default state.') @@ -269,24 +271,26 @@ def update_scheduler(loss_fn): loss = loss_fn(net=ddp, images=images, labels=labels, augment_pipe=augment_pipe) training_stats.report('Loss/loss', loss) - if enable_gradscaler: + if enable_amp: scaler.scale(loss.mean()).backward() else: - # loss.sum().mul(loss_scaling / batch_gpu_total).backward() loss.mul(loss_scaling).mean().backward() - if enable_gradscaler: - # TODO Is torch.nan_to_num(param.grad, nan=0, posinf=1e5, neginf=-1e5, out=param.grad) needed when using GradScaler? + # NOTE(aiihn & Gsunshine): This should be further tested for AMP. + for param in net.parameters(): + if param.grad is not None: + torch.nan_to_num(param.grad, nan=0, posinf=1e5, neginf=-1e5, out=param.grad) + + # LR scheduler (if needed in the future) + # for g in optimizer.param_groups: + # g['lr'] = optimizer_kwargs['lr'] * min(cur_nimg / max(lr_rampup_kimg * 1000, 1e-8), 1) + + # Update weights. + if enable_amp: scaler.step(optimizer) scaler.update() else: - # Update weights. - # for g in optimizer.param_groups: - # g['lr'] = optimizer_kwargs['lr'] * min(cur_nimg / max(lr_rampup_kimg * 1000, 1e-8), 1) - for param in net.parameters(): - if param.grad is not None: - torch.nan_to_num(param.grad, nan=0, posinf=1e5, neginf=-1e5, out=param.grad) - optimizer.step() + optimizer.step() # Update EMA. if ema_halflife_kimg is not None: @@ -341,10 +345,9 @@ def update_scheduler(loss_fn): # Save full dump of the training state. if (state_dump_ticks is not None) and (done or cur_tick % state_dump_ticks == 0) and cur_tick != 0 and dist.get_rank() == 0: - if enable_gradscaler: - data = dict(net=net, optimizer_state=optimizer.state_dict(), gradscaler_state=scaler.state_dict()) - else: - data = dict(net=net, optimizer_state=optimizer.state_dict()) + data = dict(net=net, optimizer_state=optimizer.state_dict()) + if enable_amp: + data.update(gradscaler_state=scaler.state_dict()) torch.save(data, os.path.join(run_dir, f'training-state-{cur_tick:06d}.pt')) # Save latest checkpoints @@ -363,10 +366,9 @@ def update_scheduler(loss_fn): del data # conserve memory if dist.get_rank() == 0: - if enable_gradscaler: - data = dict(net=net, optimizer_state=optimizer.state_dict(), gradscaler_state=scaler.state_dict()) - else: - data = dict(net=net, optimizer_state=optimizer.state_dict()) + data = dict(net=net, optimizer_state=optimizer.state_dict()) + if enable_amp: + data.update(gradscaler_state=scaler.state_dict()) torch.save(data, os.path.join(run_dir, f'training-state-latest.pt')) # Sample Img From 48e3f4781e4a5a4c9387da6e9363bd50d335ece2 Mon Sep 17 00:00:00 2001 From: Gsunshine Date: Mon, 23 Sep 2024 18:57:41 -0400 Subject: [PATCH 2/4] Update README. --- README.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/README.md b/README.md index 1844970..d232ef2 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ Try your own [Consistency Models](https://arxiv.org/abs/2303.01469)! You only ne Baking more in the oven. 🙃 +- **2024.09.23** - Add Gradscaler for Mixed Precision Training. To use mixed precision with GradScaler, switch to the [`amp` branch](https://github.com/locuslab/ect/tree/amp): `git checkout amp`. - **2024.04.27** - Upgrade environment to Pytorch 2.3.0. - **2024.04.12** - ECMs can now surpass SoTA GANs using 1 model step and SoTA Diffusion Models using 2 model steps on CIFAR10. Checkpoints available. @@ -46,6 +47,18 @@ bash run_ecm.sh --desc bs128.200k Replace NGPUs and PORT with the number of GPUs used for training and the port number for DDP sync. +### Half Precision Training + +In this branch, we enable fp16 training with AMP GradScaler for more stable training dynamics. +To enable fp16 and GradScaler, add the following arguments to your script: + +```bash +bash run_ecm_1hour.sh 1 --desc bs128.1hour --fp16=True --enable_amp=True +``` + +For more information, please refer to this [PR](https://github.com/locuslab/ect/pull/13). +Full support for Automatic Mixed Precision (AMP) will be added later. + ## Evaluation Run the following command to calculate FID of a pretrained checkpoint. From a2238829fa46c28272213217be7a411684f0a91f Mon Sep 17 00:00:00 2001 From: hjjjs4vbmv-netizen Date: Tue, 14 Jul 2026 23:53:59 +0800 Subject: [PATCH 3/4] fix: restore AMP CLI and metrics-none behavior --- .gitignore | 17 +++++++++++++---- ct_train.py | 4 ++-- training/ct_training_loop.py | 2 +- 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index 3857786..f3250f3 100644 --- a/.gitignore +++ b/.gitignore @@ -161,10 +161,19 @@ cython_debug/ # Project-related -datasets - -ct-runs -ct-evals +datasets/ +runs/ +checkpoints/ +ct-runs/ +ct-evals/ +wandb/ +*.pkl +*.pt +*.pth +*.ckpt +*.safetensors +*.zip +*.pyc slurm* debug.sh diff --git a/ct_train.py b/ct_train.py index f725187..1906ec4 100644 --- a/ct_train.py +++ b/ct_train.py @@ -78,7 +78,7 @@ def convert(self, value, param, ctx): @click.option('--fp16', help='Enable mixed-precision training', metavar='BOOL', type=bool, default=False, show_default=True) @click.option('--tf32', help='Enable tf32 for A100/H100 training speed', metavar='BOOL', type=bool, default=False, show_default=True) @click.option('--ls', help='Loss scaling', metavar='FLOAT', type=click.FloatRange(min=0, min_open=True), default=1, show_default=True) -@click.option('--amp', help='Enable torch.cuda.amp.GradScaler for mixed precision training, \ +@click.option('--enable_amp', '--amp', 'enable_amp', help='Enable torch.cuda.amp.GradScaler for mixed precision training, \ NOTE overwritting loss_scale set by --ls', metavar='BOOL', type=bool, default=False, show_default=True) @click.option('--bench', help='Enable cuDNN benchmarking', metavar='BOOL', type=bool, default=True, show_default=True) @click.option('--cache', help='Cache dataset in CPU memory', metavar='BOOL', type=bool, default=True, show_default=True) @@ -166,7 +166,7 @@ def main(**kwargs): c.ema_halflife_kimg = int(opts.ema * 1000) if opts.ema is not None else opts.ema c.ema_beta = opts.ema_beta c.update(batch_size=opts.batch, batch_gpu=opts.batch_gpu) - c.update(loss_scaling=opts.ls, cudnn_benchmark=opts.bench, enable_tf32=opts.tf32, enable_gradscaler=opts.enable_amp) + c.update(loss_scaling=opts.ls, cudnn_benchmark=opts.bench, enable_tf32=opts.tf32, enable_amp=opts.enable_amp) c.update(kimg_per_tick=opts.tick, snapshot_ticks=opts.snap, state_dump_ticks=opts.dump, ckpt_ticks=opts.ckpt, double_ticks=opts.double) c.update(mid_t=opts.mid_t, metrics=opts.metrics, sample_ticks=opts.sample_every, eval_ticks=opts.eval_every) diff --git a/training/ct_training_loop.py b/training/ct_training_loop.py index 10c4d4b..f89bf66 100644 --- a/training/ct_training_loop.py +++ b/training/ct_training_loop.py @@ -380,7 +380,7 @@ def update_scheduler(loss_fn): del images # Evaluation - if (eval_ticks is not None) and (done or cur_tick % eval_ticks == 0) and cur_tick > 0: + if metrics and (eval_ticks is not None) and (done or cur_tick % eval_ticks == 0) and cur_tick > 0: dist.print0('Evaluating models...') result_dict = metric_main.calc_metric(metric='fid50k_full', generator_fn=generator_fn, G=ema, G_kwargs={}, From 36dcae6f5416b2e1be770e70f6f81a20b2d2cb6b Mon Sep 17 00:00:00 2001 From: hjjjs4vbmv-netizen Date: Wed, 15 Jul 2026 11:59:52 +0800 Subject: [PATCH 4/4] chore: integrate reproducible engineering workflow --- artifacts/day2/collaborator_audit.md | 128 +++++++++++ artifacts/day2/integration_report.md | 126 +++++++++++ conda-matpool.yml | 18 ++ docs/DAY1_A.md | 202 ++++++++++++++++++ download_checkpoint.sh | 125 +++++++++++ prepare_data.sh | 147 +++++++++++++ scripts/check_environment.py | 138 ++++++++++++ scripts/smoke_engineering_100steps.sh | 295 ++++++++++++++++++++++++++ scripts/verify_assets.py | 221 +++++++++++++++++++ scripts/verify_smoke_run.py | 136 ++++++++++++ setup_env.sh | 127 +++++++++++ 11 files changed, 1663 insertions(+) create mode 100644 artifacts/day2/collaborator_audit.md create mode 100644 artifacts/day2/integration_report.md create mode 100644 conda-matpool.yml create mode 100644 docs/DAY1_A.md create mode 100755 download_checkpoint.sh create mode 100755 prepare_data.sh create mode 100755 scripts/check_environment.py create mode 100755 scripts/smoke_engineering_100steps.sh create mode 100755 scripts/verify_assets.py create mode 100755 scripts/verify_smoke_run.py create mode 100755 setup_env.sh diff --git a/artifacts/day2/collaborator_audit.md b/artifacts/day2/collaborator_audit.md new file mode 100644 index 0000000..f9e8109 --- /dev/null +++ b/artifacts/day2/collaborator_audit.md @@ -0,0 +1,128 @@ +# Day 2 collaborator audit + +Audit date: 2026-07-15 (Asia/Shanghai) + +This report records the read-only audit completed before local integration. Remote refs were +fetched and inspected, but no collaborator branch was modified, merged, rebased, reset, or +pushed. + +## Public baseline + +- `origin/main`: `4311059770f54821d151a9b0e1f76770a5f3930e` +- `origin/leader/day1-bootstrap`: `4e33194777a347ea5286b5ec1d5c29a58c792d29` +- Merge base: `origin/main` +- Relative topology: leader ahead 4, behind 0 +- Changed files versus main: 8 +- `git diff --check`: passed +- AMP CLI: `--enable_amp`, `--amp`, and `--enable_gradscaler` map to `enable_amp` +- AMP propagation: `enable_amp` is passed into `training_loop` +- GradScaler order: accumulation, `unscale_`, non-finite handling, `step`, `update` +- GradScaler state: saved in numbered/latest training states and restored on resume +- Non-AMP path: retains ordinary `optimizer.step()` +- `metrics=none`: parses to `[]`; periodic and final metric calls are skipped +- Compile and help checks: passed in the `ect` environment +- Decision: **READY_TO_MERGE** + +No formal training or FID/KID evaluation was run during the audit. + +## codex/day1-engineering + +- SHA: `3cb1c52fc56f01942c84d535dddddffd99c3af47` +- Commit: `Add reproducible Day 1 training workflow` +- Base: `origin/main`; behind the public baseline by 4 commits and ahead by 1 +- Scope: 18 changed files, 1183 insertions, 14 deletions +- Shell, Python, JSON, and whitespace syntax checks: passed + +Useful engineering deliverables: + +- `conda-matpool.yml` +- `setup_env.sh` +- `prepare_data.sh` +- `download_checkpoint.sh` +- `scripts/check_environment.py` +- `scripts/verify_assets.py` +- `scripts/verify_smoke_run.py` +- `docs/DAY1_A.md` + +Files that must not replace the public baseline: + +- `ct_train.py` +- `training/ct_training_loop.py` +- `env.yml` +- `.gitignore` +- `README.md` + +The collaborator environment pinned huggingface-hub 0.20.3, while the validated public +runtime used 0.23.4. The collaborator dataset ZIP SHA256 +`45e772cbbcb4ebb8657d383557fba2fd24cb929aeaab99fe1963b1462377da9d` differs from the +public ZIP SHA256 `2d4056e80de1a96fe16f2f58945c6c4710ecd9fc02e3cc7aa5b50513b7cdf389`. +Both reported the same byte size. `dataset_tool.py` stores variable ZIP entry timestamps, so +the converted ZIP digest is informational; source MD5, CRC, image/label content, dimensions, +color mode, and project loader behavior are authoritative. + +The collaborator smoke used global batch 10 and ran 100 fresh plus 100 resumed optimizer +updates. It explicitly used `--fp16=False`, did not enable GradScaler, and was produced from a +dirty tree whose recorded SHA was main rather than the collaborator commit. It is FP32 +engineering-connectivity evidence only. No checkpoint, training state, dataset ZIP, or large +log was committed to Git. + +Decision: **ACCEPT_SELECTED_FILES + REQUIRES_REBASE + REJECT_WHOLE_BRANCH**. + +## wk/iniBR + +- SHA: `e3158d83112a2fffb0796a515becee699c73d3fa` +- Ahead 1, behind 0 relative to main +- Raw scope: 30 files, 5375 insertions and 5375 deletions +- Ignoring end-of-line differences produces an empty diff with exit code 0 +- Typical files changed from pure LF to pure CRLF with identical logical lines +- Semantic changes: none +- Classification: **CRLF_ONLY_CHANGE** +- Decision: **DO_NOT_MERGE** + +The member should create a new branch from `origin/leader/day1-bootstrap`, set +`git config core.autocrlf input`, reapply only genuine work, and avoid cherry-picking the +line-ending conversion commit. + +## edwards365 + +- SHA: `4311059770f54821d151a9b0e1f76770a5f3930e` +- Exactly equal to `origin/main` +- Ahead 0, behind 0, changed files 0 +- Status: **NO_REMOTE_DELIVERABLE** + +No claim is made about local, unpushed work. + +## Missing collaborator + +No fourth collaborator branch was present after fetching and pruning remote-tracking refs. +The member must provide a branch name and SHA and push the deliverable before role mapping or +integration. + +## Role coverage + +| Role | Expected work | Current branch | Status | +|------|---------------|----------------|--------| +| A | Engineering and environment reproduction | `codex/day1-engineering` | Clear selective deliverable | +| B | Official fixed ECT baseline | None evidenced | Missing remote deliverable | +| C | Adaptive t-to-r schedule | None evidenced | Missing remote deliverable | +| D | Unified sampling, evaluation, visualization | None evidenced | Missing remote deliverable | + +Only Role A can be mapped from content. The other visible collaborator refs contain either no +semantic work or no commits, and one collaborator branch is absent. + +## Selective integration decision + +ACCEPT: + +- the eight engineering files listed above, subject to Day 2 path and validation hardening + +REWRITE: + +- `smoke_test.sh` as `scripts/smoke_engineering_100steps.sh` +- collaborator JSON evidence and provenance rather than importing it as frozen evidence + +REJECT: + +- direct replacement of the five protected public-baseline files +- whole-branch merge of `codex/day1-engineering` +- any merge of `wk/iniBR` diff --git a/artifacts/day2/integration_report.md b/artifacts/day2/integration_report.md new file mode 100644 index 0000000..c1ae949 --- /dev/null +++ b/artifacts/day2/integration_report.md @@ -0,0 +1,126 @@ +# Day 2 local integration report + +Integration date: 2026-07-15 (Asia/Shanghai) + +## Scope and branch + +- Local branch: `leader/day2-local-integration` +- Base: `origin/leader/day1-bootstrap@4e33194777a347ea5286b5ec1d5c29a58c792d29` +- Selective source: `origin/codex/day1-engineering@3cb1c52fc56f01942c84d535dddddffd99c3af47` +- Commit created: no +- Push or PR created: no + +The worktree versions of `ct_train.py`, `training/ct_training_loop.py`, `env.yml`, `.gitignore`, +and `README.md` were hash-checked against the public baseline immediately after extraction and +were identical. + +## Imported files + +- `conda-matpool.yml` +- `setup_env.sh` +- `prepare_data.sh` +- `download_checkpoint.sh` +- `scripts/check_environment.py` +- `scripts/verify_assets.py` +- `scripts/verify_smoke_run.py` +- `docs/DAY1_A.md` + +## Rewritten or adapted files + +- `scripts/smoke_engineering_100steps.sh` + - engineering-connectivity label, never a formal ECT baseline + - persistent `/mnt/ect_project` defaults + - `metrics=none` + - public-baseline `--enable_amp` CLI with FP16 and AMP enabled by default + - explicit statement that legacy FP32 evidence did not validate GradScaler + - separate `--check-only` and `--dry-run` paths +- `scripts/verify_assets.py` + - official CIFAR-10 tarball MD5 + - ZIP CRC + - exactly 50000 PNGs and 50000 one-to-one labels + - all images checked as 32x32 RGB + - `ImageFolderDataset` length, labels, dtype, shape, and sample reads + - ZIP SHA256 recorded without enforcing one conversion digest +- `scripts/check_environment.py` + - aligned to the protected public environment rather than collaborator `env.yml` + - records all observed package versions and validates imports + - checks the frozen core package/runtime versions and CUDA +- `setup_env.sh`, `prepare_data.sh`, and `download_checkpoint.sh` + - use `/mnt/ect_project` persistent defaults + - write optional reports under persistent `runs/day2` + - preserve check-only workflows +- `scripts/verify_smoke_run.py` + - records engineering-only provenance + - validates expected batch, FP16, AMP, and disabled formal metrics +- `docs/DAY1_A.md` + - documents persistent paths and the engineering/formal-baseline boundary + - removes dirty-main evidence claims + - does not claim a completed FP16 + GradScaler training validation + +## Validation results + +All commands below passed: + +- `bash -n setup_env.sh` +- `bash -n prepare_data.sh` +- `bash -n download_checkpoint.sh` +- `bash -n scripts/smoke_engineering_100steps.sh` +- `python -m compileall scripts training ct_train.py` with pycache redirected outside the repo +- `python ct_train.py --help` +- help checks for all imported/reworked scripts +- `bash setup_env.sh --check-only` +- `bash prepare_data.sh --check-only` +- `bash download_checkpoint.sh --check-only` +- `bash scripts/smoke_engineering_100steps.sh --check-only` +- `bash scripts/smoke_engineering_100steps.sh --dry-run` + +The dry-run configuration reported: + +- dataset size: 50000 +- global batch: 10 +- network `use_fp16`: true +- `enable_amp`: true +- `metrics`: empty list +- formal FID/KID: disabled +- output directory was not created + +Asset verification reported: + +- source tarball MD5: `c58f30108f718f92721af3b95e74349a` +- dataset ZIP CRC: passed +- PNG count: 50000 +- label count: 50000 +- image format: 32x32 RGB +- `ImageFolderDataset` length: 50000 +- dataset ZIP SHA256, informational only: + `2d4056e80de1a96fe16f2f58945c6c4710ecd9fc02e3cc7aa5b50513b7cdf389` +- official transfer checkpoint SHA256: + `4d5dcc1f1d0d41c8934ad21626eeddbdc0460182becf9fc059a0631b1eedb4da` + +No optimizer update, complete engineering smoke, formal training, FID, or KID was run. + +## Protected files rejected from import + +- `ct_train.py` +- `training/ct_training_loop.py` +- `env.yml` +- `.gitignore` +- `README.md` +- collaborator `artifacts/day1/*.json` +- collaborator `artifacts/day1/README.md` +- entire `codex/day1-engineering` branch +- entire `wk/iniBR` branch + +## Remaining issues and decisions + +1. The integration is intentionally uncommitted and unpushed. +2. No dynamic FP16 + GradScaler optimizer step or GradScaler state/resume test was run; the + successful dry run validates configuration propagation only. +3. The official fixed ECT baseline, adaptive schedule, and unified evaluation work still have + no collaborator remote deliverables. +4. `edwards365` still points to main, and one collaborator branch is still missing. +5. Public `env.yml` does not explicitly pin huggingface-hub 0.23.4 even though the validated + runtime and environment checker freeze that compatibility version. Changing `env.yml` was + explicitly outside this selective integration manifest and requires a separate decision. +6. `scripts/verify_smoke_run.py` was syntax-checked but cannot be exercised end-to-end without + running a new 100-step engineering smoke, which was explicitly prohibited in this phase. diff --git a/conda-matpool.yml b/conda-matpool.yml new file mode 100644 index 0000000..6abcf33 --- /dev/null +++ b/conda-matpool.yml @@ -0,0 +1,18 @@ +channels: + - pytorch + - nvidia + - defaults + +# Matpool's default channel alias may redirect community channels to mirror +# paths that do not exist (notably the `nvidia` channel). Keep named channels +# on Anaconda.org while retaining fast mainland mirrors for defaults. +channel_alias: https://conda.anaconda.org +default_channels: + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r + +show_channel_urls: true +default_threads: 1 +fetch_threads: 1 +verify_threads: 1 +execute_threads: 1 diff --git a/docs/DAY1_A.md b/docs/DAY1_A.md new file mode 100644 index 0000000..c2f3b57 --- /dev/null +++ b/docs/DAY1_A.md @@ -0,0 +1,202 @@ +# Day 1 A：工程环境、资产与连通性验证 + +本文档对应工程和环境复现工作线。当前整合以公共基线 +`origin/leader/day1-bootstrap@4e33194777a347ea5286b5ec1d5c29a58c792d29` +为基础，并保留该基线的 AMP、GradScaler 和 `metrics=none` 行为。 + +这里的 100-step smoke 只验证工程连通性，不是官方固定 ECT baseline， +也不能用于报告正式训练质量、FID 或 KID。 + +## 1. 持久化目录 + +脚本默认使用以下布局： + +```text +/mnt/ect_project/ +├── datasets/ +├── pretrained/ +├── runs/ +└── checkpoints/ +``` + +默认资产路径为： + +```text +/mnt/ect_project/datasets/cifar-10-python.tar.gz +/mnt/ect_project/datasets/cifar10-32x32.zip +/mnt/ect_project/pretrained/edm-cifar10-32x32-uncond-vp.pkl +``` + +如平台的持久盘挂载点不同，可统一设置： + +```bash +export ECT_PROJECT_ROOT=/path/to/persistent/ect_project +``` + +也可以分别覆盖： + +```bash +export ECT_CIFAR10_TARBALL=/path/to/cifar-10-python.tar.gz +export ECT_DATA_PATH=/path/to/cifar10-32x32.zip +export ECT_TRANSFER_PATH=/path/to/edm-cifar10-32x32-uncond-vp.pkl +export ECT_RUNS_ROOT=/path/to/persistent/runs +``` + +## 2. 环境检查 + +`env.yml` 仍由公共基线管理，本次工程整合不会覆盖它。MatrixCloud/矩池云可用 +`conda-matpool.yml` 处理 channel 映射。 + +已有 `ect` 环境时只检查，不更新： + +```bash +bash setup_env.sh --check-only +``` + +首次创建环境： + +```bash +bash setup_env.sh +``` + +只有明确决定同步依赖时才使用： + +```bash +bash setup_env.sh --update +``` + +检查器记录 Python、包版本、import、CUDA、GPU、Git SHA 和工作树状态。公共基线已 +验证的关键版本包括 Python 3.9.18、PyTorch 2.3.0、CUDA 12.1、diffusers +0.26.3、accelerate 0.27.2 和 huggingface-hub 0.23.4。 + +## 3. CIFAR-10 准备与验证 + +只检查现有持久化资产： + +```bash +bash prepare_data.sh --check-only +``` + +资产不存在时才执行下载和转换： + +```bash +bash prepare_data.sh +``` + +数据验收包括： + +1. 官方原始 tarball MD5 `c58f30108f718f92721af3b95e74349a`； +2. ZIP CRC； +3. 恰好 50000 个 PNG； +4. 恰好 50000 个 labels，且与 PNG 一一对应； +5. 所有 PNG 均为 32×32 RGB； +6. `ImageFolderDataset` 长度为 50000，且首、中、尾样本可读取； +7. 记录转换后 ZIP SHA256，但不把单一 SHA256 作为内容一致性的硬约束。 + +转换 ZIP 可能因 ZIP entry 时间戳不同而具有不同 SHA256；因此原始 tarball MD5、 +CRC、图像/标签内容和项目数据加载器检查才是主要依据。 + +## 4. 官方 EDM transfer checkpoint + +只检查现有 checkpoint： + +```bash +bash download_checkpoint.sh --check-only +``` + +文件不存在时才下载： + +```bash +bash download_checkpoint.sh +``` + +默认目标位于 `/mnt/ect_project/pretrained/`，并检查官方 SHA256： + +```text +4d5dcc1f1d0d41c8934ad21626eeddbdc0460182becf9fc059a0631b1eedb4da +``` + +## 5. 工程 smoke 与正式 baseline 的边界 + +新脚本名称为： + +```text +scripts/smoke_engineering_100steps.sh +``` + +它的职责仅包括： + +- 环境和资产可读取； +- `ct_train.py` 配置连通； +- 经单独授权后可执行 fresh 100 updates； +- 经单独授权后可执行 resume 100 updates； +- checkpoint、optimizer state 和 resume 链路可生成并检查； +- `metrics=none`，不运行正式 FID/KID。 + +它不属于官方固定 ECT baseline，不提供训练质量结论。 + +只读检查环境和资产： + +```bash +bash scripts/smoke_engineering_100steps.sh --check-only +``` + +只执行 `ct_train.py --dry_run`，不训练、不创建 run 目录： + +```bash +bash scripts/smoke_engineering_100steps.sh --dry-run +``` + +脚本默认使用公共基线的规范 AMP 参数： + +```text +--fp16=True +--enable_amp=True +``` + +公共基线也兼容 `--amp=True` 和 `--enable_gradscaler=True`。旧 collaborator smoke +明确使用 `--fp16=False`，且没有启用 GradScaler，因此旧 evidence 不能被解释为 +FP16 + GradScaler 验证。只有未来实际运行新的 AMP smoke 并检查训练日志和 state 后， +才能形成新的 GradScaler 运行证据。 + +完整 100+100-step 工程 smoke 属于训练操作，本次 Day 2 本地整合不会执行。后续只有 +得到单独授权后才能运行： + +```bash +bash scripts/smoke_engineering_100steps.sh --mode all +``` + +默认 run 路径为 `/mnt/ect_project/runs/engineering-smoke/`。总 batch 固定为 10， +所以 fresh 1 kimg 为 100 optimizer updates，resume 到 2 kimg 再增加 100 updates。 + +## 6. 正式固定 ECT baseline + +正式 baseline 必须另行定义并冻结以下信息： + +- 公共代码 SHA； +- 数据和 checkpoint 内容验证； +- batch、优化器、学习率、seed、schedule 和训练时长； +- AMP/GradScaler 是否实际启用； +- sampling 与 evaluation 协议； +- 正式结果和重复实验。 + +不能将 `smoke_engineering_100steps.sh` 的结果重命名或描述为正式 baseline。 + +## 7. 安全检查 + +提交工程改动前运行： + +```bash +bash -n setup_env.sh +bash -n prepare_data.sh +bash -n download_checkpoint.sh +bash -n scripts/smoke_engineering_100steps.sh +python -m compileall scripts training ct_train.py +python ct_train.py --help +git diff --check +git diff --stat +git status +``` + +数据或 checkpoint 验证失败时应立即停止，不得继续训练。正式训练、完整工程 smoke、 +FID 和 KID 都不属于本地整合检查。 diff --git a/download_checkpoint.sh b/download_checkpoint.sh new file mode 100755 index 0000000..ba10d45 --- /dev/null +++ b/download_checkpoint.sh @@ -0,0 +1,125 @@ +#!/usr/bin/env bash + +# Download and verify the official EDM CIFAR-10 checkpoint used for ECT transfer. + +set -Eeuo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENV_NAME="${ECT_ENV_NAME:-ect}" +PROJECT_ROOT="${ECT_PROJECT_ROOT:-/mnt/ect_project}" +URL="${ECT_EDM_CHECKPOINT_URL:-https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-uncond-vp.pkl}" +OUTPUT="${ECT_TRANSFER_PATH:-${PROJECT_ROOT}/pretrained/edm-cifar10-32x32-uncond-vp.pkl}" +REPORT_PATH="${ECT_CHECKPOINT_REPORT:-${PROJECT_ROOT}/runs/day2/checkpoint.json}" +EXPECTED_SHA256="${ECT_EDM_CHECKPOINT_SHA256:-4d5dcc1f1d0d41c8934ad21626eeddbdc0460182becf9fc059a0631b1eedb4da}" +CHECK_ONLY=0 +FORCE=0 + +fail() { + printf '[download_checkpoint] ERROR: %s\n' "$*" >&2 + exit 1 +} + +run_python() { + if [[ "${CONDA_DEFAULT_ENV:-}" == "${ENV_NAME}" ]]; then + python "$@" + elif command -v conda >/dev/null 2>&1; then + conda run --no-capture-output -n "${ENV_NAME}" python "$@" + elif command -v mamba >/dev/null 2>&1; then + mamba run --no-capture-output -n "${ENV_NAME}" python "$@" + else + fail "activate '${ENV_NAME}' or install conda/mamba first" + fi +} + +usage() { + cat <<'EOF' +Usage: bash download_checkpoint.sh [options] + + --output PATH Destination checkpoint + --url URL Download URL + --sha256 HASH Optional expected SHA-256 + --check-only Verify without downloading + --force Replace the existing checkpoint + --report PATH Validation report (default: /mnt/ect_project/runs/day2/checkpoint.json) + -h, --help Show this help +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --output) + [[ $# -ge 2 ]] || fail "--output requires a value" + OUTPUT="$2" + shift 2 + ;; + --url) + [[ $# -ge 2 ]] || fail "--url requires a value" + URL="$2" + shift 2 + ;; + --sha256) + [[ $# -ge 2 ]] || fail "--sha256 requires a value" + EXPECTED_SHA256="$2" + shift 2 + ;; + --check-only) + CHECK_ONLY=1 + shift + ;; + --force) + FORCE=1 + shift + ;; + --report) + [[ $# -ge 2 ]] || fail "--report requires a value" + REPORT_PATH="$2" + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + fail "unknown option: $1" + ;; + esac +done + +VERIFY_ARGS=( + "${ROOT_DIR}/scripts/verify_assets.py" + checkpoint + --path "${OUTPUT}" + --expected-sha256 "${EXPECTED_SHA256}" + --output "${REPORT_PATH}" +) + +if [[ "${CHECK_ONLY}" -eq 1 ]]; then + run_python "${VERIFY_ARGS[@]}" + exit 0 +fi + +if [[ -f "${OUTPUT}" && "${FORCE}" -eq 0 ]]; then + printf '[download_checkpoint] Checkpoint already exists; verifying it: %s\n' "${OUTPUT}" + run_python "${VERIFY_ARGS[@]}" + exit 0 +fi + +mkdir -p "$(dirname "${OUTPUT}")" "$(dirname "${REPORT_PATH}")" +if [[ -e "${OUTPUT}" ]]; then + [[ "${FORCE}" -eq 1 ]] || fail "output already exists: ${OUTPUT}" + rm -f "${OUTPUT}" +fi + +PARTIAL="${OUTPUT}.part" +printf '[download_checkpoint] Downloading official EDM checkpoint...\n' +if command -v curl >/dev/null 2>&1; then + curl --fail --location --retry 3 --continue-at - --output "${PARTIAL}" "${URL}" +elif command -v wget >/dev/null 2>&1; then + wget --continue --tries=3 --output-document="${PARTIAL}" "${URL}" +else + fail "curl or wget is required" +fi +mv "${PARTIAL}" "${OUTPUT}" + +run_python "${VERIFY_ARGS[@]}" +printf '[download_checkpoint] Checkpoint ready: %s\n' "${OUTPUT}" diff --git a/prepare_data.sh b/prepare_data.sh new file mode 100755 index 0000000..7d9bf25 --- /dev/null +++ b/prepare_data.sh @@ -0,0 +1,147 @@ +#!/usr/bin/env bash + +# Download CIFAR-10, verify the official archive, and convert it to EDM ZIP format. + +set -Eeuo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENV_NAME="${ECT_ENV_NAME:-ect}" +PROJECT_ROOT="${ECT_PROJECT_ROOT:-/mnt/ect_project}" +URL="${ECT_CIFAR10_URL:-https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz}" +TARBALL="${ECT_CIFAR10_TARBALL:-${PROJECT_ROOT}/datasets/cifar-10-python.tar.gz}" +OUTPUT="${ECT_DATA_PATH:-${PROJECT_ROOT}/datasets/cifar10-32x32.zip}" +REPORT_PATH="${ECT_DATA_REPORT:-${PROJECT_ROOT}/runs/day2/dataset.json}" +EXPECTED_MD5="c58f30108f718f92721af3b95e74349a" +CHECK_ONLY=0 +FORCE=0 + +fail() { + printf '[prepare_data] ERROR: %s\n' "$*" >&2 + exit 1 +} + +run_python() { + if [[ "${CONDA_DEFAULT_ENV:-}" == "${ENV_NAME}" ]]; then + python "$@" + elif command -v conda >/dev/null 2>&1; then + conda run --no-capture-output -n "${ENV_NAME}" python "$@" + elif command -v mamba >/dev/null 2>&1; then + mamba run --no-capture-output -n "${ENV_NAME}" python "$@" + else + fail "activate '${ENV_NAME}' or install conda/mamba first" + fi +} + +usage() { + cat <<'EOF' +Usage: bash prepare_data.sh [options] + + --output PATH Output dataset ZIP + --tarball PATH CIFAR-10 Python tarball (basename must stay unchanged) + --url URL Download URL + --check-only Verify the prepared dataset without downloading + --force Rebuild the output dataset + --report PATH Validation report (default: /mnt/ect_project/runs/day2/dataset.json) + -h, --help Show this help +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --output) + [[ $# -ge 2 ]] || fail "--output requires a value" + OUTPUT="$2" + shift 2 + ;; + --tarball) + [[ $# -ge 2 ]] || fail "--tarball requires a value" + TARBALL="$2" + shift 2 + ;; + --url) + [[ $# -ge 2 ]] || fail "--url requires a value" + URL="$2" + shift 2 + ;; + --check-only) + CHECK_ONLY=1 + shift + ;; + --force) + FORCE=1 + shift + ;; + --report) + [[ $# -ge 2 ]] || fail "--report requires a value" + REPORT_PATH="$2" + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + fail "unknown option: $1" + ;; + esac +done + +VERIFY_ARGS=( + "${ROOT_DIR}/scripts/verify_assets.py" + dataset + --path "${OUTPUT}" + --tarball "${TARBALL}" + --expected-md5 "${EXPECTED_MD5}" + --expected-count 50000 + --expected-labels 50000 + --expected-resolution 32 + --output "${REPORT_PATH}" +) + +if [[ "${CHECK_ONLY}" -eq 1 ]]; then + run_python "${VERIFY_ARGS[@]}" + exit 0 +fi + +if [[ -f "${OUTPUT}" && "${FORCE}" -eq 0 ]]; then + printf '[prepare_data] Dataset already exists; verifying it: %s\n' "${OUTPUT}" + run_python "${VERIFY_ARGS[@]}" + exit 0 +fi + +[[ "$(basename "${TARBALL}")" == "cifar-10-python.tar.gz" ]] || \ + fail "the tarball basename must be cifar-10-python.tar.gz for dataset_tool.py" +mkdir -p "$(dirname "${TARBALL}")" "$(dirname "${OUTPUT}")" "$(dirname "${REPORT_PATH}")" + +if [[ ! -f "${TARBALL}" ]]; then + printf '[prepare_data] Downloading CIFAR-10...\n' + PARTIAL="${TARBALL}.part" + if command -v curl >/dev/null 2>&1; then + curl --fail --location --retry 3 --continue-at - --output "${PARTIAL}" "${URL}" + elif command -v wget >/dev/null 2>&1; then + wget --continue --tries=3 --output-document="${PARTIAL}" "${URL}" + else + fail "curl or wget is required" + fi + mv "${PARTIAL}" "${TARBALL}" +fi + +if command -v md5sum >/dev/null 2>&1; then + ACTUAL_MD5="$(md5sum "${TARBALL}" | awk '{print $1}')" +elif command -v md5 >/dev/null 2>&1; then + ACTUAL_MD5="$(md5 -q "${TARBALL}")" +else + fail "md5sum or md5 is required for archive verification" +fi +[[ "${ACTUAL_MD5}" == "${EXPECTED_MD5}" ]] || \ + fail "CIFAR-10 MD5 mismatch: expected ${EXPECTED_MD5}, got ${ACTUAL_MD5}" + +if [[ -e "${OUTPUT}" ]]; then + [[ "${FORCE}" -eq 1 ]] || fail "output already exists: ${OUTPUT}" + rm -f "${OUTPUT}" +fi + +printf '[prepare_data] Converting CIFAR-10 to EDM format...\n' +run_python "${ROOT_DIR}/dataset_tool.py" --source "${TARBALL}" --dest "${OUTPUT}" +run_python "${VERIFY_ARGS[@]}" +printf '[prepare_data] Dataset ready: %s\n' "${OUTPUT}" diff --git a/scripts/check_environment.py b/scripts/check_environment.py new file mode 100755 index 0000000..4cd6dec --- /dev/null +++ b/scripts/check_environment.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Validate the frozen Day 1 ECT runtime and report the observed environment.""" + +import argparse +import importlib +import importlib.metadata +import json +import platform +import subprocess +import sys +from pathlib import Path + + +# env.yml pins these packages directly. The Day 1 validated runtime additionally +# freezes huggingface-hub because newer API removals can break diffusers 0.26.3. +EXPECTED = { + "torch": "2.3.0", + "diffusers": "0.26.3", + "accelerate": "0.27.2", + "huggingface-hub": "0.23.4", +} +REQUIRED = [ + "numpy", + "scipy", + "Pillow", + "click", + "requests", + "psutil", + "tqdm", + "imageio", + "imageio-ffmpeg", + "pyspng", +] +IMPORTS = [ + "torch", + "numpy", + "scipy", + "PIL", + "click", + "requests", + "psutil", + "tqdm", + "imageio", + "pyspng", + "diffusers", + "accelerate", + "huggingface_hub", +] + + +def git_value(repo_root: Path, *args: str): + try: + return subprocess.check_output( + ["git", *args], cwd=repo_root, text=True, stderr=subprocess.DEVNULL + ).strip() + except (OSError, subprocess.CalledProcessError): + return None + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--allow-no-cuda", action="store_true") + parser.add_argument("--output", type=Path) + args = parser.parse_args() + + versions = {} + errors = [] + for package in [*EXPECTED, *REQUIRED]: + try: + actual = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + errors.append(f"missing package: {package}") + continue + versions[package] = actual + expected = EXPECTED.get(package) + if expected is not None: + if package == "torch": + matches = actual.startswith(expected) + else: + matches = actual == expected + if not matches: + errors.append(f"{package}: expected {expected}, got {actual}") + + import_status = {} + for module in IMPORTS: + try: + importlib.import_module(module) + import_status[module] = "ok" + except Exception as exc: # Import compatibility matters, not just installation. + import_status[module] = f"failed: {type(exc).__name__}: {exc}" + errors.append(f"cannot import {module}: {type(exc).__name__}: {exc}") + + try: + import torch + except ImportError: + torch = None + + cuda_available = bool(torch and torch.cuda.is_available()) + if not args.allow_no_cuda and not cuda_available: + errors.append("CUDA is unavailable") + if sys.version_info[:3] != (3, 9, 18): + errors.append(f"expected Python 3.9.18, got {platform.python_version()}") + + repo_root = Path(__file__).resolve().parents[1] + status = git_value(repo_root, "status", "--porcelain=v1", "--untracked-files=all") + report = { + "git_commit": git_value(repo_root, "rev-parse", "HEAD"), + "git_branch": git_value(repo_root, "branch", "--show-current"), + "git_dirty": bool(status), + "python": platform.python_version(), + "platform": platform.platform(), + "packages": versions, + "imports": import_status, + "cuda": { + "available": cuda_available, + "runtime": torch.version.cuda if torch else None, + "device_count": torch.cuda.device_count() if cuda_available else 0, + "devices": [ + torch.cuda.get_device_name(index) + for index in range(torch.cuda.device_count()) + ] + if cuda_available + else [], + }, + "status": "passed" if not errors else "failed", + "errors": errors, + } + rendered = json.dumps(report, indent=2, sort_keys=True) + print(rendered) + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + if errors: + raise SystemExit("environment validation failed: " + "; ".join(errors)) + + +if __name__ == "__main__": + main() diff --git a/scripts/smoke_engineering_100steps.sh b/scripts/smoke_engineering_100steps.sh new file mode 100755 index 0000000..c219ca2 --- /dev/null +++ b/scripts/smoke_engineering_100steps.sh @@ -0,0 +1,295 @@ +#!/usr/bin/env bash + +# Engineering connectivity test only. This is not the official fixed ECT baseline. +# The old collaborator FP32 evidence did not validate FP16 or GradScaler. + +set -Eeuo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +PROJECT_ROOT="${ECT_PROJECT_ROOT:-/mnt/ect_project}" +ENV_NAME="${ECT_ENV_NAME:-ect}" +TARBALL="${ECT_CIFAR10_TARBALL:-${PROJECT_ROOT}/datasets/cifar-10-python.tar.gz}" +DATA_PATH="${ECT_DATA_PATH:-${PROJECT_ROOT}/datasets/cifar10-32x32.zip}" +TRANSFER_PATH="${ECT_TRANSFER_PATH:-${PROJECT_ROOT}/pretrained/edm-cifar10-32x32-uncond-vp.pkl}" +RUNS_ROOT="${ECT_RUNS_ROOT:-${PROJECT_ROOT}/runs}" +MODE="all" +ACTION="run" +NUM_GPUS=1 +BATCH_GPU="${ECT_SMOKE_BATCH_GPU:-2}" +FP16="${ECT_SMOKE_FP16:-True}" +ENABLE_AMP="${ECT_SMOKE_ENABLE_AMP:-True}" +PORT="${ECT_DDP_PORT:-29501}" +RUN_ROOT="" +ALLOW_DIRTY=0 + +fail() { + printf '[smoke_engineering_100steps] ERROR: %s\n' "$*" >&2 + exit 1 +} + +run_in_env() { + if [[ "${CONDA_DEFAULT_ENV:-}" == "${ENV_NAME}" ]]; then + "$@" + elif command -v conda >/dev/null 2>&1; then + conda run --no-capture-output -n "${ENV_NAME}" "$@" + elif command -v mamba >/dev/null 2>&1; then + mamba run --no-capture-output -n "${ENV_NAME}" "$@" + else + fail "activate '${ENV_NAME}' or install conda/mamba first" + fi +} + +normalize_bool() { + case "${1,,}" in + true|1|yes) printf 'True\n' ;; + false|0|no) printf 'False\n' ;; + *) fail "expected a boolean value, got: $1" ;; + esac +} + +usage() { + cat <<'EOF' +Usage: bash scripts/smoke_engineering_100steps.sh [options] + +This script is an engineering connectivity test, not the official ECT baseline. +Its training mode runs 100 fresh optimizer updates and optionally 100 resumed +updates with formal FID/KID disabled. No training is performed by --check-only +or --dry-run. + + --mode MODE fresh, resume, or all (default: all) + --gpus N GPUs used by torchrun (default: 1; must divide batch 10) + --port PORT Local DDP rendezvous port (default: 29501) + --tarball PATH Official CIFAR-10 source tarball + --data PATH Prepared CIFAR-10 EDM ZIP + --transfer PATH Official EDM transfer checkpoint + --run-root PATH Persistent output root + --batch-gpu N Microbatch per GPU (default: 2) + --fp16 BOOL Network FP16 mode (default: True) + --enable-amp BOOL Public-baseline --enable_amp value (default: True) + --check-only Validate environment and assets; do not launch ct_train.py + --dry-run Validate assets and run ct_train.py --dry_run only + --allow-dirty Permit an actual training smoke from a dirty worktree + -h, --help Show this help + +Default persistent layout: + /mnt/ect_project/datasets + /mnt/ect_project/pretrained + /mnt/ect_project/runs + /mnt/ect_project/checkpoints + +The public baseline also accepts --amp and --enable_gradscaler as aliases of +--enable_amp. This script uses the canonical --enable_amp spelling. +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --mode) + [[ $# -ge 2 ]] || fail "--mode requires a value" + MODE="$2" + shift 2 + ;; + --gpus) + [[ $# -ge 2 ]] || fail "--gpus requires a value" + NUM_GPUS="$2" + shift 2 + ;; + --port) + [[ $# -ge 2 ]] || fail "--port requires a value" + PORT="$2" + shift 2 + ;; + --tarball) + [[ $# -ge 2 ]] || fail "--tarball requires a value" + TARBALL="$2" + shift 2 + ;; + --data) + [[ $# -ge 2 ]] || fail "--data requires a value" + DATA_PATH="$2" + shift 2 + ;; + --transfer) + [[ $# -ge 2 ]] || fail "--transfer requires a value" + TRANSFER_PATH="$2" + shift 2 + ;; + --run-root) + [[ $# -ge 2 ]] || fail "--run-root requires a value" + RUN_ROOT="$2" + shift 2 + ;; + --batch-gpu) + [[ $# -ge 2 ]] || fail "--batch-gpu requires a value" + BATCH_GPU="$2" + shift 2 + ;; + --fp16) + [[ $# -ge 2 ]] || fail "--fp16 requires a value" + FP16="$2" + shift 2 + ;; + --enable-amp) + [[ $# -ge 2 ]] || fail "--enable-amp requires a value" + ENABLE_AMP="$2" + shift 2 + ;; + --check-only) + ACTION="check" + shift + ;; + --dry-run) + ACTION="dry-run" + shift + ;; + --allow-dirty) + ALLOW_DIRTY=1 + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + fail "unknown option: $1" + ;; + esac +done + +[[ "${MODE}" == "fresh" || "${MODE}" == "resume" || "${MODE}" == "all" ]] || \ + fail "--mode must be fresh, resume, or all" +[[ "${NUM_GPUS}" =~ ^[0-9]+$ && "${NUM_GPUS}" -gt 0 ]] || fail "--gpus must be positive" +(( 10 % NUM_GPUS == 0 )) || fail "--gpus must divide the total batch size 10" +[[ "${BATCH_GPU}" =~ ^[0-9]+$ && "${BATCH_GPU}" -gt 0 ]] || fail "--batch-gpu must be positive" +(( (10 / NUM_GPUS) % BATCH_GPU == 0 )) || \ + fail "--batch-gpu must divide the per-GPU batch $((10 / NUM_GPUS))" +FP16="$(normalize_bool "${FP16}")" +ENABLE_AMP="$(normalize_bool "${ENABLE_AMP}")" + +cd "${ROOT_DIR}" + +verify_prerequisites() { + run_in_env python "${ROOT_DIR}/scripts/check_environment.py" + run_in_env python "${ROOT_DIR}/scripts/verify_assets.py" dataset \ + --path "${DATA_PATH}" \ + --tarball "${TARBALL}" \ + --expected-count 50000 \ + --expected-labels 50000 \ + --expected-resolution 32 + run_in_env python "${ROOT_DIR}/scripts/verify_assets.py" checkpoint \ + --path "${TRANSFER_PATH}" \ + --expected-sha256 4d5dcc1f1d0d41c8934ad21626eeddbdc0460182becf9fc059a0631b1eedb4da +} + +verify_prerequisites +if [[ "${ACTION}" == "check" ]]; then + printf '[smoke_engineering_100steps] CHECK-ONLY PASSED; no training was launched.\n' + exit 0 +fi + +COMMON_ARGS=( + --data="${DATA_PATH}" + --cond=False + --arch=ddpmpp + --metrics=none + --batch=10 + --batch-gpu="${BATCH_GPU}" + --lr=0.0001 + --optim=RAdam + --dropout=0.2 + --augment=0.0 + --seed=2026 + --workers=1 + --cache=False + --tick=1 + --snap=1 + --dump=1 + --ckpt=1 + --double=250 + --sample_every=1000 + --eval_every=1000 + --fp16="${FP16}" + --enable_amp="${ENABLE_AMP}" + --tf32=False + --bench=True + --nosubdir +) + +launch_training() { + run_in_env torchrun \ + --nnodes=1 \ + --nproc_per_node="${NUM_GPUS}" \ + --rdzv_backend=c10d \ + --rdzv_endpoint="localhost:${PORT}" \ + "${ROOT_DIR}/ct_train.py" "$@" +} + +if [[ "${ACTION}" == "dry-run" ]]; then + DRY_RUN_ROOT="${RUN_ROOT:-${RUNS_ROOT}/engineering-smoke/dry-run}" + launch_training \ + --outdir="${DRY_RUN_ROOT}" \ + --duration=0.001 \ + --transfer="${TRANSFER_PATH}" \ + --desc=engineering-dry-run \ + --dry_run \ + "${COMMON_ARGS[@]}" + printf '[smoke_engineering_100steps] DRY-RUN PASSED; no output directory or training run was created.\n' + exit 0 +fi + +GIT_COMMIT="$(git rev-parse HEAD)" +if [[ "${ALLOW_DIRTY}" -eq 0 ]] && [[ -n "$(git status --porcelain)" ]]; then + fail "actual smoke training requires a clean worktree; use --allow-dirty only for preliminary work" +fi +if [[ -z "${RUN_ROOT}" ]]; then + [[ "${MODE}" != "resume" ]] || fail "--run-root is required with --mode resume" + RUN_ROOT="${RUNS_ROOT}/engineering-smoke/${GIT_COMMIT:0:8}-$(date -u +%Y%m%dT%H%M%SZ)" +fi +FRESH_DIR="${RUN_ROOT}/fresh-100steps" +RESUME_DIR="${RUN_ROOT}/resume-100steps" +mkdir -p "${RUN_ROOT}" + +run_in_env python "${ROOT_DIR}/scripts/check_environment.py" \ + --output "${RUN_ROOT}/environment.json" + +if [[ "${MODE}" == "fresh" || "${MODE}" == "all" ]]; then + [[ ! -e "${FRESH_DIR}" ]] || fail "fresh run directory already exists: ${FRESH_DIR}" + printf '[smoke_engineering_100steps] Fresh engineering phase: steps=100, formal_metrics=disabled\n' + launch_training \ + --outdir="${FRESH_DIR}" \ + --duration=0.001 \ + --transfer="${TRANSFER_PATH}" \ + --desc=engineering-fresh-100steps \ + "${COMMON_ARGS[@]}" +fi + +if [[ "${MODE}" == "resume" || "${MODE}" == "all" ]]; then + STATE_PATH="${FRESH_DIR}/training-state-000001.pt" + SNAPSHOT_PATH="${FRESH_DIR}/network-snapshot-000001.pkl" + [[ -f "${STATE_PATH}" ]] || fail "fresh training state not found: ${STATE_PATH}" + [[ -f "${SNAPSHOT_PATH}" ]] || fail "matching fresh snapshot not found: ${SNAPSHOT_PATH}" + [[ ! -e "${RESUME_DIR}" ]] || fail "resume run directory already exists: ${RESUME_DIR}" + printf '[smoke_engineering_100steps] Resume engineering phase: additional_steps=100, formal_metrics=disabled\n' + launch_training \ + --outdir="${RESUME_DIR}" \ + --duration=0.002 \ + --resume="${STATE_PATH}" \ + --desc=engineering-resume-100steps \ + "${COMMON_ARGS[@]}" +fi + +VERIFY_ARGS=( + "${ROOT_DIR}/scripts/verify_smoke_run.py" + --fresh "${FRESH_DIR}" + --git-commit "${GIT_COMMIT}" + --expected-batch 10 + --expected-fp16 "${FP16}" + --expected-amp "${ENABLE_AMP}" + --output "${RUN_ROOT}/smoke_report.json" +) +if [[ "${MODE}" == "resume" || "${MODE}" == "all" ]]; then + VERIFY_ARGS+=(--resume "${RESUME_DIR}") +fi +run_in_env python "${VERIFY_ARGS[@]}" + +printf '[smoke_engineering_100steps] PASSED. Engineering-only result root: %s\n' "${RUN_ROOT}" diff --git a/scripts/verify_assets.py b/scripts/verify_assets.py new file mode 100755 index 0000000..7596872 --- /dev/null +++ b/scripts/verify_assets.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +"""Verify ECT datasets and transfer checkpoints without changing the assets.""" + +import argparse +import hashlib +import io +import json +import sys +import zipfile +from pathlib import Path + +import numpy as np +from PIL import Image + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +OFFICIAL_CIFAR10_MD5 = "c58f30108f718f92721af3b95e74349a" + + +def digest(path: Path, algorithm: str) -> str: + value = hashlib.new(algorithm) + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + value.update(chunk) + return value.hexdigest() + + +def verify_tarball(path: Path, expected_md5: str) -> dict: + if not path.is_file(): + raise ValueError(f"source tarball not found: {path}") + actual_md5 = digest(path, "md5") + if actual_md5.lower() != expected_md5.lower(): + raise ValueError( + f"CIFAR-10 tarball MD5 mismatch: expected {expected_md5}, got {actual_md5}" + ) + return { + "path": str(path.resolve()), + "bytes": path.stat().st_size, + "md5": actual_md5, + "expected_md5": expected_md5.lower(), + "md5_ok": True, + } + + +def verify_dataset( + path: Path, + tarball: Path, + expected_md5: str, + expected_count: int, + expected_labels: int, + expected_resolution: int, +) -> dict: + """Validate archive integrity, every image header, labels, and project loading.""" + if not path.is_file(): + raise ValueError(f"dataset not found: {path}") + + source = verify_tarball(tarball, expected_md5) + sample_names = [] + try: + with zipfile.ZipFile(path, "r") as archive: + corrupt = archive.testzip() + if corrupt is not None: + raise ValueError(f"ZIP CRC failure: {corrupt}") + + names = archive.namelist() + images = sorted(name for name in names if name.lower().endswith(".png")) + if len(images) != expected_count: + raise ValueError(f"expected {expected_count} PNGs, found {len(images)}") + if "dataset.json" not in names: + raise ValueError("dataset.json is missing") + + metadata = json.loads(archive.read("dataset.json")) + labels = metadata.get("labels") + if not isinstance(labels, list): + raise ValueError("dataset.json must contain a labels list") + if len(labels) != expected_labels: + raise ValueError( + f"expected {expected_labels} labels, found {len(labels)}" + ) + label_names = [entry[0] for entry in labels if isinstance(entry, list) and len(entry) == 2] + if len(label_names) != expected_labels: + raise ValueError("every label must be a [filename, class] pair") + if sorted(label_names) != images: + raise ValueError("dataset labels do not map one-to-one to the PNG entries") + + expected_size = (expected_resolution, expected_resolution) + for image_name in images: + with Image.open(io.BytesIO(archive.read(image_name))) as image: + if image.size != expected_size: + raise ValueError( + f"{image_name} is {image.size}, expected {expected_size}" + ) + if image.mode != "RGB": + raise ValueError(f"{image_name} is {image.mode}, expected RGB") + + sample_names = [images[0], images[len(images) // 2], images[-1]] + except (OSError, ValueError, json.JSONDecodeError, zipfile.BadZipFile) as exc: + raise ValueError(f"dataset archive verification failed: {exc}") from exc + + # Import only after archive checks so failures clearly identify the project loader stage. + from training.dataset import ImageFolderDataset + + dataset = ImageFolderDataset(path=str(path), use_labels=True) + try: + if len(dataset) != expected_count: + raise ValueError( + f"ImageFolderDataset length is {len(dataset)}, expected {expected_count}" + ) + if not dataset.has_labels: + raise ValueError("ImageFolderDataset did not expose class labels") + + loader_samples = [] + for index in (0, len(dataset) // 2, len(dataset) - 1): + image, label = dataset[index] + if image.shape != (3, expected_resolution, expected_resolution): + raise ValueError(f"dataset sample {index} has shape {image.shape}") + if image.dtype != np.uint8: + raise ValueError(f"dataset sample {index} has dtype {image.dtype}") + loader_samples.append( + { + "index": index, + "shape": list(image.shape), + "dtype": str(image.dtype), + "label_shape": list(label.shape), + } + ) + finally: + dataset.close() + + return { + "type": "dataset", + "path": str(path.resolve()), + "bytes": path.stat().st_size, + "sha256": digest(path, "sha256"), + "sha256_policy": "record_only_not_a_unique_content_requirement", + "source_tarball": source, + "zip_crc_ok": True, + "images": expected_count, + "labels": expected_labels, + "resolution": [expected_resolution, expected_resolution], + "color_mode": "RGB", + "archive_samples": sample_names, + "imagefolderdataset": { + "length": expected_count, + "has_labels": True, + "samples": loader_samples, + }, + } + + +def verify_checkpoint(path: Path, expected_sha256: str) -> dict: + if not path.is_file(): + raise ValueError(f"checkpoint not found: {path}") + size = path.stat().st_size + if size < 10 * 1024 * 1024: + raise ValueError(f"checkpoint is unexpectedly small: {size} bytes") + with path.open("rb") as handle: + prefix = handle.read(2) + if not prefix.startswith(b"\x80"): + raise ValueError("checkpoint does not look like a binary pickle") + actual_sha256 = digest(path, "sha256") + if expected_sha256 and actual_sha256.lower() != expected_sha256.lower(): + raise ValueError( + f"checkpoint SHA-256 mismatch: expected {expected_sha256}, got {actual_sha256}" + ) + return { + "type": "checkpoint", + "path": str(path.resolve()), + "bytes": size, + "sha256": actual_sha256, + "expected_sha256": expected_sha256.lower() if expected_sha256 else None, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(dest="asset_type", required=True) + + dataset_parser = subparsers.add_parser("dataset") + dataset_parser.add_argument("--path", type=Path, required=True) + dataset_parser.add_argument("--tarball", type=Path, required=True) + dataset_parser.add_argument("--expected-md5", default=OFFICIAL_CIFAR10_MD5) + dataset_parser.add_argument("--expected-count", type=int, default=50000) + dataset_parser.add_argument("--expected-labels", type=int, default=50000) + dataset_parser.add_argument("--expected-resolution", type=int, default=32) + dataset_parser.add_argument("--output", type=Path) + + checkpoint_parser = subparsers.add_parser("checkpoint") + checkpoint_parser.add_argument("--path", type=Path, required=True) + checkpoint_parser.add_argument("--expected-sha256", default="") + checkpoint_parser.add_argument("--output", type=Path) + args = parser.parse_args() + + try: + if args.asset_type == "dataset": + report = verify_dataset( + args.path, + args.tarball, + args.expected_md5, + args.expected_count, + args.expected_labels, + args.expected_resolution, + ) + else: + report = verify_checkpoint(args.path, args.expected_sha256) + except ValueError as exc: + raise SystemExit(f"asset verification failed: {exc}") from exc + + report["status"] = "passed" + rendered = json.dumps(report, indent=2, sort_keys=True) + print(rendered) + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/scripts/verify_smoke_run.py b/scripts/verify_smoke_run.py new file mode 100755 index 0000000..afe5501 --- /dev/null +++ b/scripts/verify_smoke_run.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +"""Verify outputs from the engineering-only fresh/resume connectivity test.""" + +import argparse +import json +from pathlib import Path + + +REQUIRED_FILES = [ + "training_options.json", + "log.txt", + "stats.jsonl", + "network-snapshot-latest.pkl", + "training-state-latest.pt", + "final.png", +] + + +def parse_bool(value): + normalized = value.lower() + if normalized in {"true", "1", "yes"}: + return True + if normalized in {"false", "0", "no"}: + return False + raise argparse.ArgumentTypeError(f"expected a boolean, got {value}") + + +def inspect_run( + path: Path, + expected_kimg: int, + resumed: bool, + expected_batch: int, + expected_fp16: bool, + expected_amp: bool, +): + missing = [name for name in REQUIRED_FILES if not (path / name).is_file()] + if missing: + raise SystemExit(f"{path} is missing: {', '.join(missing)}") + + numbered_snapshots = sorted(path.glob("network-snapshot-[0-9][0-9][0-9][0-9][0-9][0-9].pkl")) + numbered_states = sorted(path.glob("training-state-[0-9][0-9][0-9][0-9][0-9][0-9].pt")) + if not numbered_snapshots or not numbered_states: + raise SystemExit(f"{path} has no numbered snapshot/state pair") + + options = json.loads((path / "training_options.json").read_text(encoding="utf-8")) + if options.get("total_kimg") != expected_kimg: + raise SystemExit( + f"{path}: expected total_kimg={expected_kimg}, got {options.get('total_kimg')}" + ) + if options.get("batch_size") != expected_batch: + raise SystemExit(f"{path}: expected batch_size={expected_batch}") + if options.get("metrics") != []: + raise SystemExit(f"{path}: smoke test must disable formal metrics") + actual_fp16 = options.get("network_kwargs", {}).get("use_fp16") + if actual_fp16 is not expected_fp16: + raise SystemExit( + f"{path}: expected network FP16={expected_fp16}, got {actual_fp16}" + ) + actual_amp = options.get("enable_amp") + if actual_amp is not expected_amp: + raise SystemExit(f"{path}: expected enable_amp={expected_amp}, got {actual_amp}") + + log_text = (path / "log.txt").read_text(encoding="utf-8", errors="replace") + if "Exiting..." not in log_text: + raise SystemExit(f"{path}: training did not exit cleanly") + if resumed and "Loading training state from" not in log_text: + raise SystemExit(f"{path}: resume state was not loaded") + + stats_lines = [ + line + for line in (path / "stats.jsonl").read_text(encoding="utf-8").splitlines() + if line.strip() + ] + if not stats_lines: + raise SystemExit(f"{path}: stats.jsonl is empty") + last_stats = json.loads(stats_lines[-1]) + last_progress = last_stats.get("Progress/kimg", {}).get("mean") + if last_progress is None or float(last_progress) < expected_kimg: + raise SystemExit( + f"{path}: expected progress >= {expected_kimg} kimg, got {last_progress}" + ) + + return { + "path": str(path.resolve()), + "resumed": resumed, + "total_kimg": expected_kimg, + "batch_size": expected_batch, + "fp16": actual_fp16, + "gradscaler": actual_amp, + "formal_metrics": False, + "last_progress_kimg": last_progress, + "numbered_snapshots": [item.name for item in numbered_snapshots], + "numbered_states": [item.name for item in numbered_states], + "latest_snapshot_bytes": (path / "network-snapshot-latest.pkl").stat().st_size, + "latest_state_bytes": (path / "training-state-latest.pt").stat().st_size, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--fresh", type=Path, required=True) + parser.add_argument("--resume", type=Path) + parser.add_argument("--git-commit", required=True) + parser.add_argument("--expected-batch", type=int, default=10) + parser.add_argument("--expected-fp16", type=parse_bool, default=True) + parser.add_argument("--expected-amp", type=parse_bool, default=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + common = { + "expected_batch": args.expected_batch, + "expected_fp16": args.expected_fp16, + "expected_amp": args.expected_amp, + } + runs = [inspect_run(args.fresh, expected_kimg=1, resumed=False, **common)] + if args.resume: + runs.append(inspect_run(args.resume, expected_kimg=2, resumed=True, **common)) + + report = { + "status": "passed", + "test_kind": "engineering_connectivity_only", + "official_ect_baseline": False, + "legacy_fp32_evidence_validates_gradscaler": False, + "git_commit": args.git_commit, + "smoke_steps_per_phase": 100, + "runs": runs, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/setup_env.sh b/setup_env.sh new file mode 100755 index 0000000..36a1732 --- /dev/null +++ b/setup_env.sh @@ -0,0 +1,127 @@ +#!/usr/bin/env bash + +# Set up the identical Conda runtime and directory layout in each MatrixCloud container. + +set -Eeuo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENV_NAME="${ECT_ENV_NAME:-ect}" +CONDA_CONFIG="${ECT_CONDA_CONFIG:-${ROOT_DIR}/conda-matpool.yml}" +PROJECT_ROOT="${ECT_PROJECT_ROOT:-/mnt/ect_project}" +REPORT_PATH="${ECT_ENV_REPORT:-${PROJECT_ROOT}/runs/day2/environment.json}" +MANAGER="" +UPDATE=0 +CHECK_ONLY=0 +ALLOW_NO_CUDA=0 + +fail() { + printf '[setup_env] ERROR: %s\n' "$*" >&2 + exit 1 +} + +usage() { + cat <<'EOF' +Usage: bash setup_env.sh [options] + + --name NAME Environment name (default: ect) + --manager CMD conda or mamba (default: auto-detect) + --update Reconcile an existing environment with env.yml + --check-only Only validate the existing environment + --allow-no-cuda Do not fail when the current machine has no visible GPU + --report PATH Environment report (default: /mnt/ect_project/runs/day2/environment.json) + -h, --help Show this help +EOF +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --name) + [[ $# -ge 2 ]] || fail "--name requires a value" + ENV_NAME="$2" + shift 2 + ;; + --manager) + [[ $# -ge 2 ]] || fail "--manager requires a value" + MANAGER="$2" + shift 2 + ;; + --update) + UPDATE=1 + shift + ;; + --check-only) + CHECK_ONLY=1 + shift + ;; + --allow-no-cuda) + ALLOW_NO_CUDA=1 + shift + ;; + --report) + [[ $# -ge 2 ]] || fail "--report requires a value" + REPORT_PATH="$2" + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + fail "unknown option: $1" + ;; + esac +done + +if [[ -z "${MANAGER}" ]]; then + if command -v conda >/dev/null 2>&1; then + MANAGER="conda" + elif command -v mamba >/dev/null 2>&1; then + MANAGER="mamba" + else + fail "conda or mamba was not found" + fi +fi +command -v "${MANAGER}" >/dev/null 2>&1 || fail "cannot find ${MANAGER}" +[[ -f "${CONDA_CONFIG}" ]] || fail "Conda configuration not found: ${CONDA_CONFIG}" + +run_manager() { + CONDARC="${CONDA_CONFIG}" "${MANAGER}" "$@" +} + +mkdir -p \ + "${PROJECT_ROOT}/datasets" \ + "${PROJECT_ROOT}/pretrained" \ + "${PROJECT_ROOT}/runs" \ + "${PROJECT_ROOT}/checkpoints" \ + "${PROJECT_ROOT}/cache" \ + "$(dirname "${REPORT_PATH}")" + +ENV_EXISTS=0 +if run_manager run -n "${ENV_NAME}" python --version >/dev/null 2>&1; then + ENV_EXISTS=1 +fi + +if [[ "${CHECK_ONLY}" -eq 0 ]]; then + if [[ "${ENV_EXISTS}" -eq 0 ]]; then + printf '[setup_env] Creating environment %s...\n' "${ENV_NAME}" + run_manager env create --name "${ENV_NAME}" --file "${ROOT_DIR}/env.yml" + elif [[ "${UPDATE}" -eq 1 ]]; then + printf '[setup_env] Updating environment %s...\n' "${ENV_NAME}" + run_manager env update --name "${ENV_NAME}" --file "${ROOT_DIR}/env.yml" --prune + else + printf '[setup_env] Environment %s already exists; validating it unchanged.\n' "${ENV_NAME}" + fi +elif [[ "${ENV_EXISTS}" -eq 0 ]]; then + fail "environment '${ENV_NAME}' does not exist" +fi + +CHECK_ARGS=( + "${ROOT_DIR}/scripts/check_environment.py" + --output "${REPORT_PATH}" +) +if [[ "${ALLOW_NO_CUDA}" -eq 1 ]]; then + CHECK_ARGS+=(--allow-no-cuda) +fi + +run_manager run --no-capture-output -n "${ENV_NAME}" python "${CHECK_ARGS[@]}" +printf '[setup_env] Environment is ready. Activate it with: %s activate %s\n' "${MANAGER}" "${ENV_NAME}"