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.