8_distributed_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training the Fraud Detection model with Ray by using Codeflare"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The example fraud detection model is very small and quickly trained. However, for many large models, training requires multiple GPUs and often multiple machines. In this notebook, you learn how to train a model by using Ray on OpenShift AI to scale out the model training. You use the Codeflare SDK to create the cluster and submit the job. You can find detailed documentation for the SDK [here](https://project-codeflare.github.io/codeflare-sdk/detailed-documentation/).\n",
    "\n",
    "For this procedure, you need to use codeflare-sdk 0.19.1 (or later).  Begin by installing the SDK if it's not already installed or up to date:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade codeflare-sdk>=0.19.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing the data\n",
    "\n",
    "Normally, the training data for your model would be available in a shared location. For this example, the data is local. You must upload it to your object storage so that you can see how data loading from a shared data source works. After you upload the data, you can work with it by using Ray Data so that it is properly shared across the worker machines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('./utils')\n",
    "\n",
    "import utils.s3\n",
    "\n",
    "utils.s3.upload_directory_to_s3(\"data\", \"data\")\n",
    "print(\"---\")\n",
    "utils.s3.list_objects(\"data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Authenticate to the cluster by using the OpenShift console login\n",
    "\n",
    "You must create the Kubernetes objects for Ray Clusters using the Codeflare SDK. In order to do so, you need access permission for your own namespace. The easiest way to set up access is by using the OpenShift CLI `oc` client. \n",
    "\n",
    "From the OpenShift web console, you can generate an `oc login` command that includes your token and server information. You can use the command to log in to the OpenShift CLI. \n",
    "\n",
    "1. To generate the command, select **Copy login command** from the username drop-down menu at the top right of the web console.\n",
    "\n",
    "    <figure>\n",
    "        <img src=\"./assets/copy-login.png\"  alt=\"copy login\"  >\n",
    "    <figure/>\n",
    "\n",
    "2. Click **Display token**.\n",
    "\n",
    "3. Below **Log in with this token**, take note of the parameters for token and server.\n",
    "   For example:\n",
    "    ```\n",
    "    oc login --token=sha256~LongString --server=https://api.your-cluster.domain.com:6443\n",
    "    ```    \n",
    "    - token: `sha256~LongString`\n",
    "    - server: `https://api.your-cluster.domain.com:6443`\n",
    "    \n",
    "4. In the following code cell, in the TokenAuthentication object, replace the token and server values with the values that you noted in Step 3.\n",
    "   For example:\n",
    "   ```\n",
    "   auth = TokenAuthentication(\n",
    "       token = \"sha256~LongString\",\n",
    "       server = \"https://api.your-cluster.domain.com:6443\",\n",
    "       skip_tls=False\n",
    "   )\n",
    "   auth.login()\n",
    "   ```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from codeflare_sdk import TokenAuthentication\n",
    "# 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",
    "auth = TokenAuthentication(\n",
    "    token = \"sha256~XXXX\",\n",
    "    server = \"https://XXXX\",\n",
    "    skip_tls=False\n",
    ")\n",
    "auth.login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Create a Ray cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure a Ray cluster\n",
    "\n",
    "CodeFlare allows you to specify parameters, such as number of workers, image, and kueue local queue name.  A full list of parameters is available [here](https://project-codeflare.github.io/codeflare-sdk/detailed-documentation/cluster/config.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from codeflare_sdk import Cluster, ClusterConfiguration\n",
    "\n",
    "cluster = Cluster(ClusterConfiguration(\n",
    "    name=\"raycluster-cpu\",\n",
    "    head_extended_resource_requests={'nvidia.com/gpu': 0},\n",
    "    worker_extended_resource_requests={'nvidia.com/gpu': 0},\n",
    "    num_workers=2,\n",
    "    worker_cpu_requests=1,\n",
    "    worker_cpu_limits=4,\n",
    "    worker_memory_requests=2,\n",
    "    worker_memory_limits=4\n",
    "))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start the cluster\n",
    "\n",
    "If you have a running cluster that you want to connect to, skip to the next cell.\n",
    "\n",
    "To start a cluster, run the following cell to create the necessary Kubernetes objects to run the Ray cluster. This step might take a few minutes to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cluster.up()\n",
    "cluster.wait_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Connect to a running cluster\n",
    "\n",
    "If you've already created a cluster, but you've restarted the Python kernel, closed the notebook, or are working in a different notebook, and you want to connect to the existing cluster, uncomment the code in the following cell and then run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from codeflare_sdk import get_cluster\n",
    "\n",
    "# cluster = get_cluster(name, namespace=namespace)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can view information about the cluster, including a link to the Ray dashboard. In the Ray dashboard, you can inspect the running jobs and logs, and see the resources being used.\n",
    "<figure>\n",
    "    <img src=\"./assets/codeflare-details.png\"  alt=\"codeflare details\" width=\"400\">\n",
    "<figure/>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cluster.details()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The link to the Ray dashboard is available in the cluster details provided as a result of running the previous cell.  It should look something like this:\n",
    "\n",
    "<figure>\n",
    "    <img src=\"./assets/ray-dashboard.png\"  alt=\"ray dashboard\" width=\"600\"\n",
    "<figure/>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ray job submission"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize the Job Submission Client\n",
    "\n",
    "If you want to submit jobs, connect to the running Ray cluster by initializing the job client that has the proper authentication and connection information.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client = cluster.job_client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After you connect to the Ray cluster, you can query the cluster to determine whether there are any existing jobs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client.list_jobs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Create a runtime environment\n",
    "\n",
    "Now you can configure the [runtime environment](https://docs.ray.io/en/latest/ray-core/handling-dependencies.html#runtime-environments) for the job. This step includes specifying the working directory, files to exclude, dependencies, and environment variables.\n",
    "\n",
    "```python\n",
    "runtime_env={\n",
    "    \"working_dir\": \"./\", # relative path to files uploaded to the job\n",
    "    \"excludes\": [\"local_data/\"], # directories and files to exclude from being uploaded to the job\n",
    "    \"pip\": [\"boto3\", \"botocore\"], # can also be a string path to a requirements.txt file\n",
    "    \"env_vars\": {\n",
    "        \"MY_ENV_VAR\": \"MY_ENV_VAR_VALUE\",\n",
    "        \"MY_ENV_VAR_2\": os.environ.get(\"MY_ENV_VAR_2\"),\n",
    "    },\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# script = \"test_data_loader.py\"\n",
    "script = \"train_tf_cpu.py\"\n",
    "runtime_env = {\n",
    "    \"working_dir\": \"./ray-scripts\",\n",
    "    \"excludes\": [],\n",
    "    \"pip\": \"./ray-scripts/requirements.txt\",\n",
    "    \"env_vars\": {\n",
    "        \"AWS_ACCESS_KEY_ID\": os.environ.get(\"AWS_ACCESS_KEY_ID\"),\n",
    "        \"AWS_SECRET_ACCESS_KEY\": os.environ.get(\"AWS_SECRET_ACCESS_KEY\"),\n",
    "        \"AWS_S3_ENDPOINT\": os.environ.get(\"AWS_S3_ENDPOINT\"),\n",
    "        \"AWS_DEFAULT_REGION\": os.environ.get(\"AWS_DEFAULT_REGION\"),\n",
    "        \"AWS_S3_BUCKET\": os.environ.get(\"AWS_S3_BUCKET\"),\n",
    "        \"NUM_WORKERS\": \"1\",\n",
    "        \"TRAIN_DATA\": \"data/train.csv\",\n",
    "        \"VALIDATE_DATA\": \"data/validate.csv\",\n",
    "        \"MODEL_OUTPUT_PREFIX\": \"models/fraud/1/\",\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Submit the configured job\n",
    "\n",
    "Now you can submit the job to the cluster. This step creates the necessary Kubernetes objects to run the job. The job runs the script with the specified runtime environment. The script for this example is located in [ray-scripts/train_tf_cpu.py](./ray-scripts/train_tf_cpu.py). The script follows the code fairly closely to the official [Ray TensorFlow example](https://docs.ray.io/en/latest/train/distributed-tensorflow-keras.html). This example uses TensorFlow, note that the [Ray site](https://docs.ray.io/en/latest/train/train.html) provides examples for PyTorch and other frameworks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "submission_id = client.submit_job(\n",
    "    entrypoint=f\"python {script}\",\n",
    "    runtime_env=runtime_env,\n",
    ")\n",
    "\n",
    "print(submission_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Query important job information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get the job's status\n",
    "print(client.get_job_status(submission_id), \"\\n\")\n",
    "\n",
    "# Get job related info\n",
    "print(client.get_job_info(submission_id), \"\\n\")\n",
    "\n",
    "# Get the job's logs\n",
    "print(client.get_job_logs(submission_id))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also tail the job logs to watch the progress of the job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "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": "markdown",
   "metadata": {},
   "source": [
    "### List jobs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client.list_jobs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Stop jobs\n",
    "\n",
    "If you want to stop a job, call `stop_job` and specify the submission ID.  In the following cell, the command lists all the jobs and stops them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for job_details in client.list_jobs():\n",
    "    print(f\"deleting {job_details.submission_id}\")\n",
    "    client.stop_job(job_details.submission_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Delete jobs\n",
    "\n",
    "You can also delete the jobs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for job_details in client.list_jobs():\n",
    "    print(f\"deleting {job_details.submission_id}\")\n",
    "    client.delete_job(job_details.submission_id)\n",
    "\n",
    "client.list_jobs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete the cluster\n",
    "\n",
    "After you complete training, you can delete the cluster. When you delete the cluster, you remove the Kubernetes objects and free up resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cluster.down()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9",
   "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": 4
}