From 15266dc3a325aa4168a60e7be470a201673aae53 Mon Sep 17 00:00:00 2001 From: danish9039 Date: Sun, 5 Jul 2026 01:33:52 +0530 Subject: [PATCH 1/2] runtimes: add restricted pss security contexts Signed-off-by: danish9039 --- .../torch-distributed-with-cache.yaml | 29 +++++++ .../files/runtimes/torch-distributed.yaml | 15 ++++ .../runtimes/torch_distributed_test.yaml | 11 +++ .../torch_distributed_with_cache_test.yaml | 38 +++++++++ cmd/data_cache/Dockerfile | 4 +- .../torch_distributed_with_cache.yaml | 29 +++++++ .../base/runtimes/torch_distributed.yaml | 15 ++++ pkg/initializers/dataset/cache.py | 18 +++++ pkg/initializers/dataset/cache_test.py | 78 +++++++++++++++++++ 9 files changed, 235 insertions(+), 2 deletions(-) diff --git a/charts/kubeflow-trainer/files/runtimes/torch-distributed-with-cache.yaml b/charts/kubeflow-trainer/files/runtimes/torch-distributed-with-cache.yaml index cb71b7960f..f5394a430d 100644 --- a/charts/kubeflow-trainer/files/runtimes/torch-distributed-with-cache.yaml +++ b/charts/kubeflow-trainer/files/runtimes/torch-distributed-with-cache.yaml @@ -38,9 +38,23 @@ spec: template: spec: serviceAccountName: kubeflow-trainer-cache-initializer + securityContext: + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault containers: - name: dataset-initializer image: {{ printf "ghcr.io/kubeflow/trainer/dataset-initializer:%s" (include "trainer.defaultImageTag" .) }} + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault env: - name: CACHE_IMAGE value: {{ include "trainer.runtimeImage" (list .Values.dataCache.cacheImage .) | quote }} @@ -60,9 +74,24 @@ spec: spec: template: spec: + securityContext: + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault containers: - name: node image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime + workingDir: /tmp + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault env: - name: TRAIN_JOB_NAME valueFrom: diff --git a/charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml b/charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml index a755e9c981..2d071bddd7 100644 --- a/charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml +++ b/charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml @@ -36,6 +36,21 @@ spec: spec: template: spec: + securityContext: + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault containers: - name: node image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime + workingDir: /tmp + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault diff --git a/charts/kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml b/charts/kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml index 60839e5600..26d21a437b 100644 --- a/charts/kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml +++ b/charts/kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml @@ -51,3 +51,14 @@ tests: - matchRegex: path: data["runtimes.yaml"] pattern: "name: torch-distributed" + + - it: Should set restricted PSS compatible security contexts + set: + runtimes: + torchDistributed: + enabled: true + asserts: + - matchRegex: + path: data["runtimes.yaml"] + pattern: |- + (?s)- name: node.*securityContext:\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault\s+containers:\s+- name: node\s+image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime\s+workingDir: /tmp\s+securityContext:\s+allowPrivilegeEscalation: false\s+capabilities:\s+drop:\s+- ALL\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault diff --git a/charts/kubeflow-trainer/tests/runtimes/torch_distributed_with_cache_test.yaml b/charts/kubeflow-trainer/tests/runtimes/torch_distributed_with_cache_test.yaml index d3b98eca59..fc9921ecbd 100644 --- a/charts/kubeflow-trainer/tests/runtimes/torch_distributed_with_cache_test.yaml +++ b/charts/kubeflow-trainer/tests/runtimes/torch_distributed_with_cache_test.yaml @@ -62,3 +62,41 @@ tests: - matchRegex: path: data["runtimes.yaml"] pattern: "ghcr.io/kubeflow/trainer/data-cache:v1.0.0" + + - it: Should have node job depend on dataset-initializer + set: + dataCache: + enabled: true + cacheImage: + registry: ghcr.io + repository: kubeflow/trainer/data-cache + tag: "v1.0.0" + runtimes: + torchDistributedWithCache: + enabled: true + asserts: + - matchRegex: + path: data["runtimes.yaml"] + pattern: |- + (?s)- name: node\s+dependsOn:\s+- name: dataset-initializer\s+status: Complete + + - it: Should set restricted PSS compatible security contexts + set: + dataCache: + enabled: true + cacheImage: + registry: ghcr.io + repository: kubeflow/trainer/data-cache + tag: "v1.0.0" + runtimes: + torchDistributedWithCache: + enabled: true + asserts: + - matchRegex: + path: data["runtimes.yaml"] + pattern: |- + (?s)- name: dataset-initializer.*securityContext:\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault\s+containers:\s+- name: dataset-initializer\s+image: ghcr.io/kubeflow/trainer/dataset-initializer:.*\s+securityContext:\s+allowPrivilegeEscalation: false\s+capabilities:\s+drop:\s+- ALL\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault + - matchRegex: + path: data["runtimes.yaml"] + pattern: |- + (?s)- name: node.*securityContext:\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault\s+containers:\s+- name: node\s+image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime\s+workingDir: /tmp\s+securityContext:\s+allowPrivilegeEscalation: false\s+capabilities:\s+drop:\s+- ALL\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault diff --git a/cmd/data_cache/Dockerfile b/cmd/data_cache/Dockerfile index e2ffc9f9d5..ebdcf5a82d 100644 --- a/cmd/data_cache/Dockerfile +++ b/cmd/data_cache/Dockerfile @@ -53,8 +53,8 @@ RUN apt-get update && \ COPY --from=builder /workspace/target/release/head /usr/local/bin/head COPY --from=builder /workspace/target/release/worker /usr/local/bin/worker -# Create non-root user -RUN groupadd -r cache_user && useradd -r -g cache_user cache_user +# Create a non-root user that matches the runtime security context. +RUN groupadd --gid 1000 cache_user && useradd --uid 1000 --gid cache_user cache_user # Change ownership and switch to non-root user RUN chown -R cache_user:cache_user /usr/local/bin/ diff --git a/manifests/base/runtimes/data-cache/torch_distributed_with_cache.yaml b/manifests/base/runtimes/data-cache/torch_distributed_with_cache.yaml index 0100479fc5..d9d6d84dc4 100644 --- a/manifests/base/runtimes/data-cache/torch_distributed_with_cache.yaml +++ b/manifests/base/runtimes/data-cache/torch_distributed_with_cache.yaml @@ -36,9 +36,23 @@ spec: template: spec: serviceAccountName: kubeflow-trainer-cache-initializer + securityContext: + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault containers: - name: dataset-initializer image: ghcr.io/kubeflow/trainer/dataset-initializer + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault env: - name: CACHE_IMAGE value: "ghcr.io/kubeflow/trainer/data-cache:latest" @@ -58,9 +72,24 @@ spec: spec: template: spec: + securityContext: + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault containers: - name: node image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime + workingDir: /tmp + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault env: - name: TRAIN_JOB_NAME valueFrom: diff --git a/manifests/base/runtimes/torch_distributed.yaml b/manifests/base/runtimes/torch_distributed.yaml index 52e7a6b68b..39cab9e426 100644 --- a/manifests/base/runtimes/torch_distributed.yaml +++ b/manifests/base/runtimes/torch_distributed.yaml @@ -33,6 +33,21 @@ spec: spec: template: spec: + securityContext: + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault containers: - name: node image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime + workingDir: /tmp + securityContext: + allowPrivilegeEscalation: false + capabilities: + drop: + - ALL + runAsNonRoot: true + runAsUser: 1000 + seccompProfile: + type: RuntimeDefault diff --git a/pkg/initializers/dataset/cache.py b/pkg/initializers/dataset/cache.py index 9eb2cad34d..e61dd0b644 100644 --- a/pkg/initializers/dataset/cache.py +++ b/pkg/initializers/dataset/cache.py @@ -21,6 +21,20 @@ import pkg.initializers.types.types as types import pkg.initializers.utils.utils as utils +RESTRICTED_POD_SECURITY_CONTEXT = { + "runAsNonRoot": True, + "runAsUser": 1000, + "seccompProfile": {"type": "RuntimeDefault"}, +} + +RESTRICTED_CONTAINER_SECURITY_CONTEXT = { + "allowPrivilegeEscalation": False, + "capabilities": {"drop": ["ALL"]}, + "runAsNonRoot": True, + "runAsUser": 1000, + "seccompProfile": {"type": "RuntimeDefault"}, +} + logging.basicConfig( format="%(asctime)s %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s", datefmt="%Y-%m-%dT%H:%M:%SZ", @@ -174,12 +188,14 @@ def download_dataset(self): }, "spec": { "serviceAccountName": service_account.metadata.name, + "securityContext": RESTRICTED_POD_SECURITY_CONTEXT, "containers": [ { "name": "head", "image": cache_image, "command": ["head"], "args": ["0.0.0.0", "50051"], + "securityContext": RESTRICTED_CONTAINER_SECURITY_CONTEXT, "resources": { "limits": { "cpu": head_cpu, @@ -212,12 +228,14 @@ def download_dataset(self): "workerTemplate": { "spec": { "serviceAccountName": f"{train_job_name}-cache", + "securityContext": RESTRICTED_POD_SECURITY_CONTEXT, "containers": [ { "name": "worker", "image": cache_image, "command": ["worker"], "args": ["0.0.0.0", "50051"], + "securityContext": RESTRICTED_CONTAINER_SECURITY_CONTEXT, "resources": { "limits": { "cpu": worker_cpu, diff --git a/pkg/initializers/dataset/cache_test.py b/pkg/initializers/dataset/cache_test.py index abd655f68c..d57bc6ebc5 100644 --- a/pkg/initializers/dataset/cache_test.py +++ b/pkg/initializers/dataset/cache_test.py @@ -248,6 +248,84 @@ def test_download_dataset(test_name, test_case): print("Test execution completed") +def test_download_dataset_creates_restricted_pss_compatible_leader_worker_set(): + cache_initializer_instance = CacheInitializer() + + config = { + "storage_uri": "cache://test_schema/test_table", + "train_job_name": "restricted-job", + "cache_image": "test-image:latest", + "iam_role": "arn:aws:iam::123456789012:role/test-role", + "metadata_loc": "s3://test-bucket/metadata", + } + + with patch.object(utils, "get_config_from_env", return_value=config): + cache_initializer_instance.load_config() + + with patch( + "pkg.initializers.dataset.cache.get_namespace", return_value="test-namespace" + ), patch("pkg.initializers.dataset.cache.config"), patch( + "pkg.initializers.dataset.cache.client" + ) as mock_client: + + mock_api_client = MagicMock() + mock_core_v1 = MagicMock() + mock_custom_api = MagicMock() + + mock_client.ApiClient.return_value = mock_api_client + mock_client.CoreV1Api.return_value = mock_core_v1 + mock_client.CustomObjectsApi.return_value = mock_custom_api + + mock_training_job = { + "apiVersion": "trainer.kubeflow.org/v1alpha1", + "kind": "TrainJob", + "metadata": {"name": "restricted-job", "uid": "test-uid"}, + } + mock_lws_ready = { + "status": {"conditions": [{"type": "Available", "status": "True"}]} + } + mock_custom_api.get_namespaced_custom_object.side_effect = [ + mock_training_job, + mock_lws_ready, + ] + + cache_initializer_instance.download_dataset() + + create_call = mock_custom_api.create_namespaced_custom_object.call_args + assert create_call.kwargs["plural"] == "leaderworkersets" + + leader_worker_set = create_call.kwargs["body"] + leader_spec = leader_worker_set["spec"]["leaderWorkerTemplate"][ + "leaderTemplate" + ]["spec"] + worker_spec = leader_worker_set["spec"]["leaderWorkerTemplate"][ + "workerTemplate" + ]["spec"] + expected_pod_security_context = { + "runAsNonRoot": True, + "runAsUser": 1000, + "seccompProfile": {"type": "RuntimeDefault"}, + } + expected_container_security_context = { + "allowPrivilegeEscalation": False, + "capabilities": {"drop": ["ALL"]}, + "runAsNonRoot": True, + "runAsUser": 1000, + "seccompProfile": {"type": "RuntimeDefault"}, + } + + assert leader_spec["securityContext"] == expected_pod_security_context + assert worker_spec["securityContext"] == expected_pod_security_context + assert ( + leader_spec["containers"][0]["securityContext"] + == expected_container_security_context + ) + assert ( + worker_spec["containers"][0]["securityContext"] + == expected_container_security_context + ) + + def test_download_dataset_service_already_exists(): """Service creation should be idempotent — a 409 AlreadyExists response must be treated as a no-op, not as a creation failure that triggers From 9bb6f2f1c2826aad9f8f5f74d38d4a6b9f396234 Mon Sep 17 00:00:00 2001 From: danish9039 Date: Wed, 15 Jul 2026 21:52:23 +0530 Subject: [PATCH 2/2] fix(runtimes): preserve trainer image working directory Signed-off-by: danish9039 --- .../files/runtimes/torch-distributed-with-cache.yaml | 1 - charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml | 1 - .../kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml | 2 +- .../tests/runtimes/torch_distributed_with_cache_test.yaml | 2 +- .../base/runtimes/data-cache/torch_distributed_with_cache.yaml | 1 - manifests/base/runtimes/torch_distributed.yaml | 1 - 6 files changed, 2 insertions(+), 6 deletions(-) diff --git a/charts/kubeflow-trainer/files/runtimes/torch-distributed-with-cache.yaml b/charts/kubeflow-trainer/files/runtimes/torch-distributed-with-cache.yaml index f5394a430d..d1b2e90f10 100644 --- a/charts/kubeflow-trainer/files/runtimes/torch-distributed-with-cache.yaml +++ b/charts/kubeflow-trainer/files/runtimes/torch-distributed-with-cache.yaml @@ -82,7 +82,6 @@ spec: containers: - name: node image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime - workingDir: /tmp securityContext: allowPrivilegeEscalation: false capabilities: diff --git a/charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml b/charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml index 2d071bddd7..52cccd027a 100644 --- a/charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml +++ b/charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml @@ -44,7 +44,6 @@ spec: containers: - name: node image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime - workingDir: /tmp securityContext: allowPrivilegeEscalation: false capabilities: diff --git a/charts/kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml b/charts/kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml index 26d21a437b..a8ab293945 100644 --- a/charts/kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml +++ b/charts/kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml @@ -61,4 +61,4 @@ tests: - matchRegex: path: data["runtimes.yaml"] pattern: |- - (?s)- name: node.*securityContext:\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault\s+containers:\s+- name: node\s+image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime\s+workingDir: /tmp\s+securityContext:\s+allowPrivilegeEscalation: false\s+capabilities:\s+drop:\s+- ALL\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault + (?s)- name: node.*securityContext:\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault\s+containers:\s+- name: node\s+image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime\s+securityContext:\s+allowPrivilegeEscalation: false\s+capabilities:\s+drop:\s+- ALL\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault diff --git a/charts/kubeflow-trainer/tests/runtimes/torch_distributed_with_cache_test.yaml b/charts/kubeflow-trainer/tests/runtimes/torch_distributed_with_cache_test.yaml index fc9921ecbd..5230b0277d 100644 --- a/charts/kubeflow-trainer/tests/runtimes/torch_distributed_with_cache_test.yaml +++ b/charts/kubeflow-trainer/tests/runtimes/torch_distributed_with_cache_test.yaml @@ -99,4 +99,4 @@ tests: - matchRegex: path: data["runtimes.yaml"] pattern: |- - (?s)- name: node.*securityContext:\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault\s+containers:\s+- name: node\s+image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime\s+workingDir: /tmp\s+securityContext:\s+allowPrivilegeEscalation: false\s+capabilities:\s+drop:\s+- ALL\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault + (?s)- name: node.*securityContext:\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault\s+containers:\s+- name: node\s+image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime\s+securityContext:\s+allowPrivilegeEscalation: false\s+capabilities:\s+drop:\s+- ALL\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault diff --git a/manifests/base/runtimes/data-cache/torch_distributed_with_cache.yaml b/manifests/base/runtimes/data-cache/torch_distributed_with_cache.yaml index d9d6d84dc4..1c5bf3ea4f 100644 --- a/manifests/base/runtimes/data-cache/torch_distributed_with_cache.yaml +++ b/manifests/base/runtimes/data-cache/torch_distributed_with_cache.yaml @@ -80,7 +80,6 @@ spec: containers: - name: node image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime - workingDir: /tmp securityContext: allowPrivilegeEscalation: false capabilities: diff --git a/manifests/base/runtimes/torch_distributed.yaml b/manifests/base/runtimes/torch_distributed.yaml index 39cab9e426..5b76b09cb5 100644 --- a/manifests/base/runtimes/torch_distributed.yaml +++ b/manifests/base/runtimes/torch_distributed.yaml @@ -41,7 +41,6 @@ spec: containers: - name: node image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime - workingDir: /tmp securityContext: allowPrivilegeEscalation: false capabilities: