examples/sagemaker/pytorch_rnn_meeshkan.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Word-level language modeling using PyTorch and Meeshkan\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install requirements\n",
    "Using `meeshkan` with sagemaker requires installing both Python SDKs. Note that if you have installed `meeshkan` previously, you may have to **restart the Jupyter kernel (`Kernel` -> `Restart`) to have the dependency updated.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade meeshkan sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import meeshkan\n",
    "import sagemaker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start the Meeshkan agent\n",
    "The Meeshkan agent is a daemonized process monitoring your jobs in the background and notifying you when events happen. The command `meeshkan init(token=YOUR_TOKEN)` starts the agent using the token you got when signing up at [meeshkan.com](https://meeshkan.com).\n",
    "The token is stored to a local file at `~/.meeshkan/credentials`, so you only need to include the token here _once_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meeshkan.init(token=\"YOUR_TOKEN_HERE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can control the agent also with the following commands:\n",
    "```python\n",
    "meeshkan.start()   # Starts the agent (assumes credentials have been setup)\n",
    "meeshkan.stop()   # Stops the agent\n",
    "meeshkan.restart()  # Restarts the agent, useful if you installed a new version of `meeshkan` and want the changes to take effect\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with SageMaker (original notebook [here](https://github.com/awslabs/amazon-sagemaker-examples/tree/master/sagemaker-python-sdk/pytorch_lstm_word_language_model))\n",
    "\n",
    "1. [Background](#Background)\n",
    "1. [Setup](#Setup)\n",
    "1. [Data](#Data)\n",
    "1. [Train](#Train)\n",
    "1. [Monitor with Meeshkan](#Monitor)\n",
    "\n",
    "---\n",
    "\n",
    "## Background\n",
    "\n",
    "This example trains a multi-layer LSTM RNN model on a language modeling task based on [PyTorch example](https://github.com/pytorch/examples/tree/master/word_language_model). By default, the training script uses the Wikitext-2 dataset. We will train a model on SageMaker, deploy it, and then use deployed model to generate new text.\n",
    "\n",
    "For more information about the PyTorch in SageMaker, please visit [sagemaker-pytorch-containers](https://github.com/aws/sagemaker-pytorch-containers) and [sagemaker-python-sdk](https://github.com/aws/sagemaker-python-sdk) github repositories.\n",
    "\n",
    "---\n",
    "\n",
    "## Setup\n",
    "\n",
    "_This notebook was created and tested on an ml.p2.xlarge notebook instance._"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a SageMaker session and specify:\n",
    "\n",
    "- The S3 bucket and prefix that you want to use for training and model data.  This should be within the same region as the Notebook Instance, training, and hosting.\n",
    "- The IAM role arn used to give training and hosting access to your data. See [the documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) for how to create these.  Note, if more than one role is required for notebook instances, training, and/or hosting, please replace the sagemaker.get_execution_role() with appropriate full IAM role arn string(s)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = sagemaker.Session()\n",
    "\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "prefix = 'sagemaker/DEMO-pytorch-rnn-lstm'\n",
    "\n",
    "role = sagemaker.get_execution_role()  # If you are running the notebook in notebook instance\n",
    "# If you are running the notebook locally, fill in the appropriate execution role ARN here.\n",
    "# role = \"EXECUTOR_ROLE_ARN_HERE\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "### Getting the data\n",
    "As mentioned above we are going to use [the wikitext-2 raw data](https://www.salesforce.com/products/einstein/ai-research/the-wikitext-dependency-language-modeling-dataset/). This data is from Wikipedia and is licensed CC-BY-SA-3.0. Before you use this data for any other purpose than this example, you should understand the data license, described at https://creativecommons.org/licenses/by-sa/3.0/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "wget http://research.metamind.io.s3.amazonaws.com/wikitext/wikitext-2-raw-v1.zip\n",
    "unzip -n wikitext-2-raw-v1.zip\n",
    "cd wikitext-2-raw\n",
    "mv wiki.test.raw test && mv wiki.train.raw train && mv wiki.valid.raw valid\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's preview what data looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head -5 wikitext-2-raw/train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Uploading the data to S3\n",
    "We are going to use the `sagemaker.Session.upload_data` function to upload our datasets to an S3 location. The return value inputs identifies the location -- we will use later when we start the training job.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = sagemaker_session.upload_data(path='wikitext-2-raw', bucket=bucket, key_prefix=prefix)\n",
    "print('input spec (in this case, just an S3 path): {}'.format(inputs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train\n",
    "### Training script\n",
    "We need to provide a training script that can run on the SageMaker platform. The training script is very similar to a training script you might run outside of SageMaker, but you can access useful properties about the training environment through various environment variables, such as:\n",
    "\n",
    "* `SM_MODEL_DIR`: A string representing the path to the directory to write model artifacts to.\n",
    "  These artifacts are uploaded to S3 for model hosting.\n",
    "* `SM_OUTPUT_DATA_DIR`: A string representing the filesystem path to write output artifacts to. Output artifacts may\n",
    "  include checkpoints, graphs, and other files to save, not including model artifacts. These artifacts are compressed\n",
    "  and uploaded to S3 to the same S3 prefix as the model artifacts.\n",
    "\n",
    "Supposing one input channel, 'training', was used in the call to the PyTorch estimator's `fit()` method,\n",
    "the following will be set, following the format `SM_CHANNEL_[channel_name]`:\n",
    "\n",
    "* `SM_CHANNEL_TRAINING`: A string representing the path to the directory containing data in the 'training' channel.\n",
    "\n",
    "A typical training script loads data from the input channels, configures training with hyperparameters, trains a model, and saves a model to `model_dir` so that it can be hosted later. Hyperparameters are passed to your script as arguments and can be retrieved with an `argparse.ArgumentParser` instance. For example, the script run by this notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize 'source/train.py'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For more information about training environment variables, please visit [SageMaker Containers](https://github.com/aws/sagemaker-containers)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the current example we also need to provide source directory since training script imports data and model classes from other modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls source"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run training in SageMaker\n",
    "The PyTorch class allows us to run our training function as a training job on SageMaker infrastructure. We need to configure it with our training script and source directory, an IAM role, the number of training instances, and the training instance type. In this case we will run our training job on ```ml.p2.xlarge``` instance. As you can see in this example you can also specify hyperparameters. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the metrics to capture\n",
    "Our training script prints, for example, rows such as 'TrainingLoss=2.2342', so tell SageMaker to capture those. **Meeshkan can only notify you of the metrics that you define here**. Note that if you use built-in Amazon algorithms, those define their own metric definitions and will be automatically watched."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "float_pattern = \"([0-9\\\\.]+)\"\n",
    "training_loss_pattern = \"TrainingLoss={}\".format(float_pattern)\n",
    "validation_loss_pattern = \"ValidationLoss={}\".format(float_pattern)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a PyTorch estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "estimator = PyTorch(entry_point='train.py',\n",
    "                    role=role,\n",
    "                    framework_version='1.0.0',\n",
    "                    train_instance_count=1,\n",
    "                    train_instance_type='ml.p2.xlarge',\n",
    "                    source_dir='source',\n",
    "                    hyperparameters={\n",
    "                        'epochs': 2,\n",
    "                        'tied': True\n",
    "                    },\n",
    "                    metric_definitions=[\n",
    "                        {'Name': 'train:loss', 'Regex': training_loss_pattern},\n",
    "                        {'Name': 'validation:loss', 'Regex': validation_loss_pattern}\n",
    "                    ]\n",
    "                   )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Launch training job\n",
    "After we've constructed our PyTorch object, we can fit it using the data we uploaded to S3. SageMaker makes sure our data is available in the local filesystem, so our training script can simply read the data from disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import gmtime, strftime\n",
    "\n",
    "job_name = 'pytorch-meeshkan-rnn-' + strftime(\"%Y-%m-%d-%H-%M-%S\", gmtime())\n",
    "\n",
    "estimator.fit({'training': inputs}, job_name=job_name, wait=False)\n",
    "print(\"Submitted job {:s}\".format(job_name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can check the status of the job with\n",
    "```python\n",
    "sagemaker_client.describe_training_job(TrainingJobName=job_name)['TrainingJobStatus']\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Monitor\n",
    "Start monitoring the job with Meeshkan, checking for updates every minute. You will be notified via Slack when new metrics are posted by the SageMaker job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meeshkan.sagemaker.monitor(job_name=job_name, poll_interval=60)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stop the job\n",
    "For demonstration purposes, you can stop the training job using the low-level `boto3` client as shown below. Stopping the job may take a few minutes, but once it's done, you'll get a Slack notification from Meeshkan of your job being finished."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "sagemaker_client = boto3.client('sagemaker')\n",
    "sagemaker_client.stop_training_job(TrainingJobName=job_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stop Meeshkan agent\n",
    "Stop the agent once you do not need monitoring anymore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meeshkan.stop()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}