adding aws_sagemaker_scheduler (#801)

Co-authored-by: Alexander Jipa <[email protected]> Co-authored-by: kiukchung <[email protected]>
pytorch · Mar 27, 2024 · dd4c61f · dd4c61f
1 parent e9a6957
commit dd4c61f
Show file tree

Hide file tree

Showing 6 changed files with 1,006 additions and 1 deletion.
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -21,6 +21,7 @@ pyre-extensions
 pyre-check
 pytest
 pytorch-lightning==1.5.10
+sagemaker>=2.149.0
 torch-model-archiver>=0.4.2
 torch==2.2.1
 torchmetrics==0.10.3

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
@@ -21,6 +21,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",
@@ -29,7 +30,8 @@
 
 class SchedulerFactory(Protocol):
     # pyre-fixme: Scheduler opts
-    def __call__(self, session_name: str, **kwargs: object) -> Scheduler: ...
+    def __call__(self, session_name: str, **kwargs: object) -> Scheduler:
+        ...
 
 
 def _defer_load_scheduler(path: str) -> SchedulerFactory: