feature: support for Sample Weights for SageMaker Autopilot (#3835)

Co-authored-by: EC2 Default User <ec2-user@ip-172-31-57-175.us-west-2.compute.internal>
Co-authored-by: Choucri Bechir <cbechir@gmail.com>
Co-authored-by: Jaeyeon <jaeyeon@amazon.com>

Author: 4 people

src/sagemaker/automl/automl.py

@@ -49,6 +49,7 @@ def __init__(
         channel_type=None,
         content_type=None,
         s3_data_type=None,
+        sample_weight_attribute_name=None,
     ):
         """Convert an S3 Uri or a list of S3 Uri to an AutoMLInput object.
 
@@ -67,13 +68,16 @@ def __init__(
                 The content type of the data from the input source.
             s3_data_type (str, PipelineVariable): The data type for S3 data source.
                 Valid values: ManifestFile or S3Prefix.
+            sample_weight_attribute_name (str, PipelineVariable):
+                the name of the dataset column representing sample weights
         """
         self.inputs = inputs
         self.target_attribute_name = target_attribute_name
         self.compression = compression
         self.channel_type = channel_type
         self.content_type = content_type
         self.s3_data_type = s3_data_type
+        self.sample_weight_attribute_name = sample_weight_attribute_name
 
     def to_request_dict(self):
         """Generates a request dictionary using the parameters provided to the class."""
@@ -96,6 +100,8 @@ def to_request_dict(self):
                 input_entry["ContentType"] = self.content_type
             if self.s3_data_type is not None:
                 input_entry["DataSource"]["S3DataSource"]["S3DataType"] = self.s3_data_type
+            if self.sample_weight_attribute_name is not None:
+                input_entry["SampleWeightAttributeName"] = self.sample_weight_attribute_name
             auto_ml_input.append(input_entry)
         return auto_ml_input
 
@@ -129,6 +135,7 @@ def __init__(
         mode: Optional[str] = None,
         auto_generate_endpoint_name: Optional[bool] = None,
         endpoint_name: Optional[str] = None,
+        sample_weight_attribute_name: str = None,
     ):
         """Initialize the an AutoML object.
 
@@ -179,6 +186,8 @@ def __init__(
                 model deployment if the endpoint name is not generated automatically.
                 Specify the endpoint_name if and only if
                 auto_generate_endpoint_name is set to False
+            sample_weight_attribute_name (str): The name of dataset column representing
+                sample weights.
 
         Returns:
             AutoML object.
@@ -234,6 +243,7 @@ def __init__(
         )
 
         self._check_problem_type_and_job_objective(self.problem_type, self.job_objective)
+        self.sample_weight_attribute_name = sample_weight_attribute_name
 
     @runnable_by_pipeline
     def fit(self, inputs=None, wait=True, logs=True, job_name=None):
@@ -342,6 +352,9 @@ def attach(cls, auto_ml_job_name, sagemaker_session=None):
                 "AutoGenerateEndpointName", False
             ),
             endpoint_name=auto_ml_job_desc.get("ModelDeployConfig", {}).get("EndpointName"),
+            sample_weight_attribute_name=auto_ml_job_desc["InputDataConfig"][0].get(
+                "SampleWeightAttributeName", None
+            ),
         )
         amlj.current_job_name = auto_ml_job_name
         amlj.latest_auto_ml_job = auto_ml_job_name  # pylint: disable=W0201
@@ -867,6 +880,7 @@ def _load_config(cls, inputs, auto_ml, expand_role=True, validate_uri=True):
                 auto_ml.target_attribute_name,
                 auto_ml.content_type,
                 auto_ml.s3_data_type,
+                auto_ml.sample_weight_attribute_name,
             )
         output_config = _Job._prepare_output_config(auto_ml.output_path, auto_ml.output_kms_key)
 
@@ -932,6 +946,7 @@ def _format_inputs_to_input_config(
         target_attribute_name=None,
         content_type=None,
         s3_data_type=None,
+        sample_weight_attribute_name=None,
     ):
         """Convert inputs to AutoML InputDataConfig.
 
@@ -961,6 +976,8 @@ def _format_inputs_to_input_config(
                 channel["ContentType"] = content_type
             if s3_data_type is not None:
                 channel["DataSource"]["S3DataSource"]["S3DataType"] = s3_data_type
+            if sample_weight_attribute_name is not None:
+                channel["SampleWeightAttributeName"] = sample_weight_attribute_name
             channels.append(channel)
         elif isinstance(inputs, list):
             for input_entry in inputs:
@@ -974,6 +991,8 @@ def _format_inputs_to_input_config(
                     channel["ContentType"] = content_type
                 if s3_data_type is not None:
                     channel["DataSource"]["S3DataSource"]["S3DataType"] = s3_data_type
+                if sample_weight_attribute_name is not None:
+                    channel["SampleWeightAttributeName"] = sample_weight_attribute_name
                 channels.append(channel)
         else:
             msg = (

src/sagemaker/session.py

@@ -2100,7 +2100,8 @@ def auto_ml(
 
         Args:
             input_config (list[dict]): A list of Channel objects. Each channel contains "DataSource"
-                and "TargetAttributeName", "CompressionType" is an optional field.
+                and "TargetAttributeName", "CompressionType" and "SampleWeightAttributeName" are
+                optional fields.
             output_config (dict): The S3 URI where you want to store the training results and
                 optional KMS key ID.
             auto_ml_job_config (dict): A dict of AutoMLJob config, containing "StoppingCondition",
@@ -2167,7 +2168,8 @@ def _get_auto_ml_request(
 
         Args:
             input_config (list[dict]): A list of Channel objects. Each channel contains "DataSource"
-                and "TargetAttributeName", "CompressionType" is an optional field.
+                and "TargetAttributeName", "CompressionType" and "SampleWeightAttributeName" are
+                optional fields.
             output_config (dict): The S3 URI where you want to store the training results and
                 optional KMS key ID.
             auto_ml_job_config (dict): A dict of AutoMLJob config, containing "StoppingCondition",

tests/unit/sagemaker/automl/test_auto_ml.py

@@ -34,6 +34,7 @@
 RESOURCE_POOLS = [{"InstanceType": INSTANCE_TYPE, "PoolSize": INSTANCE_COUNT}]
 ROLE = "DummyRole"
 TARGET_ATTRIBUTE_NAME = "target"
+SAMPLE_WEIGHT_ATTRIBUTE_NAME = "sampleWeight"
 REGION = "us-west-2"
 DEFAULT_S3_INPUT_DATA = "s3://{}/data".format(BUCKET_NAME)
 DEFAULT_S3_VALIDATION_DATA = "s3://{}/validation_data".format(BUCKET_NAME)
@@ -103,6 +104,7 @@
                 }
             },
             "TargetAttributeName": "y",
+            "SampleWeightAttributeName": "sampleWeight",
         },
         {
             "ChannelType": "validation",
@@ -115,6 +117,7 @@
                 }
             },
             "TargetAttributeName": "y",
+            "SampleWeightAttributeName": "sampleWeight",
         },
     ],
     "OutputDataConfig": {"KmsKeyId": "string", "S3OutputPath": "s3://output_prefix"},
