cherrypick #50860 (#50867)

fix 2 serve release tests. Signed-off-by: Lonnie Liu <[email protected]>
ray-project · Feb 24, 2025 · 84f2764 · 84f2764
1 parent ecd0709
commit 84f2764
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 7 deletions.
diff --git a/release/serve_tests/workloads/anyscale_service_utils.py b/release/serve_tests/workloads/anyscale_service_utils.py
@@ -15,8 +15,10 @@
 logging.basicConfig(level=logging.INFO)
 
 
-def check_service_state(service_name: str, expected_state: ServiceState):
-    state = service.status(name=service_name).state
+def check_service_state(
+    service_name: str, expected_state: ServiceState, cloud: Optional[str] = None
+):
+    state = service.status(name=service_name, cloud=cloud).state
     logger.info(
         f"Waiting for service {service_name} to be {expected_state}, currently {state}"
     )
@@ -34,6 +36,7 @@ def start_service(
     image_uri: Optional[str] = None,
     working_dir: Optional[str] = None,
     add_unique_suffix: bool = True,
+    cloud: Optional[str] = None,
 ):
     """Starts an Anyscale Service with the specified configs.
 
@@ -76,18 +79,20 @@ def start_service(
             expected_state="RUNNING",
             retry_interval_ms=10000,  # 10s
             timeout=600,
+            cloud=cloud,
         )
 
         yield service_name
 
     finally:
         logger.info(f"Terminating service {service_name}.")
-        service.terminate(name=service_name)
+        service.terminate(name=service_name, cloud=cloud)
         wait_for_condition(
             check_service_state,
             service_name=service_name,
             expected_state="TERMINATED",
             retry_interval_ms=10000,  # 10s
             timeout=600,
+            cloud=cloud,
         )
         logger.info(f"Service '{service_name}' terminated successfully.")
diff --git a/release/serve_tests/workloads/autoscaling_load_test.py b/release/serve_tests/workloads/autoscaling_load_test.py
@@ -31,6 +31,7 @@
 
 
 URI = "https://serve-resnet-benchmark-data.s3.us-west-1.amazonaws.com/000000000019.jpeg"
+CLOUD = "serve_release_tests_cloud"
 
 
 @click.command()
@@ -47,7 +48,7 @@ def main(output_path: Optional[str], image_uri: Optional[str]):
         ],
     }
     compute_config = ComputeConfig(
-        cloud="serve_release_tests_cloud",
+        cloud=CLOUD,
         head_node=HeadNodeConfig(instance_type="m5.8xlarge"),
         worker_nodes=[
             WorkerNodeGroupConfig(
@@ -67,9 +68,10 @@ def main(output_path: Optional[str], image_uri: Optional[str]):
         compute_config=compute_config,
         applications=[resnet_application],
         working_dir="workloads",
+        cloud=CLOUD,
     ) as service_name:
         ray.init(address="auto")
-        status = service.status(name=service_name)
+        status = service.status(name=service_name, cloud=CLOUD)
 
         # Start the locust workload
         num_locust_workers = int(ray.available_resources()["CPU"]) - 1

diff --git a/release/serve_tests/workloads/replica_scalability.py b/release/serve_tests/workloads/replica_scalability.py
@@ -22,6 +22,7 @@
 
 DEFAULT_FULL_TEST_NUM_REPLICA = 1000
 DEFAULT_FULL_TEST_TRIAL_LENGTH_S = 60
+CLOUD = "serve_release_tests_cloud"
 
 
 @click.command()
@@ -52,7 +53,7 @@ def main(
         ],
     }
     compute_config = ComputeConfig(
-        cloud="serve_release_tests_cloud",
+        cloud=CLOUD,
         head_node=HeadNodeConfig(instance_type="m5.8xlarge"),
         worker_nodes=[
             WorkerNodeGroupConfig(
@@ -92,9 +93,10 @@ def main(
         compute_config=compute_config,
         applications=[noop_1k_application],
         working_dir="workloads",
+        cloud=CLOUD,
     ) as service_name:
         ray.init("auto")
-        status = service.status(name=service_name)
+        status = service.status(name=service_name, cloud=CLOUD)
 
         # Start the locust workload
         num_locust_workers = int(ray.available_resources()["CPU"]) - 1