Skip to content

Commit

Permalink
feat: Add scheduled pipelines client create/get methods and unit tests.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 536465878
  • Loading branch information
vertex-sdk-bot authored and copybara-github committed May 30, 2023
1 parent 64818bf commit 4755fc7
Show file tree
Hide file tree
Showing 9 changed files with 947 additions and 5 deletions.
3 changes: 3 additions & 0 deletions google/cloud/aiplatform/compat/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
services.model_garden_service_client = services.model_garden_service_client_v1beta1
services.pipeline_service_client = services.pipeline_service_client_v1beta1
services.prediction_service_client = services.prediction_service_client_v1beta1
services.schedule_service_client = services.schedule_service_client_v1beta1
services.specialist_pool_service_client = (
services.specialist_pool_service_client_v1beta1
)
Expand Down Expand Up @@ -114,6 +115,8 @@
types.pipeline_state = types.pipeline_state_v1beta1
types.prediction_service = types.prediction_service_v1beta1
types.publisher_model = types.publisher_model_v1beta1
types.schedule = types.schedule_v1beta1
types.schedule_service = types.schedule_service_v1beta1
types.specialist_pool = types.specialist_pool_v1beta1
types.specialist_pool_service = types.specialist_pool_service_v1beta1
types.study = types.study_v1beta1
Expand Down
4 changes: 4 additions & 0 deletions google/cloud/aiplatform/compat/services/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,9 @@
from google.cloud.aiplatform_v1beta1.services.prediction_service import (
client as prediction_service_client_v1beta1,
)
from google.cloud.aiplatform_v1beta1.services.schedule_service import (
client as schedule_service_client_v1beta1,
)
from google.cloud.aiplatform_v1beta1.services.specialist_pool_service import (
client as specialist_pool_service_client_v1beta1,
)
Expand Down Expand Up @@ -140,6 +143,7 @@
model_service_client_v1beta1,
pipeline_service_client_v1beta1,
prediction_service_client_v1beta1,
schedule_service_client_v1beta1,
specialist_pool_service_client_v1beta1,
metadata_service_client_v1beta1,
tensorboard_service_client_v1beta1,
Expand Down
4 changes: 4 additions & 0 deletions google/cloud/aiplatform/compat/types/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,8 @@
pipeline_state as pipeline_state_v1beta1,
prediction_service as prediction_service_v1beta1,
publisher_model as publisher_model_v1beta1,
schedule as schedule_v1beta1,
schedule_service as schedule_service_v1beta1,
specialist_pool as specialist_pool_v1beta1,
specialist_pool_service as specialist_pool_service_v1beta1,
study as study_v1beta1,
Expand Down Expand Up @@ -283,6 +285,8 @@
pipeline_state_v1beta1,
prediction_service_v1beta1,
publisher_model_v1beta1,
schedule_v1beta1,
schedule_service_v1beta1,
specialist_pool_v1beta1,
specialist_pool_service_v1beta1,
study_v1beta1,
Expand Down
13 changes: 8 additions & 5 deletions google/cloud/aiplatform/pipeline_jobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,9 @@

import datetime
import logging
import time
import re
import tempfile
import time
from typing import Any, Callable, Dict, List, Optional, Union

from google.auth import credentials as auth_credentials
Expand All @@ -29,16 +29,16 @@
from google.cloud.aiplatform import utils
from google.cloud.aiplatform.constants import pipeline as pipeline_constants
from google.cloud.aiplatform.metadata import artifact
from google.cloud.aiplatform.metadata import constants as metadata_constants
from google.cloud.aiplatform.metadata import context
from google.cloud.aiplatform.metadata import execution
from google.cloud.aiplatform.metadata import constants as metadata_constants
from google.cloud.aiplatform.metadata import experiment_resources
from google.cloud.aiplatform.metadata import utils as metadata_utils
from google.cloud.aiplatform.utils import gcs_utils
from google.cloud.aiplatform.utils import yaml_utils
from google.cloud.aiplatform.utils import pipeline_utils
from google.protobuf import json_format
from google.cloud.aiplatform.utils import yaml_utils
from google.protobuf import field_mask_pb2 as field_mask
from google.protobuf import json_format

from google.cloud.aiplatform.compat.types import (
pipeline_job as gca_pipeline_job,
Expand Down Expand Up @@ -96,7 +96,6 @@ class PipelineJob(
),
),
):

client_class = utils.PipelineJobClientWithOverride
_resource_noun = "pipelineJobs"
_delete_method = "delete_pipeline_job"
Expand Down Expand Up @@ -443,6 +442,10 @@ def wait(self):
def pipeline_spec(self):
return self._gca_resource.pipeline_spec

@property
def runtime_config(self) -> gca_pipeline_job.PipelineJob.RuntimeConfig:
return self._gca_resource.runtime_config

@property
def state(self) -> Optional[gca_pipeline_state.PipelineState]:
"""Current pipeline state."""
Expand Down
58 changes: 58 additions & 0 deletions google/cloud/aiplatform/preview/constants/schedules.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
# -*- coding: utf-8 -*-

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from google.cloud.aiplatform.compat.types import (
schedule_v1beta1 as gca_schedule,
)
from google.cloud.aiplatform.constants import pipeline as pipeline_constants

_SCHEDULE_COMPLETE_STATES = set(
[
gca_schedule.Schedule.State.PAUSED,
gca_schedule.Schedule.State.COMPLETED,
]
)

_SCHEDULE_ERROR_STATES = set(
[
gca_schedule.Schedule.State.STATE_UNSPECIFIED,
]
)

# Pattern for valid names used as a Vertex resource name.
_VALID_NAME_PATTERN = pipeline_constants._VALID_NAME_PATTERN

# Pattern for an Artifact Registry URL.
_VALID_AR_URL = pipeline_constants._VALID_AR_URL

# Pattern for any JSON or YAML file over HTTPS.
_VALID_HTTPS_URL = pipeline_constants._VALID_HTTPS_URL

# Fields to include in returned PipelineJobSchedule when enable_simple_view=True in PipelineJobSchedule.list()
_PIPELINE_JOB_SCHEDULE_READ_MASK_FIELDS = [
"name",
"display_name",
"start_time",
"end_time",
"max_run_count",
"started_run_count",
"state",
"create_time",
"update_time",
"cron",
"catch_up",
]
Loading

0 comments on commit 4755fc7

Please sign in to comment.