---
# Source: pytorchjob-generator/templates/appwrapper.yaml
apiVersion: workload.codeflare.dev/v1beta2
kind: AppWrapper
metadata:
    name: yair-tutorial-4dot1-3b-alora-public
    namespace: security
    annotations:
        workload.codeflare.dev.mlbatch/pytorchGeneratorVersion: "1.1.12"
        workload.codeflare.dev.appwrapper/failureGracePeriodDuration: "24h"
        workload.codeflare.dev.appwrapper/retryLimit: "3"
        workload.codeflare.dev.appwrapper/deletionOnFailureGracePeriodDuration: "24h"
    labels:
        kueue.x-k8s.io/queue-name: default-queue        
spec:
    components:
        - template:
                  apiVersion: "kubeflow.org/v1"
                  kind: "PyTorchJob"
                  metadata:
                      name: yair-tutorial-4dot1-3b-alora-public
                  spec:
                      pytorchReplicaSpecs:
                          Master:
                              replicas: 1
                              restartPolicy: Never
                              template:                                  
                                  spec:                                      
                                      imagePullSecrets:
                                          - name: image-pull-secret-us                                                                            
                                      priorityClassName: default-priority
                                      affinity:
                                        nodeAffinity:
                                          requiredDuringSchedulingIgnoredDuringExecution:
                                            nodeSelectorTerms:
                                            - matchExpressions:
                                              - key: autopilot.ibm.com/gpuhealth
                                                operator: NotIn
                                                values:
                                                - ERR
                                                - TESTING
                                                - EVICT                                      
                                      volumes:
                                          - name: private-ssh-git-deploy-key
                                            secret:
                                                secretName: github-ssh-key
                                                optional: false
                                          - name: github-known-hosts
                                            configMap:
                                                name: github-known-hosts
                                          - name: dshm
                                            emptyDir:
                                                medium: Memory
                                      containers:
                                          - name: pytorch
                                            image: us.icr.io/cil15-shared-registry/sec-rl-neo:latest
                                            imagePullPolicy: IfNotPresent                                                                                        
                                            env:
                                                - name: TIMESTAMP
                                                  value: "2026-05-11_17.35"
                                                - name: NUM_PODS
                                                  value: "1"
                                                - name: GPUS_PER_POD
                                                  value: "1"
                                                - name: JOB_NAME
                                                  value: "yair-tutorial-4dot1-3b-alora-public"
                                                - name: HF_HOME
                                                  value: "/workspace/.cache/huggingface"
                                                - name: TORCH_HOME
                                                  value: "/workspace/.cache/torch"
                                                - name: XDG_CACHE_HOME
                                                  value: "/workspace/.cache"
                                                - name: PYTHONUNBUFFERED
                                                  value: "1"
                                                - name: MODEL_PATH
                                                  value: "ibm-granite/granite-switch-4.1-3b-preview"
                                                - name: VLLM_MODEL_NAME
                                                  value: "placeholder"
                                                - name: GRANITE_SWITCH_SOURCE
                                                  value: "placeholder"
                                                - name: GIT_SSH_COMMAND
                                                  value: "ssh -i /tmp/.ssh/keys/id_rsa -o UserKnownHostsFile=/tmp/.ssh/hosts/known_hosts -vv"                                            
                                            volumeMounts:
                                                - name: private-ssh-git-deploy-key
                                                  readOnly: true
                                                  mountPath: "/tmp/.ssh/keys"
                                                - name: github-known-hosts
                                                  mountPath: "/tmp/.ssh/hosts"
                                                - name: dshm
                                                  mountPath: "/dev/shm"                                            
                                            resources:
                                                requests:
                                                    cpu: 8
                                                    nvidia.com/gpu: 1
                                                    memory: 100Gi
                                                    nvidia.com/roce_gdr: 0
                                                limits:
                                                    cpu: 8
                                                    nvidia.com/gpu: 1
                                                    memory: 100Gi
                                                    nvidia.com/roce_gdr: 0                                            
                                            command:
                                                - sh
                                                - -c
                                                - |
                                                  echo "Environment variables set by the kubeflow training operator:"
                                                  echo ${MASTER_ADDR}:${MASTER_PORT}
                                                  echo "PYTHONUNBUFFERED:"${PYTHONUNBUFFERED}
                                                  echo My global rank is ${RANK} / ${WORLD_SIZE}
                                                  echo "Other injected environment variables:"
                                                  echo "NVME_MOUNT_PATH: "${NVME_MOUNT_PATH}
                                                  #
                                                  # User commands
                                                  #
                                                  nvidia-smi
                                                  echo "=== Tutorial pod (PUBLIC) for Granite Switch 4.1 3b ==="
                                                  map[echo "Model:${MODEL_PATH}"]
                                                  export VLLM_MODEL_NAME=${MODEL_PATH}
                                                  export GRANITE_SWITCH_SOURCE=${MODEL_PATH}
                                                  git clone -b main https://github.com/generative-computing/granite-switch.git
                                                  cd granite-switch
                                                  python3 -m venv .venv --system-site-packages
                                                  . .venv/bin/activate
                                                  pip install "transformers>=4.0.0,<5.0.0"
                                                  pip install nvidia-cuda-runtime-cu12
                                                  pip install -e ".[vllm]"
                                                  pip install mellea
                                                  pip install jupyter chromadb tqdm httpx python-dotenv litellm
                                                  echo "=== Starting vLLM server on port 8000 ==="
                                                  python -m vllm.entrypoints.openai.api_server --model ${MODEL_PATH} --port 8000 --host 0.0.0.0 &
                                                  echo "Waiting for vLLM to load model..."
                                                  until curl -s http://localhost:8000/v1/models > /dev/null 2>&1; do sleep 5; echo "..."; done
                                                  echo "=== vLLM ready ==="
                                                  curl -s http://localhost:8000/v1/models | python3 -m json.tool
                                                  echo "=== Starting Jupyter Lab on port 8888 ==="
                                                  cd tutorials
                                                  jupyter lab --no-browser --port=8888 --ip=0.0.0.0 --allow-root --ServerApp.token='' --ServerApp.allow_origin='*'