Upgrade to SageMaker v2 including changes:

* image -> image_uri * train_instance_count -> instance_count * train_instance_type -> instance_type * train_max_run -> max_run * train_max_run_wait -> max_run_wait * train_volume_size -> volume_size * sagemaker.session.s3_input -> sagemaker.inputs.TrainingInput

Author: brightsparc

requirements.txt

@@ -1,3 +1,3 @@
-sagemaker>=1.71.0
+sagemaker>=2.0.0
 boto3>=1.9.213
 pyyaml

src/stepfunctions/steps/sagemaker.py

@@ -36,12 +36,12 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
  data: Information about the training data. Please refer to the ``fit()`` method of the associated estimator, as this can take any of the following forms:
 
  * (str) - The S3 location where training data is saved.
- * (dict[str, str] or dict[str, sagemaker.session.s3_input]) - If using multiple
+ * (dict[str, str] or dict[str, sagemaker.inputs.TrainingInput]) - If using multiple
  channels for training data, you can specify a dict mapping channel names to
- strings or :func:`~sagemaker.session.s3_input` objects.
- * (sagemaker.session.s3_input) - Channel configuration for S3 data sources that can
+ strings or :func:`~sagemaker.inputs.TrainingInput` objects.
+ * (sagemaker.inputs.TrainingInput) - Channel configuration for S3 data sources that can
  provide additional information about the training dataset. See
- :func:`sagemaker.session.s3_input` for full details.
+ :func:`sagemaker.inputs.TrainingInput` for full details.
  * (sagemaker.amazon.amazon_estimator.RecordSet) - A collection of
  Amazon :class:`Record` objects serialized and stored in S3.
  For use with an estimator for an Amazon algorithm.
