run_train not suited for model development

[clessig]

What happened?

run_train has (recently?) been changed to require the WEATHERGEN_STAGE environment variable for the stage , at least for train_continue and inference. This is not a good workflow for model development. It also seems more consistent to use a command line argument for this.

CC @grassesi @simone99n

What are the steps to reproduce the bug?

No response

Hedgedoc link to logs and more information. This ticket is public, do not attach files directly.

No response

[grassesi]

The WEATHERGEN_STAGE env variable was introduced only for infrastructure use. For model development the direct usage of the entrypoints (train/train_continue/inference) should be preferred. I am open to discuss this approach: We could also maintain only one central entrypoint (that could be called when running run_train as a script), then we could introduce a positional argument in a now unified cli and delegate depending on the value to subparsers (run_train <train/inference/...> --some-arg <> ...)

[clessig]

The WEATHERGEN_STAGE env variable was introduced only for infrastructure use. For model development the direct usage of the entrypoints (train/train_continue/inference) should be preferred. I am open to discuss this approach: We could also maintain only one central entrypoint (that could be called when running run_train as a script), then we could introduce a positional argument in a now unified cli and delegate depending on the value to subparsers (run_train <train/inference/...> --some-arg <> ...)

Does train/train_continue/inference work with uv run torchrun --nproc-per-node=4 ...?

[clessig]

The WEATHERGEN_STAGE env variable was introduced only for infrastructure use. For model development the direct usage of the entrypoints (train/train_continue/inference) should be preferred. I am open to discuss this approach: We could also maintain only one central entrypoint (that could be called when running run_train as a script), then we could introduce a positional argument in a now unified cli and delegate depending on the value to subparsers (run_train <train/inference/...> --some-arg <> ...)

Does train/train_continue/inference work with uv run torchrun --nproc-per-node=4 ...?

With ``uv run torchrun --nproc-per-node=4 train` I get:

E0201 10:04:47.681000 126149 /iopsstor/scratch/cscs/lessig/.cache/uv/archive-v0/LCy6vPFlnveVgEnOFjmMG/torch/distributed/elastic/multiprocessing/api.py:869] failed (exitcode: 2) local_rank: 0 (pid: 126248) of binary: /users/lessig/santis/WeatherGenerator/.venv/bin/python
Traceback (most recent call last):
  File "/users/lessig/santis/WeatherGenerator/.venv/bin/torchrun", line 10, in <module>
    sys.exit(main())
             ^^^^^^
  File "/users/lessig/santis/WeatherGenerator/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/users/lessig/santis/WeatherGenerator/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 918, in main
    run(args)
  File "/users/lessig/santis/WeatherGenerator/.venv/lib/python3.12/site-packages/torch/distributed/run.py", line 909, in run
    elastic_launch(
  File "/users/lessig/santis/WeatherGenerator/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 138, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/users/lessig/santis/WeatherGenerator/.venv/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 269, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
train FAILED
------------------------------------------------------------
Failures:
[1]:
  time      : 2026-02-01_10:04:47
  host      : nid005226
  rank      : 1 (local_rank: 1)
  exitcode  : 2 (pid: 126249)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
[2]:
  time      : 2026-02-01_10:04:47
  host      : nid005226
  rank      : 2 (local_rank: 2)
  exitcode  : 2 (pid: 126250)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
[3]:
  time      : 2026-02-01_10:04:47
  host      : nid005226
  rank      : 3 (local_rank: 3)
  exitcode  : 2 (pid: 126251)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2026-02-01_10:04:47
  host      : nid005226
  rank      : 0 (local_rank: 0)
  exitcode  : 2 (pid: 126248)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================

[shmh40]

I also see this, need to define a WEATHERGEN_STAGE environment variable with torchrun. Would be good not to have to do this.

[grassesi]

I have experimented a bit with torchrun myself and did not find a good solution (Only doing the allocation already with --ntasks 4 not failed for me, but it stalled ...). I see the point of needing to run run_train.py as a module. I will implement a fix.