Update fms-hf-tuning e2e test

Signed-off-by: ted chang <htchang@us.ibm.com>
opendatahub-io · Apr 9, 2024 · 1eb5fb1 · 1eb5fb1
1 parent f46d74d
commit 1eb5fb1
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/test/e2e/config.json b/test/e2e/config.json
@@ -1,4 +1,8 @@
 {
+    "accelerate_launch_args": {
+        "num_processes": 2,
+        "cpu": true
+    },
     "model_name_or_path": "bigscience/bloom-560m",
     "training_data_path": "/etc/config/twitter_complaints_small.json",
     "output_dir": "/tmp/out",

diff --git a/test/e2e/kfto_kueue_sft_test.go b/test/e2e/kfto_kueue_sft_test.go
@@ -1,4 +1,4 @@
-//go:build ignore
+//go:build tuning_e2e
 
 /*
 Copyright 2023.
@@ -94,9 +94,9 @@ func TestPytorchjobWithSFTtrainer(t *testing.T) {
 							Containers: []corev1.Container{
 								{
 									Name:            "pytorch",
-									Image:           "quay.io/tedchang/sft-trainer:dev",
+									Image:           "quay.io/modh/fms-hf-tuning:5d8789723ec58ac1bc9c2df704395f162fed974a",
 									ImagePullPolicy: corev1.PullIfNotPresent,
-									Command:         []string{"python", "/app/launch_training.py"},
+									Command:         []string{"python", "/app/accelerate_launch.py"},
 									Env: []corev1.EnvVar{
 										{
 											Name:  "SFT_TRAINER_CONFIG_JSON_PATH",

diff --git a/test/e2e/setup.sh b/test/e2e/setup.sh
@@ -19,4 +19,4 @@ docker build -t ${KFTO_IMG} -f ${BASE_DIR}/build/images/training-operator/Docker
 
 echo "Load training operator image into cluster"
 kind load --name training-operator-cluster docker-image training-operator:dev
-KFTO_IMG=training-operator:dev make deploy
+IMG=training-operator:dev make deploy