def execute(self, context):
hook = automl_hook.AutoMLTablesHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Updating AutoML dataset %s.", self.dataset["name"])
result = hook.update_dataset(dataset=self.dataset,
location=self.location)
self.log.info("Dataset updated.")
return result
def execute(self, context):
hook = automl_hook.AutoMLTablesHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Fetch batch prediction.")
operation = hook.batch_predict(
model_id=self.model_id,
input_config=self.input_config,
output_config=self.output_config,
project_id=self.project_id,
location=self.location,
operation_timeout=self.operation_timeout,
poll_wait_time=self.poll_wait_time)
self.log.info("Batch prediction ready.")
output_info = operation["metadata"]["batchPredictDetails"][
"outputInfo"]
if "gcsOutputDirectory" in output_info:
output_loc = output_info["gcsOutputDirectory"]
elif "bigqueryOutputDataset" in output_info:
output_loc = output_info["bigqueryOutputDataset"]
else:
raise RuntimeError(
f"Output not found in prediction response:{operation}")
self.xcom_push(context,
key="prediction_result_output_loc",
value=output_loc)
return operation
def execute(self, context):
hook = automl_hook.AutoMLTablesHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Requesting table specs for %s.", self.dataset_id)
table_specs_list = hook.list_table_specs(dataset_id=self.dataset_id,
location=self.location,
project_id=self.project_id)
self.log.info(table_specs_list)
self.log.info("Table specs obtained.")
return table_specs_list
def execute(self, context):
hook = automl_hook.AutoMLTablesHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Importing dataset")
operation = hook.import_data(dataset_id=self.dataset_id,
input_config=self.input_config,
project_id=self.project_id,
location=self.location,
operation_timeout=self.operation_timeout,
poll_wait_time=self.poll_wait_time)
self.log.info("Import completed")
return operation
def execute(self, context):
hook = automl_hook.AutoMLTablesHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Creating dataset")
result = hook.create_dataset(dataset=self.dataset,
project_id=self.project_id,
location=self.location)
dataset_id = hook.extract_object_id(result)
self.log.info("Creating completed. Dataset id: %s", dataset_id)
self.xcom_push(context, key="dataset_id", value=dataset_id)
return result
def execute(self, context):
hook = automl_hook.AutoMLTablesHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Requesting column specs.")
result = hook.list_column_specs(
dataset_id=self.dataset_id,
table_spec_id=self.table_spec_id,
field_mask=self.field_mask,
location=self.location,
project_id=self.project_id,
)
self.log.info("Columns specs obtained.")
return result
def execute(self, context):
hook = automl_hook.AutoMLTablesHook(gcp_conn_id=self.gcp_conn_id)
self.log.info("Creating model.")
operation = hook.create_model(model=self.model,
project_id=self.project_id,
location=self.location,
operation_timeout=self.operation_timeout,
poll_wait_time=self.poll_wait_time)
operation_response = operation["response"]
model_id = hook.extract_object_id(operation_response)
self.log.info("Model created: %s", model_id)
self.xcom_push(context, key="model_id", value=model_id)
return operation_response