# Buggy code
def prepare_framework_container_def(model, instance_type, s3_operations):
    """Prepare the framework model container information. Specify related S3
    operations for Airflow to perform. (Upload `source_dir` )

    Args:
        model (sagemaker.model.FrameworkModel): The framework model
        instance_type (str): The EC2 instance type to deploy this Model to. For
            example, 'ml.p2.xlarge'.
        s3_operations (dict): The dict to specify S3 operations (upload
            `source_dir` ).

    Returns:
        dict: The container information of this framework model.
    """
    deploy_image = model.image
    if not deploy_image:
        region_name = model.sagemaker_session.boto_session.region_name
        deploy_image = fw_utils.create_image_uri(
            region_name,
            model.__framework_name__,
            instance_type,
            model.framework_version,
            model.py_version,
        )

    base_name = utils.base_name_from_image(deploy_image)
    model.name = model.name or utils.name_from_base(base_name)

    bucket = model.bucket or model.sagemaker_session._default_bucket
    script = os.path.basename(model.entry_point)
    key = "{}/source/sourcedir.tar.gz".format(model.name)

    if model.source_dir and model.source_dir.lower().startswith("s3://"):
        code_dir = model.source_dir
        model.uploaded_code = fw_utils.UploadedCode(
            s3_prefix=code_dir, script_name=script
        )
    else:
        code_dir = "s3://{}/{}".format(bucket, key)
        model.uploaded_code = fw_utils.UploadedCode(
            s3_prefix=code_dir, script_name=script
        )
        s3_operations["S3Upload"] = [
            {
                "Path": model.source_dir or script,
                "Bucket": bucket,
                "Key": key,
                "Tar": True,
            }
        ]

    deploy_env = dict(model.env)
    deploy_env.update(model._framework_env_vars())

    try:
        if model.model_server_workers:
            deploy_env[sagemaker.model.MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(
                model.model_server_workers
            )
    except AttributeError:
        # This applies to a FrameworkModel which is not SageMaker Deep Learning Framework Model
        pass

    return sagemaker.container_def(deploy_image, model.model_data, deploy_env)


# Traceback (AttributeError)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-42-23cc1e7f4991> in <module>()
3     role=workflow_execution_role,
4     inputs=training_data_uri,
----> 5     s3_bucket=model_output_path
6 )

~/anaconda3/envs/python3/lib/python3.6/site-packages/stepfunctions/template/pipeline/train.py in __init__(self, estimator, role, inputs, s3_bucket, client, **kwargs)
64             self.pipeline_name = 'training-pipeline-{date}'.format(date=self._generate_timestamp())
65
---> 66         self.definition = self.build_workflow_definition()
67         self.input_template = self._extract_input_template(self.definition)
68

~/anaconda3/envs/python3/lib/python3.6/site-packages/stepfunctions/template/pipeline/train.py in build_workflow_definition(self)
95             instance_type=train_instance_type,
96             model=model,
---> 97             model_name=default_name
98         )
99

~/anaconda3/envs/python3/lib/python3.6/site-packages/stepfunctions/steps/sagemaker.py in __init__(self, state_id, model, model_name, instance_type, **kwargs)
171         """
172         if isinstance(model, FrameworkModel):
--> 173             parameters = model_config(model=model, instance_type=instance_type, role=model.role, image=model.image)
174             if model_name:
175                 parameters['ModelName'] = model_name

~/anaconda3/envs/python3/lib/python3.6/site-packages/sagemaker/workflow/airflow.py in model_config(instance_type, model, role, image)
577
578     if isinstance(model, sagemaker.model.FrameworkModel):
--> 579         container_def = prepare_framework_container_def(model, instance_type, s3_operations)
580     else:
581         container_def = model.prepare_container_def(instance_type)

~/anaconda3/envs/python3/lib/python3.6/site-packages/sagemaker/workflow/airflow.py in prepare_framework_container_def(model, instance_type, s3_operations)
519         deploy_image = fw_utils.create_image_uri(
520             region_name,
--> 521             model.__framework_name__,
522             instance_type,
523             model.framework_version,

AttributeError: 'Model' object has no attribute '__framework_name__'

# Fixed code
def prepare_framework_container_def(model, instance_type, s3_operations):
    """Prepare the framework model container information. Specify related S3
    operations for Airflow to perform. (Upload `source_dir` )

    Args:
        model (sagemaker.model.FrameworkModel): The framework model
        instance_type (str): The EC2 instance type to deploy this Model to. For
            example, 'ml.p2.xlarge'.
        s3_operations (dict): The dict to specify S3 operations (upload
            `source_dir` ).

    Returns:
        dict: The container information of this framework model.
    """
    deploy_image = model.image
    if not deploy_image:
        region_name = model.sagemaker_session.boto_session.region_name
        deploy_image = model.serving_image_uri(region_name, instance_type)

    base_name = utils.base_name_from_image(deploy_image)
    model.name = model.name or utils.name_from_base(base_name)

    bucket = model.bucket or model.sagemaker_session._default_bucket
    script = os.path.basename(model.entry_point)
    key = "{}/source/sourcedir.tar.gz".format(model.name)

    if model.source_dir and model.source_dir.lower().startswith("s3://"):
        code_dir = model.source_dir
        model.uploaded_code = fw_utils.UploadedCode(
            s3_prefix=code_dir, script_name=script
        )
    else:
        code_dir = "s3://{}/{}".format(bucket, key)
        model.uploaded_code = fw_utils.UploadedCode(
            s3_prefix=code_dir, script_name=script
        )
        s3_operations["S3Upload"] = [
            {
                "Path": model.source_dir or script,
                "Bucket": bucket,
                "Key": key,
                "Tar": True,
            }
        ]

    deploy_env = dict(model.env)
    deploy_env.update(model._framework_env_vars())

    try:
        if model.model_server_workers:
            deploy_env[sagemaker.model.MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(
                model.model_server_workers
            )
    except AttributeError:
        # This applies to a FrameworkModel which is not SageMaker Deep Learning Framework Model
        pass

    return sagemaker.container_def(deploy_image, model.model_data, deploy_env)