demo-notebooks/additional-demos/ray_job_client.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this demo we will go over the basics of the RayJobClient in the SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import pieces from codeflare-sdk\n",
    "from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication, RayJobClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create authentication object for user permissions\n",
    "# IF unused, SDK will automatically check for default kubeconfig, then in-cluster config\n",
    "# KubeConfigFileAuthentication can also be used to specify kubeconfig path manually\n",
    "\n",
    "auth_token = \"XXXXX\" # The auth_token is used later for the RayJobClient\n",
    "auth = TokenAuthentication(\n",
    "    token = auth_token,\n",
    "    server = \"XXXXX\",\n",
    "    skip_tls=False\n",
    ")\n",
    "auth.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create and configure our cluster object\n",
    "cluster = Cluster(ClusterConfiguration(\n",
    "    name='jobtest',\n",
    "    namespace='default', # Update to your namespace\n",
    "    num_workers=2,\n",
    "    min_cpus=1,\n",
    "    max_cpus=1,\n",
    "    min_memory=4,\n",
    "    max_memory=4,\n",
    "    num_gpus=0,\n",
    "    image=\"quay.io/project-codeflare/ray:latest-py39-cu118\"\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bring up the cluster\n",
    "cluster.up()\n",
    "cluster.wait_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.details()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ray Job Submission - Authorized Ray Cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Submit a job using an authorized Ray dashboard and the Job Submission Client\n",
    "* Provide an entrypoint command directed to your job script\n",
    "* Set up your runtime environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gather the dashboard URL\n",
    "ray_dashboard = cluster.cluster_dashboard_uri()\n",
    "\n",
    "# Create the header for passing your bearer token\n",
    "header = {\n",
    "    'Authorization': f'Bearer {auth_token}'\n",
    "}\n",
    "\n",
    "# Initialize the RayJobClient\n",
    "client = RayJobClient(address=ray_dashboard, headers=header, verify=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit an example mnist job using the RayJobClient\n",
    "submission_id = client.submit_job(\n",
    "    entrypoint=\"python mnist.py\",\n",
    "    runtime_env={\"working_dir\": \"./\",\"pip\": \"requirements.txt\"},\n",
    ")\n",
    "print(submission_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the job's logs\n",
    "client.get_job_logs(submission_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the job's status\n",
    "client.get_job_status(submission_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get job related info\n",
    "client.get_job_info(submission_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List all existing jobs\n",
    "client.list_jobs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Iterate through the logs of a job \n",
    "async for lines in client.tail_job_logs(submission_id):\n",
    "    print(lines, end=\"\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete a job\n",
    "# Can run client.cancel_job(submission_id) first if job is still running\n",
    "client.delete_job(submission_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unauthorized Ray Cluster with the Ray Job Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Initialise the RayJobClient with the Ray Dashboard\n",
    "\"\"\"\n",
    "ray_dashboard = cluster.cluster_dashboard_uri()\n",
    "client = RayJobClient(address=ray_dashboard, verify=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit an example mnist job using the RayJobClient\n",
    "submission_id = client.submit_job(\n",
    "    entrypoint=\"python mnist.py\",\n",
    "    runtime_env={\"working_dir\": \"./\",\"pip\": \"requirements.txt\"},\n",
    ")\n",
    "print(submission_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stop the job \n",
    "client.stop_job(submission_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete the job\n",
    "client.delete_job(submission_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.down()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "auth.logout()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}