From 284617dd54cc1bd1866de4d9d41c4c5525b892ad Mon Sep 17 00:00:00 2001
From: Cuong Nguyen <128072568+can-anyscale@users.noreply.github.com>
Date: Fri, 8 Dec 2023 16:53:10 -0800
Subject: [PATCH] [ci][wins/3] add --operating-system flag to rayci (#41636)

Add --operating-system flag. Either use LinuxTesterContainer or WindowsTesterContainer based on the flag. The only if condition!

Signed-off-by: can <can@anyscale.com>
---
 ci/ray_ci/test_tester.py | 24 ++++++++++++++++++++-
 ci/ray_ci/tester.py      | 45 +++++++++++++++++++++++++++++++---------
 2 files changed, 58 insertions(+), 11 deletions(-)

diff --git a/ci/ray_ci/test_tester.py b/ci/ray_ci/test_tester.py
index d6212f959da9..b447c7a10022 100644
--- a/ci/ray_ci/test_tester.py
+++ b/ci/ray_ci/test_tester.py
@@ -7,6 +7,7 @@
 import pytest
 
 from ci.ray_ci.linux_tester_container import LinuxTesterContainer
+from ci.ray_ci.windows_tester_container import WindowsTesterContainer
 from ci.ray_ci.tester import (
     _add_default_except_tags,
     _get_container,
@@ -34,12 +35,33 @@ def test_get_container() -> None:
     with mock.patch(
         "ci.ray_ci.linux_tester_container.LinuxTesterContainer.install_ray",
         return_value=None,
+    ), mock.patch(
+        "ci.ray_ci.windows_tester_container.WindowsTesterContainer.install_ray",
+        return_value=None,
     ):
-        container = _get_container("core", 3, 1, 2, 0)
+        container = _get_container(
+            team="core",
+            operating_system="linux",
+            workers=3,
+            worker_id=1,
+            parallelism_per_worker=2,
+            gpus=0,
+        )
+        assert isinstance(container, LinuxTesterContainer)
         assert container.docker_tag == "corebuild"
         assert container.shard_count == 6
         assert container.shard_ids == [2, 3]
 
+        container = _get_container(
+            team="serve",
+            operating_system="windows",
+            workers=3,
+            worker_id=1,
+            parallelism_per_worker=2,
+            gpus=0,
+        )
+        assert isinstance(container, WindowsTesterContainer)
+
 
 def test_get_test_targets() -> None:
     _TEST_YAML = "flaky_tests: [//python/ray/tests:flaky_test_01]"
diff --git a/ci/ray_ci/tester.py b/ci/ray_ci/tester.py
index 6aed1c19f9d7..d0d79cc70dfc 100644
--- a/ci/ray_ci/tester.py
+++ b/ci/ray_ci/tester.py
@@ -13,6 +13,7 @@
     DEFAULT_ARCHITECTURE,
 )
 from ci.ray_ci.linux_tester_container import LinuxTesterContainer
+from ci.ray_ci.windows_tester_container import WindowsTesterContainer
 from ci.ray_ci.tester_container import TesterContainer
 from ci.ray_ci.utils import docker_login
 
@@ -135,12 +136,19 @@
     ),
     default="optimized",
 )
+@click.option(
+    "--operating-system",
+    default="linux",
+    type=click.Choice(["linux", "windows"]),
+    help=("Operating system to run tests on"),
+)
 def main(
     targets: List[str],
     team: str,
     workers: int,
     worker_id: int,
     parallelism_per_worker: int,
+    operating_system: str,
     except_tags: str,
     only_tags: str,
     run_flaky_tests: bool,
@@ -155,7 +163,10 @@ def main(
     if not bazel_workspace_dir:
         raise Exception("Please use `bazelisk run //ci/ray_ci`")
     os.chdir(bazel_workspace_dir)
-    docker_login(_DOCKER_ECR_REPO.split("/")[0])
+    # TODO(can): only linux uses intermediate dockers publised to ECR; remove this when
+    # we can build intermediate dockers for Windows
+    if operating_system == "linux":
+        docker_login(_DOCKER_ECR_REPO.split("/")[0])
 
     if build_type == "wheel" or build_type == "wheel-aarch64":
         # for wheel testing, we first build the wheel and then use it for running tests
@@ -163,6 +174,7 @@ def main(
         BuilderContainer(DEFAULT_PYTHON_VERSION, DEFAULT_BUILD_TYPE, architecture).run()
     container = _get_container(
         team,
+        operating_system,
         workers,
         worker_id,
         parallelism_per_worker,
@@ -195,6 +207,7 @@ def _add_default_except_tags(except_tags: str) -> str:
 
 def _get_container(
     team: str,
+    operating_system: str,
     workers: int,
     worker_id: int,
     parallelism_per_worker: int,
@@ -208,15 +221,27 @@ def _get_container(
     shard_start = worker_id * parallelism_per_worker
     shard_end = (worker_id + 1) * parallelism_per_worker
 
-    return LinuxTesterContainer(
-        build_name or f"{team}build",
-        test_envs=test_env,
-        shard_count=shard_count,
-        shard_ids=list(range(shard_start, shard_end)),
-        gpus=gpus,
-        skip_ray_installation=skip_ray_installation,
-        build_type=build_type,
-    )
+    if operating_system == "linux":
+        return LinuxTesterContainer(
+            build_name or f"{team}build",
+            test_envs=test_env,
+            shard_count=shard_count,
+            shard_ids=list(range(shard_start, shard_end)),
+            gpus=gpus,
+            skip_ray_installation=skip_ray_installation,
+            build_type=build_type,
+        )
+
+    if operating_system == "windows":
+        return WindowsTesterContainer(
+            build_name or f"{team}build",
+            test_envs=test_env,
+            shard_count=shard_count,
+            shard_ids=list(range(shard_start, shard_end)),
+            skip_ray_installation=skip_ray_installation,
+        )
+
+    assert False, f"Unsupported operating system: {operating_system}"
 
 
 def _get_tag_matcher(tag: str) -> str: