From 3fbdb3518e99caf9d4f0a5782d03273a65616a3a Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Thu, 25 Aug 2022 16:07:50 -0400 Subject: [PATCH 01/10] feat: add enable_simple_view to PipelineJob.list() --- google/cloud/aiplatform/base.py | 11 +++++ google/cloud/aiplatform/pipeline_jobs.py | 19 ++++++++ tests/unit/aiplatform/test_pipeline_jobs.py | 51 ++++++++++++++++++++- 3 files changed, 79 insertions(+), 2 deletions(-) diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index e88de02934..03a4c7febc 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -48,6 +48,7 @@ from google.cloud.aiplatform.compat.types import encryption_spec as gca_encryption_spec from google.cloud.aiplatform.constants import base as base_constants from google.protobuf import json_format +from google.protobuf import field_mask_pb2 as field_mask # This is the default retry callback to be used with get methods. _DEFAULT_RETRY = retry.Retry() @@ -1026,6 +1027,7 @@ def _list( cls_filter: Callable[[proto.Message], bool] = lambda _: True, filter: Optional[str] = None, order_by: Optional[str] = None, + read_mask: Optional[field_mask.FieldMask] = None, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, @@ -1063,6 +1065,7 @@ def _list( Returns: List[VertexAiResourceNoun] - A list of SDK resource objects """ + resource = cls._empty_constructor( project=project, location=location, credentials=credentials ) @@ -1079,6 +1082,12 @@ def _list( ), } + # We are only exposing `read_mask` PipelineJob.list() for now + if cls.__name__ == "PipelineJob": + list_request["read_mask"] = read_mask + + print(list_request) + if filter: list_request["filter"] = filter @@ -1101,6 +1110,7 @@ def _list_with_local_order( cls_filter: Callable[[proto.Message], bool] = lambda _: True, filter: Optional[str] = None, order_by: Optional[str] = None, + read_mask: Optional[field_mask.FieldMask] = None, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, @@ -1141,6 +1151,7 @@ def _list_with_local_order( cls_filter=cls_filter, filter=filter, order_by=None, # This method will handle the ordering locally + read_mask=read_mask, project=project, location=location, credentials=credentials, diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index 76992595c4..c06ab131b8 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/google/cloud/aiplatform/pipeline_jobs.py @@ -38,6 +38,7 @@ from google.cloud.aiplatform.utils import yaml_utils from google.cloud.aiplatform.utils import pipeline_utils from google.protobuf import json_format +from google.protobuf import field_mask_pb2 as field_mask from google.cloud.aiplatform.compat.types import ( pipeline_job as gca_pipeline_job, @@ -503,6 +504,7 @@ def list( cls, filter: Optional[str] = None, order_by: Optional[str] = None, + enable_simple_view: Optional[bool] = False, project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, @@ -524,6 +526,10 @@ def list( Optional. A comma-separated list of fields to order by, sorted in ascending order. Use "desc" after a field name for descending. Supported fields: `display_name`, `create_time`, `update_time` + enable_simple_view (bool): + Optional. Whether to pass the `read_mask` parameter to the list call. + This will improve the performance of calling list(). However, the + returned PipelineJob list will not include all fields for each PipelineJob. project (str): Optional. Project to retrieve list from. If not set, project set in aiplatform.init will be used. @@ -538,9 +544,22 @@ def list( List[PipelineJob] - A list of PipelineJob resource objects """ + read_mask_fields = None + + if enable_simple_view: + read_mask_fields = field_mask.FieldMask( + paths=[ + "name,state,display_name,pipeline_spec.pipeline_info,create_time,start_time,end_time,update_time,labels,template_uri,template_metadata.version,job_detail.pipeline_run_context,job_detail.pipeline_context" + ] + ) + _LOGGER.warn( + "By enabling simple view, the PipelineJob resources returned from this method will not contain all fields." + ) + return cls._list_with_local_order( filter=filter, order_by=order_by, + read_mask=read_mask_fields, project=project, location=location, credentials=credentials, diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index 9975fedd5c..57bf0d3303 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -38,6 +38,7 @@ from google.cloud.aiplatform.utils import gcs_utils from google.cloud import storage from google.protobuf import json_format +from google.protobuf import field_mask_pb2 as field_mask from google.cloud.aiplatform.compat.services import ( pipeline_service_client, @@ -62,6 +63,7 @@ _TEST_NETWORK = f"projects/{_TEST_PROJECT}/global/networks/{_TEST_PIPELINE_JOB_ID}" _TEST_PIPELINE_JOB_NAME = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/pipelineJobs/{_TEST_PIPELINE_JOB_ID}" +_TEST_PIPELINE_JOB_LIST_READ_MASK = field_mask.FieldMask(paths=["name,state,display_name,pipeline_spec.pipeline_info,create_time,start_time,end_time,update_time,labels,template_uri,template_metadata.version,job_detail.pipeline_run_context,job_detail.pipeline_context"]) _TEST_PIPELINE_PARAMETER_VALUES_LEGACY = {"string_param": "hello"} _TEST_PIPELINE_PARAMETER_VALUES = { @@ -328,6 +330,11 @@ def mock_pipeline_service_list(): with mock.patch.object( pipeline_service_client.PipelineServiceClient, "list_pipeline_jobs" ) as mock_list_pipeline_jobs: + mock_list_pipeline_jobs.return_value = [ + make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + ] yield mock_list_pipeline_jobs @@ -1286,10 +1293,50 @@ def test_list_pipeline_job( ) job.run() - job.list() + job_list = job.list() + + assert hasattr(job_list[0].gca_resource, "runtime_config") + + mock_pipeline_service_list.assert_called_once_with( + request={"parent": _TEST_PARENT, "read_mask": None} + ) + + @pytest.mark.usefixtures( + "mock_pipeline_service_create", + "mock_pipeline_service_get", + "mock_pipeline_bucket_exists", + ) + @pytest.mark.parametrize( + "job_spec", + [ + _TEST_PIPELINE_SPEC_JSON, + _TEST_PIPELINE_SPEC_YAML, + _TEST_PIPELINE_JOB, + _TEST_PIPELINE_SPEC_LEGACY_JSON, + _TEST_PIPELINE_SPEC_LEGACY_YAML, + _TEST_PIPELINE_JOB_LEGACY, + ], + ) + def test_list_pipeline_job_with_read_mask( + self, mock_pipeline_service_list, mock_load_yaml_and_json + ): + aiplatform.init( + project=_TEST_PROJECT, + staging_bucket=_TEST_GCS_BUCKET_NAME, + credentials=_TEST_CREDENTIALS, + ) + + job = pipeline_jobs.PipelineJob( + display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, + template_path=_TEST_TEMPLATE_PATH, + job_id=_TEST_PIPELINE_JOB_ID, + ) + + job.run() + job_list = job.list(enable_simple_view=True) mock_pipeline_service_list.assert_called_once_with( - request={"parent": _TEST_PARENT} + request={"parent": _TEST_PARENT, "read_mask": _TEST_PIPELINE_JOB_LIST_READ_MASK}, ) @pytest.mark.usefixtures( From 31c92e8c186bea0db6a8a441f46ba73f0b41ac1a Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Thu, 25 Aug 2022 17:58:19 -0400 Subject: [PATCH 02/10] updates to pipelinejob.list read_mask --- google/cloud/aiplatform/base.py | 2 -- google/cloud/aiplatform/pipeline_jobs.py | 4 +--- tests/unit/aiplatform/test_pipeline_jobs.py | 8 ++++++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index 03a4c7febc..593b657037 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -1086,8 +1086,6 @@ def _list( if cls.__name__ == "PipelineJob": list_request["read_mask"] = read_mask - print(list_request) - if filter: list_request["filter"] = filter diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index c06ab131b8..c1b6e81c06 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/google/cloud/aiplatform/pipeline_jobs.py @@ -548,9 +548,7 @@ def list( if enable_simple_view: read_mask_fields = field_mask.FieldMask( - paths=[ - "name,state,display_name,pipeline_spec.pipeline_info,create_time,start_time,end_time,update_time,labels,template_uri,template_metadata.version,job_detail.pipeline_run_context,job_detail.pipeline_context" - ] + paths=["name","state","display_name", "pipeline_spec.pipeline_info", "create_time", "start_time", "end_time", "update_time", "labels", "template_uri", "template_metadata.version", "job_detail.pipeline_run_context", "job_detail.pipeline_context"] ) _LOGGER.warn( "By enabling simple view, the PipelineJob resources returned from this method will not contain all fields." diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index 57bf0d3303..f917dfacdc 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -63,7 +63,7 @@ _TEST_NETWORK = f"projects/{_TEST_PROJECT}/global/networks/{_TEST_PIPELINE_JOB_ID}" _TEST_PIPELINE_JOB_NAME = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/pipelineJobs/{_TEST_PIPELINE_JOB_ID}" -_TEST_PIPELINE_JOB_LIST_READ_MASK = field_mask.FieldMask(paths=["name,state,display_name,pipeline_spec.pipeline_info,create_time,start_time,end_time,update_time,labels,template_uri,template_metadata.version,job_detail.pipeline_run_context,job_detail.pipeline_context"]) +_TEST_PIPELINE_JOB_LIST_READ_MASK = field_mask.FieldMask(paths=["name","state","display_name", "pipeline_spec.pipeline_info", "create_time", "start_time", "end_time", "update_time", "labels", "template_uri", "template_metadata.version", "job_detail.pipeline_run_context", "job_detail.pipeline_context"]) _TEST_PIPELINE_PARAMETER_VALUES_LEGACY = {"string_param": "hello"} _TEST_PIPELINE_PARAMETER_VALUES = { @@ -1295,6 +1295,7 @@ def test_list_pipeline_job( job.run() job_list = job.list() + # Confirm runtime_config is still present in a list request without enable_simple_view assert hasattr(job_list[0].gca_resource, "runtime_config") mock_pipeline_service_list.assert_called_once_with( @@ -1333,7 +1334,10 @@ def test_list_pipeline_job_with_read_mask( ) job.run() - job_list = job.list(enable_simple_view=True) + pipeline_list = job.list(enable_simple_view=True) + + # TODO: this field is still being returned + # assert not hasattr(pipeline_list[0].gca_resource, "runtime_config") mock_pipeline_service_list.assert_called_once_with( request={"parent": _TEST_PARENT, "read_mask": _TEST_PIPELINE_JOB_LIST_READ_MASK}, From afd0327a7782609e365bcac6180b39b6b697f917 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Fri, 26 Aug 2022 09:20:46 -0400 Subject: [PATCH 03/10] run linter --- google/cloud/aiplatform/pipeline_jobs.py | 16 +++++++++- tests/unit/aiplatform/test_pipeline_jobs.py | 35 ++++++++++++++++++--- 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index c1b6e81c06..632816f2ec 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/google/cloud/aiplatform/pipeline_jobs.py @@ -548,7 +548,21 @@ def list( if enable_simple_view: read_mask_fields = field_mask.FieldMask( - paths=["name","state","display_name", "pipeline_spec.pipeline_info", "create_time", "start_time", "end_time", "update_time", "labels", "template_uri", "template_metadata.version", "job_detail.pipeline_run_context", "job_detail.pipeline_context"] + paths=[ + "name", + "state", + "display_name", + "pipeline_spec.pipeline_info", + "create_time", + "start_time", + "end_time", + "update_time", + "labels", + "template_uri", + "template_metadata.version", + "job_detail.pipeline_run_context", + "job_detail.pipeline_context", + ] ) _LOGGER.warn( "By enabling simple view, the PipelineJob resources returned from this method will not contain all fields." diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index f917dfacdc..e308253ebe 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -63,7 +63,23 @@ _TEST_NETWORK = f"projects/{_TEST_PROJECT}/global/networks/{_TEST_PIPELINE_JOB_ID}" _TEST_PIPELINE_JOB_NAME = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/pipelineJobs/{_TEST_PIPELINE_JOB_ID}" -_TEST_PIPELINE_JOB_LIST_READ_MASK = field_mask.FieldMask(paths=["name","state","display_name", "pipeline_spec.pipeline_info", "create_time", "start_time", "end_time", "update_time", "labels", "template_uri", "template_metadata.version", "job_detail.pipeline_run_context", "job_detail.pipeline_context"]) +_TEST_PIPELINE_JOB_LIST_READ_MASK = field_mask.FieldMask( + paths=[ + "name", + "state", + "display_name", + "pipeline_spec.pipeline_info", + "create_time", + "start_time", + "end_time", + "update_time", + "labels", + "template_uri", + "template_metadata.version", + "job_detail.pipeline_run_context", + "job_detail.pipeline_context", + ] +) _TEST_PIPELINE_PARAMETER_VALUES_LEGACY = {"string_param": "hello"} _TEST_PIPELINE_PARAMETER_VALUES = { @@ -331,9 +347,15 @@ def mock_pipeline_service_list(): pipeline_service_client.PipelineServiceClient, "list_pipeline_jobs" ) as mock_list_pipeline_jobs: mock_list_pipeline_jobs.return_value = [ - make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), - make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), - make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED), + make_pipeline_job( + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + ), + make_pipeline_job( + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + ), + make_pipeline_job( + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + ), ] yield mock_list_pipeline_jobs @@ -1340,7 +1362,10 @@ def test_list_pipeline_job_with_read_mask( # assert not hasattr(pipeline_list[0].gca_resource, "runtime_config") mock_pipeline_service_list.assert_called_once_with( - request={"parent": _TEST_PARENT, "read_mask": _TEST_PIPELINE_JOB_LIST_READ_MASK}, + request={ + "parent": _TEST_PARENT, + "read_mask": _TEST_PIPELINE_JOB_LIST_READ_MASK, + }, ) @pytest.mark.usefixtures( From b3c02106aba64faef07151141a0a31d6a05bb1ad Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Fri, 26 Aug 2022 10:21:07 -0400 Subject: [PATCH 04/10] update to read_mask --- google/cloud/aiplatform/base.py | 4 ++-- tests/unit/aiplatform/test_pipeline_jobs.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index 593b657037..317d323f1b 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -1082,8 +1082,8 @@ def _list( ), } - # We are only exposing `read_mask` PipelineJob.list() for now - if cls.__name__ == "PipelineJob": + # We are only exposing `read_mask` on PipelineJob.list() for now + if read_mask and cls.__name__ == "PipelineJob": list_request["read_mask"] = read_mask if filter: diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index e308253ebe..5f7aa8ef4b 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -1321,7 +1321,7 @@ def test_list_pipeline_job( assert hasattr(job_list[0].gca_resource, "runtime_config") mock_pipeline_service_list.assert_called_once_with( - request={"parent": _TEST_PARENT, "read_mask": None} + request={"parent": _TEST_PARENT} ) @pytest.mark.usefixtures( @@ -1356,10 +1356,10 @@ def test_list_pipeline_job_with_read_mask( ) job.run() - pipeline_list = job.list(enable_simple_view=True) + job.list(enable_simple_view=True) # TODO: this field is still being returned - # assert not hasattr(pipeline_list[0].gca_resource, "runtime_config") + # assert not hasattr(job_list[0].gca_resource, "runtime_config") mock_pipeline_service_list.assert_called_once_with( request={ From bcdd068e820ec9057a7abb8bdf77cbc40291570d Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Fri, 26 Aug 2022 12:23:12 -0400 Subject: [PATCH 05/10] add placeholder for read_mask to system tests --- tests/system/aiplatform/test_pipeline_job.py | 5 +++++ tests/unit/aiplatform/test_pipeline_jobs.py | 7 ------- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/system/aiplatform/test_pipeline_job.py b/tests/system/aiplatform/test_pipeline_job.py index 004ad768a3..97049affa4 100644 --- a/tests/system/aiplatform/test_pipeline_job.py +++ b/tests/system/aiplatform/test_pipeline_job.py @@ -59,3 +59,8 @@ def training_pipeline(number_of_epochs: int = 10): shared_state.setdefault("resources", []).append(job) job.wait() + + # TODO: add test for `read_mask` when it is available in PipelineJob.list() + # pipeline_job_list = aiplatform.PipelineJob.list(enable_simple_view=True) + # assert not hasattr(pipeline_job_list[0].gca_resource, "runtime_config") + # assert hasattr(pipeline_job_list[0].gca_resource, "name") diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index 5f7aa8ef4b..a4d5d7902e 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -1315,10 +1315,6 @@ def test_list_pipeline_job( ) job.run() - job_list = job.list() - - # Confirm runtime_config is still present in a list request without enable_simple_view - assert hasattr(job_list[0].gca_resource, "runtime_config") mock_pipeline_service_list.assert_called_once_with( request={"parent": _TEST_PARENT} @@ -1358,9 +1354,6 @@ def test_list_pipeline_job_with_read_mask( job.run() job.list(enable_simple_view=True) - # TODO: this field is still being returned - # assert not hasattr(job_list[0].gca_resource, "runtime_config") - mock_pipeline_service_list.assert_called_once_with( request={ "parent": _TEST_PARENT, From d1de8b340926176f1342019522a6f2c8d0d69d14 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Fri, 26 Aug 2022 17:27:21 -0400 Subject: [PATCH 06/10] unit test fix --- tests/unit/aiplatform/test_pipeline_jobs.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index a4d5d7902e..0d5cc61caf 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -1315,6 +1315,7 @@ def test_list_pipeline_job( ) job.run() + job.list() mock_pipeline_service_list.assert_called_once_with( request={"parent": _TEST_PARENT} From 4772ab8e70fd0ca39dda253118bf532a32fb59ab Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 29 Aug 2022 11:20:14 -0400 Subject: [PATCH 07/10] add system test for read_mask filter --- tests/system/aiplatform/test_pipeline_job.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/system/aiplatform/test_pipeline_job.py b/tests/system/aiplatform/test_pipeline_job.py index 97049affa4..13d8207a1d 100644 --- a/tests/system/aiplatform/test_pipeline_job.py +++ b/tests/system/aiplatform/test_pipeline_job.py @@ -20,6 +20,8 @@ from google.cloud import aiplatform from tests.system.aiplatform import e2e_base +from google.protobuf.json_format import MessageToDict + @pytest.mark.usefixtures("tear_down_resources") class TestPipelineJob(e2e_base.TestEndToEnd): @@ -60,7 +62,13 @@ def training_pipeline(number_of_epochs: int = 10): job.wait() - # TODO: add test for `read_mask` when it is available in PipelineJob.list() - # pipeline_job_list = aiplatform.PipelineJob.list(enable_simple_view=True) - # assert not hasattr(pipeline_job_list[0].gca_resource, "runtime_config") - # assert hasattr(pipeline_job_list[0].gca_resource, "name") + list_with_read_mask = aiplatform.PipelineJob.list(enable_simple_view=True) + list_without_read_mask = aiplatform.PipelineJob.list() + + # enable_simple_view=True should apply the `read_mask` filter to limit PipelineJob fields returned + assert "serviceAccount" in MessageToDict( + list_without_read_mask[0].gca_resource._pb + ) + assert "serviceAccount" not in MessageToDict( + list_with_read_mask[0].gca_resource._pb + ) From d5c3ffee0bd20001df60f07f58c05e41758bcc57 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 29 Aug 2022 11:50:40 -0400 Subject: [PATCH 08/10] move read mask fields to constants file --- google/cloud/aiplatform/constants/pipeline.py | 17 ++++++++++++++++ google/cloud/aiplatform/pipeline_jobs.py | 20 +++---------------- tests/unit/aiplatform/test_pipeline_jobs.py | 17 ++-------------- 3 files changed, 22 insertions(+), 32 deletions(-) diff --git a/google/cloud/aiplatform/constants/pipeline.py b/google/cloud/aiplatform/constants/pipeline.py index d4ff2aa32e..12acf8a52e 100644 --- a/google/cloud/aiplatform/constants/pipeline.py +++ b/google/cloud/aiplatform/constants/pipeline.py @@ -37,3 +37,20 @@ # Pattern for an Artifact Registry URL. _VALID_AR_URL = re.compile(r"^https:\/\/([\w-]+)-kfp\.pkg\.dev\/.*") + +# Fields to include in returned PipelineJob when enable_simple_view=True in PipelineJob.list() +_READ_MASK_FIELDS = [ + "name", + "state", + "display_name", + "pipeline_spec.pipeline_info", + "create_time", + "start_time", + "end_time", + "update_time", + "labels", + "template_uri", + "template_metadata.version", + "job_detail.pipeline_run_context", + "job_detail.pipeline_context", +] diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index 632816f2ec..30544c3d9f 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/google/cloud/aiplatform/pipeline_jobs.py @@ -57,6 +57,8 @@ # Pattern for an Artifact Registry URL. _VALID_AR_URL = pipeline_constants._VALID_AR_URL +_READ_MASK_FIELDS = pipeline_constants._READ_MASK_FIELDS + def _get_current_time() -> datetime.datetime: """Gets the current timestamp.""" @@ -547,23 +549,7 @@ def list( read_mask_fields = None if enable_simple_view: - read_mask_fields = field_mask.FieldMask( - paths=[ - "name", - "state", - "display_name", - "pipeline_spec.pipeline_info", - "create_time", - "start_time", - "end_time", - "update_time", - "labels", - "template_uri", - "template_metadata.version", - "job_detail.pipeline_run_context", - "job_detail.pipeline_context", - ] - ) + read_mask_fields = field_mask.FieldMask(paths=_READ_MASK_FIELDS) _LOGGER.warn( "By enabling simple view, the PipelineJob resources returned from this method will not contain all fields." ) diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index 0d5cc61caf..9a737fe20f 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -29,6 +29,7 @@ from google.cloud import aiplatform from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer +from google.cloud.aiplatform.constants import pipeline as pipeline_constants from google.cloud.aiplatform_v1 import Context as GapicContext from google.cloud.aiplatform_v1 import MetadataStore as GapicMetadataStore from google.cloud.aiplatform.metadata import constants @@ -64,21 +65,7 @@ _TEST_PIPELINE_JOB_NAME = f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/pipelineJobs/{_TEST_PIPELINE_JOB_ID}" _TEST_PIPELINE_JOB_LIST_READ_MASK = field_mask.FieldMask( - paths=[ - "name", - "state", - "display_name", - "pipeline_spec.pipeline_info", - "create_time", - "start_time", - "end_time", - "update_time", - "labels", - "template_uri", - "template_metadata.version", - "job_detail.pipeline_run_context", - "job_detail.pipeline_context", - ] + paths=pipeline_constants._READ_MASK_FIELDS ) _TEST_PIPELINE_PARAMETER_VALUES_LEGACY = {"string_param": "hello"} From 6a868e40df2eb9631962ca623dff6f4eccff837e Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Mon, 29 Aug 2022 16:25:12 -0400 Subject: [PATCH 09/10] add read_mask docstrings --- google/cloud/aiplatform/base.py | 16 ++++++++++++++++ google/cloud/aiplatform/pipeline_jobs.py | 7 +++++++ 2 files changed, 23 insertions(+) diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index 317d323f1b..79ab3cd8ed 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -1050,6 +1050,14 @@ def _list( Optional. A comma-separated list of fields to order by, sorted in ascending order. Use "desc" after a field name for descending. Supported fields: `display_name`, `create_time`, `update_time` + read_mask (field_mask.FieldMask): + Optional. A FieldMask with a list of strings passed via `paths` + indicating which fields to return for each resource in the response. + For example, passing + field_mask.FieldMask(paths=["create_time", "update_time"]) + as `read_mask` would result in each returned VertexAiResourceNoun + in the result list only having the "create_time" and + "update_time" attributes. project (str): Optional. Project to retrieve list from. If not set, project set in aiplatform.init will be used. @@ -1131,6 +1139,14 @@ def _list_with_local_order( Optional. A comma-separated list of fields to order by, sorted in ascending order. Use "desc" after a field name for descending. Supported fields: `display_name`, `create_time`, `update_time` + read_mask (field_mask.FieldMask): + Optional. A FieldMask with a list of strings passed via `paths` + indicating which fields to return for each resource in the response. + For example, passing + field_mask.FieldMask(paths=["create_time", "update_time"]) + as `read_mask` would result in each returned VertexAiResourceNoun + in the result list only having the "create_time" and + "update_time" attributes. project (str): Optional. Project to retrieve list from. If not set, project set in aiplatform.init will be used. diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index 30544c3d9f..54d36d22e3 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/google/cloud/aiplatform/pipeline_jobs.py @@ -532,6 +532,13 @@ def list( Optional. Whether to pass the `read_mask` parameter to the list call. This will improve the performance of calling list(). However, the returned PipelineJob list will not include all fields for each PipelineJob. + Setting this to True will exclude the following fields in your response: + `runtime_config`, `service_account`, `network`, and some subfields of + `pipeline_spec` and `job_detail`. The following fields will be included in + each PipelineJob resource in your response: `state`, `display_name`, + `pipeline_spec.pipeline_info`, `create_time`, `start_time`, `end_time`, + `update_time`, `labels`, `template_uri`, `template_metadata.version`, + `job_detail.pipeline_run_context`, `job_detail.pipeline_context`. project (str): Optional. Project to retrieve list from. If not set, project set in aiplatform.init will be used. From 1429ea2943264be8583667213e54e6db0cc07e88 Mon Sep 17 00:00:00 2001 From: Sara Robinson Date: Wed, 31 Aug 2022 10:50:15 -0400 Subject: [PATCH 10/10] remove class name check --- google/cloud/aiplatform/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index 79ab3cd8ed..887ea77432 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -1090,8 +1090,8 @@ def _list( ), } - # We are only exposing `read_mask` on PipelineJob.list() for now - if read_mask and cls.__name__ == "PipelineJob": + # `read_mask` is only passed from PipelineJob.list() for now + if read_mask is not None: list_request["read_mask"] = read_mask if filter: