shark-job-template.yaml

# This is an advanced job template makes it easy to deploy a job that runs
# a complex inference compilation pipeline.
apiVersion: batch/v1
kind: Job
metadata:
  # Change this name with something specific to you (aka include your username)
  name: shark-test-job
  namespace: dev
spec:
  template:
    spec:
      nodeSelector:
        dev: "true"
      containers:
      - name: shark-test-container
        image: rocm/dev-ubuntu-22.04:6.3
        command: ["/bin/bash", "-c"]
        args:
        - |
          rocminfo &&
          echo "ROCm is working!" &&
          sudo apt update &&
          curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash &&
          sudo apt install git -y &&
          sudo apt install python3.11 python3.11-dev python3.11-venv -y &&
          python3.11 -m venv shark_venv && source shark_venv/bin/activate &&
          git clone https://github.com/nod-ai/SHARK-ModelDev.git -b demo &&
          git clone https://github.com/iree-org/iree-turbine.git -b demo &&
          cd SHARK-ModelDev &&
          pip install --no-compile -r ../iree-turbine/pytorch-cpu-requirements.txt &&
          pip install --pre --upgrade -r ../iree-turbine/requirements.txt &&
          pip install --pre --upgrade -r ../iree-turbine/iree-requirements.txt &&
          pip install --no-compile --pre --upgrade -e models -r models/requirements.txt &&
          python3 models/turbine_models/custom_models/sd_inference/sd_pipeline.py --device=hip --precision=fp16 --iree_target_triple=gfx942 --external_weights=safetensors --hf_model_name=stabilityai/stable-diffusion-xl-base-1.0 --width=1024 --height=1024 --batch_size=1 --use_i8_punet --attn_spec=punet --vae_decomp_attn --external_weights=safetensors --num_inference_steps=20 --verbose --benchmark=all &&
          STORAGE_KEY="V/w1DDpjcE0YgwRUT9mAH+V1LbniogpOr2tStfNtND7MPY5sFWGz9Ko3pGQAns+qu/vm3prqfWSq+AStmd7q6A==" bash upload_to_azure.sh;
        resources:
          limits:
            amd.com/gpu: 1
      restartPolicy: Never
  backoffLimit: 0