:orphan:

AWS Sagemaker Pytorch
=====================

This plugin shows an example of using Sagemaker custom training, with Pytorch distributed training.


Installation
------------

To use the Flytekit AWS Sagemaker plugin, simply run the following:

.. prompt:: bash

   pip install flytekitplugins-awssagemaker


Creating a Dockerfile for Sagemaker Custom Training [Required]
--------------------------------------------------------------

The dockerfile for Sagemaker custom training is similar to any regular dockerfile, except for the difference in using the Nvidia cuda base to use GPU's.

.. note::

    If using CPU for training, then the special dockerfile is NOT REQUIRED. If GPU or TPUs are required, the dockerfile differs only in the driver setup. The following dockerfile is enabled for GPU accelerated training using CUDA.
    The checked in version of docker file uses python:3.8-slim-buster for faster CI, but you can use the Dockerfile pasted below which uses cuda base.
    Additionally, the requirements.in uses the cpu version of pytorch. Remove the + cpu for torch and torchvision in requirements.in and make all requirements as shown below:

.. prompt:: bash

   make -C integrations/aws/sagemaker_pytorch requirements


.. code-block:: docker
    :emphasize-lines: 23-24
    :linenos:

    FROM pytorch/pytorch:1.7.0-cuda11.0-cudnn8-devel
    LABEL org.opencontainers.image.source https://github.com/flyteorg/flytesnacks

    WORKDIR /root
    ENV LANG C.UTF-8
    ENV LC_ALL C.UTF-8
    ENV PYTHONPATH /root

    # Install the AWS cli separately to prevent issues with boto being written over
    RUN pip install awscli

    ENV VENV /opt/venv
    # Virtual environment
    RUN python3 -m venv ${VENV}
    ENV PATH="${VENV}/bin:$PATH"

    # Install Python dependencies
    COPY sagemaker_pytorch/requirements.txt /root/.
    RUN pip install -r /root/requirements.txt

    # Setup Sagemaker entrypoints
    ENV SAGEMAKER_PROGRAM /opt/venv/bin/flytekit_sagemaker_runner.py

    # Copy the makefile targets to expose on the container. This makes it easier to register.
    COPY in_container.mk /root/Makefile
    COPY sagemaker_pytorch/sandbox.config /root

    # Copy the actual code
    COPY sagemaker_pytorch/ /root/sagemaker_pytorch

    # This tag is supplied by the build script and will be used to determine the version
    # when registering tasks, workflows, and launch plans
    ARG tag
    ENV FLYTE_INTERNAL_IMAGE $tag



.. raw:: html

    <div class="sphx-glr-thumbnails">


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example is adapted from `this sagemaker example &lt;https://github.com/aws/amazon-sagemaker-e...">

.. only:: html

  .. image:: /auto/integrations/aws/sagemaker_pytorch/images/thumb/sphx_glr_sagemaker_pytorch_distributed_training_thumb.png
    :alt: Distributed Pytorch on Sagemaker

  :ref:`sphx_glr_auto_integrations_aws_sagemaker_pytorch_sagemaker_pytorch_distributed_training.py`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Distributed Pytorch on Sagemaker</div>
    </div>


.. raw:: html

    </div>


.. toctree::
   :hidden:

   /auto/integrations/aws/sagemaker_pytorch/sagemaker_pytorch_distributed_training


.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-gallery

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download all examples in Python source code: sagemaker_pytorch_python.zip </auto/integrations/aws/sagemaker_pytorch/sagemaker_pytorch_python.zip>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download all examples in Jupyter notebooks: sagemaker_pytorch_jupyter.zip </auto/integrations/aws/sagemaker_pytorch/sagemaker_pytorch_jupyter.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_