apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  labels:
    trainer.kubeflow.org/framework: deepspeed
  name: deepspeed-distributed
spec:
  mlPolicy:
    mpi:
      mpiImplementation: OpenMPI
      numProcPerNode: 1
      runLauncherAsNode: true
      sshAuthMountPath: /home/mpiuser/.ssh
    numNodes: 1
  template:
    spec:
      network:
        publishNotReadyAddresses: true
      replicatedJobs:
      - name: launcher
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: trainer
          spec:
            template:
              spec:
                containers:
                - image: ghcr.io/kubeflow/trainer/deepspeed-runtime:latest
                  name: node
                  securityContext:
                    runAsUser: 1000
      - name: node
        template:
          spec:
            template:
              spec:
                containers:
                - args:
                  - -De
                  - -f
                  - /home/mpiuser/.sshd_config
                  command:
                  - /usr/sbin/sshd
                  image: ghcr.io/kubeflow/trainer/deepspeed-runtime:latest
                  name: node
                  readinessProbe:
                    initialDelaySeconds: 5
                    tcpSocket:
                      port: 2222
                  securityContext:
                    runAsUser: 1000
      successPolicy:
        operator: All
        targetReplicatedJobs:
        - launcher
# ---
# # JAX runtime - 현재 CRD에서 spec.mlPolicy.jax 미지원으로 주석처리
# apiVersion: trainer.kubeflow.org/v1alpha1
# kind: ClusterTrainingRuntime
# metadata:
#   labels:
#     trainer.kubeflow.org/framework: jax
#   name: jax-distributed
# spec:
#   mlPolicy:
#     jax: {}
#     numNodes: 1
#   template:
#     spec:
#       replicatedJobs:
#       - name: node
#         template:
#           metadata:
#             labels:
#               trainer.kubeflow.org/trainjob-ancestor-step: trainer
#           spec:
#             template:
#               spec:
#                 containers:
#                 - image: nvcr.io/nvidia/jax:25.10-py3
#                   name: node
---
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  labels:
    trainer.kubeflow.org/framework: mlx
  name: mlx-distributed
spec:
  mlPolicy:
    mpi:
      mpiImplementation: OpenMPI
      numProcPerNode: 1
      runLauncherAsNode: true
      sshAuthMountPath: /home/mpiuser/.ssh
    numNodes: 1
  template:
    spec:
      network:
        publishNotReadyAddresses: true
      replicatedJobs:
      - name: launcher
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: trainer
          spec:
            template:
              spec:
                containers:
                - image: ghcr.io/kubeflow/trainer/mlx-runtime:latest
                  name: node
                  securityContext:
                    runAsUser: 1000
      - name: node
        template:
          spec:
            template:
              spec:
                containers:
                - args:
                  - -De
                  - -f
                  - /home/mpiuser/.sshd_config
                  command:
                  - /usr/sbin/sshd
                  image: ghcr.io/kubeflow/trainer/mlx-runtime:latest
                  name: node
                  readinessProbe:
                    initialDelaySeconds: 5
                    tcpSocket:
                      port: 2222
                  securityContext:
                    runAsUser: 1000
      successPolicy:
        operator: All
        targetReplicatedJobs:
        - launcher
---
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  labels:
    trainer.kubeflow.org/framework: torch
  name: torch-distributed
spec:
  mlPolicy:
    numNodes: 1
    torch:
      numProcPerNode: auto
  template:
    spec:
      replicatedJobs:
      - name: node
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: trainer
          spec:
            template:
              spec:
                containers:
                - image: pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
                  name: node
---
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  labels:
    trainer.kubeflow.org/framework: torchtune
  name: torchtune-llama3.2-1b
spec:
  mlPolicy:
    numNodes: 1
    torch:
      numProcPerNode: auto
  template:
    spec:
      replicatedJobs:
      - name: dataset-initializer
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
          spec:
            template:
              spec:
                containers:
                - env:
                  - name: STORAGE_URI
                    value: hf://tatsu-lab/alpaca
                  image: ghcr.io/kubeflow/trainer/dataset-initializer:latest
                  name: dataset-initializer
                  volumeMounts:
                  - mountPath: /workspace
                    name: initializer
                volumes:
                - name: initializer
                  persistentVolumeClaim:
                    claimName: torchtune-llama3.2-1b
      - name: model-initializer
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: model-initializer
          spec:
            template:
              spec:
                containers:
                - env:
                  - name: STORAGE_URI
                    value: hf://meta-llama/Llama-3.2-1B-Instruct
                  image: ghcr.io/kubeflow/trainer/model-initializer:latest
                  name: model-initializer
                  volumeMounts:
                  - mountPath: /workspace
                    name: initializer
                volumes:
                - name: initializer
                  persistentVolumeClaim:
                    claimName: torchtune-llama3.2-1b
      - dependsOn:
        - name: dataset-initializer
          status: Complete
        - name: model-initializer
          status: Complete
        name: node
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: trainer
          spec:
            template:
              spec:
                containers:
                - command:
                  - tune
                  - run
                  - --rdzv_endpoint=localhost:29500
                  - full_finetune_distributed
                  - --config
                  - llama3_2/1B_full
                  - dataset=torchtune.datasets.instruct_dataset
                  - dataset.source=parquet
                  - dataset.data_dir=/workspace/dataset/data
                  - output_dir=/workspace/output
                  - tokenizer.path=/workspace/model/original/tokenizer.model
                  - checkpointer.checkpoint_dir=/workspace/model
                  image: ghcr.io/kubeflow/trainer/torchtune-trainer:latest
                  name: node
                  resources:
                    limits:
                      nvidia.com/gpu: 2
                  volumeMounts:
                  - mountPath: /workspace
                    name: initializer
                volumes:
                - name: initializer
                  persistentVolumeClaim:
                    claimName: torchtune-llama3.2-1b
