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

Signed-off-by: Kevin <kpostlet@redhat.com>

Author: KPostOffice

src/codeflare_sdk/cluster/cluster.py

@@ -21,12 +21,16 @@
 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 +45,9 @@
 import requests
 
 
+k8_client = config.new_client_from_config()
+
+
 class Cluster:
     """
     An object for requesting, bringing up, and taking down resources.
@@ -61,6 +68,21 @@ 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 +163,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 +173,9 @@ 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 +210,9 @@ 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 +281,13 @@ 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 +346,8 @@ 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 +356,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 +617,8 @@ 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,16 +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):
     def _dry_run(self, cluster: "Cluster"):
@@ -92,7 +96,9 @@ 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(
+        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,
@@ -116,7 +122,7 @@ def _dry_run(self, cluster: "Cluster"):
             scheduler=cluster.torchx_scheduler,
             cfg=cluster.torchx_config(**self.scheduler_args),
             workspace=self.workspace,
-        )
+        ), runner
 
     def _missing_spec(self, spec: str):
         raise ValueError(f"Job definition missing arg: {spec}")
@@ -125,7 +131,8 @@ 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(
+        runner = get_runner()
+        return runner.dryrun(
             app=ddp(
                 *self.script_args,
                 script=self.script,
@@ -160,7 +167,7 @@ def _dry_run_no_cluster(self):
             scheduler="kubernetes_mcad",
             cfg=self.scheduler_args,
             workspace="",
-        )
+        ), runner
 
     def submit(self, cluster: "Cluster" = None) -> "Job":
         return DDPJob(self, cluster)
@@ -171,18 +178,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)

src/codeflare_sdk/utils/generate_yaml.py

@@ -24,7 +24,14 @@
 from kubernetes import client, config
 from .kube_api_helpers import _kube_api_error_handling
 from ..cluster.auth import api_config_handler
+from os import urandom
+from base64 import b64encode
+from urllib3.util import parse_url
 
+import openshift as oc
+from kubernetes import client, config
+
+k8_client = config.new_client_from_config()
 
 def read_template(template):
     with open(template, "r") as stream:
@@ -46,12 +53,14 @@ def gen_names(name):
 
 def update_dashboard_route(route_item, cluster_name, namespace):
     metadata = route_item.get("generictemplate", {}).get("metadata")
-    metadata["name"] = f"ray-dashboard-{cluster_name}"
+    metadata["name"] = gen_dashboard_route_name(cluster_name)
     metadata["namespace"] = namespace
     metadata["labels"]["odh-ray-cluster-service"] = f"{cluster_name}-head-svc"
     spec = route_item.get("generictemplate", {}).get("spec")
     spec["to"]["name"] = f"{cluster_name}-head-svc"
 
+def gen_dashboard_route_name(cluster_name):
+    return f"ray-dashboard-{cluster_name}"
 
 # ToDo: refactor the update_x_route() functions
 def update_rayclient_route(route_item, cluster_name, namespace):
@@ -347,6 +356,63 @@ def write_user_appwrapper(user_yaml, output_file_name):
     print(f"Written to: {output_file_name}")
 
 
+def enable_openshift_oauth(user_yaml, cluster_name, namespace):
+    tls_mount_location = "/etc/tls/private"
+    oauth_port = 443
+    oauth_sa_name = f"{cluster_name}-oauth-proxy"
+    tls_secret_name = f"{cluster_name}-proxy-tls-secret"
+    tls_volume_name = "proxy-tls-secret"
+    port_name = "oauth-proxy"
+    _,_,host,_,_,_,_ = parse_url(k8_client.configuration.host)
+    host = host.replace("api.", f"{gen_dashboard_route_name(cluster_name)}-{namespace}.apps.")
+    oauth_sidecar = _create_oauth_sidecar_object(
+        namespace, tls_mount_location, oauth_port, oauth_sa_name, tls_volume_name, port_name
+    )
+    tls_secret_volume = client.V1Volume(
+        name=tls_volume_name,secret=client.V1SecretVolumeSource(secret_name=tls_secret_name)
+    )
+    # allows for setting value of Cluster object when initializing object from an existing AppWrapper on cluster
+    user_yaml["metadata"]["annotations"] = user_yaml["metadata"].get("annotations", {})
+    user_yaml["metadata"]["annotations"]["codeflare-sdk-use-oauth"] = "true"  # if the user gets an
+    ray_headgroup_pod = user_yaml["spec"]["resources"]["GenericItems"][0]["generictemplate"]["spec"]["headGroupSpec"]["template"]["spec"]
+    user_yaml["spec"]["resources"]["GenericItems"].pop(1)
+    ray_headgroup_pod["serviceAccount"] = oauth_sa_name
+    ray_headgroup_pod["volumes"] = ray_headgroup_pod.get("volumes", [])
+    ray_headgroup_pod["volumes"].append(k8_client.sanitize_for_serialization(tls_secret_volume))
+    ray_headgroup_pod["containers"].append(k8_client.sanitize_for_serialization(oauth_sidecar))
+    # add volume to headnode
+    # add sidecar container to ray object
+
+def _create_oauth_sidecar_object(
+    namespace: str,
+    tls_mount_location: str,
+    oauth_port: int,
+    oauth_sa_name: str,
+    tls_volume_name: str,
+    port_name: str
+) -> client.V1Container:
+    return client.V1Container(
+        args=[
+            f"--https-address=:{oauth_port}",
+            "--provider=openshift",
+            f"--openshift-service-account={oauth_sa_name}",
+            "--upstream=http://localhost:8265",
+            f"--tls-cert={tls_mount_location}/tls.crt",
+            f"--tls-key={tls_mount_location}/tls.key",
+            f"--cookie-secret={b64encode(urandom(64)).decode('utf-8')}",  # create random string for encrypting cookie
+            f'--openshift-delegate-urls={{"/":{{"resource":"pods","namespace":"{namespace}","verb":"get"}}}}'
+        ],
+        image="registry.redhat.io/openshift4/ose-oauth-proxy@sha256:1ea6a01bf3e63cdcf125c6064cbd4a4a270deaf0f157b3eabb78f60556840366",
+        name="oauth-proxy",
+        ports=[client.V1ContainerPort(container_port=oauth_port,name=port_name)],
+        resources = client.V1ResourceRequirements(limits=None,requests=None),
+        volume_mounts=[
+            client.V1VolumeMount(
+                mount_path=tls_mount_location,name=tls_volume_name,read_only=True
+            )
+        ],
+    )
+
 def generate_appwrapper(
     name: str,
     namespace: str,
@@ -365,6 +431,7 @@ def generate_appwrapper(
     image_pull_secrets: list,
     dispatch_priority: str,
     priority_val: int,
+    openshift_oauth: bool,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -396,6 +463,10 @@ def generate_appwrapper(
         enable_local_interactive(resources, cluster_name, namespace)
     else:
         disable_raycluster_tls(resources["resources"])
+
+    if openshift_oauth:
+        enable_openshift_oauth(user_yaml, cluster_name, namespace)
+
     outfile = appwrapper_name + ".yaml"
     write_user_appwrapper(user_yaml, outfile)
     return outfile