add functions for creating ray with oauth proxy in front of the dashboard

Signed-off-by: Kevin <[email protected]>

Author: KPostOffice

src/codeflare_sdk/cluster/cluster.py

@@ -21,12 +21,20 @@
 from time import sleep
 from typing import List, Optional, Tuple, Dict
 
+import openshift as oc
+from kubernetes import config
 from ray.job_submission import JobSubmissionClient
+import urllib3
 
 from .auth import config_check, api_config_handler
 from ..utils import pretty_print
 from ..utils.generate_yaml import generate_appwrapper
 from ..utils.kube_api_helpers import _kube_api_error_handling
+from ..utils.openshift_oauth import (
+    create_openshift_oauth_objects,
+    delete_openshift_oauth_objects,
+    download_tls_cert,
+)
 from .config import ClusterConfiguration
 from .model import (
     AppWrapper,
@@ -41,6 +49,9 @@
 import requests
 
 
+k8_client = config.new_client_from_config()
+
+
 class Cluster:
     """
     An object for requesting, bringing up, and taking down resources.
@@ -61,6 +72,25 @@ def __init__(self, config: ClusterConfiguration):
         self.config = config
         self.app_wrapper_yaml = self.create_app_wrapper()
         self.app_wrapper_name = self.app_wrapper_yaml.split(".")[0]
+        self._client = None
+
+    @property
+    def client(self):
+        if self._client:
+            return self._client
+        if self.config.openshift_oauth:
+            self._client = JobSubmissionClient(
+                self.cluster_dashboard_uri(),
+                headers={
+                    "Authorization": k8_client.configuration.auth_settings()[
+                        "BearerToken"
+                    ]["value"]
+                },
+                verify=False,
+            )
+        else:
+            self._client = JobSubmissionClient(self.cluster_dashboard_uri())
+        return self._client
 
     def evaluate_dispatch_priority(self):
         priority_class = self.config.dispatch_priority
@@ -141,6 +171,7 @@ def create_app_wrapper(self):
             image_pull_secrets=image_pull_secrets,
             dispatch_priority=dispatch_priority,
             priority_val=priority_val,
+            openshift_oauth=self.config.openshift_oauth,
         )
 
     # creates a new cluster with the provided or default spec
@@ -150,6 +181,11 @@ def up(self):
         the MCAD queue.
         """
         namespace = self.config.namespace
+        if self.config.openshift_oauth:
+            create_openshift_oauth_objects(
+                cluster_name=self.config.name, namespace=namespace
+            )
+
         try:
             config_check()
             api_instance = client.CustomObjectsApi(api_config_handler())
@@ -184,6 +220,11 @@ def down(self):
         except Exception as e:  # pragma: no cover
             return _kube_api_error_handling(e)
 
+        if self.config.openshift_oauth:
+            delete_openshift_oauth_objects(
+                cluster_name=self.config.name, namespace=namespace
+            )
+
     def status(
         self, print_to_console: bool = True
     ) -> Tuple[CodeFlareClusterStatus, bool]:
@@ -252,7 +293,16 @@ def status(
         return status, ready
 
     def is_dashboard_ready(self) -> bool:
-        response = requests.get(self.cluster_dashboard_uri(), timeout=5)
+        try:
+            response = requests.get(
+                self.cluster_dashboard_uri(),
+                headers=self.client._headers,
+                timeout=5,
+                verify=self.client._verify,
+            )
+        except requests.exceptions.SSLError:
+            # SSL exception occurs when oauth ingress has been created but cluster is not up
+            return False
         if response.status_code == 200:
             return True
         else:
@@ -311,7 +361,13 @@ def cluster_dashboard_uri(self) -> str:
             return _kube_api_error_handling(e)
 
         for route in routes["items"]:
-            if route["metadata"]["name"] == f"ray-dashboard-{self.config.name}":
+            if route["metadata"][
+                "name"
+            ] == f"ray-dashboard-{self.config.name}" or route["metadata"][
+                "name"
+            ].startswith(
+                f"{self.config.name}-ingress"
+            ):
                 protocol = "https" if route["spec"].get("tls") else "http"
                 return f"{protocol}://{route['spec']['host']}"
         return "Dashboard route not available yet, have you run cluster.up()?"
@@ -320,30 +376,24 @@ def list_jobs(self) -> List:
         """
         This method accesses the head ray node in your cluster and lists the running jobs.
         """
-        dashboard_route = self.cluster_dashboard_uri()
-        client = JobSubmissionClient(dashboard_route)
-        return client.list_jobs()
+        return self.client.list_jobs()
 
     def job_status(self, job_id: str) -> str:
         """
         This method accesses the head ray node in your cluster and returns the job status for the provided job id.
         """
-        dashboard_route = self.cluster_dashboard_uri()
-        client = JobSubmissionClient(dashboard_route)
-        return client.get_job_status(job_id)
+        return self.client.get_job_status(job_id)
 
     def job_logs(self, job_id: str) -> str:
         """
         This method accesses the head ray node in your cluster and returns the logs for the provided job id.
         """
-        dashboard_route = self.cluster_dashboard_uri()
-        client = JobSubmissionClient(dashboard_route)
-        return client.get_job_logs(job_id)
+        return self.client.get_job_logs(job_id)
 
     def torchx_config(
         self, working_dir: str = None, requirements: str = None
     ) -> Dict[str, str]:
-        dashboard_address = f"{self.cluster_dashboard_uri().lstrip('http://')}"
+        dashboard_address = urllib3.util.parse_url(self.cluster_dashboard_uri()).host
         to_return = {
             "cluster_name": self.config.name,
             "dashboard_address": dashboard_address,
@@ -587,7 +637,13 @@ def _map_to_ray_cluster(rc) -> Optional[RayCluster]:
     )
     ray_route = None
     for route in routes["items"]:
-        if route["metadata"]["name"] == f"ray-dashboard-{rc['metadata']['name']}":
+        if route["metadata"][
+            "name"
+        ] == f"ray-dashboard-{rc['metadata']['name']}" or route["metadata"][
+            "name"
+        ].startswith(
+            f"{rc['metadata']['name']}-ingress"
+        ):
             protocol = "https" if route["spec"].get("tls") else "http"
             ray_route = f"{protocol}://{route['spec']['host']}"
 

src/codeflare_sdk/cluster/config.py

@@ -48,3 +48,4 @@ class ClusterConfiguration:
     local_interactive: bool = False
     image_pull_secrets: list = field(default_factory=list)
     dispatch_priority: str = None
+    openshift_oauth: bool = False  # NOTE: to use the user must have permission to create ClusterRoleBindings

src/codeflare_sdk/job/jobs.py

@@ -18,15 +18,20 @@
 from pathlib import Path
 
 from torchx.components.dist import ddp
-from torchx.runner import get_runner
+from torchx.runner import get_runner, Runner
+from torchx.schedulers.ray_scheduler import RayScheduler
 from torchx.specs import AppHandle, parse_app_handle, AppDryRunInfo
 
+from ray.job_submission import JobSubmissionClient
+
+import openshift as oc
+
 if TYPE_CHECKING:
     from ..cluster.cluster import Cluster
 from ..cluster.cluster import get_current_namespace
+from ..utils.openshift_oauth import download_tls_cert
 
 all_jobs: List["Job"] = []
-torchx_runner = get_runner()
 
 
 class JobDefinition(metaclass=abc.ABCMeta):
@@ -92,30 +97,37 @@ def __init__(
 
     def _dry_run(self, cluster: "Cluster"):
         j = f"{cluster.config.num_workers}x{max(cluster.config.num_gpus, 1)}"  # # of proc. = # of gpus
-        return torchx_runner.dryrun(
-            app=ddp(
-                *self.script_args,
-                script=self.script,
-                m=self.m,
-                name=self.name,
-                h=self.h,
-                cpu=self.cpu if self.cpu is not None else cluster.config.max_cpus,
-                gpu=self.gpu if self.gpu is not None else cluster.config.num_gpus,
-                memMB=self.memMB
-                if self.memMB is not None
-                else cluster.config.max_memory * 1024,
-                j=self.j if self.j is not None else j,
-                env=self.env,
-                max_retries=self.max_retries,
-                rdzv_port=self.rdzv_port,
-                rdzv_backend=self.rdzv_backend
-                if self.rdzv_backend is not None
-                else "static",
-                mounts=self.mounts,
+        runner = get_runner(ray_client=cluster.client)
+        runner._scheduler_instances["ray"] = RayScheduler(
+            session_name=runner._name, ray_client=cluster.client
+        )
+        return (
+            runner.dryrun(
+                app=ddp(
+                    *self.script_args,
+                    script=self.script,
+                    m=self.m,
+                    name=self.name,
+                    h=self.h,
+                    cpu=self.cpu if self.cpu is not None else cluster.config.max_cpus,
+                    gpu=self.gpu if self.gpu is not None else cluster.config.num_gpus,
+                    memMB=self.memMB
+                    if self.memMB is not None
+                    else cluster.config.max_memory * 1024,
+                    j=self.j if self.j is not None else j,
+                    env=self.env,
+                    max_retries=self.max_retries,
+                    rdzv_port=self.rdzv_port,
+                    rdzv_backend=self.rdzv_backend
+                    if self.rdzv_backend is not None
+                    else "static",
+                    mounts=self.mounts,
+                ),
+                scheduler=cluster.torchx_scheduler,
+                cfg=cluster.torchx_config(**self.scheduler_args),
+                workspace=self.workspace,
             ),
-            scheduler=cluster.torchx_scheduler,
-            cfg=cluster.torchx_config(**self.scheduler_args),
-            workspace=self.workspace,
+            runner,
         )
 
     def _missing_spec(self, spec: str):
@@ -125,41 +137,47 @@ def _dry_run_no_cluster(self):
         if self.scheduler_args is not None:
             if self.scheduler_args.get("namespace") is None:
                 self.scheduler_args["namespace"] = get_current_namespace()
-        return torchx_runner.dryrun(
-            app=ddp(
-                *self.script_args,
-                script=self.script,
-                m=self.m,
-                name=self.name if self.name is not None else self._missing_spec("name"),
-                h=self.h,
-                cpu=self.cpu
-                if self.cpu is not None
-                else self._missing_spec("cpu (# cpus per worker)"),
-                gpu=self.gpu
-                if self.gpu is not None
-                else self._missing_spec("gpu (# gpus per worker)"),
-                memMB=self.memMB
-                if self.memMB is not None
-                else self._missing_spec("memMB (memory in MB)"),
-                j=self.j
-                if self.j is not None
-                else self._missing_spec(
-                    "j (`workers`x`procs`)"
-                ),  # # of proc. = # of gpus,
-                env=self.env,  # should this still exist?
-                max_retries=self.max_retries,
-                rdzv_port=self.rdzv_port,  # should this still exist?
-                rdzv_backend=self.rdzv_backend
-                if self.rdzv_backend is not None
-                else "c10d",
-                mounts=self.mounts,
-                image=self.image
-                if self.image is not None
-                else self._missing_spec("image"),
+        runner = get_runner()
+        return (
+            runner.dryrun(
+                app=ddp(
+                    *self.script_args,
+                    script=self.script,
+                    m=self.m,
+                    name=self.name
+                    if self.name is not None
+                    else self._missing_spec("name"),
+                    h=self.h,
+                    cpu=self.cpu
+                    if self.cpu is not None
+                    else self._missing_spec("cpu (# cpus per worker)"),
+                    gpu=self.gpu
+                    if self.gpu is not None
+                    else self._missing_spec("gpu (# gpus per worker)"),
+                    memMB=self.memMB
+                    if self.memMB is not None
+                    else self._missing_spec("memMB (memory in MB)"),
+                    j=self.j
+                    if self.j is not None
+                    else self._missing_spec(
+                        "j (`workers`x`procs`)"
+                    ),  # # of proc. = # of gpus,
+                    env=self.env,  # should this still exist?
+                    max_retries=self.max_retries,
+                    rdzv_port=self.rdzv_port,  # should this still exist?
+                    rdzv_backend=self.rdzv_backend
+                    if self.rdzv_backend is not None
+                    else "c10d",
+                    mounts=self.mounts,
+                    image=self.image
+                    if self.image is not None
+                    else self._missing_spec("image"),
+                ),
+                scheduler="kubernetes_mcad",
+                cfg=self.scheduler_args,
+                workspace="",
             ),
-            scheduler="kubernetes_mcad",
-            cfg=self.scheduler_args,
-            workspace="",
+            runner,
         )
 
     def submit(self, cluster: "Cluster" = None) -> "Job":
@@ -171,18 +189,20 @@ def __init__(self, job_definition: "DDPJobDefinition", cluster: "Cluster" = None
         self.job_definition = job_definition
         self.cluster = cluster
         if self.cluster:
-            self._app_handle = torchx_runner.schedule(job_definition._dry_run(cluster))
+            definition, runner = job_definition._dry_run(cluster)
+            self._app_handle = runner.schedule(definition)
+            self._runner = runner
         else:
-            self._app_handle = torchx_runner.schedule(
-                job_definition._dry_run_no_cluster()
-            )
+            definition, runner = job_definition._dry_run_no_cluster()
+            self._app_handle = runner.schedule(definition)
+            self._runner = runner
         all_jobs.append(self)
 
     def status(self) -> str:
-        return torchx_runner.status(self._app_handle)
+        return self._runner.status(self._app_handle)
 
     def logs(self) -> str:
-        return "".join(torchx_runner.log_lines(self._app_handle, None))
+        return "".join(self._runner.log_lines(self._app_handle, None))
 
     def cancel(self):
-        torchx_runner.cancel(self._app_handle)
+        self._runner.cancel(self._app_handle)