@@ -362,10 +365,12 @@ def test_auto_ml_validation_channel_name(sagemaker_session):
         target_attribute_name="target",
         compression="Gzip",
         channel_type="training",
+        sample_weight_attribute_name="sampleWeight",
     )
     input_validation = AutoMLInput(
         inputs=DEFAULT_S3_VALIDATION_DATA,
         target_attribute_name="target",
+        sample_weight_attribute_name="sampleWeight",
         compression="Gzip",
         channel_type="validation",
     )
@@ -384,6 +389,7 @@ def test_auto_ml_validation_channel_name(sagemaker_session):
                 }
             },
             "TargetAttributeName": TARGET_ATTRIBUTE_NAME,
+            "SampleWeightAttributeName": SAMPLE_WEIGHT_ATTRIBUTE_NAME,
         },
         {
             "ChannelType": "validation",
@@ -395,6 +401,7 @@ def test_auto_ml_validation_channel_name(sagemaker_session):
                 }
             },
             "TargetAttributeName": TARGET_ATTRIBUTE_NAME,
+            "SampleWeightAttributeName": SAMPLE_WEIGHT_ATTRIBUTE_NAME,
         },
     ]
 
@@ -617,7 +624,10 @@ def test_auto_ml_local_input(sagemaker_session):
 
 def test_auto_ml_input(sagemaker_session):
     inputs = AutoMLInput(
-        inputs=DEFAULT_S3_INPUT_DATA, target_attribute_name="target", compression="Gzip"
+        inputs=DEFAULT_S3_INPUT_DATA,
+        target_attribute_name="target",
+        compression="Gzip",
+        sample_weight_attribute_name="sampleWeight",
     )
     auto_ml = AutoML(
         role=ROLE,
@@ -636,6 +646,7 @@ def test_auto_ml_input(sagemaker_session):
                 }
             },
             "TargetAttributeName": TARGET_ATTRIBUTE_NAME,
