{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use Script Mode to train any TensorFlow script from GitHub in SageMaker\n",
    "\n",
    "In this tutorial, you train a TensorFlow script in SageMaker using the new Script Mode Tensorflow Container.\n",
    "\n",
    "For this example, you use [Multi-layer Recurrent Neural Networks (LSTM, RNN) for character-level language models in Python using Tensorflow](https://github.com/sherjilozair/char-rnn-tensorflow), but you can use the same technique for other scripts or repositories. For example, [TensorFlow Model Zoo](https://github.com/tensorflow/models) and [TensorFlow benchmark scripts](https://github.com/tensorflow/benchmarks/tree/master/scripts/tf_cnn_benchmarks)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the environment\n",
    "Let's start by creating a SageMaker session and specifying the following:\n",
    "- The S3 bucket and prefix to use for training and model data. The bucket should be in the same region as the Notebook Instance, training instance(s), and hosting instance(s). This example uses the default bucket that a SageMaker `Session` creates.\n",
    "- The IAM role that allows SageMaker services to access your data. For more information about using IAM roles in SageMaker, see  [Amazon SageMaker Roles](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "\n",
    "role = sagemaker.get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clone the repository\n",
    "Run the following command to clone the repository that contains the example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/sherjilozair/char-rnn-tensorflow > /dev/null 2>&1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This repository includes a README.md with an overview of the project, requirements, and basic usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, Markdown, Latex\n",
    "display(Markdown('char-rnn-tensorflow/README.md'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get the data\n",
    "For training data, use plain text versions of Sherlock Holmes stories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir sherlock\n",
    "!wget https://sherlock-holm.es/stories/plain-text/cnus.txt --force-directories --output-document=sherlock/input.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test locally"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Script Mode is in a development phase. We need to construct an Estimator to be able to use it with this example, see [SageMaker Python SDK](https://github.com/aws/sagemaker-python-sdk) for more information.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "from sagemaker.estimator import Framework\n",
    "from sagemaker.tensorflow import TensorFlow\n",
    "\n",
    "class ScriptModeTensorFlow(Framework):\n",
    "    \"\"\"This class is temporary until the final version of Script Mode is released.\n",
    "    \"\"\"\n",
    "    \n",
    "    __framework_name__ = \"tensorflow-scriptmode\"\n",
    "    \n",
    "    create_model = TensorFlow.create_model\n",
    "    \n",
    "    def __init__(self, py_version='py3', **kwargs):\n",
    "        super(ScriptModeTensorFlow, self).__init__(**kwargs)\n",
    "        self.py_version = py_version\n",
    "        self.image_name = None\n",
    "        self.framework_version = '1.11'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Use [Local Mode](https://github.com/aws/sagemaker-python-sdk#local-mode) to run the script locally in the notebook instance before you run a SageMaker training job:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "hyperparameters = {'num_epochs': 1, \n",
    "                   'data_dir': '/opt/ml/input/data/training',\n",
    "                   'save_dir': '/opt/ml/model'}\n",
    "\n",
    "estimator = ScriptModeTensorFlow(entry_point='train.py',\n",
    "                                 source_dir='char-rnn-tensorflow',\n",
    "                                 train_instance_type='local',      # Run in local mode\n",
    "                                 train_instance_count=1,\n",
    "                                 hyperparameters=hyperparameters,\n",
    "                                 role=role)\n",
    "\n",
    "estimator.fit({'training': 'file://%s' % os.path.join(os.getcwd(), 'sherlock')})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How Script Mode executes the script in the container\n",
    "\n",
    "The above cell downloads a Python 3 CPU container locally and simulates a SageMaker training job. When training starts, script mode installs the user script as a Python module. The module name matches the script name. In this case, **train.py** is transformed into a Python module named **train**.\n",
    "\n",
    "After that, the Python interpreter executes the user module, passing **hyperparameters** as script arguments. The example above is executed as follows:\n",
    "```bash\n",
    "python -m train --num-epochs 1 --data-dir /opt/ml/input/data/training --save-dir /opt/ml/model\n",
    "```\n",
    "\n",
    "The **train** module consumes the hyperparameters using any argument parsing library. [The example we're using](https://github.com/sherjilozair/char-rnn-tensorflow/blob/master/train.py#L11) uses the Python [argparse](https://docs.python.org/3/library/argparse.html) library:\n",
    "\n",
    "```python\n",
    "parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n",
    "# Data and model checkpoints directories\n",
    "parser.add_argument('--data_dir', type=str, default='data/tinyshakespeare', help='data directory containing input.txt with training examples')\n",
    "parser.add_argument('--save_dir', type=str, default='save', help='directory to store checkpointed models')\n",
    "...\n",
    "args = parser.parse_args()\n",
    "\n",
    "```\n",
    "\n",
    "\n",
    "Let's explain the values of **--data_dir** and **--save-dir**:\n",
    "\n",
    "- **/opt/ml/input/data/training** is the directory inside the container where the training data is downloaded. The data is downloaded to this folder because **training** is the channel name defined in ```estimator.fit({'training': inputs})```. See [training data](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html#your-algorithms-training-algo-running-container-trainingdata) for more information. \n",
    "\n",
    "- **/opt/ml/model** use this directory to save models, checkpoints, or any other data. Any data saved in this folder is saved in the S3 bucket defined for training. See [model data](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html#your-algorithms-training-algo-envvariables) for more information.\n",
    "\n",
    "### Reading additional information from the container\n",
    "\n",
    "Often, a user script needs additional information from the container that is not available in ```hyperparameters```.\n",
    "SageMaker containers write this information as **environment variables** that are available inside the script.\n",
    "\n",
    "For example, the example above can read information about the **training** channel provided in the training job request by adding the environment variable `SM_CHANNEL_TRAINING` as the default value for the `--data_dir` argument:\n",
    "\n",
    "```python\n",
    "if __name__ == '__main__':\n",
    "  parser = argparse.ArgumentParser()\n",
    "  # reads input channels training and testing from the environment variables\n",
    "  parser.add_argument('--data_dir', type=str, default=os.environ['SM_CHANNEL_TRAINING'])\n",
    "```\n",
    "\n",
    "Script mode displays the list of available environment variables in the training logs. You can find the [entire list here](https://github.com/aws/sagemaker-containers/blob/master/README.md#environment-variables-full-specification)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training in SageMaker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After you test the training job locally, upload the dataset to an S3 bucket so SageMaker can access the data during training.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = sagemaker_session.upload_data(path='sherlock', bucket=bucket, key_prefix='datasets/sherlock')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train in SageMaker, change the estimator argument **train_instance_type** to any SageMaker ml instance available for training. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = ScriptModeTensorFlow(entry_point='train.py',\n",
    "                                source_dir='char-rnn-tensorflow',\n",
    "                                train_instance_type='ml.c4.xlarge', \n",
    "                                train_instance_count=1,\n",
    "                                hyperparameters=hyperparameters,\n",
    "                                role=role)\n",
    "\n",
    "estimator.fit({'training': inputs})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Installing additional requirements"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installing pip packages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Script Mode installs the contents of your `source_dir` folder in the container as a [Python package](https://github.com/aws/sagemaker-containers/blob/master/src/sagemaker_containers/_modules.py#L100). You can include a [requirements.txt file in the root folder of your source_dir to install any pip dependencies](https://github.com/aws/sagemaker-containers/blob/master/src/sagemaker_containers/_modules.py#L111). You can, for example, install the lastest version of TensorFlow in the container:\n",
    "\n",
    "content of requirements.txt\n",
    "```\n",
    "tensorflow==1.11.0\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Installing apt-get packages and other dependencies\n",
    "You can define a `setup.py` file in your `source_dir` folder to install other dependencies. The example below installs [TensorFlow for C](https://www.tensorflow.org/install/lang_c) in the container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir tf_c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile tf_c/get-tf-c.sh\n",
    "\n",
    "wget -q -t 3 https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-cpu-linux-x86_64-1.11.0.tar.gz\n",
    "tar -xzvf libtensorflow-cpu-linux-x86_64-1.11.0.tar.gz -C /usr/local\n",
    "\n",
    "ldconfig\n",
    "\n",
    "gcc -I/usr/local/include -L/usr/local/lib hello_tf.c -ltensorflow -o hello_tf\n",
    "cp hello_tf /usr/bin/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile tf_c/hello_tf.c\n",
    "\n",
    "#include <stdio.h>\n",
    "#include <tensorflow/c/c_api.h>\n",
    "\n",
    "int main() {\n",
    "  printf(\"Hello from TensorFlow C library version %s\\n\", TF_Version());\n",
    "  return 0;\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile tf_c/setup.py\n",
    "from distutils.command.build_py import build_py as _build_py\n",
    "from distutils.core import setup\n",
    "import subprocess\n",
    "\n",
    "class build_py(_build_py):\n",
    "    def run(self):\n",
    "        subprocess.check_output(['bash', './get-tf-c.sh'])\n",
    "\n",
    "        super(build_py, self).run()\n",
    "\n",
    "\n",
    "from setuptools import setup\n",
    "setup(packages=[''],\n",
    "      name=\"test\",\n",
    "      version='1.0.0',\n",
    "      cmdclass={'build_py': build_py},\n",
    "      include_package_data=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile tf_c/train_c.py\n",
    "\n",
    "import subprocess\n",
    "\n",
    "message = subprocess.check_output('hello_tf')\n",
    "assert message == b'Hello from TensorFlow C library version 1.11.0\\n'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = ScriptModeTensorFlow(entry_point='train_c.py',\n",
    "                                 source_dir='tf_c',\n",
    "                                 train_instance_type='local', \n",
    "                                 train_instance_count=1,\n",
    "                                 role=role)\n",
    "\n",
    "estimator.fit({})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}