update module tests to use den-launched cluster (#1476)

Co-authored-by: Alexandra Belousov <[email protected]>

tests/conftest.py

@@ -260,7 +260,11 @@ def event_loop():
     unnamed_resource,  # noqa: F401
 )
 
-from tests.fixtures.static_cluster_fixtures import static_cpu_pwd_cluster  # noqa: F401
+from tests.fixtures.static_cluster_fixtures import (  # noqa: F401
+    static_cpu_pwd_cluster,
+    static_cpu_pwd_cluster_den_launcher,
+    static_gpu_pwd_cluster_den_launcher,
+)
 
 
 # ----------------- Folders -----------------

tests/fixtures/static_cluster_fixtures.py

@@ -1,54 +1,79 @@
+from enum import Enum
+from typing import Union
+
 import pytest
 
 import runhouse as rh
+from runhouse.resources.hardware.utils import LauncherType
 
 from tests.conftest import init_args
 from tests.utils import setup_test_base
 
 
-@pytest.fixture(scope="session")
-def static_cpu_pwd_cluster(test_rns_folder):
+class computeType(str, Enum):
+    cpu = "cpu"
+    gpu = "gpu"
+
+
+def setup_static_cluster(
+    launcher: Union[LauncherType, str] = None,
+    compute_type: computeType = computeType.cpu,
+):
+    from tests.fixtures.resource_fixtures import create_folder_path
+
     rh.constants.SSH_SKY_SECRET_NAME = (
-        f"{test_rns_folder}-{rh.constants.SSH_SKY_SECRET_NAME}"
+        f"{create_folder_path()}-{rh.constants.SSH_SKY_SECRET_NAME}"
     )
-    sky_cluster = rh.cluster(
-        "aws-cpu-password", instance_type="CPU:4", provider="aws"
+    instance_type = "CPU:4" if compute_type == computeType.cpu else "g5.xlarge"
+    cluster = rh.cluster(
+        f"aws-{compute_type}-password",
+        instance_type=instance_type,
+        provider="aws",
+        launcher=launcher,
     ).save()
-    if not sky_cluster.is_up():
-        sky_cluster.up()
+    if not cluster.is_up():
+        cluster.up()
 
         # set up password on remote
-        sky_cluster.run_bash(
+        cluster.run_bash(
             [
-                [
-                    'sudo sed -i "/^[^#]*PasswordAuthentication[[:space:]]no/c\PasswordAuthentication yes" '
-                    "/etc/ssh/sshd_config"
-                ]
+                'sudo sed -i "/^[^#]*PasswordAuthentication[[:space:]]no/c\PasswordAuthentication yes" /etc/ssh/sshd_config'
             ]
         )
-        sky_cluster.run_bash(["sudo /etc/init.d/ssh force-reload"])
-        sky_cluster.run_bash(["sudo /etc/init.d/ssh restart"])
-        sky_cluster.run_bash(
+        cluster.run_bash(["sudo /etc/init.d/ssh force-reload"])
+        cluster.run_bash(["sudo /etc/init.d/ssh restart"])
+        cluster.run_bash(
             ["(echo 'cluster-pass' && echo 'cluster-pass') | sudo passwd ubuntu"]
         )
-        sky_cluster.run_bash(
-            ["pip uninstall skypilot runhouse -y", "pip install pytest"]
-        )
-        sky_cluster.run_bash(["rm -rf runhouse/"])
+        cluster.run_bash(["pip uninstall skypilot runhouse -y", "pip install pytest"])
+        cluster.run_bash(["rm -rf runhouse/"])
 
     # instantiate byo cluster with password
     ssh_creds = {
         "ssh_user": "ubuntu",
         "ssh_private_key": "~/.ssh/sky-key",
         "password": "cluster-pass",
     }
-    args = dict(
-        name="static-cpu-password", host=[sky_cluster.head_ip], ssh_creds=ssh_creds
-    )
+    args = dict(name="static-cpu-password", host=[cluster.head_ip], ssh_creds=ssh_creds)
     c = rh.cluster(**args).save()
     c.restart_server(resync_rh=True)
     init_args[id(c)] = args
 
     setup_test_base(c)
 
     return c
+
+
+@pytest.fixture(scope="session")
+def static_cpu_pwd_cluster():
+    return setup_static_cluster()
+
+
+@pytest.fixture(scope="session")
+def static_cpu_pwd_cluster_den_launcher():
+    return setup_static_cluster(launcher=LauncherType.DEN)
+
+
+@pytest.fixture(scope="session")
+def static_gpu_pwd_cluster_den_launcher():
+    return setup_static_cluster(launcher=LauncherType.DEN, compute_type=computeType.gpu)

tests/test_resources/test_modules/test_module.py

@@ -253,6 +253,40 @@ def nested_call_logs_stream_helper(slow_numpy_array):
 @pytest.mark.moduletest
 class TestModule:
 
+    MAP_FIXTURES = {"resource": "cluster"}
+
+    UNIT = {"cluster": ["named_cluster"]}
+    LOCAL = {
+        "cluster": [
+            "docker_cluster_pk_ssh_no_auth",  # Represents private dev use case
+            "docker_cluster_pk_ssh_den_auth",  # Helps isolate Auth issues
+            "docker_cluster_pk_http_exposed",  # Represents within VPC use case
+            "docker_cluster_pwd_ssh_no_auth",
+        ],
+    }
+    MINIMAL = {
+        "cluster": [
+            "static_cpu_pwd_cluster",
+            "static_cpu_pwd_cluster_den_launcher",
+        ]
+    }
+    RELEASE = {
+        "cluster": [
+            "static_cpu_pwd_cluster_den_launcher",  # tests modules on a den-launched cpu cluster
+            "static_gpu_pwd_cluster_den_launcher",  # tests modules on a den-launched gpu cluster
+            "static_cpu_pwd_cluster",  # tests modules on a local (sky) -launched cpu cluster
+        ]
+    }
+    MAXIMAL = {
+        "cluster": [
+            "docker_cluster_pk_ssh_no_auth",
+            "docker_cluster_pk_ssh_den_auth",
+            "docker_cluster_pwd_ssh_no_auth",
+            "static_cpu_pwd_cluster",
+            "multinode_cpu_docker_conda_cluster",
+        ]
+    }
+
     # --------- integration tests ---------
     @pytest.mark.level("local")
     def test_call_module_method(self, cluster):