{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7ee3764b",
   "metadata": {},
   "source": [
    "# Using PyFireCREST in a Jupyter notebook\n",
    "\n",
    "This page implements the tutorial from the CSC User Guide [PyFirecREST tutorial](https://docs.csc.fi/support/tutorial/pyfirecrest). For more detailed explanations of each step read the tutorial.\n",
    "\n",
    "We will train a Classifier model on the [Iris dataset](https://archive.ics.uci.edu/dataset/53/iris), and output the confusion matrix.\n",
    "\n",
    "This workflow enables easy modification of your unprocessed data within a notebook or Python script, while still using HPC resources for the heavy computations. This way you use BU:s only on the heavy computation, not on the parts you can do locally.\n",
    "\n",
    "We will be using Roihu for this tutorial. Any differences to Lumi can be found in the [tutorial](https://docs.csc.fi/support/tutorial/pyfirecrest).\n",
    "## Setup\n",
    "Install PyFireCREST in your environment. You can do this with `pip install pyfirecrest`. Then you can import firecrest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7793e91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mandatory imports:\n",
    "import firecrest as fc\n",
    "\n",
    "# Whatever you need for your own code.\n",
    "import jwt\n",
    "import pandas as pd\n",
    "import os\n",
    "import time\n",
    "\n",
    "# Set constants:\n",
    "RAW_DATA_PATH = \"~/Downloads/iris.csv\"\n",
    "ROIHU_PROJ_DIR = \"/scratch/project_1234567/<username>/jupyter-dir/\"\n",
    "OUTPUT_FILENAME = \"confusion_matrix.png\"\n",
    "FIRECREST_URL = \"https://api.roihu.csc.fi/v1\"\n",
    "ACCOUNT = \"project_1234567\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c95f754",
   "metadata": {},
   "source": [
    "Retrieve your personal access token, and store it in a .env-file in your project directory. Instructions for retrieving your token and the exact API endpoint are found in the [Connecting to Roihu FirecREST HPC API](../../computing/firecrest/connecting.md) and for Lumi in the [Lumi Documentation](https://docs.lumi-supercomputer.eu/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9299108",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7306917",
   "metadata": {},
   "source": [
    "We must implement the TokentAuth-class, read the CSC guide on [FirecREST](https://docs.csc.fi/computing/firecrest/pyfirecrest/) for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "9df8d27d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TokenAuth:\n",
    "  def __init__(self):\n",
    "    pass\n",
    "\n",
    "  # Use PyJWT to decode the token and verify expiration time.\n",
    "  # Return False if decoding fails (input is not valid JWT) or if the token has expired\n",
    "  def _is_token_valid(self, token: str) -> bool:\n",
    "    try:\n",
    "      payload = jwt.decode(token, options={\"verify_signature\": False, \"verify_exp\": False, \"verify_aud\": False})\n",
    "      return time.time() <= payload[\"exp\"]\n",
    "    except Exception:\n",
    "      return False\n",
    "\n",
    "  # A PyFirecREST Authorization object is required to have method get_access_token(),\n",
    "  # which, when called, will return a valid JWT access token.\n",
    "  def get_access_token(self):\n",
    "    token = os.getenv('FIRECREST_TOKEN', None)\n",
    "    if not token:\n",
    "      raise RuntimeError(\"Environment variable FIRECREST_TOKEN is not defined.\")\n",
    "    if not self._is_token_valid(token):\n",
    "      raise RuntimeError(\"Token is invalid or has expired.\")\n",
    "    return token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "988e73af",
   "metadata": {},
   "outputs": [],
   "source": [
    "firecrest = fc.v2.Firecrest(firecrest_url=FIRECREST_URL, authorization=TokenAuth())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5f33723",
   "metadata": {},
   "source": [
    "## Data preprocessing\n",
    "\n",
    "Load your data and do whatever you need to do. Here we will use the [Iris dataset](https://archive.ics.uci.edu/dataset/53/iris), and do some filtering of outliers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "4a4d8b51",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   sepal_length  sepal_width  petal_length  petal_width species\n",
      "0           5.1          3.5           1.4          0.2  setosa\n",
      "1           4.9          3.0           1.4          0.2  setosa\n",
      "2           4.7          3.2           1.3          0.2  setosa\n",
      "3           4.6          3.1           1.5          0.2  setosa\n",
      "4           5.0          3.6           1.4          0.2  setosa\n"
     ]
    }
   ],
   "source": [
    "df = pd.read_csv(RAW_DATA_PATH)\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "d0b05f6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   sepal_length  petal_length  petal_width species\n",
      "0           5.1           1.4          0.2  setosa\n",
      "1           4.9           1.4          0.2  setosa\n",
      "2           4.7           1.3          0.2  setosa\n",
      "3           4.6           1.5          0.2  setosa\n",
      "4           5.0           1.4          0.2  setosa\n",
      "Original row count: 150, filtered row count: 141\n"
     ]
    }
   ],
   "source": [
    "df = df.dropna()\n",
    "q_low = df[\"sepal_length\"].quantile(0.01)\n",
    "q_high = df[\"sepal_length\"].quantile(0.99)\n",
    "\n",
    "df_filtered = df[(df[\"sepal_length\"] < q_high) & (df[\"sepal_length\"] > q_low)]\n",
    "df_filtered = df_filtered.drop(columns=\"sepal_width\")\n",
    "print(df_filtered.head())\n",
    "print(f\"Original row count: {len(df)}, filtered row count: {len(df_filtered)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e085f3a9",
   "metadata": {},
   "source": [
    "Now we are ready to train our model. We save the data we have processed to a csv file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "00629c1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "upload_file = \"/tmp/processed_data.csv\"\n",
    "df_filtered.to_csv(upload_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52db6b2a",
   "metadata": {},
   "source": [
    "## Upload files to Roihu\n",
    "\n",
    "Now we upload the file using `firecrest.upload()`, but first we'll make sure the directory exists with `firecrest.mkdir()`.\n",
    "When using the firecrest methods, all of them require `system_name` as an input. This distinguishes the different node types, \"cpu\" and \"gpu\". As Roihu uses a shared filesystem, the only command this has an effect on is the `firecrest.submit()`.\n",
    "\n",
    "If you are using Lumi, use `system_name=\"lumi\"` on all firecrest commands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49bbdb75",
   "metadata": {},
   "outputs": [],
   "source": [
    "firecrest.mkdir(system_name=\"cpu\", path=ROIHU_PROJ_DIR, create_parents=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cc0dde0",
   "metadata": {},
   "outputs": [],
   "source": [
    "filename_on_roihu = \"training_data.csv\"\n",
    "upload = firecrest.upload(system_name=\"cpu\", local_file=upload_file, directory=ROIHU_PROJ_DIR, filename=filename_on_roihu, account=ACCOUNT)\n",
    "\n",
    "# Check if upload is done as a batch job or not:\n",
    "if upload != None:\n",
    "  print(\"Upload as a batch job, may take a while.\")\n",
    "  # Wait for job to finish\n",
    "  upload.wait_for_transfer_job()\n",
    "\n",
    "print(f\"Upload complete for file {upload_file}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b231ebe0",
   "metadata": {},
   "source": [
    "## Submit the job\n",
    "\n",
    "The environment variable `CSC_ENV_INIT_NON_INTERACTIVE=yes` must be passed to the slurm job,\n",
    "else the environment won't be set up properly, and among other things the modules will not work properly.    \n",
    "\n",
    "To pass environment variables we use a dictionary. In addition to `CSC_ENV_INIT_NON_INTERACTIVE=yes` we will pass the `OUTPUT_FILENAME` and `DATA_FILE` variables which we can then access in the Slurm script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "08f216ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "env_vars = dict()\n",
    "env_vars[\"CSC_ENV_INIT_NON_INTERACTIVE\"] = \"yes\"\n",
    "env_vars[\"OUTPUT_FILENAME\"] = OUTPUT_FILENAME\n",
    "env_vars[\"DATA_FILE\"] = filename_on_roihu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a86f0c0",
   "metadata": {},
   "source": [
    "To submit a job, you need a Slurm script. Our script is `iris_slurm_script.sh` and it includes the Python script we want to run within it. See the [tutorial](https://docs.csc.fi/support/tutorial/pyfirecrest) for details on the script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd05ca77",
   "metadata": {},
   "outputs": [],
   "source": [
    "job = firecrest.submit(system_name=\"cpu\", working_dir=ROIHU_PROJ_DIR, script_local_path=\"iris_slurm_script.sh\", env_vars=env_vars)\n",
    "jobid = job[\"jobId\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c58cacef",
   "metadata": {},
   "source": [
    "## Download results\n",
    "Wait for the job to finish using `firecrest.wait_for_job()`. When it is, we can download the results, which could be any file. In this case is a png image of the confusion matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da628e71",
   "metadata": {},
   "outputs": [],
   "source": [
    "firecrest.wait_for_job(system_name=\"cpu\", job_id=jobid, timeout=None, not_found_timeout=80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f9521be",
   "metadata": {},
   "outputs": [],
   "source": [
    "download = firecrest.download(system_name=\"cpu\", source_path=os.path.join(ROIHU_PROJ_DIR, OUTPUT_FILENAME), target_path=OUTPUT_FILENAME, account=ACCOUNT)\n",
    "if download != None:\n",
    "    print(\"Download is done as a batch job, waiting for it to finish.\")\n",
    "    download.wait_for_transfer_job()\n",
    "print(f\"Results downloaded successfully to {OUTPUT_FILENAME}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da613169",
   "metadata": {},
   "source": [
    "Now you can analyse the results locally with whatever tools you have installed on your machine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e51846a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_results = your_post_processing(results)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.11.15.final.0)",
   "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.11.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}