Ensure Environment Secret Creation in Kubeflow

Add functionality to create and manage Kubernetes Secrets for
environment variables in the KubeflowExecutor class. The new
_method_ `_ensure_env_secret` verifies the existence of the
secret and creates or updates it as necessary. This allows
for better handling of sensitive environment configurations.

Update the template rendering to include secrets as environment
variables, enhancing the flexibility of environment variable
management.

- Implemented `_ensure_env_secret` to manage Secrets
- Updated YAML template to include `env_from_secrets`
- Added tests for secret creation and conflict handling

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

Author: jskswamy

nemo_run/core/execution/kubeflow.py

@@ -373,6 +373,9 @@ def _create_cluster_training_runtime(self, configmap_name: str, sha: str) -> str
         # Ensure storage objects exist prior to runtime creation
         self._ensure_storage()
 
+        # Ensure env secret exists prior to runtime creation
+        env_from_secrets: list[str] = self._ensure_env_secret(sha)
+
         template_vars = {
             "runtime_name": runtime_name,
             "namespace": self.namespace,
@@ -385,6 +388,7 @@ def _create_cluster_training_runtime(self, configmap_name: str, sha: str) -> str
             "gpus": self.gpus,
             "enable_tcpxo": self.enable_tcpxo,
             "storage_pvc_mounts": self._get_normalized_storage_mounts(),
+            "env_from_secrets": env_from_secrets,
         }
         rendered = fill_template(
             template_name="kubeflow_clustertrainingruntime.yaml.j2",
@@ -507,6 +511,37 @@ def _get_additional_files(self, task) -> dict[str, tuple[str, str]]:
 
         return files_to_stage
 
+    def _ensure_env_secret(self, sha: str) -> list[str]:
+        """Ensure a Secret exists when env_vars are configured; return list of envFrom names."""
+        if not self.env_vars:
+            return []
+        generated_secret_name = self._env_secret_name(sha)
+        try:
+            core_client = client.CoreV1Api()
+            body = client.V1Secret(
+                metadata=client.V1ObjectMeta(name=generated_secret_name, namespace=self.namespace),
+                string_data=self.env_vars,
+                type="Opaque",
+            )
+            core_client.create_namespaced_secret(namespace=self.namespace, body=body)
+            logger.info(f"Created Secret {generated_secret_name} in {self.namespace}")
+        except ApiException as e:
+            if e.status == 409:
+                # Secret exists; patch to ensure latest env_vars are reflected
+                try:
+                    patch_body = {"stringData": self.env_vars, "type": "Opaque"}
+                    core_client.patch_namespaced_secret(
+                        name=generated_secret_name, namespace=self.namespace, body=patch_body
+                    )
+                    logger.info(
+                        f"Patched Secret {generated_secret_name} with updated stringData in {self.namespace}"
+                    )
+                except Exception as patch_err:
+                    logger.warning(f"Failed to patch Secret {generated_secret_name}: {patch_err}")
+            else:
+                logger.warning(f"Failed to create Secret {generated_secret_name}: {e}")
+        return [generated_secret_name]
+
     def stage_files(self, task_dir: str, task=None) -> tuple[str, str]:
         """Stage files using the packager.
 
@@ -700,6 +735,11 @@ def _runtime_name(self, sha: str) -> str:
         identifier = self._get_experiment_identifier()
         return sanitize_kubernetes_name(f"nemo-runtime-{identifier}-{sha}")
 
+    def _env_secret_name(self, sha: str) -> str:
+        """Return a deterministic Secret name for env vars derived from experiment+sha."""
+        identifier = self._get_experiment_identifier()
+        return sanitize_kubernetes_name(f"nemo-env-{identifier}-{sha}")
+
     def _get_staged_file_path(self, filename: str) -> str:
         """Return path where a staged file would be mounted inside the container.
 

nemo_run/core/execution/templates/kubeflow_clustertrainingruntime.yaml.j2

@@ -125,6 +125,15 @@ spec:
                           value: /usr/local/nvidia/lib64
                         - name: NCCL_FASTRAK_LLCM_DEVICE_DIRECTORY
                           value: /dev/aperture_devices
+                        {% if secret_env_vars and secret_env_vars|length > 0 %}
+                        {% for sev in secret_env_vars %}
+                        - name: {{ sev.name }}
+                          valueFrom:
+                            secretKeyRef:
+                              name: {{ sev.secret_name }}
+                              key: {{ sev.key }}
+                        {% endfor %}
+                        {% endif %}
                       volumeMounts:
                         - name: workspace
                           mountPath: {{ workspace_mount_path }}

test/core/execution/test_kubeflow.py

@@ -138,6 +138,28 @@ def test_crt_template_renders_storage_pvc(self):
         assert "mountPath: /mnt/a" in rendered
         assert "readOnly: true" in rendered
 
+    def test_crt_template_renders_envfrom_secret(self):
+        rendered = fill_template(
+            template_name="kubeflow_clustertrainingruntime.yaml.j2",
+            variables={
+                "runtime_name": "rt",
+                "namespace": "ns",
+                "nodes": 1,
+                "image": "img",
+                "workspace_mount_path": "/src",
+                "configmap_name": "cfg",
+                "cpu_limit": None,
+                "memory_limit": None,
+                "gpus": None,
+                "enable_tcpxo": False,
+                "storage_pvc_mounts": [],
+                "env_from_secrets": ["my-secret"],
+            },
+        )
+
+        assert "envFrom:" in rendered
+        assert "name: my-secret" in rendered
+
     def test_pvc_creation_when_missing(self, mocker):
         # Configure an executor with a PVC that should be created
 
@@ -374,6 +396,66 @@ def test_kubeflow_executor_get_custom_trainer_fallback():
         assert mounted_path in " ".join(call_args.get("args", []))
 
 
+class TestEnvSecretHandling:
+    def test_secret_creation_without_conflict(self, mocker):
+        executor = KubeflowExecutor(namespace="default")
+        executor.packager = ConfigMapPackager()
+        executor.assign("exp-abc", "/tmp/exp", "task-1", "task_dir")
+
+        executor.env_vars = {"CONFIG_KEY1": "xyz", "FOO": "bar"}
+
+        mock_core = mocker.patch("kubernetes.client.CoreV1Api")
+        api = mock_core.return_value
+        # No exception on create (no conflict)
+        api.create_namespaced_secret.return_value = None
+
+        with patch("nemo_run.core.execution.kubeflow.fill_template") as ft:
+            ft.return_value = "apiVersion: v1\nkind: ClusterTrainingRuntime\nmetadata: {}"
+            with patch("kubernetes.client.CustomObjectsApi") as mock_coa:
+                coa = mock_coa.return_value
+                coa.create_cluster_custom_object.return_value = {}
+
+                executor._create_cluster_training_runtime(configmap_name="cfg", sha="beadfeed")
+
+        # Ensure create was called, and patch was NOT called
+        assert api.create_namespaced_secret.called
+        assert not api.patch_namespaced_secret.called
+
+        # Capture variables passed to template and assert env_from_secrets includes our secret
+        called_vars = ft.call_args[1]["variables"]
+        assert "env_from_secrets" in called_vars
+        assert isinstance(called_vars["env_from_secrets"], list)
+        assert len(called_vars["env_from_secrets"]) == 1
+
+    def test_secret_creation_and_patch_on_conflict(self, mocker):
+        executor = KubeflowExecutor(namespace="default")
+        executor.packager = ConfigMapPackager()
+        # Simulate assignment to set experiment identifier used in secret name
+        executor.assign("exp-xyz", "/tmp/exp", "task-1", "task_dir")
+
+        # Set env vars that should be converted to a Secret
+        executor.env_vars = {"CONFIG_KEY1": "abc", "OTHER": "val"}
+
+        # Mock k8s CoreV1Api to simulate create 409 then patch
+        mock_core = mocker.patch("kubernetes.client.CoreV1Api")
+        api = mock_core.return_value
+        from kubernetes.client.exceptions import ApiException
+
+        # First call: create raises 409 (already exists)
+        api.create_namespaced_secret.side_effect = ApiException(status=409)
+
+        # Run ensure function indirectly via _create_cluster_training_runtime
+        with patch("nemo_run.core.execution.kubeflow.fill_template") as ft:
+            ft.return_value = "apiVersion: v1\nkind: ClusterTrainingRuntime\nmetadata: {}"
+            with patch("kubernetes.client.CustomObjectsApi") as mock_coa:
+                coa = mock_coa.return_value
+                coa.create_cluster_custom_object.return_value = {}
+                # Should call patch on conflict
+                executor._create_cluster_training_runtime(configmap_name="cfg", sha="deadbeef")
+
+        assert api.patch_namespaced_secret.called
+
+
 def test_kubeflow_executor_create_trainjob():
     """Test create_trainjob method."""
     executor = KubeflowExecutor(nodes=1)