feat: adding aws_sagemaker_scheduler (#785)

pytorch · Feb 6, 2024 · 78c5286 · 78c5286
1 parent 11b3fa3
commit 78c5286
Show file tree

Hide file tree

Showing 6 changed files with 999 additions and 0 deletions.
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -2,6 +2,7 @@ aiobotocore
 ax-platform[mysql]==0.2.3
 black==23.3.0
 boto3
+sagemaker>=2.149.0
 captum>=0.4.0
 docker
 flake8==3.9.0

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -77,6 +77,7 @@ Works With
    schedulers/slurm
    schedulers/ray
    schedulers/aws_batch
+   schedulers/aws_sagemaker
    schedulers/lsf
    schedulers/gcp_batch
 

diff --git a/docs/source/schedulers/aws_sagemaker.rst b/docs/source/schedulers/aws_sagemaker.rst
@@ -0,0 +1,19 @@
+AWS SageMaker
+=================
+
+.. automodule:: torchx.schedulers.aws_sagemaker_scheduler
+
+.. currentmodule:: torchx.schedulers.aws_sagemaker_scheduler
+
+.. autoclass:: AWSSageMakerScheduler
+   :members:
+   :show-inheritance:
+
+.. autoclass:: AWSSageMakerJob
+   :members:
+
+Reference
+~~~~~~~~~~~~
+
+.. autofunction:: create_scheduler
+
diff --git a/torchx/schedulers/__init__.py b/torchx/schedulers/__init__.py
@@ -19,6 +19,7 @@
     "kubernetes": "torchx.schedulers.kubernetes_scheduler",
     "kubernetes_mcad": "torchx.schedulers.kubernetes_mcad_scheduler",
     "aws_batch": "torchx.schedulers.aws_batch_scheduler",
+    "aws_sagemaker": "torchx.schedulers.aws_sagemaker_scheduler",
     "gcp_batch": "torchx.schedulers.gcp_batch_scheduler",
     "ray": "torchx.schedulers.ray_scheduler",
     "lsf": "torchx.schedulers.lsf_scheduler",