demo-notebooks/guided-demos/3_pytorch_lightning_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we are going to run a Ray Train & Pytorch Lightning script using the CodeFlare SDK and Ray Job Submission.\n",
    "\n",
    "NOTE: For distributed training an external persistent storage option should be set in the `run_config`.\n",
    "You can find examples in the `pytorch_lightning.py` script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import pieces from codeflare-sdk\n",
    "from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication"
   ]
  },
  {
   "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",
    "auth = TokenAuthentication(\n",
    "    token = \"XXXXX\",\n",
    "    server = \"XXXXX\",\n",
    "    skip_tls=False\n",
    ")\n",
    "auth.login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once again, let's start by running through the same cluster setup as before:\n",
    "\n",
    "NOTE: We must specify the `image` which will be used in our RayCluster, we recommend you bring your own image which suits your purposes. \n",
    "The example here is a community image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create and configure our cluster object\n",
    "# The SDK will try to find the name of your default local queue based on the annotation \"kueue.x-k8s.io/default-queue\": \"true\" unless you specify the local queue manually below\n",
    "cluster = Cluster(ClusterConfiguration(\n",
    "    name='raytest',\n",
    "    namespace='default', # Update to your namespace\n",
    "    num_workers=2,\n",
    "    min_cpus=2,\n",
    "    max_cpus=2,\n",
    "    min_memory=8,\n",
    "    max_memory=8,\n",
    "    num_gpus=1,\n",
    "    head_gpus=1,\n",
    "    image=\"quay.io/project-codeflare/ray:2.20.0-py39-cu118\",\n",
    "    write_to_file=True, # When enabled Ray Cluster yaml files are written to /HOME/.codeflare/resources \n",
    "    # local_queue=\"local-queue-name\" # Specify the local queue manually\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": [
    "Note: For this example external S3 compatible storage is required. Please refer to our [documentation](https://github.com/project-codeflare/codeflare-sdk/blob/main/docs/s3-compatible-storage.md) for steps on how to configure this training script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Job Submission Client\n",
    "\"\"\"\n",
    "The SDK will automatically gather the dashboard address and authenticate using the Ray Job Submission Client\n",
    "\"\"\"\n",
    "client = cluster.job_client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit an example mnist job using the Job Submission Client\n",
    "submission_id = client.submit_job(\n",
    "    entrypoint=\"python pytorch_lightning.py\",\n",
    "    runtime_env={\"working_dir\": \"./\",\"pip\": \"requirements_lightning.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": [
    "cluster.down()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "auth.logout()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}