trainer/charts/kubeflow-trainer/templates/runtimes/data-cache/torch-distributed-with-cache.yaml

{{- /*
Copyright 2026 The Kubeflow authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/ -}}

{{- if and .Values.dataCache.runtimes.torchDistributedWithCache.enabled (not .Values.dataCache.enabled) }}
{{- fail "dataCache.runtimes.torchDistributedWithCache.enabled requires dataCache.enabled to be true" }}
{{- end }}

{{- if and .Values.dataCache.enabled .Values.dataCache.runtimes.torchDistributedWithCache.enabled }}
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-distributed-with-cache
  labels:
    trainer.kubeflow.org/framework: torch
    {{- with .Values.runtimes.commonLabels }}
    {{- toYaml . | nindent 4 }}
    {{- end }}
    {{- include "trainer.labels" . | nindent 4 }}
spec:
  mlPolicy:
    numNodes: 1
    torch: {}
  template:
    spec:
      replicatedJobs:
      - name: dataset-initializer
        replicas: 1
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
          spec:
            template:
              spec:
                serviceAccountName: kubeflow-trainer-cache-initializer
                containers:
                  - name: dataset-initializer
                    image: {{ printf "ghcr.io/kubeflow/trainer/dataset-initializer:%s" (include "trainer.defaultImageTag" .) }}
                    env:
                      - name: CACHE_IMAGE
                        value: {{ include "trainer.runtimeImage" (list .Values.dataCache.cacheImage .) | quote }}
                      - name: TRAIN_JOB_NAME
                        valueFrom:
                          fieldRef:
                            apiVersion: v1
                            fieldPath: metadata.labels['jobset.sigs.k8s.io/jobset-name']
      - name: node
        dependsOn:
          - name: dataset-initializer
            status: Complete
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: trainer
          spec:
            template:
              spec:
                containers:
                  - name: node
                    image: pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime
                    env:
                      - name: TRAIN_JOB_NAME
                        valueFrom:
                          fieldRef:
                            apiVersion: v1
                            fieldPath: metadata.labels['jobset.sigs.k8s.io/jobset-name']
{{- end }}