[autoscaler][kubernetes] Ray client setup, example config simplification, example scripts. (#13920)

python/ray/tests/test_k8s_operator_examples.py

@@ -20,7 +20,7 @@ NAMESPACE = "test-k8s-operator-examples"
 def retry_until_true(f):
     # Retry 60 times with 1 second delay between attempts.
     def f_with_retries(*args, **kwargs):
-        for _ in range(60):
+        for _ in range(120):
             if f(*args, **kwargs):
                 return
             else:
@@ -47,25 +47,38 @@ def wait_for_logs():
     cmd = f"kubectl -n {NAMESPACE} logs ray-operator-pod"\
         "| grep ^example-cluster: | tail -n 100"
     log_tail = subprocess.check_output(cmd, shell=True).decode()
-    return ("head-node" in log_tail) and ("worker-nodes" in log_tail)
+    return ("head-node" in log_tail) and ("worker-node" in log_tail)
 
 
-def operator_configs_directory():
+@retry_until_true
+def wait_for_job(job_pod):
+    cmd = f"kubectl -n {NAMESPACE} logs {job_pod}"
+    out = subprocess.check_output(cmd, shell=True).decode()
+    return ("success" in out.lower())
+
+
+def kubernetes_configs_directory():
     here = os.path.realpath(__file__)
     ray_python_root = os.path.dirname(os.path.dirname(here))
-    relative_path = "autoscaler/kubernetes/operator_configs"
+    relative_path = "autoscaler/kubernetes"
     return os.path.join(ray_python_root, relative_path)
 
 
+def get_kubernetes_config_path(name):
+    return os.path.join(kubernetes_configs_directory(), name)
+
+
 def get_operator_config_path(file_name):
-    return os.path.join(operator_configs_directory(), file_name)
+    operator_configs = get_kubernetes_config_path("operator_configs")
+    return os.path.join(operator_configs, file_name)
 
 
 class KubernetesOperatorTest(unittest.TestCase):
     def test_examples(self):
         with tempfile.NamedTemporaryFile("w+") as example_cluster_file, \
                 tempfile.NamedTemporaryFile("w+") as example_cluster2_file,\
-                tempfile.NamedTemporaryFile("w+") as operator_file:
+                tempfile.NamedTemporaryFile("w+") as operator_file,\
+                tempfile.NamedTemporaryFile("w+") as job_file:
 
             # Get paths to operator configs
             example_cluster_config_path = get_operator_config_path(
@@ -73,6 +86,7 @@ class KubernetesOperatorTest(unittest.TestCase):
             example_cluster2_config_path = get_operator_config_path(
                 "example_cluster2.yaml")
             operator_config_path = get_operator_config_path("operator.yaml")
+            job_path = get_kubernetes_config_path("job-example.yaml")
             self.crd_path = get_operator_config_path("cluster_crd.yaml")
 
             # Load operator configs
@@ -82,19 +96,23 @@ class KubernetesOperatorTest(unittest.TestCase):
                 open(example_cluster2_config_path).read())
             operator_config = list(
                 yaml.safe_load_all(open(operator_config_path).read()))
+            job_config = yaml.safe_load(open(job_path).read())
 
             # Fill image fields
             podTypes = example_cluster_config["spec"]["podTypes"]
             podTypes2 = example_cluster2_config["spec"]["podTypes"]
-            pod_configs = ([operator_config[-1]] + [
-                podType["podConfig"] for podType in podTypes
-            ] + [podType["podConfig"] for podType in podTypes2])
-            for pod_config in pod_configs:
-                pod_config["spec"]["containers"][0]["image"] = IMAGE
+            pod_specs = ([operator_config[-1]["spec"]] + [
+                job_config["spec"]["template"]["spec"]
+            ] + [podType["podConfig"]["spec"] for podType in podTypes
+                 ] + [podType["podConfig"]["spec"] for podType in podTypes2])
+            for pod_spec in pod_specs:
+                pod_spec["containers"][0]["image"] = IMAGE
+                pod_spec["containers"][0]["imagePullPolicy"] = "IfNotPresent"
 
             # Dump to temporary files
             yaml.dump(example_cluster_config, example_cluster_file)
             yaml.dump(example_cluster2_config, example_cluster2_file)
+            yaml.dump(job_config, job_file)
             yaml.dump_all(operator_config, operator_file)
             files = [
                 example_cluster_file, example_cluster2_file, operator_file
@@ -131,6 +149,19 @@ class KubernetesOperatorTest(unittest.TestCase):
             # Four pods remain
             wait_for_pods(4)
 
+            # Check job submission
+            cmd = f"kubectl -n {NAMESPACE} create -f {job_file.name}"
+            subprocess.check_call(cmd, shell=True)
+
+            cmd = f"kubectl -n {NAMESPACE} get pods --no-headers -o"\
+                " custom-columns=\":metadata.name\""
+            pods = subprocess.check_output(cmd, shell=True).decode().split()
+            job_pod = [pod for pod in pods if "job" in pod].pop()
+            time.sleep(10)
+            wait_for_job(job_pod)
+            cmd = f"kubectl -n {NAMESPACE} delete jobs --all"
+            subprocess.check_call(cmd, shell=True)
+
             # Check that cluster updates work: increase minWorkers to 3
             # and check that one worker is created.
             example_cluster_edit = copy.deepcopy(example_cluster_config)