From 5b15433834d688dfdbaa32227ea68cac6ff65d65 Mon Sep 17 00:00:00 2001 From: Michael Tarnawa <18899420+mtar@users.noreply.github.com> Date: Thu, 22 Jan 2026 10:43:19 +0100 Subject: [PATCH 1/7] add argument --- src/weathergen/run_train.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py index 2af04edfe..4d886edb3 100644 --- a/src/weathergen/run_train.py +++ b/src/weathergen/run_train.py @@ -190,9 +190,15 @@ def train_with_args(argl: list[str], stream_dir: str | None): if __name__ == "__main__": - # Entry point for slurm script. - # Check whether --from_run_id passed as argument. - if any("--from_run_id" in arg for arg in sys.argv): - train_continue() - else: - train() + + stage = sys.argv[1] + + if stage == "train": + # Entry point for slurm script. + # Check whether --from_run_id passed as argument. + if any("--from_run_id" in arg for arg in sys.argv): + train_continue() + else: + train() + elif stage == "inference": + inference() From 8cfd4edd79ed0cdebfd8ac2ae8ddc830b7c7f722 Mon Sep 17 00:00:00 2001 From: Michael Tarnawa <18899420+mtar@users.noreply.github.com> Date: Thu, 22 Jan 2026 11:17:01 +0100 Subject: [PATCH 2/7] check stage argument --- src/weathergen/run_train.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py index 4d886edb3..1e88bb14a 100644 --- a/src/weathergen/run_train.py +++ b/src/weathergen/run_train.py @@ -28,7 +28,7 @@ def inference(): # By default, arguments from the command line are read. - inference_from_args(sys.argv[1:]) + inference_from_args(sys.argv[2:]) def inference_from_args(argl: list[str]): @@ -95,7 +95,7 @@ def train_continue() -> None: continue training. Defaults to None. Note: All model configurations are set in the function body. """ - train_continue_from_args(sys.argv[1:]) + train_continue_from_args(sys.argv[2:]) def train_continue_from_args(argl: list[str]): @@ -146,7 +146,7 @@ def train() -> None: continue training. Defaults to None. Note: All model configurations are set in the function body. """ - train_with_args(sys.argv[1:], None) + train_with_args(sys.argv[2:], None) def train_with_args(argl: list[str], stream_dir: str | None): @@ -191,6 +191,12 @@ def train_with_args(argl: list[str], stream_dir: str | None): if __name__ == "__main__": + if (args_count := len(sys.argv)) < 1: + logger.error(f"At least one argument expected, got {args_count - 1}") + raise SystemExit(2) + + stage = 'STAGE_NOT_SET' + stage = sys.argv[1] if stage == "train": @@ -202,3 +208,5 @@ def train_with_args(argl: list[str], stream_dir: str | None): train() elif stage == "inference": inference() + else: + logger.error(f'The selected stage is unknown, got {stage}.') From 54bc44f13c5289815ca18efeff18385274d63f99 Mon Sep 17 00:00:00 2001 From: Michael Tarnawa <18899420+mtar@users.noreply.github.com> Date: Thu, 22 Jan 2026 11:32:25 +0100 Subject: [PATCH 3/7] removed unnecessary code --- src/weathergen/run_train.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py index 1e88bb14a..3ebe041a3 100644 --- a/src/weathergen/run_train.py +++ b/src/weathergen/run_train.py @@ -195,8 +195,6 @@ def train_with_args(argl: list[str], stream_dir: str | None): logger.error(f"At least one argument expected, got {args_count - 1}") raise SystemExit(2) - stage = 'STAGE_NOT_SET' - stage = sys.argv[1] if stage == "train": From c38ed2da5ce9de0affd85c5f472392e34a8384c7 Mon Sep 17 00:00:00 2001 From: Michael Tarnawa <18899420+mtar@users.noreply.github.com> Date: Thu, 22 Jan 2026 15:14:06 +0100 Subject: [PATCH 4/7] arbitrary position arguments --- src/weathergen/run_train.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py index 3ebe041a3..daf133e01 100644 --- a/src/weathergen/run_train.py +++ b/src/weathergen/run_train.py @@ -28,7 +28,7 @@ def inference(): # By default, arguments from the command line are read. - inference_from_args(sys.argv[2:]) + inference_from_args(sys.argv[1:]) def inference_from_args(argl: list[str]): @@ -95,7 +95,7 @@ def train_continue() -> None: continue training. Defaults to None. Note: All model configurations are set in the function body. """ - train_continue_from_args(sys.argv[2:]) + train_continue_from_args(sys.argv[1:]) def train_continue_from_args(argl: list[str]): @@ -146,7 +146,7 @@ def train() -> None: continue training. Defaults to None. Note: All model configurations are set in the function body. """ - train_with_args(sys.argv[2:], None) + train_with_args(sys.argv[1:], None) def train_with_args(argl: list[str], stream_dir: str | None): @@ -191,20 +191,14 @@ def train_with_args(argl: list[str], stream_dir: str | None): if __name__ == "__main__": - if (args_count := len(sys.argv)) < 1: - logger.error(f"At least one argument expected, got {args_count - 1}") - raise SystemExit(2) - - stage = sys.argv[1] - - if stage == "train": + if any("train" in arg for arg in sys.argv): # Entry point for slurm script. # Check whether --from_run_id passed as argument. if any("--from_run_id" in arg for arg in sys.argv): train_continue() else: train() - elif stage == "inference": + elif any("inference" in arg for arg in sys.argv): inference() else: logger.error(f'The selected stage is unknown, got {stage}.') From 1d174d82f18ae5efe1af15ca584906f584027146 Mon Sep 17 00:00:00 2001 From: Michael Tarnawa <18899420+mtar@users.noreply.github.com> Date: Thu, 22 Jan 2026 15:52:33 +0100 Subject: [PATCH 5/7] Fix error text --- src/weathergen/run_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py index daf133e01..5c1ded4f1 100644 --- a/src/weathergen/run_train.py +++ b/src/weathergen/run_train.py @@ -201,4 +201,4 @@ def train_with_args(argl: list[str], stream_dir: str | None): elif any("inference" in arg for arg in sys.argv): inference() else: - logger.error(f'The selected stage is unknown, got {stage}.') + logger.error("No stage was found.") From a66c9c738ff2f5639071bbca64c04a73c69f467d Mon Sep 17 00:00:00 2001 From: Simon Grasse Date: Mon, 26 Jan 2026 02:15:27 +0100 Subject: [PATCH 6/7] get stage info from environment variable. --- src/weathergen/run_train.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py index 5c1ded4f1..e2d789767 100644 --- a/src/weathergen/run_train.py +++ b/src/weathergen/run_train.py @@ -12,6 +12,7 @@ """ import logging +import os import pdb import sys import time @@ -190,15 +191,20 @@ def train_with_args(argl: list[str], stream_dir: str | None): if __name__ == "__main__": + try: + stage = os.environ.get("WEATHERGEN_STAGE") + except KeyError as e: + msg = f"missing environment variable 'WEATHERGEN_STAGE'" + raise ValueError(msg) from e - if any("train" in arg for arg in sys.argv): + if stage == "train": # Entry point for slurm script. # Check whether --from_run_id passed as argument. if any("--from_run_id" in arg for arg in sys.argv): train_continue() else: train() - elif any("inference" in arg for arg in sys.argv): + elif stage == "inference": inference() else: logger.error("No stage was found.") From 57827cd48ce69a00eaed5a380948026c4c2dd3d8 Mon Sep 17 00:00:00 2001 From: Michael Tarnawa <18899420+mtar@users.noreply.github.com> Date: Mon, 26 Jan 2026 14:49:56 +0100 Subject: [PATCH 7/7] Update run_train.py --- src/weathergen/run_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py index e2d789767..a4ffb1f5b 100644 --- a/src/weathergen/run_train.py +++ b/src/weathergen/run_train.py @@ -194,7 +194,7 @@ def train_with_args(argl: list[str], stream_dir: str | None): try: stage = os.environ.get("WEATHERGEN_STAGE") except KeyError as e: - msg = f"missing environment variable 'WEATHERGEN_STAGE'" + msg = "missing environment variable 'WEATHERGEN_STAGE'" raise ValueError(msg) from e if stage == "train":