def do_create_model(**kwargs):
""" Creates a model only if one with the same name did not exist.
It leverages the check from the previous task pushed using xcom.
"""
model_params = {
'name': kwargs['dag_run'].conf.get('model_name'),
'description': 'A custom DNN regressor model',
'regions': [REGION]
}
ti = kwargs['ti']
is_model = ti.xcom_pull(key='is_project', task_ids='check_model')
if not is_model:
mle = MLEngineHook()
mle.create_model(PROJECT, model_params)
def execute(self, context):
hook = MLEngineHook(gcp_conn_id=self._gcp_conn_id,
delegate_to=self._delegate_to)
if self._operation == 'create':
return hook.create_model(self._project_id, self._model)
elif self._operation == 'get':
return hook.get_model(self._project_id, self._model['name'])
else:
raise ValueError('Unknown operation: {}'.format(self._operation))
def execute(self, context):
hook = MLEngineHook(
gcp_conn_id=self._gcp_conn_id, delegate_to=self._delegate_to)
if self._operation == 'create':
return hook.create_model(self._project_id, self._model)
elif self._operation == 'get':
return hook.get_model(self._project_id, self._model['name'])
else:
raise ValueError('Unknown operation: {}'.format(self._operation))
