From cf5268f2eca93b1f027b7d9cc51a2eabedf3f3bd Mon Sep 17 00:00:00 2001 From: Kiuk Chung Date: Mon, 4 Aug 2025 10:31:24 -0700 Subject: [PATCH] (bugfix)(torchx/runner) properly remove TORCHX_ prefix from env var before setting scheduler_params Summary: kwargs to `create_scheduler()` scheduler factory method can be specified via the env var `TORCHX_{ARGUMENT_NAME}=arg_value`. For example: `TORCHX_TIER=prod` would call the factory method as `create_scheduler(session_name, tier="prod")`. Python's `str.strip("torchx_")` doesn't remove the prefix `torchx_` but rather strips any characters `t`, `o`, `r`, `c`, `h`, `x`, `_` from the front and back of the string. Change to using `str.removeprefix("torchx_")` which does what we intended it to. Reviewed By: ethanbwaite Differential Revision: D79473894 --- torchx/runner/api.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/torchx/runner/api.py b/torchx/runner/api.py index 32f9a5314..67b0cbe75 100644 --- a/torchx/runner/api.py +++ b/torchx/runner/api.py @@ -129,9 +129,9 @@ def __init__( def _get_scheduler_params_from_env(self) -> Dict[str, str]: scheduler_params = {} for key, value in os.environ.items(): - lower_case_key = key.lower() - if lower_case_key.startswith("torchx_"): - scheduler_params[lower_case_key.strip("torchx_")] = value + key = key.lower() + if key.startswith("torchx_"): + scheduler_params[key.removeprefix("torchx_")] = value return scheduler_params def __enter__(self) -> "Self":