---
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  labels:
    trainer.kubeflow.org/framework: torchtune
  name: torchtune-llama3.2-3b
spec:
  mlPolicy:
    numNodes: 1
    torch:
      numProcPerNode: auto
  template:
    spec:
      replicatedJobs:
      - name: dataset-initializer
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
          spec:
            template:
              spec:
                containers:
                - env:
                  - name: STORAGE_URI
                    value: hf://tatsu-lab/alpaca
                  image: ghcr.io/kubeflow/trainer/dataset-initializer:latest
                  name: dataset-initializer
                  volumeMounts:
                  - mountPath: /workspace
                    name: initializer
                volumes:
                - name: initializer
                  persistentVolumeClaim:
                    claimName: torchtune-llama3.2-3b
      - name: model-initializer
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: model-initializer
          spec:
            template:
              spec:
                containers:
                - env:
                  - name: STORAGE_URI
                    value: hf://meta-llama/Llama-3.2-3B-Instruct
                  image: ghcr.io/kubeflow/trainer/model-initializer:latest
                  name: model-initializer
                  volumeMounts:
                  - mountPath: /workspace
                    name: initializer
                volumes:
                - name: initializer
                  persistentVolumeClaim:
                    claimName: torchtune-llama3.2-3b
      - dependsOn:
        - name: dataset-initializer
          status: Complete
        - name: model-initializer
          status: Complete
        name: node
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: trainer
          spec:
            template:
              spec:
                containers:
                - command:
                  - tune
                  - run
                  - --rdzv_endpoint=localhost:29500
                  - full_finetune_distributed
                  - --config
                  - llama3_2/3B_full
                  - dataset=torchtune.datasets.instruct_dataset
                  - dataset.source=parquet
                  - dataset.data_dir=/workspace/dataset/data
                  - output_dir=/workspace/output
                  - tokenizer.path=/workspace/model/original/tokenizer.model
                  - checkpointer.checkpoint_dir=/workspace/model
                  image: ghcr.io/kubeflow/trainer/torchtune-trainer:latest
                  name: node
                  resources:
                    limits:
                      nvidia.com/gpu: 2
                  volumeMounts:
                  - mountPath: /workspace
                    name: initializer
                volumes:
                - name: initializer
                  persistentVolumeClaim:
                    claimName: torchtune-llama3.2-3b
---
apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  labels:
    trainer.kubeflow.org/framework: torchtune
  name: torchtune-qwen2.5-1.5b
spec:
  mlPolicy:
    numNodes: 1
    torch:
      numProcPerNode: auto
  template:
    spec:
      replicatedJobs:
      - name: dataset-initializer
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: dataset-initializer
          spec:
            template:
              spec:
                containers:
                - env:
                  - name: STORAGE_URI
                    value: hf://tatsu-lab/alpaca
                  image: ghcr.io/kubeflow/trainer/dataset-initializer:latest
                  name: dataset-initializer
                  volumeMounts:
                  - mountPath: /workspace
                    name: initializer
                volumes:
                - name: initializer
                  persistentVolumeClaim:
                    claimName: torchtune-qwen2.5-1.5b
      - name: model-initializer
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: model-initializer
          spec:
            template:
              spec:
                containers:
                - env:
                  - name: STORAGE_URI
                    value: hf://Qwen/Qwen2.5-1.5B-Instruct
                  image: ghcr.io/kubeflow/trainer/model-initializer:latest
                  name: model-initializer
                  volumeMounts:
                  - mountPath: /workspace
                    name: initializer
                volumes:
                - name: initializer
                  persistentVolumeClaim:
                    claimName: torchtune-qwen2.5-1.5b
      - dependsOn:
        - name: dataset-initializer
          status: Complete
        - name: model-initializer
          status: Complete
        name: node
        template:
          metadata:
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: trainer
          spec:
            template:
              spec:
                containers:
                - command:
                  - tune
                  - run
                  - --rdzv_endpoint=localhost:29500
                  - full_finetune_distributed
                  - --config
                  - qwen2_5/1.5B_full
                  - dataset=torchtune.datasets.instruct_dataset
                  - dataset.source=parquet
                  - dataset.data_dir=/workspace/dataset/data
                  - output_dir=/workspace/output
                  - tokenizer.path=/workspace/model/vocab.json
                  - tokenizer.merges_file=/workspace/model/merges.txt
                  - checkpointer.checkpoint_dir=/workspace/model
                  image: ghcr.io/kubeflow/trainer/torchtune-trainer:latest
                  name: node
                  resources:
                    limits:
                      nvidia.com/gpu: 2
                  volumeMounts:
                  - mountPath: /workspace
                    name: initializer
                volumes:
                - name: initializer
                  persistentVolumeClaim:
                    claimName: torchtune-qwen2.5-1.5b