Add Command-Line Argument for Specifying Wandb Directory

[nqhq-lou]

Issue: This PR addresses issue #433, which involves adding a command-line argument to specify the wandb directory.

Changes:

1. Added --wandb_dir in mace.tools.arg_parser.build_default_arg_parser(). The arg default is None to preserve the default wandb directory settings.
2. Added keyword directory in mace.tools.torch_tools.init_wandb(). All associated function calls are adapted to the new keyword.

Code behavior:

with --wandb_dir:

• command:

  mace_run_train --name=dbg_MACE_train_11267359_0 \
  --log_dir=./runs/dbg_MACE_train_11267359_0 --model_dir=./runs/dbg_MACE_train_11267359_0 \
  --checkpoints_dir=./runs/dbg_MACE_train_11267359_0 --results_dir=./runs/dbg_MACE_train_11267359_0 \
  ... \
  --wandb --wandb_dir=./runs/dbg_MACE_train_11267359_0 --wandb_project=[project_name] --wandb_entity=[entity_name] --wandb_name=[wandb_name]

• directory structure:

  .
  ├── logs
  │   ├── dbg_MACE_train.11267359_0.err
  │   └── dbg_MACE_train.11267359_0.out
  ├── runs
  │   └── dbg_MACE_train_11267359_0
  │       ├── dbg_MACE_train_11267359_0_run-42_epoch-26.pt
  │       ├── dbg_MACE_train_11267359_0_run-42.log
  │       ├── dbg_MACE_train_11267359_0_run-42_train.txt
  │       ├── node_usage.11267359_0.log
  │       └── wandb
  │           ├── debug-internal.log -> run-20240617_152910-sw6ffqpr/logs/debug-internal.log
  │           ├── debug.log -> run-20240617_152910-sw6ffqpr/logs/debug.log
  │           ├── latest-run -> run-20240617_152910-sw6ffqpr
  │           └── run-20240617_152910-sw6ffqpr
  │               ├── files
  │               │   ├── conda-environment.yaml
  │               │   ├── config.yaml
  │               │   ├── output.log
  │               │   ├── requirements.txt
  │               │   ├── wandb-metadata.json
  │               │   └── wandb-summary.json
  │               ├── logs
  │               │   ├── debug-internal.log
  │               │   └── debug.log
  │               ├── run-sw6ffqpr.wandb
  │               └── tmp
  │                   └── code
  ├── train.sbatch.sh
  └── wandb
  └── debug-cli.zklou.log

• As expected the wandb metadata dir is relocated as specified by --wandb_dir

without --wandb_dir:

• command:

  mace_run_train --name=dbg_MACE_train_11267416_0 \
  --log_dir=./runs/dbg_MACE_train_11267416_0 --model_dir=./runs/dbg_MACE_train_11267416_0 \
  --checkpoints_dir=./runs/dbg_MACE_train_11267416_0 --results_dir=./runs/dbg_MACE_train_11267416_0 \
  ... \
  --wandb --wandb_project=[project_name] --wandb_entity=[entity_name] --wandb_name=[wandb_name]

• directory structure:

  .
  ├── logs
  │   ├── dbg_MACE_train.11267416_0.err
  │   └── dbg_MACE_train.11267416_0.out
  ├── runs
  │   └── dbg_MACE_train_11267416_0
  │       ├── dbg_MACE_train_11267416_0_run-42_epoch-60.pt
  │       ├── dbg_MACE_train_11267416_0_run-42.log
  │       ├── dbg_MACE_train_11267416_0_run-42_train.txt
  │       └── node_usage.11267416_0.log
  ├── train.sbatch.sh
  └── wandb
  ├── debug-cli.zklou.log
  ├── debug-internal.log -> run-20240617_153307-sfclv2b3/logs/debug-internal.log
  ├── debug.log -> run-20240617_153307-sfclv2b3/logs/debug.log
  ├── latest-run -> run-20240617_153307-sfclv2b3
  └── run-20240617_153307-sfclv2b3
      ├── files
      │   ├── conda-environment.yaml
      │   ├── config.yaml
      │   ├── output.log
      │   ├── requirements.txt
      │   ├── wandb-metadata.json
      │   └── wandb-summary.json
      ├── logs
      │   ├── debug-internal.log
      │   └── debug.log
      ├── run-sfclv2b3.wandb
      └── tmp
          └── code

• There is not impact on the default wandb logging directory, ensuring compatibility.

[nqhq-lou]

@ilyes319 Hi Ilyes, how does this PR look like? For the failed testing I checked the output logs and it tells some package errors. This PR is actually quite simple so I think maybe it would be safe to merge.

[ilyes319]

yep can you just update your fork with the latest develop, it should make the test pass.

[nqhq-lou]

Just merged branch develop into this branch! Please approve the testing workflow.