Export Airflow model config from a SageMaker estimator

    Args:
        instance_type (str): The EC2 instance type to deploy this Model to. For example, 'ml.p2.xlarge'
        estimator (sagemaker.model.EstimatorBase): The SageMaker estimator to export Airflow config from.
            It has to be an estimator associated with a training job.
        task_id (str): The task id of any airflow.contrib.operators.SageMakerTrainingOperator or
            airflow.contrib.operators.SageMakerTuningOperator that generates training jobs in the DAG. The model config
            is built based on the training job generated in this operator.
        task_type (str): Whether the task is from SageMakerTrainingOperator or SageMakerTuningOperator. Values can be
            'training', 'tuning' or None (which means training job is not from any task).
        role (str): The ``ExecutionRoleArn`` IAM Role ARN for the model
        image (str): An container image to use for deploying the model
        name (str): Name of the model
        model_server_workers (int): The number of worker processes used by the inference server.
            If None, server will use one worker per vCPU. Only effective when estimator is a
            SageMaker framework.
        vpc_config_override (dict[str, list[str]]): Override for VpcConfig set on the model.
            Default: use subnets and security groups from this Estimator.
            * 'Subnets' (list[str]): List of subnet ids.
            * 'SecurityGroupIds' (list[str]): List of security group ids.

    Returns:
        dict: Model config that can be directly used by SageMakerModelOperator in Airflow. It can also be part
            of the config used by SageMakerEndpointOperator in Airflow.

def model_config_from_estimator(instance_type, estimator, task_id, task_type, role=None, image=None, name=None,
                                model_server_workers=None, vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT):
    """Export Airflow model config from a SageMaker estimator

    Args:
        instance_type (str): The EC2 instance type to deploy this Model to. For example, 'ml.p2.xlarge'
        estimator (sagemaker.model.EstimatorBase): The SageMaker estimator to export Airflow config from.
            It has to be an estimator associated with a training job.
        task_id (str): The task id of any airflow.contrib.operators.SageMakerTrainingOperator or
            airflow.contrib.operators.SageMakerTuningOperator that generates training jobs in the DAG. The model config
            is built based on the training job generated in this operator.
        task_type (str): Whether the task is from SageMakerTrainingOperator or SageMakerTuningOperator. Values can be
            'training', 'tuning' or None (which means training job is not from any task).
        role (str): The ``ExecutionRoleArn`` IAM Role ARN for the model
        image (str): An container image to use for deploying the model
        name (str): Name of the model
        model_server_workers (int): The number of worker processes used by the inference server.
            If None, server will use one worker per vCPU. Only effective when estimator is a
            SageMaker framework.
        vpc_config_override (dict[str, list[str]]): Override for VpcConfig set on the model.
            Default: use subnets and security groups from this Estimator.
            * 'Subnets' (list[str]): List of subnet ids.
            * 'SecurityGroupIds' (list[str]): List of security group ids.

    Returns:
        dict: Model config that can be directly used by SageMakerModelOperator in Airflow. It can also be part
            of the config used by SageMakerEndpointOperator in Airflow.
    """
    update_estimator_from_task(estimator, task_id, task_type)
    if isinstance(estimator, sagemaker.estimator.Estimator):
        model = estimator.create_model(role=role, image=image, vpc_config_override=vpc_config_override)
    elif isinstance(estimator, sagemaker.amazon.amazon_estimator.AmazonAlgorithmEstimatorBase):
        model = estimator.create_model(vpc_config_override=vpc_config_override)
    elif isinstance(estimator, sagemaker.estimator.Framework):
        model = estimator.create_model(model_server_workers=model_server_workers, role=role,
                                       vpc_config_override=vpc_config_override)
    else:
        raise TypeError('Estimator must be one of sagemaker.estimator.Estimator, sagemaker.estimator.Framework'
                        ' or sagemaker.amazon.amazon_estimator.AmazonAlgorithmEstimatorBase.')
    model.name = name

    return model_config(instance_type, model, role, image)