from sagemaker import AutoML, AutoMLInput, get_execution_role
from sagemaker.workflow.automl_step import AutoMLStep
from sagemaker.workflow.pipeline import Pipeline
from sagemaker.workflow.pipeline_context import PipelineSession

execution_role = get_execution_role()
pipeline_session = PipelineSession()
target_attribute_name = "target"

automl = AutoML(
    role=execution_role,
    target_attribute_name=target_attribute_name,
    sagemaker_session=pipeline_session,
    total_job_runtime_in_seconds=3600,
    mode="ENSEMBLING",  # only ensembling mode is supported for native AutoML step integration in SageMaker Pipelines
    problem_type="BinaryClassification",
    job_objective="F1",
)

train_args = automl.fit(
    inputs=[
        AutoMLInput(
            inputs="s3://sagemaker-us-east-1-862774760132/loan-data/data.csv",
            target_attribute_name=target_attribute_name,
            channel_type="training",
        )
    ]
)

step_auto_ml_training = AutoMLStep(
    name="AutoMLTrainingStep",
    step_args=train_args,
)

pipeline = Pipeline(
    name="TrainingPipeline",
    steps=[
        step_auto_ml_training,
    ],
    sagemaker_session=pipeline_session,
)

pipeline.upsert(role_arn=execution_role)