diff --git a/components/vertex-components/src/vertex_components/custom_train_job.py b/components/vertex-components/src/vertex_components/custom_train_job.py index 0a2773bd..9cc51b17 100644 --- a/components/vertex-components/src/vertex_components/custom_train_job.py +++ b/components/vertex-components/src/vertex_components/custom_train_job.py @@ -1,5 +1,5 @@ # Copyright 2022 Google LLC -from typing import List, Dict +from typing import List, Dict, NamedTuple # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,7 +18,10 @@ @component( base_image="python:3.7", - packages_to_install=["google-cloud-aiplatform==1.24.1"], + packages_to_install=[ + "google-cloud-aiplatform==1.24.1", + "google-cloud-pipeline-components==1.0.42", + ], ) def custom_train_job( train_script_uri: str, @@ -41,7 +44,7 @@ def custom_train_job( accelerator_type: str = "ACCELERATOR_TYPE_UNSPECIFIED", accelerator_count: int = 0, parent_model: str = None, -): +) -> NamedTuple("Outputs", [("gcp_resources", str)]): """Run a custom training job using a training script. The provided script will be invoked by passing the following command-line arguments: @@ -87,12 +90,15 @@ def custom_train_job( Returns: parent_model (str): Resource URI of the parent model (empty string if the trained model is the first model version of its kind). + NamedTuple: gcp_resources for Vertex AI UI integration. """ import json import logging import os.path import time import google.cloud.aiplatform as aip + from google_cloud_pipeline_components.proto.gcp_resources_pb2 import GcpResources + from google.protobuf.json_format import MessageToJson logging.info(f"Using train script: {train_script_uri}") script_path = "/gcs/" + train_script_uri[5:] @@ -143,3 +149,12 @@ def custom_train_job( for k, v in parsed_metrics.items(): if type(v) is float: metrics.log_metric(k, v) + + # return GCP resource for Vertex AI UI integration + custom_job_name = job.to_dict()["trainingTaskMetadata"]["backingCustomJob"] + custom_train_job_resources = GcpResources() + ctr = custom_train_job_resources.resources.add() + ctr.resource_type = "CustomJob" + ctr.resource_uri = custom_job_name + gcp_resources = MessageToJson(custom_train_job_resources) + return (gcp_resources,) diff --git a/components/vertex-components/tests/test_custom_training_job.py b/components/vertex-components/tests/test_custom_training_job.py index 46ec46ca..ac7675c4 100644 --- a/components/vertex-components/tests/test_custom_training_job.py +++ b/components/vertex-components/tests/test_custom_training_job.py @@ -2,7 +2,7 @@ from kfp.v2.dsl import Dataset, Metrics, Artifact from unittest import mock import pytest - +import json import vertex_components @@ -24,7 +24,12 @@ def test_custom_train_job(mock_open, mock_exists, mock_job, tmpdir): mock_model = Artifact(uri=tmpdir, metadata={"resourceName": ""}) mock_metrics = Metrics(uri=tmpdir) - custom_train_job( + mock_job_instance = mock_job.return_value + mock_job_instance.to_dict.return_value = { + "trainingTaskMetadata": {"backingCustomJob": "mock_custom_job_name"} + } + + (gcp_resources,) = custom_train_job( train_script_uri="gs://my-bucket/train_script.py", train_data=mock_train_data, valid_data=mock_valid_data, @@ -53,6 +58,11 @@ def test_custom_train_job(mock_open, mock_exists, mock_job, tmpdir): # Assert metrics loading mock_open.assert_called_once_with(tmpdir, "r") + # Assert gcp_resources contains the expected value + assert ( + json.loads(gcp_resources)["resources"][0]["resourceUri"] + == "mock_custom_job_name" + ) @mock.patch("google.cloud.aiplatform.CustomTrainingJob") @@ -72,7 +82,7 @@ def test_custom_train_script_not_found(mock_open, mock_exists, mock_job, tmpdir) mock_metrics = Metrics(uri=tmpdir) with pytest.raises(ValueError): - custom_train_job( + (gcp_resources,) = custom_train_job( train_script_uri="gs://my-bucket/train_script.py", train_data=mock_train_data, valid_data=mock_valid_data,