add multi-user jobs smoke test

Author: cg505

tests/smoke_tests/test_api_server.py

@@ -2,9 +2,18 @@
 
 from smoke_tests import smoke_tests_utils
 
+import sky
 from sky.skylet import constants
 
 
+def set_user(user_id: str, user_name: str, commands: List[str]) -> List[str]:
+    return [
+        f'export {constants.USER_ID_ENV_VAR}="{user_id}"; '
+        f'export {constants.USER_ENV_VAR}="{user_name}"; ' + cmd
+        for cmd in commands
+    ]
+
+
 # ---------- Test multi-tenant ----------
 def test_multi_tenant(generic_cloud: str):
     name = smoke_tests_utils.get_cluster_name()
@@ -13,14 +22,6 @@ def test_multi_tenant(generic_cloud: str):
     user_2 = 'abcdef13'
     user_2_name = 'user2'
 
-    def set_user(user_id: str, user_name: str,
-                 commands: List[str]) -> List[str]:
-        return [
-            f'export {constants.USER_ID_ENV_VAR}="{user_id}"; '
-            f'export {constants.USER_ENV_VAR}="{user_name}"; ' + cmd
-            for cmd in commands
-        ]
-
     stop_test_cmds = [
         'echo "==== Test multi-tenant cluster stop ===="',
         *set_user(
@@ -94,3 +95,49 @@ def set_user(user_id: str, user_name: str,
         f'sky down -y {name}-1 {name}-2',
     )
     smoke_tests_utils.run_one_test(test)
+
+
+def test_multi_tenant_managed_jobs(generic_cloud: str):
+    name = smoke_tests_utils.get_cluster_name()
+    user_1 = 'abcdef12'
+    user_1_name = 'user1'
+    user_2 = 'abcdef13'
+    user_2_name = 'user2'
+
+    test = smoke_tests_utils.Test(
+        'test_multi_tenant_managed_jobs',
+        [
+            'echo "==== Test multi-tenant managed jobs ===="',
+            *set_user(user_1, user_1_name, [
+                f'sky jobs launch -n {name}-1 --cloud {generic_cloud} --cpus 2+ tests/test_yamls/minimal.yaml -y',
+                f's=$(sky jobs queue) && echo "$s" && echo "$s" | grep {name}-1 | grep SUCCEEDED',
+                f's=$(sky jobs queue -u) && echo "$s" && echo "$s" | grep {user_1_name} | grep {name}-1 | grep SUCCEEDED',
+            ]), *set_user(user_2, user_2_name, [
+                f's=$(sky jobs queue) && echo "$s" && echo "$s" | grep {name}-1 && exit 1 || true',
+                f's=$(sky jobs queue -u) && echo "$s" && echo "$s" | grep {user_1_name} | grep {name}-1 | grep SUCCEEDED',
+                f'sky jobs launch -n {name}-2 --cloud {generic_cloud} --cpus 2+ tests/test_yamls/minimal.yaml -y',
+                f's=$(sky jobs queue) && echo "$s" && echo "$s" | grep {name}-2 | grep SUCCEEDED',
+                f's=$(sky jobs queue -u) && echo "$s" && echo "$s" | grep {user_2_name} | grep {name}-2 | grep SUCCEEDED',
+            ]), 'echo "==== Test jobs controller cluster user ===="',
+            f's=$(sky status -u) && echo "$s" && echo "$s" | grep sky-jobs-controller- | grep -v {user_1_name} | grep -v {user_2_name}'
+            'echo "==== Test controller down blocked by other users ===="',
+            *set_user(user_1, user_1_name, [
+                f'sky jobs launch -n {name}-3 --cloud {generic_cloud} --cpus 2+ -y -d sleep 1000',
+                smoke_tests_utils.
+                get_cmd_wait_until_managed_job_status_contains_matching_job_name(
+                    job_name=f'{name}-3',
+                    job_status=[
+                        sky.ManagedJobStatus.PENDING,
+                        sky.ManagedJobStatus.SUBMITTED,
+                        sky.ManagedJobStatus.STARTING,
+                        sky.ManagedJobStatus.RUNNING
+                    ],
+                    timeout=60),
+            ]), *set_user(user_2, user_2_name, [
+                f'sky down sky-jobs-controller-\'*\' -y -p && exit 1 || true',
+                f'sky jobs cancel -y -n {name}-3',
+            ])
+        ],
+        f'sky jobs cancel -y -n {name}-1; sky jobs cancel -y -n {name}-2; sky jobs cancel -y -n {name}-3',
+    )
+    smoke_tests_utils.run_one_test(test)