Skip to content

fix(runtimes): add restricted PSS security contexts #3702

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
danish9039 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

fix(runtimes): add restricted PSS security contexts #3702

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
15266dc
runtimes: add restricted pss security contexts
danish9039 Jul 4, 2026
9bb6f2f
fix(runtimes): preserve trainer image working directory
danish9039 Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions charts/kubeflow-trainer/files/runtimes/torch-distributed-with-cache.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Comment thread
danish9039 marked this conversation as resolved.
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 }}
Expand All @@ -60,9 +74,23 @@ 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
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop:
- ALL
runAsNonRoot: true
runAsUser: 1000
seccompProfile:
type: RuntimeDefault
env:
- name: TRAIN_JOB_NAME
valueFrom:
Expand Down
14 changes: 14 additions & 0 deletions charts/kubeflow-trainer/files/runtimes/torch-distributed.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,20 @@ 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
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop:
- ALL
runAsNonRoot: true
runAsUser: 1000
seccompProfile:
type: RuntimeDefault
11 changes: 11 additions & 0 deletions charts/kubeflow-trainer/tests/runtimes/torch_distributed_test.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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+securityContext:\s+allowPrivilegeEscalation: false\s+capabilities:\s+drop:\s+- ALL\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault
38 changes: 38 additions & 0 deletions charts/kubeflow-trainer/tests/runtimes/torch_distributed_with_cache_test.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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+securityContext:\s+allowPrivilegeEscalation: false\s+capabilities:\s+drop:\s+- ALL\s+runAsNonRoot: true\s+runAsUser: 1000\s+seccompProfile:\s+type: RuntimeDefault
4 changes: 2 additions & 2 deletions cmd/data_cache/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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/
Expand Down
28 changes: 28 additions & 0 deletions manifests/base/runtimes/data-cache/torch_distributed_with_cache.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Comment thread
danish9039 marked this conversation as resolved.
allowPrivilegeEscalation: false
capabilities:
drop:
- ALL
runAsNonRoot: true
runAsUser: 1000
seccompProfile:
type: RuntimeDefault
env:
- name: CACHE_IMAGE
value: "ghcr.io/kubeflow/trainer/data-cache:latest"
Expand All @@ -58,9 +72,23 @@ 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
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop:
- ALL
runAsNonRoot: true
runAsUser: 1000
seccompProfile:
type: RuntimeDefault
env:
- name: TRAIN_JOB_NAME
valueFrom:
Expand Down
14 changes: 14 additions & 0 deletions manifests/base/runtimes/torch_distributed.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,20 @@ 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
securityContext:
Comment thread
danish9039 marked this conversation as resolved.
allowPrivilegeEscalation: false
capabilities:
drop:
- ALL
runAsNonRoot: true
runAsUser: 1000
seccompProfile:
type: RuntimeDefault
18 changes: 18 additions & 0 deletions pkg/initializers/dataset/cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
78 changes: 78 additions & 0 deletions pkg/initializers/dataset/cache_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading