examples/training_without_local_storage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🪣 Training without Local Storage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Composer natively supports training in environments where servers lack local persistent storage.\n",
    "\n",
    "Without local persistent storage, checkpoints and datasets will need to be downloaded from the cloud, and all checkpoints, logs, metrics, and other artifacts will need to be backed up directly to the cloud.\n",
    "\n",
    "Composer can automatically load checkpoints from cloud storage, convert traditional datasets into a format that can be streamed in when training, and asynchronously back up checkpoints and other artifacts without blocking the training loop. In this tutorial, we'll walk through how to convert a local training workflow for an MNIST classifier into one that does not require persistent disks.\n",
    "\n",
    "_In reality, small models with datasets that converge quickly - like MNIST - would likely not require streaming datasets, since it's usually fast enough to download all data at the start of training. However, it's a simple example that highlights the steps involved to use these features in Composer._\n",
    "\n",
    "1. [Prerequisites](#1-Prerequisites)\n",
    "1. [The local training workflow](#2-the-local-training-workflow)\n",
    "1. [Storing and loading checkpoints and logs with the cloud](#3-storing-and-loading-checkpoints-and-logs-with-the-cloud)\n",
    "1. [Switching to streaming datasets](#4-switching-to-streaming-datasets)\n",
    "1. [Putting it all together](#5-putting-it-all-together)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Prerequisites\n",
    "\n",
    "This tutorial requires access to an AWS S3 Bucket. If AWS credentials are not already available in your environment, then you will need to obtain an `AWS_ACCESS_KEY_ID` and an `AWS_SECRET_ACCESS_KEY` that have permission to upload to and download from a S3 Bucket.\n",
    "\n",
    "Here, we'll define all configuration variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_bucket_name = 'my-bucket'  # The S3 bucket to use\n",
    "# Give all objects in the bucket a prefix, allowing the\n",
    "# bucket to be shared across training runs\n",
    "bucket_prefix = 'composer-diskless-training-tutorial'\n",
    "\n",
    "# If necessary, uncomment the following lines to set AWS credentials\n",
    "# Do NOT include quotes\n",
    "# %env AWS_ACCESS_KEY_ID ***\n",
    "# %env AWS_SECRET_ACCESS_KEY ***\n",
    "# %env AWS_DEFAULT_REGION us-west-2\n",
    "\n",
    "# Also define local (temporary) folders that will be used for\n",
    "# staging datasets, checkpoints, and log files before uploading\n",
    "data_dir = '/tmp/data'\n",
    "tensorboard_log_dir = '/tmp/tb_logs'\n",
    "checkpoint_dir = '/tmp/checkpoints'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, if you didn't already, install Composer with streaming support:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install 'mosaicml[streaming,tensorboard]'\n",
    "\n",
    "# Or, to install the latest composer directly from GitHub\n",
    "%pip install 'mosaicml[streaming,tensorboard] @ git+https://github.com/mosaicml/composer.git@dev'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. The Local Training Workflow\n",
    "\n",
    "Let's first define our local training code for MNIST. This code downloads all data before training starts. We include the Tensorboard logger, which we will use to visualize results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "import os\n",
    "\n",
    "import torch.utils.data\n",
    "from torch.optim import SGD\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "from composer import Trainer\n",
    "from composer.loggers import TensorboardLogger\n",
    "from composer.utils.reproducibility import seed_all\n",
    "from composer.models.classify_mnist import mnist_model\n",
    "\n",
    "\n",
    "# Configure the trainer\n",
    "\n",
    "# Model and optimizer\n",
    "def get_model_and_optimizer():\n",
    "    # Set the seed before creating the model\n",
    "    # for consistent initialization\n",
    "    seed_all(42)\n",
    "    model = mnist_model(num_classes=10)\n",
    "    optimizer = SGD(model.parameters(), lr=0.01)\n",
    "    return model, optimizer\n",
    "\n",
    "model, optimizer = get_model_and_optimizer()\n",
    "\n",
    "# Datasets\n",
    "batch_size = 2048\n",
    "train_dataset = MNIST(\n",
    "    root=os.path.join(data_dir, 'imagefolder'),\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset=train_dataset,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "eval_dataset = MNIST(\n",
    "    root=os.path.join(data_dir, 'imagefolder'),\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")\n",
    "eval_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset=eval_dataset,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "\n",
    "# Tensorboard Logger (for visualizing results)\n",
    "def get_tensorboard_logger():\n",
    "    shutil.rmtree(tensorboard_log_dir, ignore_errors=True)\n",
    "    return TensorboardLogger(log_dir=tensorboard_log_dir, flush_interval=1)\n",
    "\n",
    "\n",
    "# Clean up the checkpoint directory (if it already exists)\n",
    "shutil.rmtree(checkpoint_dir, ignore_errors=True)\n",
    "\n",
    "# Create the trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    max_duration='2ep',\n",
    "    # Make training fast: Terminate each epoch after 5 batches\n",
    "    train_subset_num_batches=5,\n",
    "    optimizers=optimizer,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=eval_dataloader,\n",
    "    # Make evaluation fast: Evaluate on only five batches\n",
    "    eval_subset_num_batches=5,\n",
    "    # Flush log files every batch\n",
    "    loggers=get_tensorboard_logger(),\n",
    "    save_folder=checkpoint_dir,\n",
    "    save_interval=\"1ba\",\n",
    "    run_name='local_training_run',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train!\n",
    "trainer.fit()\n",
    "\n",
    "# Close the trainer\n",
    "trainer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We trained our model. Now let's verify that we have our checkpoints and can visualize our Tensorboard logs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Checkpoint Files\")\n",
    "!ls -al {checkpoint_dir}\n",
    "\n",
    "# Visualize Tensorboard Logs\n",
    "%load_ext tensorboard\n",
    "import tensorboard.notebook\n",
    "\n",
    "# NOTE: Tensorboard can take a few moments to appear\n",
    "# If it does not show up after ~30 seconds, run this cell again\n",
    "%tensorboard --logdir {tensorboard_log_dir}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Storing and loading checkpoints and logs with the cloud"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we verified our local training workflow works as intended, let's configure Composer to back up our checkpoints and Tensorboard TF Event files to the cloud. \n",
    "\n",
    "First, a brief overview of Composer architecture:\n",
    "\n",
    "* [Logger][Logger]: Composer includes a centralized logger, which passes logged data to each `LoggerDestination` (more on that below). Logged data can be either metrics or artifacts. The logger is similar to Python's built-in `logging.getLogger(...)` but is designed to log structured metrics and artifacts in addition to just text.\n",
    "\n",
    "[Logger]: https://docs.mosaicml.com/en/stable/api_reference/composer.loggers.logger.html#composer.loggers.logger.Logger\n",
    "\n",
    "* [LoggerDestination][LoggerDestination]: Where logs are sent is specified via the `loggers` argument of the [Trainer constructor][trainer]. The centralized `Logger` (above) passes all metrics and artifacts to each `LoggerDestination`, which is responsible for handling and storing the data. For example, Composer includes `LoggerDestination`s for logging to [files][FileLogger], [Tensorboard][TensorboardLogger], [Weights & Biases][wandb], and [Object Stores][ObjectStoreLogger] like S3. Not all `LoggerDestination`s support storing all types of data; for example, the [ObjectStoreLogger][ObjectStoreLogger] only supports logging artifacts, whereas the [FileLogger][FileLogger] only supports logging files. Others, such as the [WandBLogger][wandb], support both.\n",
    "\n",
    "[LoggerDestination]: https://docs.mosaicml.com/en/stable/api_reference/composer.loggers.logger_destination.html#composer.loggers.logger_destination.LoggerDestination\n",
    "[trainer]: https://docs.mosaicml.com/en/stable/api_reference/composer.trainer.trainer.html#composer.trainer.trainer.Trainer\n",
    "[FileLogger]: https://docs.mosaicml.com/en/stable/api_reference/composer.loggers.file_logger.html#composer.loggers.file_logger.FileLogger\n",
    "[TensorboardLogger]: https://docs.mosaicml.com/en/stable/api_reference/composer.loggers.tensorboard_logger.html#composer.loggers.tensorboard_logger.TensorboardLogger\n",
    "[wandb]: https://docs.mosaicml.com/en/stable/api_reference/composer.loggers.wandb_logger.html#composer.loggers.wandb_logger.WandBLogger\n",
    "[ObjectStoreLogger]: https://docs.mosaicml.com/en/stable/api_reference/composer.loggers.object_store_logger.html#composer.loggers.object_store_logger.ObjectStoreLogger\n",
    "[FileLogger]: https://docs.mosaicml.com/en/stable/api_reference/composer.loggers.file_logger.html#composer.loggers.file_logger.FileLogger\n",
    "\n",
    "* Metrics: A metric is a scalar, such as accuracy, that can be logged. Usually you would want to plot metrics over time (e.g., to see how accuracy improves over batches).\n",
    "\n",
    "* [Artifacts][Artifacts]: An artifact is a file generated throughout training, such as a checkpoint or log file. Each file is a separate artifact.\n",
    "\n",
    "* [ObjectStore][ObjectStore]: The abstract `ObjectStore` class provides an API for uploading and downloading checkpoints. Composer includes object store implementations for [S3][S3], [SFTP][SFTP], and [Libcloud][Libcloud]. You can also write your own implementation by extending the [base class][ObjectStore] if you are using a custom backend.\n",
    "\n",
    "[Artifacts]: https://docs.mosaicml.com/en/latest/trainer/artifact_logging.html\n",
    "[ObjectStore]: https://docs.mosaicml.com/en/stable/api_reference/composer.utils.object_store.object_store.html#composer.utils.object_store.object_store.ObjectStore\n",
    "[S3]: https://docs.mosaicml.com/en/stable/api_reference/composer.utils.object_store.s3_object_store.html#composer.utils.object_store.s3_object_store.S3ObjectStore\n",
    "[SFTP]: https://docs.mosaicml.com/en/stable/api_reference/composer.utils.object_store.sftp_object_store.html#composer.utils.object_store.sftp_object_store.SFTPObjectStore\n",
    "[Libcloud]: https://docs.mosaicml.com/en/stable/api_reference/composer.utils.object_store.libcloud_object_store.html#composer.utils.object_store.libcloud_object_store.LibcloudObjectStore\n",
    "\n",
    "We'll use these components together to back up our checkpoints and Tensorboard TF Event files to the cloud. Internally, the [CheckpointSaver][CheckpointSaver] callback and [TensorboardLogger][TensorboardLogger] pass all generated files to the [Logger][Logger] as artifacts by calling [Logger.file_artifact][Logger_file_artifact]. The Logger then passes these files to each [LoggerDestination][LoggerDestination] (specified via the `loggers` argument of the [Trainer constructor][trainer]). A `LoggerDestination`, which can implement the `log_file_artifact` method, is responsible for uploading the file to the cloud.\n",
    "\n",
    "[CheckpointSaver]: https://docs.mosaicml.com/en/stable/api_reference/composer.callbacks.checkpoint_saver.html#composer.callbacks.checkpoint_saver.CheckpointSaver\n",
    "[Logger_file_artifact]: https://docs.mosaicml.com/en/stable/api_reference/composer.loggers.logger.html#composer.loggers.logger.Logger.file_artifact\n",
    "\n",
    "\n",
    "Here, our \"cloud\" will be an S3 bucket. To upload checkpoints and Tensorboard TF Event files to the bucket, we'll add the [ObjectStoreLogger][ObjectStoreLogger] with the [S3ObjectStore][S3] backend to our list of logger destinations. This class asynchronously uploads artifacts to an object store without blocking the training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.loggers import ObjectStoreLogger\n",
    "from composer.utils.object_store import S3ObjectStore\n",
    "\n",
    "# Clean the directories from the previous training run\n",
    "shutil.rmtree(checkpoint_dir, ignore_errors=True)\n",
    "\n",
    "model, optimizer = get_model_and_optimizer()\n",
    "\n",
    "def get_object_store_logger():\n",
    "    return ObjectStoreLogger(\n",
    "        object_store_cls=S3ObjectStore,\n",
    "        # Keyword arguments passed to the S3ObjectStore constructor\n",
    "        object_store_kwargs={\n",
    "            'bucket': s3_bucket_name,\n",
    "            'prefix': bucket_prefix,\n",
    "        },\n",
    "        # In Jupyter, we set use_procs to False, since subprocess do not work\n",
    "        # well within notebooks. Outside of Jupyter, it is recommended to let\n",
    "        # use_procs default to True for performance\n",
    "        use_procs=False,\n",
    "    )\n",
    "\n",
    "# Create the trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    max_duration='2ep',\n",
    "    train_subset_num_batches=5,\n",
    "    optimizers=optimizer,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=eval_dataloader,\n",
    "    eval_subset_num_batches=5,\n",
    "    loggers=[\n",
    "        get_object_store_logger(),\n",
    "        get_tensorboard_logger(),\n",
    "    ],\n",
    "    save_folder=checkpoint_dir,\n",
    "    save_interval=\"1ba\",\n",
    "    # Because we are uploading checkpoints to the cloud, we set\n",
    "    # save_num_checkpoints_to_keep to 0 to delete them locally\n",
    "    # after uploading to save disk space!\n",
    "    save_num_checkpoints_to_keep=0,\n",
    "    run_name='cloud_training_run',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train!\n",
    "trainer.fit()\n",
    "\n",
    "# Close the trainer, which will block on until all files have been uploaded\n",
    "trainer.close()\n",
    "\n",
    "# Remove all local Tensorboard traces, since they were uploaded to the cloud\n",
    "shutil.rmtree(tensorboard_log_dir, ignore_errors=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! We trained our model again, and this time, uploaded the checkpoints and Tensorboard logs to our S3 bucket. Let's verify that our files exist in the bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "s3 = boto3.client('s3')\n",
    "\n",
    "def print_objects_in_bucket(bucket: str, prefix: str):\n",
    "    response = s3.list_objects(Bucket=bucket, Prefix=prefix)\n",
    "    keys = [obj['Key'] for obj in response['Contents']]\n",
    "    keys.sort()\n",
    "    for k in keys:\n",
    "        print(k)\n",
    "\n",
    "print_objects_in_bucket(s3_bucket_name, bucket_prefix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also take a look at visualizing tensorboard from our S3 Bucket:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorboard.notebook\n",
    "\n",
    "%tensorboard --logdir s3://{s3_bucket_name}/{bucket_prefix}/tensorboard_logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also resume training from a cloud checkpoint, without having to first download it. To do so, we'll need to set the `load_object_store` argument of the Trainer constructor to our `cloud_logger`, and the `load_path` to the object name of our checkpoint file within the bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the cloud logger\n",
    "cloud_logger = get_object_store_logger()\n",
    "tensorboard_logger = get_tensorboard_logger()\n",
    "\n",
    "model, optimizer = get_model_and_optimizer()\n",
    "\n",
    "# Create the trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    max_duration='4ep',\n",
    "    # Load the latest checkpoint\n",
    "    load_path='cloud_training_run/checkpoints/latest-rank0',\n",
    "    # Load from the cloud logger\n",
    "    load_object_store=cloud_logger,\n",
    "    train_subset_num_batches=5,\n",
    "    optimizers=optimizer,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=eval_dataloader,\n",
    "    eval_subset_num_batches=5,\n",
    "    loggers=[\n",
    "        cloud_logger,\n",
    "        tensorboard_logger,\n",
    "    ],\n",
    "    save_folder=checkpoint_dir,\n",
    "    save_interval=\"1ba\",\n",
    "    save_num_checkpoints_to_keep=0,\n",
    "    run_name='cloud_training_run',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train!\n",
    "trainer.fit()\n",
    "\n",
    "# Close the trainer, which will block until all files have been uploaded\n",
    "trainer.close()\n",
    "\n",
    "# Remove all local Tensorboard traces, since they were uploaded to the cloud\n",
    "shutil.rmtree(tensorboard_log_dir, ignore_errors=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take another look at Tensorboard, which should show another two epochs (four total):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorboard.notebook\n",
    "\n",
    "%tensorboard --logdir s3://{s3_bucket_name}/{bucket_prefix}/tensorboard_logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an alternative to setting the `load_path` and `load_object_store`, we could instead set `autoresume=True`. See our [tutorial on autoresumption](https://docs.mosaicml.com/en/latest/examples/checkpoint_autoresume.html) for more information about how this feature works."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Switching to Streaming Datasets\n",
    "\n",
    "So far, we covered how to store and load checkpoints and other artifacts with cloud storage. But, we relied on all of our training data being available at the start. While it works to download all data to a local folder before training, this technique isn't scalable for large datasets. Training wouldn't begin until _all_ data is downloaded, causing accelerators to idle.\n",
    "\n",
    "To overcome this bottleneck, we can switch to streaming datasets. With streaming datasets, the dataset must be split into multiple shard files, each of which contains a subset of the samples. Shards are downloaded asynchronously (in separate subprocesses) while we train. Instead of waiting until all data is downloaded, we can begin training as soon as the first shard is downloaded. This feature lets training begin instantly!\n",
    "\n",
    "For an in-depth walk-through on how to use streaming datasets, see our tutorial on training [FaceSynthetics with a Streaming Dataloader](https://docs.mosaicml.com/en/latest/examples/streaming_dataloader_facesynthetics.html). Here, we'll encode the MNIST dataset in the Composer streaming format and upload it to our S3 bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write the streaming dataset\n",
    "import os\n",
    "import torch\n",
    "import io\n",
    "from typing import Dict, cast, Iterable\n",
    "import struct\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from composer.datasets.streaming import StreamingDatasetWriter\n",
    "\n",
    "train_out_folder = os.path.join(data_dir, 'streaming', 'train')\n",
    "eval_out_folder = os.path.join(data_dir, 'streaming', 'eval')\n",
    "\n",
    "train_remote = S3ObjectStore(\n",
    "    bucket=s3_bucket_name,\n",
    "    prefix=bucket_prefix + '/train',\n",
    ")\n",
    "\n",
    "eval_remote = S3ObjectStore(\n",
    "    bucket=s3_bucket_name,\n",
    "    prefix=bucket_prefix + '/eval',\n",
    ")\n",
    "\n",
    "\n",
    "shutil.rmtree(train_out_folder, ignore_errors=True)\n",
    "shutil.rmtree(eval_out_folder, ignore_errors=True)\n",
    "os.makedirs(train_out_folder, exist_ok=True)\n",
    "os.makedirs(eval_out_folder, exist_ok=True)\n",
    "\n",
    "fields = ['i', 'x', 'y']\n",
    "\n",
    "def encode_sample(i: int, x: torch.Tensor, y: int) -> Dict[str, bytes]:\n",
    "    \"\"\"Encode a (x,y) sample into a dictionary.\"\"\"\n",
    "    x_buffer = io.BytesIO()\n",
    "    torch.save(x, x_buffer)\n",
    "    x_buffer.seek(0)\n",
    "    # See https://docs.python.org/3/library/struct.html#format-characters for\n",
    "    # struct format characters\n",
    "    return {\n",
    "        # The sample index, encoded as an uint64 (Q format code)\n",
    "        'i': struct.pack('Q', i),\n",
    "        # The sample input, in bytes\n",
    "        'x': x_buffer.read(),\n",
    "        # The class index, encoded as a uint64 (Q format code)\n",
    "        'y': struct.pack('Q', y),\n",
    "    }\n",
    "\n",
    "print(\"Writing training dataset\")\n",
    "with StreamingDatasetWriter(train_out_folder, fields, remote=train_remote) as out:\n",
    "    for i, (x, y) in tqdm(\n",
    "        enumerate(cast(Iterable, train_dataset)),\n",
    "        total=len(train_dataset),\n",
    "    ):\n",
    "        out.write_sample(encode_sample(i, x, y))\n",
    "print(f\"Uploaded training dataset to s3://{s3_bucket_name}/{bucket_prefix}/train\")\n",
    "\n",
    "print(\"Writing evaluation dataset\")\n",
    "with StreamingDatasetWriter(eval_out_folder, fields, remote=eval_remote) as out:\n",
    "    for i, (x, y) in tqdm(\n",
    "        enumerate(cast(Iterable, eval_dataset)),\n",
    "        total=len(eval_dataset),\n",
    "    ):\n",
    "        out.write_sample(encode_sample(i, x, y))\n",
    "print(f\"Uploaded evaluation dataset to s3://{s3_bucket_name}/{bucket_prefix}/eval\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a look at the shard files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Training shard files\")\n",
    "print_objects_in_bucket(s3_bucket_name, f'{bucket_prefix}/train')\n",
    "\n",
    "print(\"Evaluation shard files\")\n",
    "print_objects_in_bucket(s3_bucket_name, f'{bucket_prefix}/eval')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's define our `MNISTStreamingDataset` subclass so we can decode the samples from bytes into an `(x, y)` tuple similar to a Torchvision-style dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from composer.datasets.streaming import StreamingDataset\n",
    "from typing import Tuple\n",
    "\n",
    "# Define the decoders for the dataset\n",
    "def decode_x(x: bytes):\n",
    "    x_buffer = io.BytesIO(x)\n",
    "    x = torch.load(x_buffer)\n",
    "    return x\n",
    "\n",
    "def decode_y(y: bytes) -> int:\n",
    "    y, = struct.unpack('Q', y)\n",
    "    assert isinstance(y, int)\n",
    "    return y\n",
    "\n",
    "decoders = {\n",
    "    'x': decode_x,\n",
    "    'y': decode_y, \n",
    "}\n",
    "\n",
    "# Create a custom subclass of StreamingDataset to automatically unpack\n",
    "# samples into tuples \n",
    "class MNISTStreamingDataset(StreamingDataset):\n",
    "    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:\n",
    "        # Overriding __getitem__ to unpack the dictionary\n",
    "        # from encode_sample into an (x, y) tuple\n",
    "        sample_dict = super().__getitem__(idx)\n",
    "        return sample_dict['x'], sample_dict['y']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's train using our streaming version of the MNIST dataset!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.utils.data\n",
    "\n",
    "model, optimizer = get_model_and_optimizer()\n",
    "\n",
    "# Redefine the dataset and dataloader to use the streaming format\n",
    "train_dataset = MNISTStreamingDataset(\n",
    "    remote=f's3://{s3_bucket_name}/{bucket_prefix}/train',\n",
    "    local=os.path.join(data_dir, 'streaming_cache', 'train'),\n",
    "    shuffle=True,\n",
    "    decoders=decoders,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "\n",
    "eval_dataset = MNISTStreamingDataset(\n",
    "    remote=f's3://{s3_bucket_name}/{bucket_prefix}/eval',\n",
    "    local=os.path.join(data_dir, 'streaming_cache', 'eval'),\n",
    "    shuffle=True,\n",
    "    decoders=decoders,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset=train_dataset,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "eval_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset=eval_dataset,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "\n",
    "# Create the trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    max_duration='2ep',\n",
    "    train_subset_num_batches=5,\n",
    "    optimizers=optimizer,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=eval_dataloader,\n",
    "    eval_subset_num_batches=5,\n",
    "    loggers=[\n",
    "        get_object_store_logger(),\n",
    "        get_tensorboard_logger(),\n",
    "    ],\n",
    "    save_folder=checkpoint_dir,\n",
    "    save_interval=\"1ba\",\n",
    "    save_num_checkpoints_to_keep=0,\n",
    "    run_name='streaming_training_run',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train!\n",
    "trainer.fit()\n",
    "\n",
    "# Close the trainer, which will block on until all files have been uploaded\n",
    "trainer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take another look at our Tensorboard plots. Note that the loss and accuracy will slightly differ from the local dataset run because the streaming dataset shuffling is nondeterministic -- it depends on the order in which shards are downloaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorboard.notebook\n",
    "\n",
    "%tensorboard --logdir s3://{s3_bucket_name}/{bucket_prefix}/tensorboard_logs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations! You trained your first model without relying on a persistent local disk!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Putting it all together\n",
    "\n",
    "In this tutorial, we walked through storing and loading checkpoints with the cloud, converting existing datasets into a streaming format, and training a model using both of these features.\n",
    "\n",
    "Below is a complete example, showing everything we did in one cell. Feel free to use this as a reference for your own training workflows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import os\n",
    "import struct\n",
    "import shutil\n",
    "import time\n",
    "from typing import Dict, cast, Iterable, Tuple\n",
    "\n",
    "import torch\n",
    "from torch.optim import SGD\n",
    "import torch.utils.data\n",
    "from torchvision.datasets import MNIST\n",
    "from torchvision.transforms import ToTensor\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from composer import Trainer\n",
    "from composer.models import mnist_model\n",
    "from composer.utils.reproducibility import seed_all\n",
    "from composer.utils.object_store import S3ObjectStore\n",
    "from composer.loggers import ObjectStoreLogger, TensorboardLogger\n",
    "from composer.datasets.streaming import StreamingDatasetWriter, StreamingDataset\n",
    "\n",
    "\n",
    "###############\n",
    "# Configuration\n",
    "###############\n",
    "\n",
    "# # The S3 bucket to use\n",
    "# s3_bucket_name = 'my-bucket'\n",
    "\n",
    "# # Prefix all objects in the bucket with key,\n",
    "# # allowing the bucket can be shared across training runs\n",
    "# bucket_prefix = 'composer-diskless-training-tutorial'\n",
    "\n",
    "# # If necessary, uncomment the following lines to set AWS credentials\n",
    "# # Do NOT include quotes\n",
    "# %env AWS_ACCESS_KEY_ID ***\n",
    "# %env AWS_SECRET_ACCESS_KEY ***\n",
    "# %env AWS_DEFAULT_REGION us-west-2\n",
    "\n",
    "# # Also define local (temporary) folders that will be used for\n",
    "# # staging datasets, checkpoints, and log files before uploading\n",
    "# data_dir = '/tmp/data'\n",
    "# tensorboard_log_dir = '/tmp/tb_logs'\n",
    "# checkpoint_dir = '/tmp/checkpoints'\n",
    "\n",
    "############################\n",
    "# Writing Streaming Datasets\n",
    "############################\n",
    "\n",
    "print(\"Writing the streaming datasets\")\n",
    "\n",
    "batch_size = 2048\n",
    "\n",
    "local_train_dataset = MNIST(\n",
    "    root=os.path.join(data_dir, 'imagefolder'),\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")\n",
    "local_eval_dataset = MNIST(\n",
    "    root=os.path.join(data_dir, 'imagefolder'),\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")\n",
    "\n",
    "train_out_folder = os.path.join(data_dir, 'streaming', 'train')\n",
    "eval_out_folder = os.path.join(data_dir, 'streaming', 'eval')\n",
    "\n",
    "train_remote = S3ObjectStore(\n",
    "    bucket=s3_bucket_name,\n",
    "    prefix=bucket_prefix + '/train',\n",
    ")\n",
    "\n",
    "eval_remote = S3ObjectStore(\n",
    "    bucket=s3_bucket_name,\n",
    "    prefix=bucket_prefix + '/eval',\n",
    ")\n",
    "\n",
    "shutil.rmtree(train_out_folder, ignore_errors=True)\n",
    "shutil.rmtree(eval_out_folder, ignore_errors=True)\n",
    "os.makedirs(train_out_folder, exist_ok=True)\n",
    "os.makedirs(eval_out_folder, exist_ok=True)\n",
    "\n",
    "fields = ['i', 'x', 'y']\n",
    "\n",
    "def encode_sample(i: int, x: torch.Tensor, y: int) -> Dict[str, bytes]:\n",
    "    \"\"\"Encode a (x,y) sample into a dictionary.\"\"\"\n",
    "    x_buffer = io.BytesIO()\n",
    "    torch.save(x, x_buffer)\n",
    "    x_buffer.seek(0)\n",
    "    # See https://docs.python.org/3/library/struct.html#format-characters for\n",
    "    # struct format characters\n",
    "    return {\n",
    "        # The sample index, encoded as an uint64 (Q format code)\n",
    "        'i': struct.pack('Q', i),\n",
    "        # The sample input, in bytes\n",
    "        'x': x_buffer.read(),\n",
    "        # The class index, encoded as a uint64 (Q format code)\n",
    "        'y': struct.pack('Q', y),\n",
    "    }\n",
    "\n",
    "print(\"Writing training dataset\")\n",
    "with StreamingDatasetWriter(train_out_folder, fields, remote=train_remote) as out:\n",
    "    for i, (x, y) in tqdm(\n",
    "        enumerate(cast(Iterable, local_train_dataset)),\n",
    "        total=len(local_train_dataset),\n",
    "    ):\n",
    "        out.write_sample(encode_sample(i, x, y))\n",
    "print(f\"Uploaded training dataset to s3://{s3_bucket_name}/{bucket_prefix}/train\")\n",
    "\n",
    "print(\"Writing evaluation dataset\")\n",
    "with StreamingDatasetWriter(eval_out_folder, fields, remote=eval_remote) as out:\n",
    "    for i, (x, y) in tqdm(\n",
    "        enumerate(cast(Iterable, local_eval_dataset)),\n",
    "        total=len(local_eval_dataset),\n",
    "    ):\n",
    "        out.write_sample(encode_sample(i, x, y))\n",
    "print(f\"Uploaded evaluation dataset to s3://{s3_bucket_name}/{bucket_prefix}/eval\")\n",
    "\n",
    "\n",
    "#################################\n",
    "# Loading from Streaming Datasets\n",
    "#################################\n",
    "\n",
    "def decode_x(x: bytes):\n",
    "    x_buffer = io.BytesIO(x)\n",
    "    x = torch.load(x_buffer)\n",
    "    return x\n",
    "\n",
    "def decode_y(y: bytes) -> int:\n",
    "    y, = struct.unpack('Q', y)\n",
    "    assert isinstance(y, int)\n",
    "    return y\n",
    "\n",
    "decoders = {\n",
    "    'x': decode_x,\n",
    "    'y': decode_y, \n",
    "}\n",
    "\n",
    "class MNISTStreamingDataset(StreamingDataset):\n",
    "    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:\n",
    "        # Overriding __getitem__ to unpack the dictionary from encode_sample\n",
    "        # into an (x, y) tuple\n",
    "        sample_dict = super().__getitem__(idx)\n",
    "        return sample_dict['x'], sample_dict['y']\n",
    "\n",
    "\n",
    "#######################\n",
    "# Trainer Configuration\n",
    "#######################\n",
    "\n",
    "run_name = f'{int(time.time())}-final-training-run'\n",
    "\n",
    "seed_all(42)  # Set the seed before creating the model\n",
    "model = mnist_model(num_classes=10)\n",
    "optimizer = SGD(model.parameters(), lr=0.01)\n",
    "\n",
    "shutil.rmtree(tensorboard_log_dir, ignore_errors=True)\n",
    "shutil.rmtree(checkpoint_dir, ignore_errors=True)\n",
    "\n",
    "\n",
    "train_dataset = MNISTStreamingDataset(\n",
    "    remote=f's3://{s3_bucket_name}/{bucket_prefix}/train',\n",
    "    local=os.path.join(data_dir, 'streaming_cache', 'train'),\n",
    "    shuffle=True,\n",
    "    decoders=decoders,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "\n",
    "eval_dataset = MNISTStreamingDataset(\n",
    "    remote=f's3://{s3_bucket_name}/{bucket_prefix}/eval',\n",
    "    local=os.path.join(data_dir, 'streaming_cache', 'eval'),\n",
    "    shuffle=True,\n",
    "    decoders=decoders,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset=train_dataset,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "eval_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset=eval_dataset,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "\n",
    "cloud_logger = ObjectStoreLogger(\n",
    "    object_store_cls=S3ObjectStore,\n",
    "    # Keyword arguments passed to the S3ObjectStore constructor\n",
    "    object_store_kwargs={\n",
    "        'bucket': s3_bucket_name,\n",
    "        'prefix': bucket_prefix,\n",
    "    },\n",
    "    use_procs=False,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    max_duration='2ep',\n",
    "    # Make training fast: Terminate each epoch after 5 batches\n",
    "    train_subset_num_batches=5,\n",
    "    optimizers=optimizer,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=eval_dataloader,\n",
    "    # Make evaluation fast: Evaluate only on five batches\n",
    "    eval_subset_num_batches=5,\n",
    "    loggers=[\n",
    "        cloud_logger,\n",
    "        TensorboardLogger(log_dir=tensorboard_log_dir, flush_interval=1),\n",
    "    ],\n",
    "    save_folder=checkpoint_dir,\n",
    "    save_interval=\"1ba\",\n",
    "    # Because we are uploading checkpoints to the cloud,\n",
    "    # delete them locally to save disk space!\n",
    "    save_num_checkpoints_to_keep=0,\n",
    "    run_name=run_name,\n",
    ")\n",
    "\n",
    "########\n",
    "# Train!\n",
    "########\n",
    "\n",
    "print(\"Training from the beginning\")\n",
    "trainer.fit()\n",
    "trainer.close()\n",
    "\n",
    "\n",
    "#################################\n",
    "# Train again (from a checkpoint)\n",
    "#################################\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    max_duration='4ep',\n",
    "    load_path=f'{run_name}/checkpoints/latest-rank0',\n",
    "    load_object_store=cloud_logger,\n",
    "    train_subset_num_batches=5,\n",
    "    optimizers=optimizer,\n",
    "    train_dataloader=train_dataloader,\n",
    "    eval_dataloader=eval_dataloader,\n",
    "    eval_subset_num_batches=5,\n",
    "    loggers=[\n",
    "        cloud_logger,\n",
    "        TensorboardLogger(log_dir=tensorboard_log_dir, flush_interval=1),\n",
    "    ],\n",
    "    save_folder=checkpoint_dir,\n",
    "    save_interval=\"1ba\",\n",
    "    save_num_checkpoints_to_keep=0,\n",
    "    run_name=run_name,\n",
    ")\n",
    "\n",
    "print(\"Training from a checkpoint\")\n",
    "trainer.fit()\n",
    "trainer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "###################\n",
    "# Visualize Results\n",
    "###################\n",
    "\n",
    "import tensorboard.notebook\n",
    "%load_ext tensorboard\n",
    "\n",
    "%tensorboard --logdir s3://{s3_bucket_name}/{bucket_prefix}/tensorboard_logs"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}