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

flytekit/models/sagemaker/training_job.py

@@ -163,15 +163,15 @@ def __init__(
             self,
             algorithm_name: int,
             algorithm_version: str,
-            metric_definitions: List[MetricDefinition],
             input_mode: int,
+            metric_definitions: List[MetricDefinition] = None,
             input_content_type: int = InputContentType.TEXT_CSV,
     ):
         self._input_mode = input_mode
         self._input_content_type = input_content_type
         self._algorithm_name = algorithm_name
         self._algorithm_version = algorithm_version
-        self._metric_definitions = metric_definitions
+        self._metric_definitions = metric_definitions or []
 
     @property
     def input_mode(self) -> int:
@@ -213,7 +213,12 @@ def metric_definitions(self) -> List[MetricDefinition]:
         """
         A list of metric definitions for SageMaker to evaluate/track on the progress of the training job
         See this: https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_AlgorithmSpecification.html
-        and this: https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-define-metrics.html
+
+        Note that, when you use one of the Amazon SageMaker built-in algorithms, you cannot define custom metrics.
+        If you are doing hyperparameter tuning, built-in algorithms automatically send metrics to hyperparameter tuning.
+        When using hyperparameter tuning, you do need to choose one of the metrics that the built-in algorithm emits as
+        the objective metric for the tuning job.
+        See this: https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-define-metrics.html
         :rtype: List[MetricDefinition]
         """
         return self._metric_definitions

tests/flytekit/unit/sdk/tasks/test_sagemaker_tasks.py

@@ -53,7 +53,6 @@
         input_content_type=InputContentType.TEXT_CSV,
         algorithm_name=AlgorithmName.XGBOOST,
         algorithm_version="0.72",
-        metric_definitions=[MetricDefinition(name="Validation error", regex="validation:error")]
     ),
 )
 
@@ -96,6 +95,8 @@ def test_simple_training_job_task():
     assert simple_training_job_task.metadata.discoverable is False
     assert simple_training_job_task.metadata.discovery_version == ''
     assert simple_training_job_task.metadata.retries.retries == 0
+    assert "metricDefinitions" not in simple_training_job_task.custom["algorithmSpecification"].keys()
+
     ParseDict(simple_training_job_task.custom['trainingJobResourceConfig'],
               _pb2_TrainingJobResourceConfig)  # fails the test if it cannot be parsed
 
@@ -129,6 +130,8 @@ def test_simple_training_job_task():
 
 
 def test_simple_hpo_job_task():
+    print(simple_xgboost_hpo_job_task.custom["trainingJob"])
+
     assert isinstance(simple_xgboost_hpo_job_task, SdkSimpleHyperparameterTuningJobTask)
     assert isinstance(simple_xgboost_hpo_job_task, _sdk_task.SdkTask)
     # Checking if the input of the underlying SdkTrainingJobTask has been embedded
@@ -176,7 +179,8 @@ def test_simple_hpo_job_task():
     assert simple_xgboost_hpo_job_task.metadata.retries.retries == 2
 
     assert simple_xgboost_hpo_job_task.metadata.deprecated_error_message == ''
-
+    assert "metricDefinitions" in simple_xgboost_hpo_job_task.custom["trainingJob"]["algorithmSpecification"].keys()
+    assert len(simple_xgboost_hpo_job_task.custom["trainingJob"]["algorithmSpecification"]["metricDefinitions"]) == 1
     """ These are attributes for SdkRunnable. We will need these when supporting CustomTrainingJobTask and CustomHPOJobTask 
     assert simple_xgboost_hpo_job_task.task_module == __name__
     assert simple_xgboost_hpo_job_task._get_container_definition().args[0] == 'pyflyte-execute'