test

Signed-off-by: Cody Yu <[email protected]>

.buildkite/llm.rayci.yml

@@ -5,6 +5,19 @@ depends_on:
 steps:
   - name: llmbuild
     wanda: ci/docker/llm.build.wanda.yaml
+    depends_on: oss-ci-base_ml
+    env:
+      IMAGE_FROM: cr.ray.io/rayproject/oss-ci-base_build-py3.11
+      IMAGE_TO: llmbuild
+      RAYCI_IS_GPU_BUILD: "false"
+
+  - name: llmgpubuild
+    wanda: ci/docker/llm.build.wanda.yaml
+    depends_on: oss-ci-base_gpu
+    env:
+      IMAGE_FROM: cr.ray.io/rayproject/oss-ci-base_gpu
+      IMAGE_TO: llmgpubuild
+      RAYCI_IS_GPU_BUILD: "true"
 
   - label: "llm batch tests cpu"
     key: "llm-batch-tests-cpu"

ci/docker/llm.build.wanda.yaml

@@ -1,10 +1,14 @@
-name: "llmbuild"
-froms: ["cr.ray.io/rayproject/oss-ci-base_build-py3.11"]
+name: "$IMAGE_TO"
+froms: ["$IMAGE_FROM"]
 dockerfile: ci/docker/llm.build.Dockerfile
 srcs:
   - ci/env/install-dependencies.sh
   - ci/env/install-llvm-binaries.sh
   - ci/suppress_output
   - python/requirements_compiled_rayllm_test_py311.txt
+build_args:
+  - DOCKER_IMAGE_BASE_BUILD=$IMAGE_FROM
+  - RAYCI_IS_GPU_BUILD
+  - PYTHON
 tags:
-  - cr.ray.io/rayproject/llmbuild
+  - cr.ray.io/rayproject/$IMAGE_TO