# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# 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.

apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-distributed
  labels:
    trainer.kubeflow.org/framework: torch
  annotations: {}
spec:
  mlPolicy:
    numNodes: 1
    # numProcPerNode was removed from mlPolicy.torch in kubeflow-trainer v2.2.0
    # (kubeflow/trainer#3239) — Torch now infers parallelism from
    # the container's nvidia.com/gpu resource limit. mlPolicy.mpi.numProcPerNode
    # is unaffected.
    torch: {}
  template:
    spec:
      replicatedJobs:
        - name: node
          template:
            metadata:
              labels:
                trainer.kubeflow.org/trainjob-ancestor-step: trainer
            spec:
              template:
                spec:
                  # nodeSelector and tolerations are injected by the AICR bundler
                  # from --accelerated-node-selector / --accelerated-node-toleration
                  # flags via the registry's nodeScheduling.accelerated paths
                  # (see recipes/registry.yaml). This lets users submit a bare
                  # TrainJob with no podTemplateOverrides / runtimePatches — the
                  # runtime carries the per-cluster scheduling vocabulary baked
                  # in at bundle time.
                  nodeSelector:
                    nvidia.com/gpu.present: "true"
                  tolerations:
                    - operator: Exists
                  containers:
                    - name: node
                      image: pytorch/pytorch:2.11.0-cuda12.8-cudnn9-runtime@sha256:eee11b3b3872a8c838e35ef48f08b2d5def2080902c7f666831310ca1a0ef2be