Vertex AI job operators in deferrable mode assume job uses Managed Model

[vignesh-sc]

Apache Airflow Provider(s)

google

Versions of Apache Airflow Providers

10.18.0

Apache Airflow version

2.8.4

Operating System

Linux

Deployment

Official Apache Airflow Helm Chart

Deployment details

No response

What happened

When using the CreateCustomContainerTrainingJobOperator without exporting a managed model, the operator failed when fetching the result with the following error:

  File "/opt/python3.11/lib/python3.11/site-packages/airflow/models/baseoperator.py", line 1629, in resume_execution
    return execute_callable(context)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/python3.11/lib/python3.11/site-packages/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py", line 574, in execute_complete
    model_id = self.hook.extract_model_id_from_training_pipeline(result)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/python3.11/lib/python3.11/site-packages/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py", line 266, in extract_model_id_from_training_pipeline
    return training_pipeline["model_to_upload"]["name"].rpartition("/")[-1]

On further exploration, this seems to be because the execute_complete function assumes the model parameter is returned. The execute function actually checks if a model is returned or not.

What you think should happen instead

Wether the operator is deferrable mode or not, it should check if a model is exported or not

How to reproduce

In order to reproduce:

1. Create a task using CreateCustomContainerTrainingJobOperator in deferrable mode
2. Dont set the model parameters
3. Execute the task

Anything else

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[eladkal]

cc @e-galan @VladaZakharova can you take a look?

[e-galan]

Hi @eladkal , I will check it

[e-galan]

Hi @vignesh-sc , I have just submitted a community PR to address the issue. You are welcome to check it.