Skip to content

Commit

Permalink
Adding more cluster status tests (#992)
Browse files Browse the repository at this point in the history
Co-authored-by: Alexandra Belousov <sashabelousovrh@Alexandras-MacBook-Pro.local>
  • Loading branch information
BelSasha and Alexandra Belousov committed Jul 15, 2024
1 parent f2a2ed2 commit 923e9e6
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 0 deletions.
1 change: 1 addition & 0 deletions tests/test_resources/test_clusters/test_cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -680,6 +680,7 @@ def test_rh_status_stopped(self, cluster):
@pytest.mark.level("local")
@pytest.mark.clustertest
def test_send_status_to_db(self, cluster):

import json

cluster.save()
Expand Down
50 changes: 50 additions & 0 deletions tests/test_resources/test_clusters/test_on_demand_cluster.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import asyncio
import subprocess
import time

import pytest
Expand Down Expand Up @@ -181,6 +182,9 @@ def test_fn_to_docker_container(self, ondemand_aws_cluster):
remote_torch_exists = rh.function(torch_exists).to(ondemand_aws_cluster)
assert remote_torch_exists()

####################################################################################################
# Status tests
####################################################################################################
@pytest.mark.level("minimal")
@pytest.mark.skip("Test requires terminating the cluster")
def test_set_status_after_teardown(self, cluster):
Expand All @@ -202,6 +206,52 @@ def test_set_status_after_teardown(self, cluster):
assert get_status_data["resource_type"] == cluster_config.get("resource_type")
assert get_status_data["status"] == ResourceServerStatus.terminated

@pytest.mark.level("minimal")
def test_status_autostop_cluster(self, cluster):
cluster_config = cluster.config()
cluster_uri = rns_client.format_rns_address(cluster.rns_address)
api_server_url = cluster_config.get("api_server_url", rns_client.api_server_url)
cluster_name_no_owner = cluster.rns_address.split("/")[-1]

# Mocking autostop by running sky down
result_teardown = subprocess.run(
["sky", "down", "-y", cluster_name_no_owner], capture_output=True, text=True
)
assert result_teardown.returncode == 0

get_status_data_resp = requests.get(
f"{api_server_url}/resource/{cluster_uri}/cluster/status",
headers=rns_client.request_headers(),
)
assert get_status_data_resp.status_code == 200
# For UI displaying purposes, the cluster/status endpoint returns cluster status history.
# The latest status info is the first element in the list returned by the endpoint.
get_status_data = get_status_data_resp.json()["data"][0]
assert get_status_data["resource_type"] == cluster_config.get("resource_type")
assert get_status_data["status"] == ResourceServerStatus.terminated

@pytest.mark.level("minimal")
def test_status_cluster_rh_daemon_stopped(self, cluster):
cluster_config = cluster.config()
cluster_uri = rns_client.format_rns_address(cluster.rns_address)
api_server_url = cluster_config.get("api_server_url", rns_client.api_server_url)

cluster.run(["runhouse stop"])

get_status_data_resp = requests.get(
f"{api_server_url}/resource/{cluster_uri}/cluster/status",
headers=rns_client.request_headers(),
)
assert get_status_data_resp.status_code == 200
# For UI displaying purposes, the cluster/status endpoint returns cluster status history.
# The latest status info is the first element in the list returned by the endpoint.
get_status_data = get_status_data_resp.json()["data"][0]
assert get_status_data["resource_type"] == cluster_config.get("resource_type")
assert get_status_data["status"] == ResourceServerStatus.runhouse_daemon_down

####################################################################################################
# Logs surfacing tests
####################################################################################################
@pytest.mark.level("minimal")
def test_logs_surfacing_scheduler_basic_flow(self, cluster):

Expand Down

0 comments on commit 923e9e6

Please sign in to comment.