raycluster.yaml

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: rayml-efa
  labels:
    controller-tools.k8s.io: "1.0"
  annotations:
    karpenter.sh/do-not-disrupt: "true"
spec:
  # Ray head pod template
  headGroupSpec:
    # The `rayStartParams` are used to configure the `ray start` command.
    # See https://github.com/ray-project/kuberay/blob/master/docs/guidance/rayStartParams.md for the default settings of `rayStartParams` in KubeRay.
    # See https://docs.ray.io/en/latest/cluster/cli.html#ray-start for all available options in `rayStartParams`.
    rayStartParams:
      dashboard-host: '0.0.0.0'
      metrics-export-port: '8080'  # Explicitly set metrics port
    #pod template
    template:
      spec:
        nodeSelector:
          node.kubernetes.io/instance-type: PLACEHOLDER_INSTANCE_TYPE
          sagemaker.amazonaws.com/node-health-status: Schedulable
        securityContext:
          runAsUser: 0
          runAsGroup: 0
          fsGroup: 0
        containers:
        - name: ray-head
          image: PLACEHOLDER_IMAGE     ## IMAGE: Here you may choose which image your head pod will run
          env:                                ## ENV: Here is where you can send stuff to the head pod
            ## PROMETHEUS AND GRAFANA - AWS MANAGED SERVICES
            - name: RAY_GRAFANA_IFRAME_HOST
              value: http://localhost:3000
            - name: RAY_GRAFANA_HOST
              value: http://prometheus-grafana.prometheus-system.svc:80
            - name: RAY_PROMETHEUS_HOST
              value: http://prometheus-kube-prometheus-prometheus.prometheus-system.svc:9090
            ## EFA AND NCCL CONFIGURATION
            - name: FI_PROVIDER
              value: "efa"
            - name: FI_EFA_USE_DEVICE_RDMA
              value: "1"
            - name: FI_EFA_FORK_SAFE
              value: "1"
            - name: NCCL_PROTO
              value: "simple"
            - name: NCCL_SOCKET_IFNAME
              value: "^docker,lo,veth"
            - name: NCCL_DEBUG
              value: "INFO"
            - name: TORCH_NCCL_DUMP_ON_TIMEOUT
              value: "1"
            - name: TORCH_NCCL_ASYNC_ERROR_HANDLING
              value: "1"
            - name: HF_TOKEN
              value: PLACEHOLDER_HF_TOKEN
          lifecycle:
            preStop:
              exec:
                command: ["/bin/sh","-c","ray stop"]
          resources:
            limits:                                    ## LIMITS: Set resource limits for your head pod
              cpu: 8
              memory: 32Gi
            requests:                                    ## REQUESTS: Set resource requests for your head pod
              cpu: 8
              memory: 32Gi
          ports:
          - containerPort: 6379
            name: gcs-server
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          - containerPort: 8000
            name: serve
          - containerPort: 8080
            name: metrics
          volumeMounts:                                    ## VOLUMEMOUNTS: Mount your S3 CSI EKS Add-On to head pod
          - name: fsx-storage
            mountPath: /fsx
          - name: ray-logs
            mountPath: /tmp/ray
          # - name: checkpoint-logs
          #   mountPath: /var/log/sagemaker_checkpointing
        volumes:
          - name: ray-logs
            emptyDir: {}
          - name: fsx-storage
            persistentVolumeClaim:
              claimName: fsx-claim
          # - name: checkpoint-logs
          #   hostPath:
          #     path: /var/logs/sagemaker_checkpointing
          #     type: DirectoryOrCreate
  workerGroupSpecs:
  # the pod replicas in this group typed worker
  - replicas: PLACEHOLDER_NUM_NODES                                    ## REPLICAS: How many worker pods you want 
    minReplicas: 1
    maxReplicas: 10
    # logical group name, for this called small-group, also can be functional
    groupName: gpu-group
    rayStartParams:
      num-gpus: "PLACEHOLDER_NUM_GPU_PER_NODE"
      metrics-export-port: '8080'  # Explicitly set metrics port for workers
    #pod template
    template:
      spec:
        nodeSelector:
          node.kubernetes.io/instance-type: PLACEHOLDER_INSTANCE_TYPE
          sagemaker.amazonaws.com/node-health-status: Schedulable
        securityContext:
          runAsUser: 0
          runAsGroup: 0
          fsGroup: 0
        containers:
        - name: ray-worker
          image: PLACEHOLDER_IMAGE             ## IMAGE: Here you may choose which image your head node will run
          env:
            ## EFA CONFIGURATION
            - name: FI_PROVIDER
              value: "efa"
            - name: FI_EFA_USE_DEVICE_RDMA
              value: "1"
            - name: FI_EFA_FORK_SAFE
              value: "1"
            ## NCCL CONFIGURATION
            - name: NCCL_SOCKET_IFNAME
              value: "^docker,lo,veth"
            - name: NCCL_BUFFSIZE
              value: "8388608"
            - name: NCCL_P2P_NET_CHUNKSIZE
              value: "8388608"
            - name: NCCL_TUNER_PLUGIN
              value: "/nccl-links/lib64/libnccl-tuner-aws-ofi.so"
            - name: NCCL_DEBUG
              value: "INFO"
            - name: LD_LIBRARY_PATH
              value: "/nccl-links/lib64:/host-ofi-nccl/lib64:/host-efa/lib64:/usr/local/lib/python3.12/dist-packages/torch/lib:/usr/local/lib:/usr/lib/x86_64-linux-gnu"
            ## VLLM AND PYTORCH
            - name: VLLM_USE_V1
              value: "1"
            - name: TOKENIZERS_PARALLELISM
              value: "true"
            - name: PYTORCH_CUDA_ALLOC_CONF
              value: "expandable_segments:True"
            ## TORCH NCCL
            - name: TORCH_NCCL_DUMP_ON_TIMEOUT
              value: "1"
            - name: TORCH_NCCL_ASYNC_ERROR_HANDLING
              value: "1"
            - name: HF_TOKEN
              value: PLACEHOLDER_HF_TOKEN
          lifecycle:
            preStop:
              exec:
                command: ["/bin/sh","-c","ray stop"]
          resources:
            limits:                                    ## LIMITS: Set resource limits for your worker pods
              cpu: PLACEHOLDER_WORKER_CPU
              memory: PLACEHOLDER_WORKER_MEMORY
              nvidia.com/gpu: PLACEHOLDER_NUM_GPU_PER_NODE
              vpc.amazonaws.com/efa: PLACEHOLDER_NUM_EFA_PER_NODE
            requests:                                    ## REQUESTS: Set resource requests for your worker pods
              cpu: PLACEHOLDER_WORKER_CPU
              memory: PLACEHOLDER_WORKER_MEMORY
              nvidia.com/gpu: PLACEHOLDER_NUM_GPU_PER_NODE
              vpc.amazonaws.com/efa: PLACEHOLDER_NUM_EFA_PER_NODE
          ports:
          - containerPort: 8080
            name: metrics
          volumeMounts:                                    ## VOLUMEMOUNTS: Mount EFA/NCCL host libs and shared storage
          - name: host-efa
            mountPath: /host-efa
            readOnly: true
          - name: host-ofi-nccl
            mountPath: /host-ofi-nccl
            readOnly: true
          - name: nccl-links
            mountPath: /nccl-links
          - name: ray-logs
            mountPath: /tmp/ray
          - name: fsx-storage
            mountPath: /fsx
          - name: dshm
            mountPath: /dev/shm
        ## INIT CONTAINER: Create NCCL symlinks from host OFI-NCCL libs
        ## Required when the Docker image does not bundle EFA/OFI-NCCL libraries.
        ## The host nodes (e.g. p5en) have these at /opt/amazon/ofi-nccl but with
        ## different filenames than what NCCL expects. The init container creates
        ## the expected symlinks in a shared emptyDir volume.
        initContainers:
        - name: nccl-symlinks
          image: busybox:latest
          command: ["sh", "-c"]
          args:
          - |
            mkdir -p /nccl-links/lib64
            ln -sf /host-ofi-nccl/lib64/libnccl-net-ofi.so /nccl-links/lib64/libnccl-net-aws-ofi.so
            ln -sf /host-ofi-nccl/lib64/libnccl-ofi-tuner.so /nccl-links/lib64/libnccl-tuner-aws-ofi.so
            ls -la /nccl-links/lib64/
            echo "NCCL symlinks created"
          volumeMounts:
          - name: nccl-links
            mountPath: /nccl-links
          - name: host-ofi-nccl
            mountPath: /host-ofi-nccl
            readOnly: true
        volumes:
        ## HOST EFA AND OFI-NCCL LIBRARIES
        ## These are pre-installed on EFA-enabled nodes (p5, p5e, p5en, etc.)
        ## at /opt/amazon/efa and /opt/amazon/ofi-nccl. Mount them into the pod
        ## so NCCL can use EFA for inter-node communication.
        - name: host-efa
          hostPath:
            path: /opt/amazon/efa
            type: Directory
        - name: host-ofi-nccl
          hostPath:
            path: /opt/amazon/ofi-nccl
            type: Directory
        - name: nccl-links
          emptyDir: {}
        - name: fsx-storage
          persistentVolumeClaim:
            claimName: fsx-claim
        - name: ray-logs
          emptyDir: {}
        ## SHARED MEMORY: Required for PyTorch DataLoader shared memory and
        ## NCCL shared-memory transport. Size should be generous for large models.
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 200Gi