airflow icon indicating copy to clipboard operation
airflow copied to clipboard

Remove parent_model version suffix if passed to Vertex AI operators

Open e-galan opened this issue 9 months ago • 1 comments

Check the parent_model parameter for model version suffix (looks like <model_id>@1 or <model_name>@1) and remove it from the parent_model string if it is present there. Keep the same value if the suffix is not present.

This is needed because sometimes operators can be passed with the parent_model value which contains such a suffix, for example when the user wants to extract a model id / model name from a serialized result of CreateCustomTrainingJobOperator run and then pass it to another operator as the parent_model. Vertex AI is not able to find a resource if it contains this suffix, and will return an error that will prevent the user from submitting a new training job.

This PR partially reverts #38417 which removed a similar functionality because it caused a bug in the cases when the parent_model did not contain any suffixes. This PR should cover both problems.


^ Add meaningful description above Read the Pull Request Guidelines for more information. In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed. In case of a new dependency, check compliance with the ASF 3rd Party License Policy. In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

e-galan avatar May 15 '24 13:05 e-galan

Hi @eladkal ! Can we check this PR and include it to next google-provider release? Thanks!

VladaZakharova avatar May 16 '24 09:05 VladaZakharova