SageMaker on Flyte: TrainingJob for training with built-in algorithms and basic HPOJob support [Alpha]

[bnsblue]

TL;DR

This PR adds the necessary definitions for basic support of SageMaker TrainingJob and HPOJob (Hyperparameter Optimization)

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

This PR adds the basic support for users to invoke SageMaker TrainingJob (built-in algorithm mode) and SageMaker HPOJob from within Flyte. The follow paragraphs demonstrates the supported use cases.

Defining a Simple Training Job

Users can leverage SageMaker's powerful built-in algorithms easily without needing to write any function or logic. They can simply define a SdkBuildinAlgorithmTrainingJobTask in Flytekit and supplies the settings and the spec of the built-in algorithm as follows:

alg_spec = training_job_models.AlgorithmSpecification(
    input_mode=_sdk_sagemaker_types.InputMode.FILE,
    algorithm_name=_sdk_sagemaker_types.AlgorithmName.XGBOOST,
    algorithm_version="0.72",
    metric_definitions=[training_job_models.MetricDefinition(name="Minimize", regex="validation:error")]
)

# Definition of a SageMaker training job using SageMaker's built-in algorithm mode of XGBoost
xgboost_simple_train_task = training_job_task.SdkBuiltinAlgorithmTrainingJobTask(
    training_job_config=training_job_models.TrainingJobConfig(
        instance_type="ml.m4.xlarge",
        instance_count=1,
        volume_size_in_gb=25,
    ),
    algorithm_specification=alg_spec,
    cache_version='v1',
    cacheable=True,
)

Defining a simple Hyperparameter Tuning Job

SageMaker-on-Flyte also supports easy chaining between a TrainingJob task and a hpo job. After users define a TrainingJob task, he/she may want to apply hyperparameter tuning to the the training, while also maintaining the flexibility to run the TrainingJob task standalone. This should be easily doable by using the SdkSimpleHPOJobTask in our SDK. SdkSimpleHPOJobTask accepts the definition of a TrainingJob task as a part of the spec.

# Definition of a SageMaker hyperparameter-tuning job.
# Note that is chained behind the simple TrainingJob task defined above
xgboost_simple_hpo_task = hpo_job_task.SdkSimpleHPOJobTask(
    training_job=xgboost_simple_train_task,
    max_number_of_training_jobs=10,
    max_parallel_training_jobs=5,
    cache_version='2',
    retries=2,
    cacheable=True,
)

Invoking Training Jobs Task and Hyperparameter Tuning Jobs Task

Invoking Training Job Tasks and HPO Job Tasks from inside a Flyte workflow is pretty much the same as invoking other types of tasks. You should be able to achieve this by simply supplying the required inputs to the task definition. For Training Job Tasks and HPO Job Tasks , required inputs are pre-defined inputs that we think is needed for every training job. That's why you don't see the declaration of these inputs in the task definition -- we added them for you in our SDK.

@workflow_class
class SageMakerSimpleWorkflow(object):
    static_hyperparameters = Input(Types.Generic, required=True, help="Sample hyperparameter input")

    my_simple_trianing_task_exec = xgboost_simple_train_task(
        train="s3://path/to/train/data",
        validation="s3://path/to/validation/data",
        static_hyperparameters=static_hyperparameters,
        stopping_condition=StoppingCondition(
            max_runtime_in_seconds=43200,
        ).to_flyte_idl(),
    )
    ...

Flyte's Single Task Execution capability also makes it easy to invoke a SimpleTrainingJobTask and SimpleHPOJobTask. That is, users do not need a workflow to launch the SageMaker tasks; instead, they can simply define the tasks, and then register and launch the tasks standalone, which enables fast iterations.

xgboost_simple_train_task = training_job_task.SdkBuiltinAlgorithmTrainingJobTask( ... )

xgboost_simple_hpo_task = hpo_job_task.SdkSimpleHPOJobTask(
    training_job=xgboost_simple_train_task,
    ...
)

Tracking Issue

flyteorg/flyte#255

Follow-up issue

flyteorg/flyte#431

[codecov-commenter]

Codecov Report

Merging #120 into master will increase coverage by 0.02%.
The diff coverage is 81.42%.

@@            Coverage Diff             @@
##           master     #120      +/-   ##
==========================================
+ Coverage   80.83%   80.86%   +0.02%     
==========================================
  Files         219      225       +6     
  Lines       14313    14678     +365     
  Branches     1195     1205      +10     
==========================================
+ Hits        11570    11869     +299     
- Misses       2460     2526      +66     
  Partials      283      283              

Impacted Files	Coverage Δ
flytekit/plugins/__init__.py	100.00% <ø> (ø)
flytekit/sdk/tasks.py	71.59% <28.57%> (-3.72%)	⬇️
flytekit/models/sagemaker/parameter_ranges.py	53.75% <53.75%> (ø)
flytekit/models/sagemaker/hpo_job.py	71.42% <71.42%> (ø)
...ytekit/common/tasks/sagemaker/training_job_task.py	92.85% <92.85%> (ø)
flytekit/models/sagemaker/training_job.py	95.45% <95.45%> (ø)
flytekit/__init__.py	100.00% <100.00%> (ø)
flytekit/common/constants.py	100.00% <100.00%> (ø)
flytekit/common/tasks/sagemaker/hpo_job_task.py	100.00% <100.00%> (ø)
...ts/flytekit/unit/sdk/tasks/test_sagemaker_tasks.py	100.00% <100.00%> (ø)
... and 7 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update a32dd2d...7dbac57. Read the comment docs.

[wild-endeavor]

+1 - will approve after the final version is in.

[kumare3]

What does Simple mean?

[kumare3]

is this meant to be built in?

[bnsblue]

Yes -- meaning using built-in algorithm mode where users don't write his/her own decorated function.

[kumare3]

This looks awesome, two major comments

1. Support for other formats (currently you can keep it at CSV)
2. Docs to be added

[kumare3]

Looks good to me, but we should get plugin merged before this