@@ -202,7 +202,7 @@ def __init__(self, state_id, model, model_name=None, instance_type=None, tags=No
  tags (list[dict], optional): `List to tags <https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html>`_ to associate with the resource.
  """
  if isinstance(model, FrameworkModel):
- parameters = model_config(model=model, instance_type=instance_type, role=model.role, image=model.image)
+ parameters = model_config(model=model, instance_type=instance_type, role=model.role, image_uri=model.image_uri)
  if model_name:
  parameters['ModelName'] = model_name
  elif isinstance(model, Model):
@@ -211,7 +211,7 @@ def __init__(self, state_id, model, model_name=None, instance_type=None, tags=No
  'ModelName': model_name or model.name,
  'PrimaryContainer': {
  'Environment': {},
- 'Image': model.image,
+ 'Image': model.image_uri,
  'ModelDataUrl': model.model_data
  }
  }
@@ -322,12 +322,12 @@ def __init__(self, state_id, tuner, job_name, data, wait_for_completion=True, ta
  data: Information about the training data. Please refer to the ``fit()`` method of the associated estimator in the tuner, as this can take any of the following forms:
 
  * (str) - The S3 location where training data is saved.
- * (dict[str, str] or dict[str, sagemaker.session.s3_input]) - If using multiple
+ * (dict[str, str] or dict[str, sagemaker.inputs.TrainingInput]) - If using multiple
  channels for training data, you can specify a dict mapping channel names to
- strings or :func:`~sagemaker.session.s3_input` objects.
- * (sagemaker.session.s3_input) - Channel configuration for S3 data sources that can
+ strings or :func:`~sagemaker.inputs.TrainingInput` objects.
+ * (sagemaker.inputs.TrainingInput) - Channel configuration for S3 data sources that can
  provide additional information about the training dataset. See
- :func:`sagemaker.session.s3_input` for full details.
+ :func:`sagemaker.inputs.TrainingInput` for full details.
  * (sagemaker.amazon.amazon_estimator.RecordSet) - A collection of
  Amazon :class:`Record` objects serialized and stored in S3.
  For use with an estimator for an Amazon algorithm.

src/stepfunctions/template/pipeline/inference.py

@@ -48,8 +48,8 @@ def __init__(self, preprocessor, estimator, inputs, s3_bucket, role, client=None
  inputs: Information about the training data. Please refer to the `fit()` method of the associated estimator, as this can take any of the following forms:
 
  * (str) - The S3 location where training data is saved.
- * (dict[str, str] or dict[str, `sagemaker.session.s3_input`]) - If using multiple channels for training data, you can specify a dict mapping channel names to strings or `sagemaker.session.s3_input` objects.
- * (`sagemaker.session.s3_input`) - Channel configuration for S3 data sources that can provide additional information about the training dataset. See `sagemaker.session.s3_input` for full details.
+ * (dict[str, str] or dict[str, `sagemaker.inputs.TrainingInput`]) - If using multiple channels for training data, you can specify a dict mapping channel names to strings or `sagemaker.inputs.TrainingInput` objects.
+ * (`sagemaker.inputs.TrainingInput`) - Channel configuration for S3 data sources that can provide additional information about the training dataset. See `sagemaker.inputs.TrainingInput` for full details.
  * (`sagemaker.amazon.amazon_estimator.RecordSet`) - A collection of Amazon `Record` objects serialized and stored in S3. For use with an estimator for an Amazon algorithm.
  * (list[`sagemaker.amazon.amazon_estimator.RecordSet`]) - A list of `sagemaker.amazon.amazon_estimator.RecordSet` objects, where each instance is a different channel of training data.
  s3_bucket (str): S3 bucket under which the output artifacts from the training job will be stored. The parent path used is built using the format: ``s3://{s3_bucket}/{pipeline_name}/models/{job_name}/``. In this format, `pipeline_name` refers to the keyword argument provided for TrainingPipeline. If a `pipeline_name` argument was not provided, one is auto-generated by the pipeline as `training-pipeline-<timestamp>`. Also, in the format, `job_name` refers to the job name provided when calling the :meth:`TrainingPipeline.run()` method.
@@ -87,8 +87,8 @@ def build_workflow_definition(self):
  """
  default_name = self.pipeline_name
 
- train_instance_type = self.preprocessor.train_instance_type
- train_instance_count = self.preprocessor.train_instance_count
+ instance_type = self.preprocessor.instance_type
+ instance_count = self.preprocessor.instance_count
 
  # Preprocessor for feature transformation
  preprocessor_train_step = TrainingStep(
@@ -100,13 +100,13 @@ def build_workflow_definition(self):
  preprocessor_model = self.preprocessor.create_model()
  preprocessor_model_step = ModelStep(
  StepId.CreatePreprocessorModel.value,
- instance_type=train_instance_type,
+ instance_type=instance_type,
  model=preprocessor_model,
  model_name=default_name
  )
  preprocessor_transform_step = TransformStep(
  StepId.TransformInput.value,
- transformer=self.preprocessor.transformer(instance_count=train_instance_count, instance_type=train_instance_type, max_payload=20),
+ transformer=self.preprocessor.transformer(instance_count=instance_count, instance_type=instance_type, max_payload=20),
  job_name=default_name,
  model_name=default_name,
  data=self.inputs['train'],
@@ -115,8 +115,8 @@ def build_workflow_definition(self):
  )
 
  # Training
- train_instance_type = self.estimator.train_instance_type
- train_instance_count = self.estimator.train_instance_count
+ instance_type = self.estimator.instance_type
+ instance_count = self.estimator.instance_count
 
  training_step = TrainingStep(
  StepId.Train.value,
@@ -135,21 +135,21 @@ def build_workflow_definition(self):
  )
  pipeline_model_step = ModelStep(
  StepId.CreatePipelineModel.value,
- instance_type=train_instance_type,
+ instance_type=instance_type,
  model=preprocessor_model,
  model_name=default_name
  )
- pipeline_model_step.parameters = self.pipeline_model_config(train_instance_type, pipeline_model)
+ pipeline_model_step.parameters = self.pipeline_model_config(instance_type, pipeline_model)
 
- deployable_model = Model(model_data='', image='')
+ deployable_model = Model(model_data='', image_uri='')
 
  # Deployment
  endpoint_config_step = EndpointConfigStep(
  StepId.ConfigureEndpoint.value,
  endpoint_config_name=default_name,
  model_name=default_name,
- initial_instance_count=train_instance_count,
- instance_type=train_instance_type
+ initial_instance_count=instance_count,
+ instance_type=instance_type
  )
 
  deploy_step = EndpointStep(

src/stepfunctions/template/pipeline/train.py

@@ -43,8 +43,8 @@ def __init__(self, estimator, role, inputs, s3_bucket, client=None, **kwargs):
  inputs: Information about the training data. Please refer to the `fit()` method of the associated estimator, as this can take any of the following forms:
 
  * (str) - The S3 location where training data is saved.
- * (dict[str, str] or dict[str, `sagemaker.session.s3_input`]) - If using multiple channels for training data, you can specify a dict mapping channel names to strings or `sagemaker.session.s3_input` objects.
- * (`sagemaker.session.s3_input`) - Channel configuration for S3 data sources that can provide additional information about the training dataset. See `sagemaker.session.s3_input` for full details.
+ * (dict[str, str] or dict[str, `sagemaker.inputs.TrainingInput`]) - If using multiple channels for training data, you can specify a dict mapping channel names to strings or `sagemaker.inputs.TrainingInput` objects.
+ * (`sagemaker.inputs.TrainingInput`) - Channel configuration for S3 data sources that can provide additional information about the training dataset. See `sagemaker.inputs.TrainingInput` for full details.
  * (`sagemaker.amazon.amazon_estimator.RecordSet`) - A collection of Amazon `Record` objects serialized and stored in S3. For use with an estimator for an Amazon algorithm.
  * (list[`sagemaker.amazon.amazon_estimator.RecordSet`]) - A list of `sagemaker.amazon.amazon_estimator.RecordSet` objects, where each instance is a different channel of training data.
  s3_bucket (str): S3 bucket under which the output artifacts from the training job will be stored. The parent path used is built using the format: ``s3://{s3_bucket}/{pipeline_name}/models/{job_name}/``. In this format, `pipeline_name` refers to the keyword argument provided for TrainingPipeline. If a `pipeline_name` argument was not provided, one is auto-generated by the pipeline as `training-pipeline-<timestamp>`. Also, in the format, `job_name` refers to the job name provided when calling the :meth:`TrainingPipeline.run()` method.
@@ -79,8 +79,8 @@ def build_workflow_definition(self):
  """
  default_name = self.pipeline_name
 
- train_instance_type = self.estimator.train_instance_type
- train_instance_count = self.estimator.train_instance_count
+ instance_type = self.estimator.instance_type
+ instance_count = self.estimator.instance_count
 
  training_step = TrainingStep(
  StepId.Train.value,
@@ -92,7 +92,7 @@ def build_workflow_definition(self):
  model = self.estimator.create_model()
  model_step = ModelStep(
  StepId.CreateModel.value,
- instance_type=train_instance_type,
+ instance_type=instance_type,
  model=model,
  model_name=default_name
  )
@@ -101,8 +101,8 @@ def build_workflow_definition(self):
  StepId.ConfigureEndpoint.value,
  endpoint_config_name=default_name,
  model_name=default_name,
- initial_instance_count=train_instance_count,
- instance_type=train_instance_type
+ initial_instance_count=instance_count,
+ instance_type=instance_type
  )
  deploy_step = EndpointStep(
  StepId.Deploy.value,

tests/integ/conftest.py

@@ -53,8 +53,8 @@ def sagemaker_role_arn(aws_account_id):
 def pca_estimator_fixture(sagemaker_role_arn):
  estimator = pca.PCA(
  role=sagemaker_role_arn,
- train_instance_count=1,
- train_instance_type="ml.m5.large",
+ instance_count=1,
+ instance_type="ml.m5.large",
  num_components=48
  )
  return estimator

tests/integ/test_inference_pipeline.py

@@ -45,7 +45,7 @@ def sklearn_preprocessor(sagemaker_role_arn, sagemaker_session):
  sklearn_preprocessor = SKLearn(
  entry_point=script_path,
  role=sagemaker_role_arn,
- train_instance_type="ml.m5.large",
+ instance_type="ml.m5.large",
  sagemaker_session=sagemaker_session,
  hyperparameters={"epochs": 1},
  )
@@ -60,7 +60,7 @@ def sklearn_estimator(sagemaker_role_arn, sagemaker_session):
  sklearn_estimator = SKLearn(
  entry_point=script_path,
  role=sagemaker_role_arn,
- train_instance_type="ml.m5.large",
+ instance_type="ml.m5.large",
  sagemaker_session=sagemaker_session,
  hyperparameters={"epochs": 1},
  input_mode='File'

tests/integ/test_sagemaker_steps.py

@@ -254,8 +254,8 @@ def test_tuning_step(sfn_client, record_set_for_hyperparameter_tuning, sagemaker
 
  kmeans = KMeans(
  role=sagemaker_role_arn,
- train_instance_count=1,
- train_instance_type=INSTANCE_TYPE,
+ instance_count=1,
+ instance_type=INSTANCE_TYPE,
  k=10
  )
 

tests/integ/test_training_pipeline_estimators.py

@@ -50,8 +50,8 @@ def pca_estimator(sagemaker_role_arn):
  pca_estimator = PCA(
  role=sagemaker_role_arn,
  num_components=1,
- train_instance_count=1,
- train_instance_type='ml.m5.large',
+ instance_count=1,
+ instance_type='ml.m5.large',
  )
 
  pca_estimator.feature_dim=500

tests/integ/test_training_pipeline_framework_estimator.py

@@ -36,8 +36,8 @@ def torch_estimator(sagemaker_role_arn):
  entry_point=script_path,
  role=sagemaker_role_arn,
  framework_version='1.2.0',
- train_instance_count=1,
- train_instance_type='ml.m5.large',
+ instance_count=1,
+ instance_type='ml.m5.large',
  hyperparameters={
  'epochs': 6,
  'backend': 'gloo'
@@ -50,8 +50,8 @@ def sklearn_estimator(sagemaker_role_arn):
  return SKLearn(
  entry_point=script_path,
  role=sagemaker_role_arn,
- train_instance_count=1,
- train_instance_type='ml.m5.large',
+ instance_count=1,
+ instance_type='ml.m5.large',
  framework_version='0.20.0',
  hyperparameters={
  "epochs": 1

tests/unit/test_pipeline.py

@@ -36,8 +36,8 @@ def pca_estimator():
  pca = sagemaker.estimator.Estimator(
  PCA_IMAGE,
  role=SAGEMAKER_EXECUTION_ROLE,
- train_instance_count=1,
- train_instance_type='ml.c4.xlarge',
+ instance_count=1,
+ instance_type='ml.c4.xlarge',
  output_path=s3_output_location,
  sagemaker_session=sagemaker_session
  )
@@ -62,7 +62,7 @@ def sklearn_preprocessor():
  sklearn_preprocessor = SKLearn(
  entry_point=script_path,
  role=SAGEMAKER_EXECUTION_ROLE,
- train_instance_type="ml.c4.xlarge",
+ instance_type="ml.c4.xlarge",
  source_dir=source_dir,
  sagemaker_session=sagemaker_session
  )
@@ -82,10 +82,10 @@ def linear_learner_estimator():
  ll_estimator = sagemaker.estimator.Estimator(
  LINEAR_LEARNER_IMAGE,
  SAGEMAKER_EXECUTION_ROLE, 
- train_instance_count=1, 
- train_instance_type='ml.c4.xlarge',
- train_volume_size=20,
- train_max_run=3600,
+ instance_count=1, 
+ instance_type='ml.c4.xlarge',
+ volume_size=20,
+ max_run=3600,
  input_mode='File',
  output_path=s3_output_location,
  sagemaker_session=sagemaker_session