Implement Additional Packages Configuration for Executor

Add support for specifying additional package installations in the
KubeflowExecutor. This enhancement allows users to configure packages
to be installed within the training container, improving flexibility
for various training requirements.

- Introduced an AdditionalPackages dataclass to encapsulate
  installation parameters.
- Updated the CommandTrainer instantiation to accept additional
  packages if configured.
- Modified get_volume_name and get_pvc_claim_name to ensure
  lowercase names for Kubernetes compatibility.

This change enables more robust customization of the training
environment, facilitating the inclusion of necessary dependencies
directly through executor configuration.

Signed-off-by: Krishnaswamy Subramanian <subramk@thoughtworks.com>

Author: jskswamy

nemo_run/core/execution/kubeflow.py

@@ -16,7 +16,7 @@
 import logging
 import os
 import re
-from dataclasses import dataclass, field
+from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, Optional, Union
 
 import yaml
@@ -139,13 +139,35 @@ def to_template_fragment(self, index: int) -> dict[str, Any]:
     def get_volume_name(self, index: int) -> str:
         """Return a DNS-1123 safe volume name, defaulting to pvc-{index}."""
         base = self.name or f"pvc-{index}"
-        return sanitize_kubernetes_name(base)
+        return sanitize_kubernetes_name(base).lower()
 
     def get_pvc_claim_name(self) -> Optional[str]:
         """Return a DNS-1123 safe PVC claim name or None if unset."""
         if not self.pvc_claim_name:
             return None
-        return sanitize_kubernetes_name(self.pvc_claim_name)
+        return sanitize_kubernetes_name(self.pvc_claim_name).lower()
+
+
+@dataclass
+class AdditionalPackages:
+    """Optional package installation configuration for the training container.
+
+    Fields map directly to SDK `CommandTrainer` parameters.
+    """
+
+    packages_to_install: Optional[list[str]] = None
+    pip_index_urls: Optional[list[str]] = None
+    pip_extra_args: Optional[list[str]] = None
+
+    def as_trainer_kwargs(self) -> Dict[str, Any]:
+        """Return subset of kwargs for CommandTrainer based on configured fields."""
+        allowed = {"packages_to_install", "pip_index_urls", "pip_extra_args"}
+        return asdict(
+            self,
+            dict_factory=lambda items: {
+                k: (list(v) if isinstance(v, list) else v) for k, v in items if k in allowed and v
+            },
+        )
 
 
 @dataclass(kw_only=True)
@@ -224,6 +246,9 @@ class KubeflowExecutor(Executor):
 
     storage_mounts: list["StorageMount"] = field(default_factory=list)
 
+    #: Optional package installation configuration
+    additional_packages: Optional[AdditionalPackages] = None
+
     def __post_init__(self):
         """Validate executor configuration and setup Kubernetes access."""
         if self.nodes < 1:
@@ -533,12 +558,16 @@ def _get_custom_trainer(self, task) -> CommandTrainer:
         mounted_path = f"{self.workspace_mount_path}/{self.training_entry}"
         command, args = _build_trainer_command(task, mounted_path)
 
-        trainer = CommandTrainer(
-            command=command,
-            args=args,
-            num_nodes=self.nodes,
-            resources_per_node=resources_per_node,
-        )
+        trainer_kwargs: Dict[str, Any] = {
+            "command": command,
+            "args": args,
+            "num_nodes": self.nodes,
+            "resources_per_node": resources_per_node,
+        }
+        if self.additional_packages:
+            trainer_kwargs.update(self.additional_packages.as_trainer_kwargs())
+
+        trainer = CommandTrainer(**trainer_kwargs)
 
         logger.info(
             f"CommandTrainer created with command={trainer.command}, args={trainer.args}, "

test/core/execution/test_kubeflow.py

@@ -21,6 +21,7 @@
 
 from nemo_run.config import Partial, Script
 from nemo_run.core.execution.kubeflow import (
+    AdditionalPackages,
     KubeflowExecutor,
     StorageMount,
 )
@@ -787,3 +788,28 @@ def _dummy(x, y=2):
         assert "--nproc_per_node ${PET_NPROC_PER_NODE}" in args_joined
         assert "--rdzv_backend c10d" in args_joined
         assert "--rdzv_endpoint ${PET_MASTER_ADDR}:${PET_MASTER_PORT}" in args_joined
+
+
+def test_executor_additional_packages_forwarding():
+    script_task = Script(inline="python train.py")
+    executor = KubeflowExecutor(nodes=1, ntasks_per_node=4)
+    executor.packager = ConfigMapPackager()
+    executor.assign("exp-abc123", "/tmp/exp", "task-1", "task_dir")
+
+    executor.additional_packages = AdditionalPackages(
+        packages_to_install=["nemo==2.0.0", "deepspeed>=0.14.0"],
+        pip_index_urls=["https://pypi.org/simple", "https://extra/simple"],
+        pip_extra_args=["--no-cache-dir", "--find-links", "/wheels"],
+    )
+
+    with patch("nemo_run.core.execution.kubeflow.CommandTrainer") as mock_trainer:
+        instance = MagicMock()
+        mock_trainer.return_value = instance
+
+        res = executor._get_custom_trainer(script_task)
+
+        assert res == instance
+        kwargs = mock_trainer.call_args[1]
+        assert kwargs["packages_to_install"] == ["nemo==2.0.0", "deepspeed>=0.14.0"]
+        assert kwargs["pip_index_urls"] == ["https://pypi.org/simple", "https://extra/simple"]
+        assert kwargs["pip_extra_args"] == ["--no-cache-dir", "--find-links", "/wheels"]