demo-notebooks/guided-demos/2_basic_interactive.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bbc21043",
   "metadata": {},
   "source": [
    "In this notebook, we will go over how to leverage the SDK to directly work interactively with a Ray Cluster during development."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b55bc3ea-4ce3-49bf-bb1f-e209de8ca47a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import pieces from codeflare-sdk\n",
    "from codeflare_sdk import Cluster, ClusterConfiguration, TokenAuthentication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "614daa0c",
   "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",
   "id": "bc27f84c",
   "metadata": {},
   "source": [
    "Once again, let's start by running through the same cluster setup as before:\n",
    "\n",
    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
    "\n",
    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
    "\n",
    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f4bc870-091f-4e11-9642-cba145710159",
   "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_name = \"interactivetest\"\n",
    "cluster = Cluster(ClusterConfiguration(\n",
    "    name=cluster_name,\n",
    "    head_cpu_requests=1,\n",
    "    head_cpu_limits=1,\n",
    "    head_memory_requests=6,\n",
    "    head_memory_limits=6,\n",
    "    head_extended_resource_requests={'nvidia.com/gpu':1}, # For GPU enabled workloads set the head_extended_resource_requests and worker_extended_resource_requests\n",
    "    worker_extended_resource_requests={'nvidia.com/gpu':1},\n",
    "    num_workers=2,\n",
    "    worker_cpu_requests='250m',\n",
    "    worker_cpu_limits=1,\n",
    "    worker_memory_requests=4,\n",
    "    worker_memory_limits=6,\n",
    "    # image=\"\", # Optional Field \n",
    "    write_to_file=False, # 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,
   "id": "f0884bbc-c224-4ca0-98a0-02dfa09c2200",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bring up the cluster\n",
    "cluster.up()\n",
    "cluster.wait_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df71c1ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.details()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33663f47",
   "metadata": {},
   "source": [
    "This time we will demonstrate another potential method of use: working with the Ray cluster interactively.\n",
    "\n",
    "Using the SDK, we can get both the Ray cluster URI and dashboard URI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1719bca",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray_dashboard_uri = cluster.cluster_dashboard_uri()\n",
    "ray_cluster_uri = cluster.cluster_uri()\n",
    "print(ray_dashboard_uri)\n",
    "print(ray_cluster_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a2aca6a",
   "metadata": {},
   "source": [
    "Now we can connect directly to our Ray cluster via the Ray python client:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9436436",
   "metadata": {},
   "outputs": [],
   "source": [
    "from codeflare_sdk import generate_cert\n",
    "# Create required TLS cert and export the environment variables to enable TLS\n",
    "generate_cert.generate_tls_cert(cluster_name, cluster.config.namespace)\n",
    "generate_cert.export_env(cluster_name, cluster.config.namespace)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "300146dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# before proceeding make sure the cluster exists and the uri is not empty\n",
    "assert ray_cluster_uri, \"Ray cluster needs to be started and set before proceeding\"\n",
    "\n",
    "import ray\n",
    "\n",
    "# reset the ray context in case there's already one. \n",
    "ray.shutdown()\n",
    "# establish connection to ray cluster\n",
    "\n",
    "# install additional libraries that will be required for model training\n",
    "runtime_env = {\"pip\": [\"transformers==4.41.2\", \"datasets==2.17.0\", \"accelerate==0.31.0\", \"scikit-learn==1.5.0\"]}\n",
    "# NOTE: This will work for in-cluster notebook servers (RHODS/ODH), but not for local machines\n",
    "# To see how to connect from your laptop, go to demo-notebooks/additional-demos/local_interactive.ipynb\n",
    "ray.init(address=ray_cluster_uri, runtime_env=runtime_env)\n",
    "\n",
    "print(\"Ray cluster is up and running: \", ray.is_initialized())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9711030b",
   "metadata": {},
   "source": [
    "Now that we are connected (and have passed in some package requirements), let's try writing some training code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b36e0d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "def train_fn():\n",
    "    import os\n",
    "    import numpy as np\n",
    "    from datasets import load_dataset, load_metric\n",
    "    import transformers\n",
    "    from transformers import (\n",
    "        Trainer,\n",
    "        TrainingArguments,\n",
    "        AutoTokenizer,\n",
    "        AutoModelForSequenceClassification,\n",
    "    )\n",
    "    import ray.train.huggingface.transformers\n",
    "    from ray.train import ScalingConfig\n",
    "    from ray.train.torch import TorchTrainer\n",
    "\n",
    "    # When running in a multi-node cluster you will need persistent storage that is accessible across all worker nodes. \n",
    "    # See www.github.com/project-codeflare/codeflare-sdk/tree/main/docs/s3-compatible-storage.md for more information.\n",
    "    \n",
    "    def train_func():\n",
    "        # Datasets\n",
    "        dataset = load_dataset(\"imdb\")\n",
    "        tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "\n",
    "        def tokenize_function(examples):\n",
    "            return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True)\n",
    "\n",
    "        small_train_dataset = (\n",
    "            dataset[\"train\"].select(range(100)).map(tokenize_function, batched=True)\n",
    "        )\n",
    "        small_eval_dataset = (\n",
    "            dataset[\"test\"].select(range(100)).map(tokenize_function, batched=True)\n",
    "        )\n",
    "\n",
    "        # Model\n",
    "        model = AutoModelForSequenceClassification.from_pretrained(\n",
    "            \"distilbert-base-uncased\", num_labels=2\n",
    "        )\n",
    "\n",
    "        def compute_metrics(eval_pred):\n",
    "            metric = load_metric(\"accuracy\")\n",
    "            logits, labels = eval_pred\n",
    "            predictions = np.argmax(logits, axis=-1)\n",
    "            return metric.compute(predictions=predictions, references=labels)\n",
    "\n",
    "        # Hugging Face Trainer\n",
    "        training_args = TrainingArguments(\n",
    "            output_dir=\"test_trainer\",\n",
    "            evaluation_strategy=\"epoch\",\n",
    "            save_strategy=\"epoch\",\n",
    "            report_to=\"none\",\n",
    "        )\n",
    "\n",
    "        trainer = Trainer(\n",
    "            model=model,\n",
    "            args=training_args,\n",
    "            train_dataset=small_train_dataset,\n",
    "            eval_dataset=small_eval_dataset,\n",
    "            compute_metrics=compute_metrics,\n",
    "        )\n",
    "\n",
    "\n",
    "        callback = ray.train.huggingface.transformers.RayTrainReportCallback()\n",
    "        trainer.add_callback(callback)\n",
    "\n",
    "        trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)\n",
    "\n",
    "        trainer.train()\n",
    "\n",
    "\n",
    "    ray_trainer = TorchTrainer(\n",
    "        train_func,\n",
    "        scaling_config=ScalingConfig(\n",
    "            # num_workers = number of worker nodes with the ray head node included\n",
    "            num_workers=3,\n",
    "            use_gpu=True,\n",
    "            resources_per_worker={\n",
    "                \"CPU\": 1,\n",
    "            },\n",
    "            trainer_resources={\n",
    "                \"CPU\": 0,\n",
    "            }\n",
    "        )\n",
    "        # Configure persistent storage that is accessible across \n",
    "        # all worker nodes.\n",
    "        # Uncomment and update the RunConfig below to include your storage details.\n",
    "        # run_config=ray.train.RunConfig(storage_path=\"storage path\"),\n",
    "    )\n",
    "    result: ray.train.Result = ray_trainer.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4d8fd65",
   "metadata": {},
   "source": [
    "Once we want to test our code out, we can run the training function we defined above remotely on our Ray cluster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5901d958",
   "metadata": {},
   "outputs": [],
   "source": [
    "#call the above cell as a remote ray function\n",
    "ray.get(train_fn.remote())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5af8cd32",
   "metadata": {},
   "source": [
    "Once complete, we can bring our Ray cluster down and clean up:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f36db0f-31f6-4373-9503-dc3c1c4c3f57",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster.down()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d41b90e",
   "metadata": {},
   "outputs": [],
   "source": [
    "auth.logout()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "f9f85f796d01129d0dd105a088854619f454435301f6ffec2fea96ecbd9be4ac"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}