diff --git a/configs/production_training.yaml b/configs/production_training.yaml new file mode 100644 index 0000000..215b2be --- /dev/null +++ b/configs/production_training.yaml @@ -0,0 +1,24 @@ +# Full PPO training config for Brittle Star (HPC Production) +exp_name: "production_training" +seed: 1 +track: true +capture_video: true +save_model: true +checkpoint_frequency: 100 # not yet implemented in train.py but here for future use + +# Scaling for HPC (using A100 GPU slices) +num_envs: 128 +total_timesteps: 10000000 +num_steps: 128 +num_minibatches: 4 +update_epochs: 4 + +# Algorithm +learning_rate: 2.5e-4 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +clip_coef: 0.1 +ent_coef: 0.01 +vf_coef: 0.5 +cuda: true diff --git a/docs/HPC.md b/docs/HPC.md index e5d0794..2f23704 100644 --- a/docs/HPC.md +++ b/docs/HPC.md @@ -6,39 +6,15 @@ Full documentation: Choose the appropriate cluster before submitting a job with `module swap cluster/`. The default login cluster is **doduo**. -| Cluster | Type | Use case | -|-----------|-------------------------|-----------------------------------------------| -| `donphan` | Interactive / debug GPU | First-time setup, interactive debugging | -| `doduo` | CPU (default login) | Rapid iteration, CPU-only smoke tests | -| `joltik` | GPU (A100 ¼-slice) | Standard training runs | -| `accelgor`| GPU (A100 full) | Large-scale / long experiments | -| `litleo` | GPU | Alternative GPU option | - -> **Rule:** use at most **1 GPU per group at a time** on shared GPU clusters. > Check the current queue load at . ---- - ## Storage Overview -The HPC provides three filesystems for different purposes. Understanding this is critical to avoid filling up your home directory. - -| Variable | Typical size | Purpose | -|-----------------|-------------|------------------------------------------------------| -| `$VSC_HOME` | ~3 GB | Config files, SSH keys, project source code only | -| `$VSC_DATA` | ~25 GB | Persistent outputs: trained models, final logs | -| `$VSC_SCRATCH` | Large | Fast I/O during jobs: caches, intermediate files | - -**Important:** -- Clone the repository into `$VSC_HOME` — it is small in size and accessible from all clusters. -- All caches (pip, uv, matplotlib) must be redirected to `$VSC_SCRATCH` to avoid filling `$VSC_HOME`. - Run outputs are written to `$VSC_SCRATCH` during the job (fast I/O) and copied to `$VSC_DATA` at the end for persistence. - `$VSC_SCRATCH` may be purged periodically — do not use it as long-term storage. Check your quota: (Usage section). ---- - ## Initial Environment Setup Run **once** from a login shell on `donphan` after cloning the repository: @@ -56,21 +32,10 @@ cd 2026SEL3-project-BrittleStar bash scripts/hpc/install.sh ``` -This uses the official [`vsc-venv`](https://docs.hpc.ugent.be/Linux/setting_up_python_virtual_environments/#vsc-venv-python-virtual-environment-wrapper-script) wrapper to: -- Redirect caches to `$VSC_SCRATCH` (to preserve your `$VSC_HOME` quota) -- Load the EasyBuild modules listed in `env/hpc/modules.txt` (JAX, Flax, WandB, …) -- Create a per-cluster virtual environment in `$VSC_DATA` -- Pip-install the remaining packages from `env/hpc/requirements.txt` -- Register a Jupyter kernel named `SEL3 ()` - > **Note:** virtual environments are cluster-specific. Re-run the script when switching to a new cluster. ---- - ## Interactive Debugging on donphan -The `donphan` cluster provides quick access and is ideal for verifying your environment before submitting batch jobs. - ### Option A — Interactive shell session ```bash @@ -121,8 +86,6 @@ python src/train.py --config configs/hpc/smoke_test.yaml > **Warning:** JAX can only be loaded by one kernel at a time. Shut down other kernels before switching notebooks. ---- - ## Submitting Batch Training Jobs ```bash @@ -135,10 +98,9 @@ qsub scripts/hpc/train.pbs ``` The job script automatically: -- Redirects all caches to `$VSC_SCRATCH` -- Writes run outputs to `$VSC_SCRATCH/runs/` during the run -- Copies final results to `$VSC_DATA/runs/` on completion -- Writes PBS stdout/stderr to `$VSC_DATA/runs//job.out` / `job.err` +- Writes run outputs to `$VSC_SCRATCH/runs/` during the run using the `--run_dir` argument. This ensures that frequent I/O (like tensorboard logs and checkpoints) happens on the fastest available filesystem. +- Copies the final results to `$VSC_DATA/runs/` on completion for long-term persistence. +- Writes PBS stdout/stderr to `runs/brittlestar-ppo.o` / `.e` (standard PBS convention, relative to the project root). Monitor your jobs: @@ -148,8 +110,6 @@ qstat -f # detailed info for a specific job qdel # cancel a job ``` ---- - ## Managing Dependencies `env/hpc/requirements.txt` is auto-generated by CI whenever `pyproject.toml` changes. To regenerate locally: diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index 03c220e..edaf318 100644 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ b/src/brittle_star_project/dataclasses/PPOArgs.py @@ -13,6 +13,12 @@ class PPOArgs: # the name of this experiment exp_name: str = "brittle_star_ppo" + # the directory to save the experiment results + run_dir: str | None = None + + # how often to save checkpoints (0 to disable) + checkpoint_frequency: int = 0 + # seed of the experiment seed: int = 1 diff --git a/src/main.py b/src/main.py deleted file mode 100644 index ef7c36e..0000000 --- a/src/main.py +++ /dev/null @@ -1,4 +0,0 @@ -import jax - -if __name__ == "__main__": - print(jax.devices()) diff --git a/src/train.py b/src/train.py index c1a646d..4f049f2 100644 --- a/src/train.py +++ b/src/train.py @@ -46,6 +46,12 @@ def train(args: PPOArgs): run_name = f"{args.exp_name}__seed_{args.seed}__{int(time.time())}" print(f"running name: {run_name}") + if args.run_dir is None: + args.run_dir = f"runs/{run_name}" + + import os + os.makedirs(args.run_dir, exist_ok=True) + if args.track: import wandb @@ -58,7 +64,7 @@ def train(args: PPOArgs): save_code=True, ) - writer = SummaryWriter(f"runs/{run_name}") + writer = SummaryWriter(args.run_dir) writer.add_text( "hyperparameters", "|param|value|\n|---|---|\n" + "\n".join(f"|{k}|{v}|" for k, v in vars(args).items()), @@ -297,7 +303,7 @@ def train(args: PPOArgs): ) if args.save_model: - model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model" + model_path = f"{args.run_dir}/{args.exp_name}.cleanrl_model" with open(model_path, "wb") as f: f.write( flax.serialization.to_bytes( @@ -319,7 +325,7 @@ def train(args: PPOArgs): print("Saving loss plot...") plt.plot(losses) plt.title("PPO Loss, mean over minibatches") - plt.savefig(f"runs/{run_name}/{args.exp_name}_losses.png") + plt.savefig(f"{args.run_dir}/{args.exp_name}_losses.png") plt.close()