[ci] Separate Windows Serve tests into their own build (ray-project#40057)

The windows tests are quite flaky and take a long time to run. When I'm merging Serve PRs, I want to be able to quickly find and evaluate Serve-specific tests on Windows. This will help avoid introducing failures/flakiness.

Author: edoakes

.buildkite/pipeline.windows.yml

@@ -28,8 +28,8 @@ upload_wheels_if_needed: &upload_wheels_if_needed |-
 
 
 steps:
-- label: ":windows: Build & Test"
-  parallelism: 6
+- label: ":windows: Core + Python tests (excluding serve)"
+  parallelism: 5
   commands:
     - *prelude_commands
     # conda init should be moved to the docker image setup
@@ -41,6 +41,18 @@ steps:
     # The next command will be sharded into $parallelism shards.
     - ./ci/ci.sh test_python
 
+- label: ":windows: Serve tests"
+  parallelism: 2
+  commands:
+    - *prelude_commands
+    # conda init should be moved to the docker image setup
+    - conda init
+    - . ./ci/ci.sh init
+    - ./ci/ci.sh build
+    - export RAY_ENABLE_WINDOWS_OR_OSX_CLUSTER="1"
+    # The next command will be sharded into $parallelism shards.
+    - ./ci/ci.sh test_serve
+
 - label: ":windows: Wheels"
   commands:
     - *prelude_commands

ci/ci.sh

@@ -206,15 +206,13 @@ prepare_docker() {
     popd
 }
 
-# For running Python tests on Windows.
-test_python() {
+# For running Serve tests on Windows.
+test_serve() {
   local pathsep=":" args=()
   if [ "${OSTYPE}" = msys ]; then
     pathsep=";"
     args+=(
       python/ray/serve/...
-      python/ray/tests/...
-      python/ray/train:test_windows
       -python/ray/serve/tests:test_cross_language # Ray java not built on Windows yet.
       -python/ray/serve/tests:test_gcs_failure # Fork not supported in windows
       -python/ray/serve/tests:test_standalone_2 # Multinode not supported on Windows
@@ -224,6 +222,41 @@ test_python() {
       -python/ray/serve/tests:test_fastapi
       -python/ray/serve/tests:test_get_deployment # address violation
       -python/ray/serve/tests:test_proxy_state # skip running on Windows until we refactor the test
+    )
+  fi
+  if [ 0 -lt "${#args[@]}" ]; then  # Any targets to test?
+    install_ray
+
+    # Shard the args.
+    BUILDKITE_PARALLEL_JOB=${BUILDKITE_PARALLEL_JOB:-'0'}
+    BUILDKITE_PARALLEL_JOB_COUNT=${BUILDKITE_PARALLEL_JOB_COUNT:-'1'}
+    test_shard_selection=$(python ./ci/ray_ci/bazel_sharding.py --exclude_manual --index "${BUILDKITE_PARALLEL_JOB}" --count "${BUILDKITE_PARALLEL_JOB_COUNT}" "${args[@]}")
+
+    # TODO(mehrdadn): We set PYTHONPATH here to let Python find our pickle5 under pip install -e.
+    # It's unclear to me if this should be necessary, but this is to make tests run for now.
+    # Check why this issue doesn't arise on Linux/Mac.
+    # Ideally importing ray.cloudpickle should import pickle5 automatically.
+    # shellcheck disable=SC2046,SC2086
+    bazel test --config=ci \
+      --build_tests_only $(./ci/run/bazel_export_options) \
+      --test_env=PYTHONPATH="${PYTHONPATH-}${pathsep}${WORKSPACE_DIR}/python/ray/pickle5_files" \
+      --test_env=CI="1" \
+      --test_env=RAY_CI_POST_WHEEL_TESTS="1" \
+      --test_env=USERPROFILE="${USERPROFILE}" \
+      --test_output=streamed \
+      -- \
+      ${test_shard_selection};
+  fi
+}
+
+# For running Python tests on Windows (excluding Serve).
+test_python() {
+  local pathsep=":" args=()
+  if [ "${OSTYPE}" = msys ]; then
+    pathsep=";"
+    args+=(
+      python/ray/tests/...
+      python/ray/train:test_windows
       -python/ray/tests:test_actor_advanced  # crashes in shutdown
       -python/ray/tests:test_autoscaler # We don't support Autoscaler on Windows
       -python/ray/tests:test_autoscaler_aws