{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7033a0a9",
   "metadata": {},
   "source": [
    "# Using cloudknot to run pyAFQ on AWS batch:\n",
    "\n",
    "One of the purposes of `pyAFQ` is to analyze large-scale openly-available\n",
    "datasets, such as those in the\n",
    "[Human Connectome Project](https://www.humanconnectome.org/).\n",
    "\n",
    "To analyze these datasets, large amounts of compute are needed.\n",
    "One way to gain access to massive computational power is by using\n",
    "cloud computing. Here, we will demonstrate\n",
    "how to use `pyAFQ` in the Amazon Web Services cloud.\n",
    "\n",
    "We will rely on the [AWS Batch Service](https://aws.amazon.com/batch/),\n",
    "and we will submit work into AWS Batch using software that our group\n",
    "developed called [Cloudknot](https://nrdg.github.io/cloudknot/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fd19b7b",
   "metadata": {},
   "source": [
    "Import cloudknot and set the AWS region within which computations will take\n",
    "place. Setting a region is important, because if the data that you are\n",
    "analyzing is stored in `AWS S3 <https://aws.amazon.com/s3/>`_ in a\n",
    "particular region, it is best to run the computation in that region as well.\n",
    "That is because AWS charges for inter-region transfer of data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9481b3ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cloudknot as ck\n",
    "ck.set_region('us-east-1')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5d45193",
   "metadata": {},
   "source": [
    "## Define the function to use\n",
    "\n",
    "`Cloudknot` uses the single program multiple data paradigm of computing.\n",
    "This means that the same function will be run on multiple different inputs.\n",
    "For example, a `pyAFQ` processing function run\n",
    "on multiple different subjects in a dataset.\n",
    "Below, we define the function that we will use. Notice that\n",
    "`Cloudknot` functions include the import statements of the dependencies\n",
    "used. This is necessary so that `Cloudknot` knows\n",
    "what dependencies to install into AWS Batch to run this function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41dffb63",
   "metadata": {},
   "outputs": [],
   "source": [
    "def afq_process_subject(subject):\n",
    "    # define a function that each job will run\n",
    "    # In this case, each process does a single subject\n",
    "    import s3fs\n",
    "    # all imports must be at the top of the function\n",
    "    # cloudknot installs the appropriate packages from pip\n",
    "    from s3bids.utils import S3BIDSStudy\n",
    "    from AFQ.api.group import GroupAFQ\n",
    "    import AFQ.definitions.image as afm\n",
    "\n",
    "    # Download the given subject to your local machine from s3\n",
    "    # Can find subjects more easily if they are specified in a\n",
    "    # BIDS participants.tsv file, even if it is sparse\n",
    "    study_ixi = S3BIDSStudy(\n",
    "        \"my_study\",\n",
    "        \"my_study_bucket\",\n",
    "        \"my_study_prefix\",\n",
    "        subjects=[subject],\n",
    "        use_participants_tsv=True,\n",
    "        anon=False)\n",
    "    study_ixi.download(\n",
    "        \"local_bids_dir\",\n",
    "        include_derivs=[\"pipeline_name\"])\n",
    "\n",
    "    # define the api AFQ object\n",
    "    myafq = GroupAFQ(\n",
    "        \"local_bids_dir\",\n",
    "        dwi_preproc_pipeline=\"pipeline_name\",\n",
    "        viz_backend_spec='plotly')  # this will generate both interactive html and MP4s # noqa\n",
    "\n",
    "    # export_all runs the entire pipeline and creates many useful derivates\n",
    "    myafq.export_all()\n",
    "\n",
    "    # upload the results to some location on s3\n",
    "    myafq.upload_to_s3(\n",
    "        s3fs.S3FileSystem(),\n",
    "        \"my_study_bucket/my_study_prefix/derivatives/afq\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdcbb6b3",
   "metadata": {},
   "source": [
    "Here we provide a list of subjects that we have selected to process\n",
    "to randomly select 3 subjects without replacement, instead do:\n",
    "subjects = [[1], [2], [3]]\n",
    "see the docstring for S3BIDSStudy.__init__ for more information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be7b3f04",
   "metadata": {},
   "outputs": [],
   "source": [
    "subjects = [\"123456\", \"123457\", \"123458\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de6a870a",
   "metadata": {},
   "source": [
    "## Defining a `Knot` instance\n",
    "\n",
    "We instantiate a class instance of the `ck.Knot` class.\n",
    "This object will be used to run your jobs.\n",
    "The object is instantiated with the `'AmazonS3FullAccess'` policy,\n",
    "so that it can write the results\n",
    "out to S3, into a bucket that you have write permissions on.\n",
    "Setting the `bid_percentage` key-word makes AWS Batch use\n",
    "[spot EC2 instances](https://aws.amazon.com/ec2/spot/) for the\n",
    "computation. This can result in substantial cost-savings, as spot compute\n",
    "instances can cost much less than on-demand instances.\n",
    "However, not that spot instances can also\n",
    "be evicted, so if completing all of the work is very time-sensitive,\n",
    "do not set this key-word argument. Using the `image_github_installs`\n",
    "key-word argument will install pyAFQ from GitHub.\n",
    "You can also specify other forks and branches to install from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e50903ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "knot = ck.Knot(\n",
    "    name='afq-process-subject-201009-0',\n",
    "    func=afq_process_subject,\n",
    "    base_image='python:3.11',\n",
    "    image_github_installs=\"https://github.com/tractometry/pyAFQ.git\",\n",
    "    pars_policies=('AmazonS3FullAccess',),\n",
    "    bid_percentage=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cfa3ddd",
   "metadata": {},
   "source": [
    "## Launching the computation\n",
    "\n",
    "The `map` method of the `Knot` object maps each of the inputs\n",
    "provided as a sequence onto the function and executes the function on each\n",
    "one of them in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5341b35f",
   "metadata": {},
   "outputs": [],
   "source": [
    "result_futures = knot.map(subjects)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17ef9987",
   "metadata": {},
   "source": [
    "Once computations have started, you can call the following\n",
    "function to view the progress of jobs:\n",
    "\n",
    "```python\n",
    "knot.view_jobs()\n",
    "```\n",
    "\n",
    "You can also view the status of a specific job:\n",
    "\n",
    "```python\n",
    "knot.jobs[0].status\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f873473f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "caefef48",
   "metadata": {},
   "source": [
    "When all jobs are finished, remember to use the `clobber` method to\n",
    "destroy all of the AWS resources created by the `Knot`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "599f6df8",
   "metadata": {},
   "outputs": [],
   "source": [
    "result_futures.result()\n",
    "knot.clobber(clobber_pars=True, clobber_repo=True, clobber_image=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0ec5d23",
   "metadata": {},
   "source": [
    "In a second `Knot` object, we use a function that takes the\n",
    "resulting profiles of each subject and combines them into one csv file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "458259d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def afq_combine_profiles(dummy_argument):\n",
    "    from AFQ.api import download_and_combine_afq_profiles\n",
    "    download_and_combine_afq_profiles(\n",
    "        \"my_study_bucket\", \"my_study_prefix\")\n",
    "\n",
    "\n",
    "knot2 = ck.Knot(\n",
    "    name='afq_combine_subjects-201009-0',\n",
    "    func=afq_combine_profiles,\n",
    "    base_image='python:3.11',\n",
    "    image_github_installs=\"https://github.com/tractometry/pyAFQ.git\",\n",
    "    pars_policies=('AmazonS3FullAccess',),\n",
    "    bid_percentage=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "424be449",
   "metadata": {},
   "source": [
    "This knot is called with a dummy argument, which is not used within the\n",
    "function itself. The `job_type` key-word argument is used to signal to\n",
    "`Cloudknot` that only one job is submitted rather than the default\n",
    "array of jobs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef9c028d",
   "metadata": {},
   "outputs": [],
   "source": [
    "result_futures2 = knot2.map([\"dummy_argument\"], job_type=\"independent\")\n",
    "result_futures2.result()\n",
    "knot2.clobber(clobber_pars=True, clobber_repo=True, clobber_image=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7aab31b",
   "metadata": {},
   "source": [
    ":::{only} html\n",
    "{download}`Download as Jupyter Notebook <cloudknot_example.ipynb>`\n",
    ":::"
   ]
  }
 ],
 "metadata": {
  "file_format": "mystnb",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  },
  "mystnb": {
   "execution_mode": "off"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}