Repository: allenai/tango (language: python)

Issue:
Create a manually triggered Actions workflow to run integration tests
Requires #183 first, since these would also have to run on Beaker.

We'd also need to:
- [x] Set up a 4xA100 Beaker cluster: https://beaker.org/cl/ai2/tango-integration-tests/details
- [ ] #188 


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 5ad8b7a0b..766a951f8 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -219,10 +219,12 @@ jobs:
   gpu_tests:
     name: GPU Tests
     runs-on: [ubuntu-latest]
+    timeout-minutes: 20
     env:
       BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
       BEAKER_WORKSPACE: ai2/tango-testing
-      IMAGE_NAME: tango-testing
+      BEAKER_CLUSTER: ai2/tango-gpu-tests
+      IMAGE_NAME: petew/tango-testing
     steps:
       - uses: actions/checkout@v2
 
@@ -263,7 +265,7 @@ jobs:
           tasks:
             - name: tests
               image:
-                beaker: petew/tango-testing
+                beaker: ${{ env.IMAGE_NAME }}
               command: ["/entrypoint.sh", "pytest", "-v", "-m", "gpu", "tests/"]
               envVars:
                 - name: COMMIT_SHA
@@ -273,14 +275,15 @@ jobs:
               resources:
                 gpuCount: 2
               context:
-                cluster: ai2/tango-gpu-tests
+                cluster: ${{ env.BEAKER_CLUSTER }}
                 priority: normal
           EOL
           cat beaker_config.yml
 
       - name: Submit beaker job
         run: |
-          EXPERIMENT=$(beaker experiment create beaker_config.yml --workspace $BEAKER_WORKSPACE --name gpu-tests-$COMMIT_SHA | awk '{print $2}')
+          TIMESTAMP=$(date +%H%M%S)
+          EXPERIMENT=$(beaker experiment create beaker_config.yml --workspace $BEAKER_WORKSPACE --name "gpu-tests-${COMMIT_SHA}-${TIMESTAMP}" | awk '{print $2}')
           if [ -z "$EXPERIMENT" ]; then
             exit 1
           else
@@ -297,13 +300,31 @@ jobs:
       - name: Get logs
         if: always()
         run: |
+          # EXPERIMENT could be empty if the submission step failed.
+          # We'll exit right away if that's the case.
           if [ -z "$EXPERIMENT" ]; then
             echo "No logs to show"
             exit 0
           fi
+
+          # Download logs from beaker.
           beaker experiment results $EXPERIMENT --prefix out.log --output results
+
+          # If the experiment failed during startup, there might not be any logs.
           if [ -f results/tests/out.log ]; then
+            echo ""
+            echo ">>> Logs:"
+            echo ""
             cat results/tests/out.log
+          else
+            echo "No logs to show"
+          fi
+
+      - name: Stop job
+        if: cancelled()
+        run: |
+          if [ ! -z "$EXPERIMENT" ]; then
+            beaker experiment stop $EXPERIMENT
           fi
 
   release:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 119dab985..5ad5e2b70 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed a small bug `LocalWorkspace` would fail to capture the conda environment in our Docker image.
+- Fixed activation of `FILE_FRIENDLY_LOGGING` when set from the corresponding environment variable.
 - Use relative paths within the `work_dir` for symbolic links to the latest and the best checkpoints in `TorchTrainStep`.
 
 
diff --git a/scripts/entrypoint.sh b/scripts/entrypoint.sh
index b20ab0bcb..51546d104 100644
--- a/scripts/entrypoint.sh
+++ b/scripts/entrypoint.sh
@@ -1,7 +1,8 @@
-#!/bin/sh
+#!/bin/bash
 
 # Exit script if any commands fail.
 set -e
+set -o pipefail
 
 # Check that the environment variable has been set correctly
 if [ -z "$COMMIT_SHA" ]; then
@@ -19,4 +20,4 @@ git checkout "$COMMIT_SHA"
 mkdir -p /results
 
 # Execute the arguments to this script as commands themselves, piping output into a log file.