+            "SampleWeightAttributeName": SAMPLE_WEIGHT_ATTRIBUTE_NAME,
         }
     ]
 

tests/unit/sagemaker/workflow/test_automl_step.py

@@ -56,12 +56,14 @@ def test_single_automl_step(pipeline_session):
         target_attribute_name="target",
         compression="Gzip",
         channel_type="training",
+        sample_weight_attribute_name="sampleWeight",
     )
     input_validation = AutoMLInput(
         inputs="s3://bucket/validation_data",
         target_attribute_name="target",
         compression="Gzip",
         channel_type="validation",
+        sample_weight_attribute_name="sampleWeight",
     )
     inputs = [input_training, input_validation]
 
@@ -114,6 +116,7 @@ def test_single_automl_step(pipeline_session):
                         }
                     },
                     "TargetAttributeName": "target",
+                    "SampleWeightAttributeName": "sampleWeight",
                 },
                 {
                     "ChannelType": "validation",
@@ -125,6 +128,7 @@ def test_single_automl_step(pipeline_session):
                         }
                     },
                     "TargetAttributeName": "target",
+                    "SampleWeightAttributeName": "sampleWeight",
                 },
             ],
             "OutputDataConfig": {

tests/unit/test_session.py

@@ -3175,6 +3175,7 @@ def test_wait_until_fail_access_denied_after_5_mins(patched_sleep):
         {
             "DataSource": {"S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": S3_INPUT_URI}},
             "TargetAttributeName": "y",
+            "SampleWeightAttributeName": "sampleWeight",
         }
     ],
     "OutputDataConfig": {"S3OutputPath": S3_OUTPUT},
@@ -3202,6 +3203,7 @@ def test_wait_until_fail_access_denied_after_5_mins(patched_sleep):
                 }
             },
             "TargetAttributeName": "y",
+            "SampleWeightAttributeName": "sampleWeight",
         },
         {
             "ChannelType": "validation",
@@ -3213,6 +3215,7 @@ def test_wait_until_fail_access_denied_after_5_mins(patched_sleep):
                 }
             },
             "TargetAttributeName": "y",
+            "SampleWeightAttributeName": "sampleWeight",
         },
     ],
     "OutputDataConfig": {"S3OutputPath": S3_OUTPUT},
@@ -3254,6 +3257,7 @@ def test_auto_ml_pack_to_request(sagemaker_session):
         {
             "DataSource": {"S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": S3_INPUT_URI}},
             "TargetAttributeName": "y",
+            "SampleWeightAttributeName": "sampleWeight",
         }
     ]
 
@@ -3288,6 +3292,7 @@ def test_create_auto_ml_with_sagemaker_config_injection(sagemaker_session):
         {
             "DataSource": {"S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": S3_INPUT_URI}},
             "TargetAttributeName": "y",
+            "SampleWeightAttributeName": "sampleWeight",
         }
     ]
 
@@ -3350,6 +3355,7 @@ def test_auto_ml_pack_to_request_with_optional_args(sagemaker_session):
                 }
             },
             "TargetAttributeName": "y",
+            "SampleWeightAttributeName": "sampleWeight",
         },
         {
             "ChannelType": "validation",
@@ -3361,6 +3367,7 @@ def test_auto_ml_pack_to_request_with_optional_args(sagemaker_session):
                 }
             },
             "TargetAttributeName": "y",
+            "SampleWeightAttributeName": "sampleWeight",
         },
     ]