schedulers/aws_batch: add a scheduler to launch jobs directly on aws_batch

Author: d4l3k

.github/workflows/aws-batch-integration-tests.yaml

@@ -0,0 +1,45 @@
+name: AWS Batch Integration Tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  awsbatch:
+    runs-on: ubuntu-18.04
+    permissions:
+      id-token: write
+      contents: read
+    steps:
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.9
+          architecture: x64
+      - name: Checkout TorchX
+        uses: actions/checkout@v2
+      - name: Configure AWS
+        env:
+          AWS_ROLE_ARN: ${{ secrets.AWS_ROLE_ARN }}
+        run: |
+          if [ -n "$AWS_ROLE_ARN" ]; then
+            export AWS_WEB_IDENTITY_TOKEN_FILE=/tmp/awscreds
+            export AWS_DEFAULT_REGION=us-west-2
+
+            echo AWS_WEB_IDENTITY_TOKEN_FILE=$AWS_WEB_IDENTITY_TOKEN_FILE >> $GITHUB_ENV
+            echo AWS_ROLE_ARN=$AWS_ROLE_ARN >> $GITHUB_ENV
+            echo AWS_DEFAULT_REGION=$AWS_DEFAULT_REGION >> $GITHUB_ENV
+
+            curl -H "Authorization: bearer $ACTIONS_ID_TOKEN_REQUEST_TOKEN" "$ACTIONS_ID_TOKEN_REQUEST_URL" | jq -r '.value' > $AWS_WEB_IDENTITY_TOKEN_FILE
+          fi
+      - name: Install dependencies
+        run: |
+          set -eux
+          pip install -e .[dev]
+      - name: Run AWS Batch Integration Tests
+        run: |
+          set -ex
+
+          scripts/awsbatchint.sh

dev-requirements.txt

@@ -21,3 +21,4 @@ ts==0.5.1
 usort==0.6.4
 ipython
 pytest
+boto3==1.20.49

docs/source/index.rst

@@ -97,6 +97,7 @@ Works With
    schedulers/kubernetes
    schedulers/slurm
    schedulers/ray
+   schedulers/aws_batch
 
 .. _Pipelines:
 .. toctree::

docs/source/schedulers/aws_batch.rst

@@ -0,0 +1,8 @@
+AWS Batch
+=================
+
+.. automodule:: torchx.schedulers.aws_batch_scheduler
+.. currentmodule:: torchx.schedulers.aws_batch_scheduler
+
+.. autoclass:: AWSBatchScheduler
+   :members:

scripts/awsbatchint.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -ex
+
+APP_ID="$(torchx run --wait --scheduler aws_batch -c queue=torchx utils.echo)"
+torchx status "$APP_ID"
+torchx describe "$APP_ID"
+torchx log "$APP_ID"
+LINES="$(torchx log "$APP_ID" | wc -l)"
+
+if [ "$LINES" -ne 1 ]
+then
+    echo "expected 1 log lines"
+    exit 1
+fi

torchx/schedulers/__init__.py

@@ -7,6 +7,7 @@
 
 from typing import Dict, Optional
 
+import torchx.schedulers.aws_batch_scheduler as aws_batch_scheduler
 import torchx.schedulers.docker_scheduler as docker_scheduler
 import torchx.schedulers.kubernetes_scheduler as kubernetes_scheduler
 import torchx.schedulers.local_scheduler as local_scheduler
@@ -48,6 +49,7 @@ def get_scheduler_factories() -> Dict[str, SchedulerFactory]:
         "local_cwd": local_scheduler.create_cwd_scheduler,
         "slurm": slurm_scheduler.create_scheduler,
         "kubernetes": kubernetes_scheduler.create_scheduler,
+        "aws_batch": aws_batch_scheduler.create_scheduler,
     }
 
     ray_scheduler_creator = try_get_ray_scheduler()