Refactor CI workflows

Author: aybchan

.github/gke-workflow/jax-vllm-offloading/deploy-transfer.sh

@@ -1,7 +1,7 @@
-kubectl apply -f deployment/gateway-pod.yml
-kubectl apply -f deployment/gateway-svc.yml
+kubectl apply -f transfer/deployment/gateway-pod.yml
+kubectl apply -f transfer/deployment/gateway-svc.yml
 
 kubectl apply -f huggingface-secret.yml
 
-kubectl apply -f deployment/rollout.yml
-kubectl apply -f deploymeny/trainer.yml
+kubectl apply -f transfer/deployment/rollout.yml
+kubectl apply -f transfer/deployment/trainer.yml

.github/workflows/jax-vllm-offloading/ci.yml

@@ -0,0 +1,29 @@
+name: JAX-vLLM offloading
+
+on:
+  workflow_call:
+    inputs:
+      JAX_VLLM_OFFLOADING_IMAGE:
+        type: string
+        description: MaxText image from ghcr.io/nvidia
+        default: ghcr.io/nvidia/jax-toolbox-internal:19461214142-jio-amd64
+        required: false
+  pull_request:
+    types:
+      - opened
+      - reopened
+      - ready_for_review
+      - synchronize
+    paths:
+      - 'jax-inference-offloading/**'
+      - '.github/gke-workflow/jax-vllm-offloading/**'
+      - '.github/workflows/jax-vllm-offloading/**'
+
+jobs:
+  jax-vllm-offloading-transfer-gke-xpk:
+    uses: .github/workflows/jax-vllm-offloading/gke_transfer.yml
+    # needs: build # .github/workflows/jio.yaml
+    with:
+      # JAX_VLLM_OFFLOADING_IMAGE: ${{ needs.build.outputs.DOCKER_TAG_FINAL }}
+      JAX_VLLM_OFFLOADING_IMAGE: ghcr.io/nvidia/jax-toolbox-internal:19461214142-jio-amd64
+    secrets: inherit

…/workflows/_test_jax_vllm_offloading.yml …ows/jax-vllm-offloading/gke_transfer.yml.github/workflows/_test_jax_vllm_offloading.yml renamed to .github/workflows/jax-vllm-offloading/gke_transfer.yml .github/workflows/_test_jax_vllm_offloading.yml renamed to .github/workflows/jax-vllm-offloading/gke_transfer.yml

@@ -1,4 +1,4 @@
-name: ~Test JAX-vLLM offloading transfer (GKE, XPK)
+name: JAX-vLLM offloading transfer (GKE, XPK)
 
 on:
   workflow_call:
@@ -8,25 +8,14 @@ on:
         description: MaxText image from ghcr.io/nvidia
         default: ghcr.io/nvidia/jax-toolbox-internal:19461214142-jio-amd64
         required: false
-  pull_request:
-    types:
-      - opened
-      - reopened
-      - ready_for_review
-      - synchronize
-    paths:
-      - '.github/workflows/_test_jax_vllm_offloading.yaml'
-      - '.github/gke-workflow/jax-vllm-offloading/**'
-      - 'jax-inference-offloading/**'
 
 jobs:
   jax-vllm-offloading-transfer-gke-xpk:
     runs-on: gke-a3mega
 
     env:
       WORKLOAD_NAME_PREFIX: gke-jax-vllm-transfer
-      #JAX_VLLM_OFFLOADING_IMAGE: ${{ inputs.JAX_VLLM_OFFLOADING_IMAGE }}
-      JAX_VLLM_OFFLOADING_IMAGE: ghcr.io/nvidia/jax-toolbox-internal:19461214142-jio-amd64
+      JAX_VLLM_OFFLOADING_IMAGE: ${{ inputs.JAX_VLLM_OFFLOADING_IMAGE }}
 
       NUM_NODES: 2
       ENV_FILE: ../../.github/gke-workflow/jax-vllm-offloading/transfer/jobset.env