fix: Fix artifact registry link not showing in ui when creating sched…

…ules with SDK. PiperOrigin-RevId: 547572327
googleapis · Jul 12, 2023 · 203cb47 · 203cb47
1 parent 1f55b05
commit 203cb47
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/google/cloud/aiplatform/preview/pipelinejobschedule/pipeline_job_schedules.py b/google/cloud/aiplatform/preview/pipelinejobschedule/pipeline_job_schedules.py
@@ -107,6 +107,10 @@ def __init__(
                 "pipeline_spec": pipeline_job.pipeline_spec,
             },
         }
+        if "template_uri" in pipeline_job._gca_resource:
+            create_pipeline_job_request["pipeline_job"][
+                "template_uri"
+            ] = pipeline_job._gca_resource.template_uri
         pipeline_job_schedule_args = {
             "display_name": display_name,
             "create_pipeline_job_request": create_pipeline_job_request,

diff --git a/tests/unit/aiplatform/test_pipeline_job_schedules.py b/tests/unit/aiplatform/test_pipeline_job_schedules.py
@@ -658,6 +658,7 @@ def test_call_schedule_service_create_artifact_registry(
                     "pipeline_spec": dict_to_struct(pipeline_spec),
                     "service_account": _TEST_SERVICE_ACCOUNT,
                     "network": _TEST_NETWORK,
+                    "template_uri": _TEST_AR_TEMPLATE_PATH,
                 },
             },
         )
@@ -740,6 +741,7 @@ def test_call_schedule_service_create_https(
                     "pipeline_spec": dict_to_struct(pipeline_spec),
                     "service_account": _TEST_SERVICE_ACCOUNT,
                     "network": _TEST_NETWORK,
+                    "template_uri": _TEST_HTTPS_TEMPLATE_PATH,
                 },
             },
         )