-exec "$@" | tee /results/out.log
+exec "$@" 2>&1 | tee /results/out.log
diff --git a/tango/__main__.py b/tango/__main__.py
index 2c137955c..11a042d64 100644
--- a/tango/__main__.py
+++ b/tango/__main__.py
@@ -110,12 +110,10 @@ class TangoGlobalSettings(FromParams):
 
     """
 
-    file_friendly_logging: bool = False
+    file_friendly_logging: Optional[bool] = None
     """
     If this flag is set to ``True``, we add newlines to tqdm output, even on an interactive terminal, and we slow
     down tqdm's output to only once every 10 seconds.
-
-    By default, it is set to ``False``.
     """
 
     multiprocessing_start_method: str = "spawn"
@@ -216,7 +214,7 @@ def main(
     if log_level is not None:
         config.log_level = log_level
 
-    if file_friendly_logging is not None:
+    if file_friendly_logging:
         config.file_friendly_logging = file_friendly_logging
 
     initialize_logging(
diff --git a/tango/local_workspace.py b/tango/local_workspace.py
index 302fbd7e8..14990dc30 100644
--- a/tango/local_workspace.py
+++ b/tango/local_workspace.py
@@ -187,8 +187,24 @@ def _save_conda(self, run_dir: Path):
             import subprocess
 
             try:
-                with (run_dir / "conda-environment.yaml").open("w") as f:
-                    subprocess.call(["conda", "env", "export"], stdout=f)
+                result = subprocess.run(["conda", "env", "export"], capture_output=True)
+                if (
+                    result.returncode != 0
+                    and result.stderr is not None
+                    and "Unable to determine environment" in result.stderr.decode()
+                ):
+                    result = subprocess.run(
+                        ["conda", "env", "export", "-n", "base"], capture_output=True
+                    )
+
+                if result.returncode != 0:
+                    if result.stderr is not None:
+                        logger.exception("Error saving conda packages: %s", result.stderr.decode())
+                    else:
+                        result.check_returncode()
+                elif result.stdout is not None:
+                    with (run_dir / "conda-environment.yaml").open("w") as f:
+                        f.write(result.stdout.decode())
             except Exception as exc:
                 logger.exception("Error saving conda packages: %s", exc)
 


Test patch:
diff --git a/.github/workflows/docker_testing.yml b/.github/workflows/docker_testing.yml
index 26d8af242..7bab0bf67 100644
--- a/.github/workflows/docker_testing.yml
+++ b/.github/workflows/docker_testing.yml
@@ -13,12 +13,14 @@ on:
     paths:
       - 'Dockerfile.test'
       - '.dockerignore'
+      - 'scripts/entrypoint.sh'
   push:
     branches:
       - main
     paths:
       - 'Dockerfile.test'
       - '.dockerignore'
+      - 'scripts/entrypoint.sh'
 
 jobs:
   build:
diff --git a/.github/workflows/integration_tests.yml b/.github/workflows/integration_tests.yml
new file mode 100644
index 000000000..47e0253d5
--- /dev/null
+++ b/.github/workflows/integration_tests.yml
@@ -0,0 +1,154 @@
+name: Integration tests
+
+on:
+  workflow_dispatch:
+    inputs:
+      test:
+        description: the integration test to run
+        default: fairscale_benchmarks
+        required: true
+        type: choice
+        options:
+          - fairscale_benchmarks
+      cluster:
+        description: the beaker cluster to run the test on
+        default: ai2/tango-integration-tests
+        required: true
+        type: choice
+        options:
+          - ai2/tango-integration-tests
+          - ai2/allennlp-cirrascale
+  # Uncomment this trigger to test changes on a pull request.
+  # You also have to uncomment the lines below that mention 'for pull request checks'
+  # pull_request:
+  #   branches:
+  #     - '*'
+
+jobs:
+  run_test:
+    name: ${{ github.event.inputs.test }}
+    # name: fairscale_benchmarks  # for pull request checks
+    runs-on: [ubuntu-latest]
+    timeout-minutes: 60
+    env:
+      TEST_NAME: ${{ github.event.inputs.test }}
+      # TEST_NAME: fairscale_benchmarks  # for pull request checks
+      BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
+      BEAKER_WORKSPACE: ai2/tango-integration-tests
+      BEAKER_CLUSTER: ${{ github.event.inputs.cluster }}
+      # BEAKER_CLUSTER: ai2/allennlp-cirrascale  # for pull request checks
+      IMAGE_NAME: petew/tango-testing
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Validate inputs
+        run: |
+          # The 'test' input should be a directory in `integration_tests/`
+          test -d "integration_tests/${TEST_NAME}"
+
+      - name: Determine current commit SHA (pull request)
+        if: github.event_name == 'pull_request'
+        run: |
+          echo "COMMIT_SHA=${{ github.event.pull_request.head.sha }}" >> $GITHUB_ENV
+
+      - name: Determine current commit SHA (push)
+        if: github.event_name != 'pull_request'
+        run: |
+          echo "COMMIT_SHA=$GITHUB_SHA" >> $GITHUB_ENV
+
+      - name: Install beaker client
+        shell: bash
+        run: |
+          mkdir -p "$HOME/bin"
+
+          # Download and install from latest GitHub release.
+          curl -s https://api.github.com/repos/allenai/beaker/releases/latest \
+            | grep 'browser_download_url.*linux' \
+            | cut -d '"' -f 4 \
+            | wget -qi - \
+          && tar -xvzf beaker_linux.tar.gz -C "$HOME/bin"
+
+          # Add to path.
+          echo "$HOME/bin" >> "$GITHUB_PATH"
+
+      - name: Verify beaker install
+        run: |
+          beaker account whoami
+
+      - name: Create beaker experiment config
+        run: |
+          cat >beaker_config.yml << EOL
+          version: v2-alpha
+          description: ${{ env.TEST_NAME }}
+          tasks:
+            - name: test
+              image:
+                beaker: ${{ env.IMAGE_NAME }}
+              command: ["/entrypoint.sh", "integration_tests/${{ env.TEST_NAME }}/run.sh"]
+              envVars:
+                - name: COMMIT_SHA
+                  value: $COMMIT_SHA
+                - name: WANDB_API_KEY
+                  secret: WANDB_API_KEY
+                - name: FILE_FRIENDLY_LOGGING
+                  value: "true"
+                - name: TOKENIZERS_PARALLELISM  # set this to avoid warnings
+                  value: "true"
+                - name: PYTHONUNBUFFERED
+                  value: "true"
+              result:
+                path: '/results'
+              resources:
+                gpuCount: 4
+              context:
+                cluster: ${{ env.BEAKER_CLUSTER }}
+                priority: normal
+          EOL
+          cat beaker_config.yml
+
+      - name: Submit beaker job
+        run: |
+          TIMESTAMP=$(date +%H%M%S)
+          EXPERIMENT=$(beaker experiment create beaker_config.yml --workspace $BEAKER_WORKSPACE --name "${TEST_NAME}-${{ github.run_number }}-${TIMESTAMP}" | awk '{print $2}')
+          if [ -z "$EXPERIMENT" ]; then
+            exit 1
+          else
+            echo "EXPERIMENT=$EXPERIMENT" >> $GITHUB_ENV
+            echo "Experiment $EXPERIMENT submitted. See progress at https://beaker.org/ex/$EXPERIMENT"
+          fi
+
+      - name: Wait for job to finish
+        run: |
+          beaker experiment await $EXPERIMENT test finalized --timeout 60m
+          # Check the job's exit code.
+          test $(beaker experiment get $EXPERIMENT --format=json | jq '.[0].jobs[0].status.exitCode') -eq 0
+
+      - name: Get logs
+        if: always()
+        run: |
+          # EXPERIMENT could be empty if the submission step failed.
+          # We'll exit right away if that's the case.
+          if [ -z "$EXPERIMENT" ]; then
+            echo "No logs to show"
+            exit 0
+          fi
+
+          # Download logs from beaker.
+          beaker experiment results $EXPERIMENT --prefix out.log --output results
+
+          # If the experiment failed during startup, there might not be any logs.
+          if [ -f results/test/out.log ]; then
+            echo ""
+            echo ">>> Logs:"
+            echo ""
+            cat results/test/out.log
+          else
+            echo "No logs to show"
+          fi
+
+      - name: Stop job
+        if: cancelled()
+        run: |
+          if [ ! -z "$EXPERIMENT" ]; then
+            beaker experiment stop $EXPERIMENT
+          fi
diff --git a/integration_tests/README.md b/integration_tests/README.md
index f4ddb7ece..d6f74e19c 100644
--- a/integration_tests/README.md
+++ b/integration_tests/README.md
@@ -4,4 +4,5 @@ These are a collection of longer running end-to-end tests of various parts of th
 At the moment we are not running these in CI since they are computationally expensive and/or require special hardware,
 so they need to be run manually.
 
-The README in each test's folder should explain when and how that particular test should be ran.
+The README in each test's folder should explain where and when that particular test should be run.
+Each test should have a `run.sh` file in its folder that will run the relevant tango command.
diff --git a/integration_tests/fairscale_benchmarks/README.md b/integration_tests/fairscale_benchmarks/README.md
index 2ce3bd38d..d59d7caa8 100644
--- a/integration_tests/fairscale_benchmarks/README.md
+++ b/integration_tests/fairscale_benchmarks/README.md
@@ -8,8 +8,11 @@ This integration test is for checking the performance of the `FairScaleTrainingE
 
 **How to run it:** From the root directory of this repository, run:
 ```
-tango run integration_tests/fairscale_benchmarks/config.jsonnet -i examples/train_lm/tokenize_step.py
+integration_tests/fairscale_benchmarks/run.sh
 ```
 
+By default, not all configurations are run. If you want to run change which configurations are run, open `config.jsonnet`
+are search for "enabled". Then toggle this `enabled` field to `true` or `false` for each configuration.
+
 **What to look for:** The training jobs shouldn't fail, for one. After `tango run` completes, check the corresponding Weights & Biases
 dashboard and inspect the results. Compare the various "fsdp" training runs with the baseline to ensure you see memory savings.
diff --git a/integration_tests/fairscale_benchmarks/config.jsonnet b/integration_tests/fairscale_benchmarks/config.jsonnet
index 41c786b84..109010fe8 100644
--- a/integration_tests/fairscale_benchmarks/config.jsonnet
+++ b/integration_tests/fairscale_benchmarks/config.jsonnet
@@ -51,6 +51,17 @@ local single_device_dataloader = {
 local TrainStep(options) =
     local training_engine = {
         type: if options.fsdp_config != null then "fairscale" else "torch",
+        optimizer: {
+            type: "torch::AdamW",
+            lr: learning_rate,
+            betas: [0.9, 0.95],
+            eps: 1e-6,
+        },
+        lr_scheduler: {
+            type: "transformers::linear",
+            num_warmup_steps: warmup_steps,
+            num_training_steps: training_steps,
+        },
         amp: options.amp,
         [if options.fsdp_config != null then "fsdp_config" else null]: options.fsdp_config,
     };
@@ -71,17 +82,6 @@ local TrainStep(options) =
         dataset_dict: { type: "ref", ref: "tokenized_data" },
         train_dataloader: distributed_dataloader,
         validation_split: "validation",
-        optimizer: {
-            type: "torch::AdamW",
-            lr: learning_rate,
-            betas: [0.9, 0.95],
-            eps: 1e-6,
-        },
-        lr_scheduler: {
-            type: "transformers::linear",
-            num_warmup_steps: warmup_steps,
-            num_training_steps: training_steps,
-        },
         grad_accum: grad_accum,
         train_steps: training_steps,
         validate_every: validate_every,
@@ -117,33 +117,39 @@ local TrainStep(options) =
     } + {
         ["trained_model_" + options.name]: TrainStep(options)
         for options in [
-            # With 6B model, baseline will fail with CUDA OOM
+            # NOTE: With 6B model, baseline and many others will fail with CUDA OOM.
+            # FSDP and activation checkpointing will be required for a 6B model.
             {
                 name: "baseline",
+                enabled: false,
                 amp: false,
                 fsdp_config: null,
                 activation_checkpointing: false,
             },
             {
                 name: "amp",
+                enabled: false,
                 amp: true,
                 fsdp_config: null,
                 activation_checkpointing: false,
             },
             {
                 name: "checkpointing",
+                enabled: false,
                 amp: false,
                 fsdp_config: null,
                 activation_checkpointing: true,
             },
             {
                 name: "amp_and_checkpointing",
+                enabled: false,
                 amp: true,
                 fsdp_config: null,
                 activation_checkpointing: true,
             },
             {
                 name: "fsdp",
+                enabled: false,
                 amp: false,
                 activation_checkpointing: false,
                 fsdp_config: {
@@ -155,6 +161,7 @@ local TrainStep(options) =
             },
             {
                 name: "fsdp_no_reshard",
+                enabled: false,
                 amp: false,
                 activation_checkpointing: false,
                 fsdp_config: {
@@ -166,6 +173,7 @@ local TrainStep(options) =
             },
             {
                 name: "amp_and_fsdp",
+                enabled: false,
                 amp: true,
                 activation_checkpointing: false,
                 fsdp_config: {
@@ -177,6 +185,7 @@ local TrainStep(options) =
             },
             {
                 name: "amp_and_fsdp_no_reshard",
+                enabled: false,
                 amp: true,
                 activation_checkpointing: false,
                 fsdp_config: {
@@ -188,6 +197,7 @@ local TrainStep(options) =
             },
             {
                 name: "amp_and_fsdp_mp",
+                enabled: false,
                 amp: true,
                 activation_checkpointing: false,
                 fsdp_config: {
@@ -199,6 +209,7 @@ local TrainStep(options) =
             },
             {
                 name: "amp_and_fsdp_mp_no_reshard",
+                enabled: false,
                 amp: true,
                 activation_checkpointing: false,
                 fsdp_config: {
@@ -210,6 +221,7 @@ local TrainStep(options) =
             },
             {
                 name: "checkpointing_and_fsdp",
+                enabled: false,
                 amp: false,
                 activation_checkpointing: true,
                 fsdp_config: {
@@ -221,6 +233,7 @@ local TrainStep(options) =
             },
             {
                 name: "amp_and_checkpointing_and_fsdp",
+                enabled: false,
                 amp: true,
                 activation_checkpointing: true,
                 fsdp_config: {
@@ -232,6 +245,7 @@ local TrainStep(options) =
             },
             {
                 name: "amp_and_checkpointing_and_fsdp_mp",
+                enabled: true,
                 amp: true,
                 activation_checkpointing: true,
                 fsdp_config: {
@@ -243,6 +257,7 @@ local TrainStep(options) =
             },
             {
                 name: "checkpointing_and_fsdp_mp",
+                enabled: false,
                 amp: false,
                 activation_checkpointing: true,
                 fsdp_config: {
@@ -252,20 +267,21 @@ local TrainStep(options) =
                     mixed_precision: true,
                 },
             },
-            # Currently does not work. Tracking https://github.com/facebookresearch/fairscale/issues/918
-            # {
-            #     name: "amp_and_checkpointing_and_fsdp_mp_with_partial_offloading",
-            #     amp: true,
-            #     activation_checkpointing: true,
-            #     fsdp_config: {
-            #         reshard_after_forward: true,
-            #         move_params_to_cpu: true,
-            #         move_grads_to_cpu: false,
-            #         mixed_precision: true,
-            #     },
-            # },
+            {  # This configuration currently does not work. Tracking https://github.com/facebookresearch/fairscale/issues/918
+                name: "amp_and_checkpointing_and_fsdp_mp_with_partial_offloading",
+                enabled: false,
+                amp: true,
+                activation_checkpointing: true,
+                fsdp_config: {
+                    reshard_after_forward: true,
+                    move_params_to_cpu: true,
+                    move_grads_to_cpu: false,
+                    mixed_precision: true,
+                },
+            },
             {
                 name: "amp_and_checkpointing_and_fsdp_mp_with_full_offloading",
+                enabled: false,
                 amp: true,
                 activation_checkpointing: true,
                 fsdp_config: {
@@ -275,6 +291,6 @@ local TrainStep(options) =
                     mixed_precision: true,
                 },
             },
-        ]
+        ] if options.enabled
     }
 }
diff --git a/integration_tests/fairscale_benchmarks/run.sh b/integration_tests/fairscale_benchmarks/run.sh
new file mode 100755
index 000000000..0b9ae1906
--- /dev/null
+++ b/integration_tests/fairscale_benchmarks/run.sh
@@ -0,0 +1,3 @@
+#!/bin/sh
+
+tango run integration_tests/fairscale_benchmarks/config.jsonnet -i examples/train_lm/tokenize_step.py