e-galan commented on code in PR #40685:
URL: https://github.com/apache/airflow/pull/40685#discussion_r1675976036
##########
airflow/providers/google/cloud/operators/vertex_ai/custom_job.py:
##########
@@ -180,6 +180,40 @@ def execute(self, context: Context) -> None:
stacklevel=2,
)
+ def execute_complete(self, context: Context, event: dict[str, Any]) ->
dict[str, Any] | None:
+ if event["status"] == "error":
+ raise AirflowException(event["message"])
+ training_pipeline = event["job"]
+ custom_job_id =
self.hook.extract_custom_job_id_from_training_pipeline(training_pipeline)
+ self.xcom_push(context, key="custom_job_id", value=custom_job_id)
+ try:
+ model = training_pipeline["model_to_upload"]
+ model_id = self.hook.extract_model_id(model)
+ self.xcom_push(context, key="model_id", value=model_id)
+ VertexAIModelLink.persist(context=context, task_instance=self,
model_id=model_id)
+ return model
+ except KeyError:
+ self.log.warning(
+ "It is impossible to get the Model. "
+ "The Training Pipeline did not produce a Managed Model because
it was not "
+ "configured to upload a Model. Please ensure that the
'model_serving_container_image_uri' "
+ "and 'model_display_name' parameters are passed in when
creating a Training Pipeline, "
+ "and check that your training script saves the model to
os.environ['AIP_MODEL_DIR']."
+ )
+ return None
Review Comment:
It's hard to assess if every possible task can take an empty model as an
input, but if we're speaking about the Custom Training Job operators, they can
accept an existing model as a parent to train its next version via the
`parent_model` parameter.
And if the supplied `parent_model` turns out to be `None`, it will just mean
that no parent model will be used and a completely new model will be trained.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
