fix(kfpytorch): add environment property to pass elastic config via env vars

This fixes bug #1 where overriding nproc_per_node via with_overrides would not
change the number of processes for single-node elastic tasks.

For single-node elastic tasks (task_type='python-task'), the _execute method
reads PET_NPROC_PER_NODE, PET_NNODES, PET_MAX_RESTARTS, and PET_MONITOR_INTERVAL
from environment variables. However, these env vars were never being set in the
task template during serialization.

The fix adds an environment property override to PytorchElasticFunctionTask that
includes the elastic config as environment variables. This ensures that when
task_config is modified via with_overrides, the elastic configuration is
correctly passed to the pod via environment variables.

Combined with the previous fix (dynamic task_type property), this now fully
supports:
- Bug #1: single-node (1 proc) -> single-node (multiple procs) override
- Bug #2: single-node (1 proc) -> multi-node (multiple procs) override

Co-Authored-By: carlos@exa.ai <carlosmarques.personal@gmail.com>

Author: devin-ai-integration[bot]

plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py

@@ -364,6 +364,33 @@ def task_type(self) -> str:
         """
         return self._ELASTIC_TASK_TYPE_STANDALONE if self._task_config.nnodes == 1 else self._ELASTIC_TASK_TYPE
 
+    @property
+    def environment(self) -> Dict[str, str]:
+        """
+        Dynamically compute environment variables based on current _task_config.
+
+        This property overrides the base class's environment to include elastic-specific
+        configuration as environment variables. This ensures that when task_config is
+        modified via with_overrides (e.g., in a dynamic task), the elastic configuration
+        is correctly passed to the pod via environment variables.
+
+        For single-node elastic tasks (task_type="python-task"), the _execute method reads
+        these environment variables to configure torch elastic's LaunchConfig. Without this
+        override, the environment variables would not be set, and _execute would fall back
+        to the original _task_config values from the decorator, ignoring any overrides.
+
+        This fixes a bug where overriding nproc_per_node via with_overrides would not
+        change the number of processes for single-node elastic tasks.
+        """
+        env = self._environment.copy() if self._environment else {}
+        # Always include elastic config in environment variables so that _execute
+        # can read the current (potentially overridden) values
+        env["PET_NNODES"] = str(self._task_config.nnodes)
+        env["PET_NPROC_PER_NODE"] = str(self._task_config.nproc_per_node)
+        env["PET_MAX_RESTARTS"] = str(self._task_config.max_restarts)
+        env["PET_MONITOR_INTERVAL"] = str(self._task_config.monitor_interval)
+        return env
+
     def _execute(self, **kwargs) -> Any:
         """
         Execute the task function using torch distributed's `elastic_launch`.

plugins/flytekit-kf-pytorch/tests/test_elastic_task.py

@@ -361,3 +361,52 @@ def elastic_range_task():
     # Update to single-node
     elastic_range_task._task_config = Elastic(nnodes=1, nproc_per_node=1)
     assert elastic_range_task.task_type == "python-task"
+
+
+def test_environment_includes_elastic_config() -> None:
+    """Test that environment property includes elastic config as environment variables.
+
+    This test verifies the fix for a bug where overriding nproc_per_node via with_overrides
+    would not change the number of processes for single-node elastic tasks. The fix adds
+    elastic config to environment variables so that _execute can read the current
+    (potentially overridden) values.
+    """
+    # Create a task with single-node config
+    @task(task_config=Elastic(nnodes=1, nproc_per_node=1, max_restarts=3, monitor_interval=10))
+    def single_node_task():
+        pass
+
+    # Check that environment includes elastic config
+    env = single_node_task.environment
+    assert env["PET_NNODES"] == "1"
+    assert env["PET_NPROC_PER_NODE"] == "1"
+    assert env["PET_MAX_RESTARTS"] == "3"
+    assert env["PET_MONITOR_INTERVAL"] == "10"
+
+    # Simulate what with_overrides does: update _task_config
+    single_node_task._task_config = Elastic(nnodes=1, nproc_per_node=4, max_restarts=5, monitor_interval=20)
+
+    # After updating _task_config, environment should reflect the new values
+    env = single_node_task.environment
+    assert env["PET_NNODES"] == "1"
+    assert env["PET_NPROC_PER_NODE"] == "4"
+    assert env["PET_MAX_RESTARTS"] == "5"
+    assert env["PET_MONITOR_INTERVAL"] == "20"
+
+
+def test_environment_preserves_existing_env_vars() -> None:
+    """Test that environment property preserves existing environment variables."""
+    # Create a task with custom environment variables
+    @task(
+        task_config=Elastic(nnodes=1, nproc_per_node=2),
+        environment={"CUSTOM_VAR": "custom_value", "ANOTHER_VAR": "another_value"}
+    )
+    def task_with_env():
+        pass
+
+    # Check that both custom and elastic env vars are present
+    env = task_with_env.environment
+    assert env["CUSTOM_VAR"] == "custom_value"
+    assert env["ANOTHER_VAR"] == "another_value"
+    assert env["PET_NNODES"] == "1"
+    assert env["PET_NPROC_PER_NODE"] == "2"