Add SkyPilot Benchmark doc (skypilot-org#1066)

* Add benchmark doc

* Address comments

* Address comments

Author: WoosukKwon

docs/source/index.rst

@@ -49,6 +49,7 @@ Key features:
    reference/job-queue
    reference/auto-stop
    examples/spot-jobs
+   reference/benchmark/index
 
 .. toctree::
    :maxdepth: 1

docs/source/reference/benchmark/callback.rst

@@ -0,0 +1,134 @@
+.. _benchmark-skycallback:
+
+SkyCallback
+===========
+
+SkyCallback is a simple Python library that works in conjunction with SkyPilot Benchmark.
+It enables SkyPilot to provide a more detailed benchmark report without the need to wait until the task finishes.
+
+What SkyCallback is for
+--------------------------------------------
+
+SkyCallback is designed for **machine learning tasks** which have a loop iterating many `steps`.
+SkyCallback measures the average time taken by each step, and extrapolates it to the total execution time of the task.
+
+Installing SkyCallback
+--------------------------------------------
+
+Unlike SkyPilot, SkyCallback must be installed and imported `in your program`.
+To install it, add the following line in the ``setup`` section of your task YAML.
+
+.. code-block:: yaml
+
+    setup:
+        # Activate conda or virtualenv if you use one
+        # Then, install SkyCallback
+        pip install "git+https://github.com/skypilot-org/skypilot.git#egg=sky-callback&subdirectory=sky/callbacks/"
+
+
+Using SkyCallback generic APIs
+--------------------------------------------
+
+The SkyCallback generic APIs are for **PyTorch, TensorFlow, and JAX** programs where training loops are exposed to the users.
+Below we provide the instructions for using the APIs.
+
+First, import the SkyCallback package and initialize it using ``init``.
+
+.. code-block:: python
+
+    import sky_callback
+    sky_callback.init()
+
+Next, mark the beginning and end of each step using one of the three equivalent methods.
+
+.. code-block:: python
+
+    # Method 1: wrap your iterable (e.g., dataloader) with `step_iterator`.
+    from sky_callback import step_iterator
+    for batch in step_iterator(train_dataloader):
+        ...
+
+    # Method 2: wrap your loop body with the `step` context manager.
+    for batch in train_dataloader:
+        with sky_callback.step():
+            ...
+
+    # Method 3: call `step_begin` and `step_end` directly.
+    for batch in train_dataloader:
+        sky_callback.step_begin()
+        ...
+        sky_callback.step_end()
+
+That's it.
+Now you can launch your task and get a detailed benchmark report using SkyPilot Benchmark CLI.
+`Here <https://github.com/skypilot-org/skypilot/blob/master/examples/benchmark/timm/callback.patch>`_ we provide an example of applying SkyCallback to Pytorch ImageNet training.
+
+.. note::
+    
+    Optionally in ``sky_callback.init``, you can specify the total number of steps that the task will iterate through.
+    This information is needed to estimate the total execution time/cost of your task.
+
+    .. code-block:: python
+    
+        sky_callback.init(
+            total_steps=num_epochs * len(train_dataloader), # Optional
+        )
+
+.. note::
+    In distributed training, ``global_rank`` should be additionally passed to ``sky_callback.init`` as follows:
+
+    .. code-block:: python
+
+        # PyTorch DDP users
+        global_rank = torch.distributed.get_rank()
+
+        # Horovod users
+        global_rank = hvd.rank()
+
+        sky_callback.init(
+            global_rank=global_rank,
+            total_steps=num_epochs * len(train_dataloader), # Optional
+        )
+
+Integrations with ML frameworks
+----------------------------------------------------------
+
+Using SkyCallback is even easier for **Keras, PytorchLightning, and HuggingFace Transformers** programs where trainer APIs are used.
+SkyCallback natively supports these frameworks with simple interface.
+
+* Keras example
+
+.. code-block:: python
+
+    from sky_callback import SkyKerasCallback
+
+    # Add the callback to your Keras model.
+    model.fit(..., callbacks=[SkyKerasCallback()])
+
+`Here <https://github.com/skypilot-org/skypilot/blob/master/examples/benchmark/keras_asr/callback.patch>`_ you can find an example of applying SkyCallback to Keras ASR model training.
+
+* PytorchLightning example
+
+.. code-block:: python
+
+    from sky_callback import SkyLightningCallback
+
+    # Add the callback to your trainer.
+    trainer = pl.Trainer(..., callbacks=[SkyLightningCallback()])
+
+`Here <https://github.com/skypilot-org/skypilot/blob/master/examples/benchmark/lightning_gan/callback.patch>`_ you can find an example of applying SkyCallback to PyTorchLightning GAN model training.
+
+* HuggingFace Transformers example
+
+.. code-block:: python
+
+    from sky_callback import SkyTransformersCallback
+
+    # Add the callback to your trainer.
+    trainer = transformers.Trainer(..., callbacks=[SkyTransformersCallback()])
+
+`Here <https://github.com/skypilot-org/skypilot/blob/master/examples/benchmark/transformers_qa/callback.patch>`_ you can find an example of applying SkyCallback to HuggingFace BERT fine-tuning.
+
+.. note::
+    When using the framework-integrated callbacks, do not call ``sky_callback.init`` for initialization.
+    The callbacks will do it for you.

docs/source/reference/benchmark/cli.rst

@@ -0,0 +1,78 @@
+.. _benchmark-cli:
+
+CLI
+=============
+
+Workflow
+--------------------------------
+
+You can use SkyPilot Benchmark by simply replacing your ``sky launch`` command with ``sky bench launch``:
+
+.. code-block:: bash
+
+    # Launch mytask on a V100 VM and a T4 VM
+    $ sky bench launch mytask.yaml --gpus V100,T4 --benchmark mybench
+
+The second command will launch ``mytask.yaml`` on a V100 VM and a T4 VM simultaneously, with a benchmark name ``mybench``.
+After the task finishes, you can check the benchmark results using ``sky bench show``:
+
+.. code-block:: bash
+
+    # Show the benchmark report on `mybench`
+    $ sky bench show mybench
+
+    CLUSTER              RESOURCES                          STATUS    DURATION  SPENT($)  STEPS  SEC/STEP  $/STEP  EST(hr)  EST($)  
+    sky-bench-mybench-0  1x GCP(n1-highmem-8, {'V100': 1})  FINISHED  12m 51s   0.6317    -       -         -       -        -       
+    sky-bench-mybench-1  1x AWS(g4dn.xlarge, {'T4': 1})     FINISHED  16m 19s   0.1430    -       -         -       -        -     
+
+In the report, SkyPilot shows the duration and cost of ``mybench`` on each VM.
+The VMs can be terminated by either ``sky bench down`` or ``sky down``:
+
+.. code-block:: bash
+
+    # Terminate all the clusters used for `mybench`
+    $ sky bench down mybench
+
+    # Terminate all the clusters used for `mybench` except `sky-bench-mybench-0`
+    $ sky bench down mybench --exclude sky-bench-mybench-0
+
+    # Terminate individual clusters as usual
+    $ sky down sky-bench-mybench-0
+
+.. note::
+
+    Each cluster launched by ``sky bench launch`` will automatically **stop** itself 5 minutes after the task is finished.
+    However, you don't have to restart those clusters.
+    Regardless of the status of the clusters, ``sky bench show`` will provide the benchmark results.
+
+.. note::
+
+    SkyPilot Benchmark does not consider the time/cost of provisioning and setup.
+    The columns (such as ``DURATION`` and ``SPENT($)``) in the report indicate the time/cost spent in executing the ``run`` section of your task YAML.
+
+.. note::
+
+    Here, the columns other than ``DURATION`` and ``SPENT($)`` are empty.
+    To get a complete benchmark report, please refer to :ref:`SkyCallback`.
+
+
+Managing benchmark reports
+---------------------------
+
+``sky bench ls`` shows the list of the benchmark reports you have:
+
+.. code-block:: bash
+
+    # List all the benchmark reports
+    $ sky bench ls
+
+    BENCHMARK  TASK         LAUNCHED             CANDIDATE 1                    CANDIDATE 2            CANDIDATE 3            CANDIDATE 4               
+    bert       bert_qa      2022-08-10 10:07:27  1x Standard_NC6_Promo (K80:1)  1x g4dn.xlarge (T4:1)  1x g5.xlarge (A10G:1)  1x n1-highmem-8 (V100:1)  
+    mybench    mytask       2022-08-10 11:24:27  1x n1-highmem-8 (V100:1)       1x g4dn.xlarge (T4:1)
+
+To delete a benchmark report, use ``sky bench delete``:
+
+.. code-block:: bash
+
+    # Delete the benchmark report on `mybench`
+    $ sky bench delete mybench

docs/source/reference/benchmark/config.rst

@@ -0,0 +1,48 @@
+.. _benchmark-yaml:
+
+YAML Configuration
+===================
+
+The resources to benchmark can be configured in the SkyPilot YAML interface.
+Below we provide an example:
+
+.. code-block:: yaml
+
+    # Only shows `resources` as other fields do not change.
+    resources:
+        cloud: gcp  # Works as a default value for `cloud`.
+
+        # Added only for SkyPilot Benchmark.
+        candidates:
+        - {accelerators: A100}
+        - {accelerators: V100, instance_type: n1-highmem-16}
+        - {accelerators: T4, cloud: aws}  # Overrides `cloud` to `aws`.
+
+For SkyPilot Benchmark, ``candidates`` is newly added under the ``resources`` field.
+``candidates`` is the list of dictionaries that configure the resources to benchmark.
+Any subfield of ``resources`` (``accelerators``, ``instance_type``, etc.) can be re-defined in the dictionaries.
+Subfields defined outside ``candidates`` (e.g. ``cloud`` in this example) are used as default values and are overriden by those defined in the dictionaries.
+Thus, the above example can be interpreted as follows:
+
+.. code-block:: yaml
+
+    # Configuration of the first candidate.
+    resources:
+        cloud: gcp
+        accelerators: A100
+
+    # Configuration of the second candidate.
+    resources:
+        cloud: gcp
+        accelerators: V100
+        instance_type: n1-highmem-16
+
+    # Configuration of the third candidate.
+    resources:
+        cloud: aws
+        accelerators: T4
+
+.. note::
+
+    Currently, SkyPilot Benchmark does not support on-prem jobs and managed spot jobs.
+    While you can set ``use_spot: True`` to benchmark spot VMs, automatic recovery will not be provided when preemption occurs.

docs/source/reference/benchmark/index.rst

@@ -0,0 +1,43 @@
+.. _benchmark-overview:
+
+Benchmark
+================================================
+
+SkyPilot allows **easy measurement of performance and cost of different kinds of cloud resources** through the benchmark feature.
+With minimal effort, you can find the right cloud resource for your task that fits your performance goals and budget constraints.
+
+For example, say you want to fine-tune a BERT model and you do not know which GPU type is the best for you.
+With SkyPilot Benchmark, you can quickly run your task on different types of VMs and get a benchmark report like the following:
+
+.. code-block:: bash
+
+    Legend:
+    - STEPS: Number of steps taken.
+    - SEC/STEP, $/STEP: Average time (cost) per step.
+    - EST(hr), EST($): Estimated total time (cost) to complete the benchmark.
+
+    CLUSTER            RESOURCES                                STATUS      DURATION  SPENT($)  STEPS   SEC/STEP  $/STEP    EST(hr)  EST($)  
+    sky-bench-bert-0  1x Azure(Standard_NC6_Promo, {'K80': 1})  TERMINATED  12m 48s   0.0384    1415    1.1548    0.000058  10.60    1.91    
+    sky-bench-bert-1  1x AWS(g4dn.xlarge, {'T4': 1})            TERMINATED  14m 2s    0.1230    2387    0.6429    0.000094  5.92     3.11    
+    sky-bench-bert-2  1x AWS(g5.xlarge, {'A10G': 1})            TERMINATED  13m 57s   0.2339    7423    0.1859    0.000052  1.75     1.76      
+    sky-bench-bert-3  1x GCP(n1-highmem-8, {'V100': 1})         TERMINATED  13m 45s   0.6768    7306    0.2005    0.000165  1.87     5.51
+
+The report shows the benchmarking results of 4 VMs each with a different GPU type.
+Based on the report, you can pick the VM with either the lowest cost (``EST($)``) or the fastest execution time (``EST(hr)``), or find a sweet spot between them.
+In this example, AWS g5.xlarge (NVIDIA A10G GPU) turns out to be the best choice in terms of both cost and time.
+
+Using SkyPilot Benchmark
+------------------------
+
+A part of the SkyPilot Benchmark report relies on the :ref:`SkyCallback` library instrumented in the training code to report step completion.
+Depending on the level of detail required by you in the benchmark report, SkyPilot Benchmark can be used in two modes:
+
+1. Without SkyCallback - You can get a basic benchmark report using SkyPilot Benchmark :ref:`benchmark-cli`. **This requires zero changes in your code**.
+2. With SkyCallback - You can get a more detailed benchmark report **by a few lines of code changes**. Please refer to :ref:`SkyCallback`.
+
+Table of Contents
+-------------------
+.. toctree::
+   cli
+   config
+   callback

examples/benchmark/keras_asr.yaml

@@ -12,8 +12,7 @@ setup: |
   conda activate keras
 
   # Install SkyCallback
-  git clone [email protected]:skypilot-org/skypilot.git
-  pip install skypilot/sky/callbacks/
+  pip install "git+https://github.com/skypilot-org/skypilot.git#egg=sky-callback&subdirectory=sky/callbacks/"
 
   # User setup
   pip install numpy pandas tensorflow

examples/benchmark/lightning_gan.yaml

@@ -12,8 +12,7 @@ setup: |
   conda activate pl
 
   # Install SkyCallback
-  git clone [email protected]:skypilot-org/skypilot.git
-  pip install skypilot/sky/callbacks/
+  pip install "git+https://github.com/skypilot-org/skypilot.git#egg=sky-callback&subdirectory=sky/callbacks/"
 
   # User setup
   pip install "torchvision" "pytorch-lightning>=1.4" "torch>=1.6, <1.9"

examples/benchmark/timm.yaml

@@ -12,8 +12,7 @@ setup: |
   conda activate timm
 
   # Install SkyCallback
-  git clone [email protected]:skypilot-org/skypilot.git
-  pip install skypilot/sky/callbacks/
+  pip install "git+https://github.com/skypilot-org/skypilot.git#egg=sky-callback&subdirectory=sky/callbacks/"
 
   # User setup
   git clone https://github.com/rwightman/pytorch-image-models.git timm

examples/benchmark/transformers_qa.yaml

@@ -12,8 +12,7 @@ setup: |
   conda activate hf
 
   # Install SkyCallback
-  git clone [email protected]:skypilot-org/skypilot.git
-  pip install skypilot/sky/callbacks/
+  pip install "git+https://github.com/skypilot-org/skypilot.git#egg=sky-callback&subdirectory=sky/callbacks/"
 
   # User setup
   pip install transformers