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

@@ -18,15 +18,21 @@
 cluster setup queue, a list of all existing clusters, and the user's working namespace.
 """
 
+<<<<<<< HEAD
 from time import sleep
 from typing import List, Optional, Tuple, Dict
 
+import openshift as oc
+from kubernetes import config
+>>>>>>> bb0a0a7 (add functions for creating ray with oauth proxy in front of the dashboard)
 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 +47,9 @@
 import requests
 
 
+k8_client = config.new_client_from_config()
+
+
 class Cluster:
     """
     An object for requesting, bringing up, and taking down resources.
@@ -61,6 +70,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 evaluate_dispatch_priority(self):
         priority_class = self.config.dispatch_priority
@@ -141,6 +165,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 +175,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 +212,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]:
@@ -322,14 +353,14 @@ def list_jobs(self) -> List:
         """
         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)
+        client = JobSubmissionClient(dashboard_route,)
         return client.get_job_status(job_id)
 
     def job_logs(self, job_id: str) -> str:
@@ -343,7 +374,7 @@ def job_logs(self, job_id: str) -> str:
     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,

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

src/codeflare_sdk/job/jobs.py

@@ -21,21 +21,35 @@
 from torchx.runner import get_runner
 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
 
 all_jobs: List["Job"] = []
 torchx_runner = get_runner()
 
-
 class JobDefinition(metaclass=abc.ABCMeta):
     def _dry_run(self, cluster: "Cluster"):
         pass
 
     def submit(self, cluster: "Cluster"):
         pass
 
+    def _get_torchx_runner(self, cluster: "Cluster"):
+        return get_runner(
+            scheduler_params={
+                "ray_client": JobSubmissionClient(
+                    address=cluster.cluster_dashboard_uri(),
+                    headers={"Authorization": f"Bearer {oc.get_auth_token()}"},
+                    verify=cluster.config.openshift_oauth,
+                )
+            }
+        )
+
 
 class Job(metaclass=abc.ABCMeta):
     def status(self):
@@ -171,9 +185,11 @@ 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))
+            runner = job_definition._get_torchx_runner(cluster=cluster)
+            self._app_handle = runner.schedule(job_definition._dry_run(cluster))
         else:
-            self._app_handle = torchx_runner.schedule(
+            runner = get_runner()
+            self._app_handle = runner.schedule(
                 job_definition._dry_run_no_cluster()
             )
         all_jobs.append(self)

src/codeflare_sdk/utils/generate_yaml.py

@@ -21,10 +21,20 @@
 import sys
 import argparse
 import uuid
+<<<<<<< HEAD
 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
+>>>>>>> bb0a0a7 (add functions for creating ray with oauth proxy in front of the dashboard)
 
+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 +56,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 +359,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,
@@ -365,6 +435,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 +467,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