How to pass the experiment configuration to a SagemakerTrainingOperator while training?

Question

Idea:

• To use experiments and trials to log the training parameters and artifacts in sagemaker while using MWAA as the pipeline orchestrator

I am using the training_config to create the dict to pass the training configuration to the Tensorflow estimator, but there is no parameter to pass the experiment configuration

tf_estimator = TensorFlow(entry_point='train_model.py',
                                      source_dir= source
                                      role=sagemaker.get_execution_role(),
                                      instance_count=1,
                                      framework_version='2.3.0',
                                      instance_type=instance_type,
                                      py_version='py37',
                                      script_mode=True,
                                      enable_sagemaker_metrics = True,
                                      metric_definitions=metric_definitions,
                                      output_path=output

model_training_config = training_config(
                    estimator=tf_estimator,
                    inputs=input
                    job_name=training_jobname,
                )
    



training_task = SageMakerTrainingOperator(
                    task_id=test_id,
                    config=model_training_config,
                    aws_conn_id="airflow-sagemaker",  
                    print_log=True,
                    wait_for_completion=True,
                    check_interval=60  
                )

Answer 1

You can use the experiment_config in estimator.fit. More detailed example can be found here