adding aws_sagemaker_scheduler

[clumsy]

Providing an implementation of torchx scheduler for AWS SageMaker.

Several considerations that influenced this particular implementation:

1. SageMaker's torch_distributed entry_point does not allow specifying custom torchrun command, it constructs the command itself meaning the component that one should use for aws_sagemaker scheduler is utils.python. That also means that instance count and instance type have to be passed as scheduler args.
2. Current SageMaker does not allow using modules, only scripts are allowed. This should go away once these two changes are adopted: feature: python module support to torch_distributed aws/sagemaker-python-sdk#4324 and feature: add python module entrypoint type, add python module support to torch_distributed aws/sagemaker-training-toolkit#205. This PR assumes they get adopted and allows specifying the module.
3. In order to make torchx one stop for submitting jobs to SageMaker I added all currently supported scheduler parameters with their documentation from AWS docs.
4. Some parameters had to be controlled by torchx (e.g. from role) - failing if one attempts to set those.
5. Unfortunately special care had to be taken when processing hyperparameters (script arguments), as SageMakers expect those in a key value format and would prepend them with --.

Unit test is to be added later, sending this PR ahead to initiate the discussion and gather feedback.

Test plan:

• tested by submitting a job to SageMaker in an AWS account: works for a script, fails for a module at runtime (see the notes above)

[clumsy]

Please find this first version and let me know what you think, @kiukchung

[kiukchung]

Could you include unittests? Thanks!

[clumsy]

Added doc files as requested, @kiukchung

[clumsy]

Do we need to add anything else to this PR, @kiukchung?

[kiukchung]

LGTM thanks for adding this (and making the docs changes)!

[kiukchung]

@clumsy it looks like there are unittest failures. Can you take a look? (and also rebase onto HEAD)

[clumsy]

I fixed the mypy issue, @kiukchung. Looks like we need a permission to rerun all the workflows.

[goswamig]

Thanks for adding this.

[kiukchung]

@clumsy looks like there are unittest failures. Do the tests pass locally?

[clumsy]

Weird, they were passing locally, @kiukchung. Let me have a look.

[clumsy]

The test should pass now, @kiukchung

[kiukchung]

The test should pass now, @kiukchung

Thanks, it looks like lint is still failing. Can you try running the linter locally (essentially running usort, black, and flake8)? This is the lint action: https://github.com/pytorch/torchx/blob/main/scripts/lint.sh#L58-L60

[clumsy]

@kiukchung looks like black rules might have been changed recently. rerun lint.sh and submitted a new revision.

[kiukchung]

@clumsy looks like pyre (typecheck) is failing on a few sagemaker related files:

 torchx/schedulers/aws_sagemaker_scheduler.py:28:0 Unused ignore [0]: The `pyre-ignore[21]` or `pyre-fixme[21]` comment is not suppressing type errors, please remove it.
torchx/schedulers/test/aws_sagemaker_scheduler_test.py:225:28 Missing parameter annotation [2]: Parameter `mock_pytorch_estimator` has no type specified.
torchx/schedulers/test/aws_sagemaker_scheduler_test.py:259:25 Undefined attribute [16]: `Optional` has no attribute `app_id`.
torchx/schedulers/test/aws_sagemaker_scheduler_test.py:260:25 Undefined attribute [16]: `Optional` has no attribute `state`.
torchx/schedulers/test/aws_sagemaker_scheduler_test.py:262:12 Undefined attribute [16]: `Optional` has no attribute `ui_url`.

[clumsy]

Fixed pyre and docs builds, @kiukchung

[clumsy]

Please let me know if additional changes are needed, @kiukchung

[clumsy]

Fixed pyre issue by upgrading local version of pyre dependencies, the errors I see are not related to this PR:

torchx/examples/apps/lightning/profiler.py:22:0 Undefined import [21]: Could not find a module corresponding to import `pytorch_lightning.loggers.base`.
torchx/examples/apps/lightning/profiler.py:23:0 Undefined import [21]: Could not find a module corresponding to import `pytorch_lightning.profiler.base`.
torchx/examples/apps/lightning/profiler.py:26:28 Undefined or invalid type [11]: Annotation `BaseProfiler` is not defined as a type.
torchx/examples/apps/lightning/profiler.py:33:31 Undefined or invalid type [11]: Annotation `LightningLoggerBase` is not defined as a type.
torchx/examples/apps/lightning/profiler.py:37:8 Missing attribute annotation [4]: Attribute `logger` of class `SimpleLoggingProfiler` has no type specified.
torchx/examples/apps/lightning/train.py:172:12 Incompatible parameter type [6]: In call `pl.trainer.trainer.Trainer.__init__`, for argument `profiler`, expected `Union[None, Profiler, str]` but got `SimpleLoggingProfiler`.

link failure is not related to this PR:

Running linters ...
Checking torchx/schedulers/__init__.py
torchx/schedulers/__init__.py:33:5: E704 multiple statements on one line (def)
Checking torchx/schedulers/aws_sagemaker_scheduler.py
Checking torchx/schedulers/test/aws_sagemaker_scheduler_test.py
One of the linters returned an error. See above output.

Unit test failure is not related to this PR:

=========================== short test summary info ============================
FAILED torchx/schedulers/test/local_scheduler_test.py::LocalDirectorySchedulerTest::test_no_orphan_process_function - AssertionError: OSError not raised

Docs build fine locally if I disable nbsphinx.

[kiukchung]

Thanks @clumsy for fixing! There was a merge conflict on dev-requirements.txt which I just resolved. Will let the CI run again and merge once everything is green.

[d4l3k]

@kiukchung looks like this PR introduced a issue in the doctest build

https://github.com/pytorch/torchx/actions/runs/8456379306/job/23166027094