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

@@ -23,10 +23,12 @@
 from typing import List, Optional, Tuple, Dict
 
 import openshift as oc
+from kubernetes import config
 from ray.job_submission import JobSubmissionClient
 
 from ..utils import pretty_print
 from ..utils.generate_yaml import generate_appwrapper
+from ..utils.openshift_oauth import create_openshift_oauth_objects, delete_openshift_oauth_objects
 from .config import ClusterConfiguration
 from .model import (
     AppWrapper,
@@ -37,6 +39,9 @@
 )
 
 
+k8_client = config.new_client_from_config()
+
+
 class Cluster:
     """
     An object for requesting, bringing up, and taking down resources.
@@ -57,6 +62,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:
+            # user must be logged in to OpenShift
+            self._client = JobSubmissionClient(
+                self.cluster_dashboard_uri, 
+                headers={"Authorization": k8_client.configuration.auth_settings()["BearerToken"]["value"]}
+            )
+        else:
+            self._client = JobSubmissionClient(self.cluster_dashboard_uri)
+        return self._client
 
     def create_app_wrapper(self):
         """
@@ -102,6 +122,7 @@ def create_app_wrapper(self):
             env=env,
             local_interactive=local_interactive,
             image_pull_secrets=image_pull_secrets,
+            openshift_oauth=self.config.openshift_oauth,
         )
 
     # creates a new cluster with the provided or default spec
@@ -111,6 +132,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:
             with oc.project(namespace):
                 oc.invoke("apply", ["-f", self.app_wrapper_yaml])
@@ -146,6 +170,8 @@ def down(self):
                 print("Cluster not found, have you run cluster.up() yet?")
             else:
                 raise osp
+        if self.config.openshift_oauth:
+            delete_openshift_oauth_objects(cluster_name=self.config.name, namespace=namespace)
 
     def status(
         self, print_to_console: bool = True
@@ -254,7 +280,7 @@ def cluster_dashboard_uri(self) -> str:
                 route = route.out().split(" ")
                 route = [x for x in route if f"ray-dashboard-{self.config.name}" in x]
                 route = route[0].strip().strip("'")
-            return f"http://{route}"
+            return f"https://{route}"
         except:
             return "Dashboard route not available yet, have you run cluster.up()?"
 

src/codeflare_sdk/cluster/config.py

@@ -50,3 +50,4 @@ class ClusterConfiguration:
     image: str = "quay.io/project-codeflare/ray:2.5.0-py38-cu116"
     local_interactive: bool = False
     image_pull_secrets: list = field(default_factory=list)
+    openshift_oauth: bool = False

src/codeflare_sdk/utils/generate_yaml.py

@@ -21,8 +21,14 @@
 import sys
 import argparse
 import uuid
+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:
@@ -44,12 +50,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):
@@ -289,6 +297,64 @@ 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",
+            "--cookie-secret=SECRET",
+            # 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,
@@ -305,6 +371,7 @@ def generate_appwrapper(
     env,
     local_interactive: bool,
     image_pull_secrets: list,
+    openshift_oauth: bool,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -335,6 +402,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

src/codeflare_sdk/utils/openshift_oauth.py

@@ -0,0 +1,112 @@
+from urllib3.util import parse_url
+from .generate_yaml import gen_dashboard_route_name
+
+from kubernetes import config, client
+
+k8_client = config.new_client_from_config()
+core_api = client.CoreV1Api(k8_client)
+rbac_auth_api = client.RbacAuthorizationV1Api(k8_client)
+networking_api = client.NetworkingV1Api(k8_client)
+
+def create_openshift_oauth_objects(cluster_name, namespace):
+    oauth_port = 443
+    oauth_sa_name = f"{cluster_name}-oauth-proxy"
+    tls_secret_name = f"{cluster_name}-proxy-tls-secret"
+    service_name = f"{cluster_name}-oauth"
+    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_crb = client.V1ClusterRoleBinding(
+        api_version="rbac.authorization.k8s.io/v1", kind="ClusterRoleBinding",
+        metadata=client.V1ObjectMeta(name=f"{cluster_name}-rb"),
+        role_ref=client.V1RoleRef(
+            api_group="rbac.authorization.k8s.io",
+            kind="ClusterRole",
+            name="system:auth-delegator",
+        ),
+        subjects=[client.V1Subject(kind="ServiceAccount", name=oauth_sa_name, namespace=namespace)]
+    )
+    oauth_sa = client.V1ServiceAccount(
+        api_version="v1",
+        kind="ServiceAccount",
+        metadata=client.V1ObjectMeta(
+            name=oauth_sa_name,
+            namespace=namespace,
+            annotations={"serviceaccounts.openshift.io/oauth-redirecturi.first": f"https://{host}"}
+        )
+    )
+    oauth_service = _create_oauth_service_obj(
+        cluster_name, namespace, oauth_port, tls_secret_name, service_name, port_name
+    )
+    ingress = _create_oauth_ingress_object(cluster_name, namespace, service_name, port_name, host)
+    core_api.create_namespaced_service_account(namespace=namespace, body=oauth_sa)
+    core_api.create_namespaced_service(namespace=namespace, body=oauth_service)
+    networking_api.create_namespaced_ingress(namespace=namespace, body=ingress)
+    rbac_auth_api.create_cluster_role_binding(body=oauth_crb)
+
+def delete_openshift_oauth_objects(cluster_name, namespace):
+    oauth_sa_name = f"{cluster_name}-oauth-proxy"
+    service_name = f"{cluster_name}-oauth"
+    core_api.delete_namespaced_service_account(name=oauth_sa_name, namespace=namespace)
+    core_api.delete_namespaced_service(name=service_name, namespace=namespace)
+    networking_api.delete_namespaced_ingress(name=f"{cluster_name}-ingress", namespace=namespace)
+    rbac_auth_api.delete_cluster_role_binding(name= f"{cluster_name}-rb")
+
+
+def _create_oauth_service_obj(
+    cluster_name: str,
+    namespace: str,
+    oauth_port: int,
+    tls_secret_name: str,
+    service_name: str,
+    port_name: str,
+) -> client.V1Service:
+    return client.V1Service(
+        api_version="v1",
+        kind="Service",
+        metadata=client.V1ObjectMeta(
+            annotations={"service.beta.openshift.io/serving-cert-secret-name": tls_secret_name},
+            name=service_name,
+            namespace=namespace
+        ),
+        spec=client.V1ServiceSpec(
+            ports=[client.V1ServicePort(name=port_name, protocol="TCP", port=oauth_port, target_port=oauth_port)],
+            selector={
+                "app.kubernetes.io/created-by": "kuberay-operator",
+                "app.kubernetes.io/name": "kuberay",
+                "ray.io/cluster": cluster_name,
+                "ray.io/identifier": f"{cluster_name}-head",
+                "ray.io/node-type": "head",
+            }
+        )
+    )
+
+def _create_oauth_ingress_object(
+    cluster_name: str,
+    namespace: str,
+    service_name: str,
+    port_name: str,
+    host: str,
+) -> client.V1Ingress:
+    return client.V1Ingress(
+        api_version="networking.k8s.io/v1",
+        kind="Ingress",
+        metadata=client.V1ObjectMeta(
+            annotations={"route.openshift.io/termination": "passthrough"},
+            name=f"{cluster_name}-ingress",
+            namespace=namespace
+        ),
+        spec=client.V1IngressSpec(rules=[client.V1IngressRule(
+            host=host,
+            http=client.V1HTTPIngressRuleValue(paths=[
+                client.V1HTTPIngressPath(
+                    backend=client.V1IngressBackend(
+                        service=client.V1IngressServiceBackend(
+                            name=service_name,port=client.V1ServiceBackendPort(name=port_name)
+                        )
+                    ),
+                    path_type="ImplementationSpecific"
+                )
+            ])
+        )]),
+    )