yangaws commented on a change in pull request #4091: [AIRFLOW-2524] Update SageMaker hook, operator and sensor for training, tuning and transform URL: https://github.com/apache/incubator-airflow/pull/4091#discussion_r228372702
########## File path: airflow/contrib/sensors/sagemaker_training_sensor.py ########## @@ -37,12 +41,31 @@ class SageMakerTrainingSensor(SageMakerBaseSensor): @apply_defaults def __init__(self, job_name, - region_name=None, + print_log=True, *args, **kwargs): super(SageMakerTrainingSensor, self).__init__(*args, **kwargs) self.job_name = job_name - self.region_name = region_name + self.print_log = print_log + self.hook = SageMakerHook(aws_conn_id=self.aws_conn_id) + self.positions = {} + self.stream_names = [] + self.instance_count = None + self.state = None + self.last_description = None + self.last_describe_job_call = None + self.log_resource_inited = False + + def init_log_resource(self): + description = self.hook.describe_training_job(self.job_name) + self.instance_count = description['ResourceConfig']['InstanceCount'] + + status = description['TrainingJobStatus'] + job_already_completed = False if status in self.non_terminal_states() else True Review comment: Yep that's much better style! But I guess we should do `job_already_completed = status not in self.non_terminal_states()` ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services