chore: refactored code

Signed-off-by: Akash Jaiswal <akashjaiswal3846@gmail.com>

Author: jaiakash

.github/workflows/test-e2e-gpu.yaml

@@ -65,17 +65,17 @@ jobs:
         run: |
           make test-e2e-delete-gpu-cluster
 
-      - name: Setup cluster with GPU
+      - name: Setup cluster with GPU support using nvidia/kind
         if: steps.check-label.outputs.skip == 'false'
         run: |
           make test-e2e-setup-gpu-cluster K8S_VERSION=${{ matrix.kubernetes-version }}
 
-      # - name: Run e2e with Go
-      #   if: steps.check-label.outputs.skip == 'false'
-      #   run: |
-      #     make test-e2e || (kubectl logs -n kubeflow-system -l app.kubernetes.io/name=trainer && exit 1)
+      - name: Run e2e with Go
+        if: steps.check-label.outputs.skip == 'false'
+        run: |
+          make test-e2e || (kubectl logs -n kubeflow-system -l app.kubernetes.io/name=trainer && exit 1)
 
-      - name: Run e2e test for example Notebooks
+      - name: Run e2e test for torchtrainer notebook on GPU cluster
         if: steps.check-label.outputs.skip == 'false'
         run: |
           mkdir -p artifacts/notebooks

Makefile

@@ -179,7 +179,7 @@ test-e2e-setup-cluster: kind ## Setup Kind cluster for e2e test.
 	KIND=$(KIND) K8S_VERSION=$(K8S_VERSION) ./hack/e2e-setup-cluster.sh
 
 .PHONY: test-e2e-setup-gpu-cluster
-test-e2e-setup-gpu-cluster: kind ## Setup Kind cluster with GPU e2e test.
+test-e2e-setup-gpu-cluster: kind ## Setup Kind cluster for GPU e2e test.
 	KIND=$(KIND) K8S_VERSION=$(K8S_VERSION) ./hack/e2e-setup-gpu-cluster.sh
 
 .PHONY: test-e2e-delete-cluster

examples/torchtune/llama3_2/alpaca-trainjob-yaml.ipynb

@@ -199,7 +199,7 @@
     "        ),\n",
     "        model=HuggingFaceModelInitializer(\n",
     "            storage_uri=\"hf://meta-llama/Llama-3.2-1B-Instruct\",\n",
-    "            access_token=os.environ[\"HF_TOKEN_AKASH\"] # Replace with your Hugging Face token,\n",
+    "            access_token=os.environ[\"HF_TOKEN\"] # Replace with your Hugging Face token,\n",
     "        )\n",
     "    ),\n",
     "    trainer=BuiltinTrainer(\n",

hack/e2e-run-notebook.sh

@@ -42,7 +42,7 @@ print_results() {
     kubectl describe trainjob
     kubectl logs -n kubeflow-system -l app.kubernetes.io/name=trainer
     kubectl logs -l jobset.sigs.k8s.io/replicatedjob-name=trainer-node,batch.kubernetes.io/job-completion-index=0 --tail -1
-    kubectl wait trainjob --for=condition=Complete --all --timeout 900s
+    kubectl wait trainjob --for=condition=Complete --all --timeout 600s
 }
 
 (papermill "${NOTEBOOK_INPUT}" "${NOTEBOOK_OUTPUT}" --execution-timeout "${PAPERMILL_TIMEOUT}" && print_results) ||