diff --git a/google/cloud/aiplatform/__init__.py b/google/cloud/aiplatform/__init__.py index 1ad69b2a54..31f459d3f7 100644 --- a/google/cloud/aiplatform/__init__.py +++ b/google/cloud/aiplatform/__init__.py @@ -42,7 +42,7 @@ MatchingEngineIndex, MatchingEngineIndexEndpoint, ) -from google.cloud.aiplatform.metadata import metadata +from google.cloud.aiplatform import metadata from google.cloud.aiplatform.models import Endpoint from google.cloud.aiplatform.models import Model from google.cloud.aiplatform.model_evaluation import ModelEvaluation @@ -56,6 +56,7 @@ Tensorboard, TensorboardExperiment, TensorboardRun, + TensorboardTimeSeries, ) from google.cloud.aiplatform.training_jobs import ( CustomTrainingJob, @@ -78,24 +79,39 @@ """ init = initializer.global_config.init -log_params = metadata.metadata_service.log_params -log_metrics = metadata.metadata_service.log_metrics -get_experiment_df = metadata.metadata_service.get_experiment_df -get_pipeline_df = metadata.metadata_service.get_pipeline_df -start_run = metadata.metadata_service.start_run +get_pipeline_df = metadata.metadata._LegacyExperimentService.get_pipeline_df + +log_params = metadata.metadata._experiment_tracker.log_params +log_metrics = metadata.metadata._experiment_tracker.log_metrics +get_experiment_df = metadata.metadata._experiment_tracker.get_experiment_df +start_run = metadata.metadata._experiment_tracker.start_run +start_execution = metadata.metadata._experiment_tracker.start_execution +log = metadata.metadata._experiment_tracker.log +log_time_series_metrics = metadata.metadata._experiment_tracker.log_time_series_metrics +end_run = metadata.metadata._experiment_tracker.end_run + +Experiment = metadata.experiment_resources.Experiment +ExperimentRun = metadata.experiment_run_resource.ExperimentRun +Artifact = metadata.artifact.Artifact +Execution = metadata.execution.Execution __all__ = ( + "end_run", "explain", "gapic", "init", "helpers", "hyperparameter_tuning", + "log", "log_params", "log_metrics", + "log_time_series_metrics", "get_experiment_df", "get_pipeline_df", "start_run", + "start_execution", + "Artifact", "AutoMLImageTrainingJob", "AutoMLTabularTrainingJob", "AutoMLForecastingTrainingJob", @@ -108,6 +124,9 @@ "CustomPythonPackageTrainingJob", "Endpoint", "EntityType", + "Execution", + "Experiment", + "ExperimentRun", "Feature", "Featurestore", "MatchingEngineIndex", @@ -122,6 +141,7 @@ "Tensorboard", "TensorboardExperiment", "TensorboardRun", + "TensorboardTimeSeries", "TextDataset", "TimeSeriesDataset", "VideoDataset", diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index e0b0ae6312..ceb9287322 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -610,6 +610,25 @@ def name(self) -> str: self._assert_gca_resource_is_available() return self._gca_resource.name.split("/")[-1] + @property + def _project_tuple(self) -> Tuple[Optional[str], Optional[str]]: + """Returns the tuple of project id and project inferred from the local instance. + + Another option is to use resource_manager_utils but requires the caller have resource manager + get role. + """ + # we may not have the project if project inferred from the resource name + maybe_project_id = self.project + if self._gca_resource is not None and self._gca_resource.name: + project_no = self._parse_resource_name(self._gca_resource.name)["project"] + else: + project_no = None + + if maybe_project_id == project_no: + return (None, project_no) + else: + return (maybe_project_id, project_no) + @property def resource_name(self) -> str: """Full qualified resource name.""" diff --git a/google/cloud/aiplatform/compat/__init__.py b/google/cloud/aiplatform/compat/__init__.py index 6aea51d133..00e213505d 100644 --- a/google/cloud/aiplatform/compat/__init__.py +++ b/google/cloud/aiplatform/compat/__init__.py @@ -80,6 +80,7 @@ types.io = types.io_v1beta1 types.job_service = types.job_service_v1beta1 types.job_state = types.job_state_v1beta1 + types.lineage_subgraph = types.lineage_subgraph_v1beta1 types.machine_resources = types.machine_resources_v1beta1 types.manual_batch_tuning_parameters = types.manual_batch_tuning_parameters_v1beta1 types.matching_engine_deployed_index_ref = ( @@ -88,6 +89,7 @@ types.matching_engine_index = types.index_v1beta1 types.matching_engine_index_endpoint = types.index_endpoint_v1beta1 types.metadata_service = types.metadata_service_v1beta1 + types.metadata_schema = types.metadata_schema_v1beta1 types.metadata_store = types.metadata_store_v1beta1 types.model = types.model_v1beta1 types.model_evaluation = types.model_evaluation_v1beta1 @@ -162,6 +164,7 @@ types.io = types.io_v1 types.job_service = types.job_service_v1 types.job_state = types.job_state_v1 + types.lineage_subgraph = types.lineage_subgraph_v1 types.machine_resources = types.machine_resources_v1 types.manual_batch_tuning_parameters = types.manual_batch_tuning_parameters_v1 types.matching_engine_deployed_index_ref = ( @@ -170,6 +173,7 @@ types.matching_engine_index = types.index_v1 types.matching_engine_index_endpoint = types.index_endpoint_v1 types.metadata_service = types.metadata_service_v1 + types.metadata_schema = types.metadata_schema_v1 types.metadata_store = types.metadata_store_v1 types.model = types.model_v1 types.model_evaluation = types.model_evaluation_v1 diff --git a/google/cloud/aiplatform/compat/types/__init__.py b/google/cloud/aiplatform/compat/types/__init__.py index 14ff93f011..a3f6f96147 100644 --- a/google/cloud/aiplatform/compat/types/__init__.py +++ b/google/cloud/aiplatform/compat/types/__init__.py @@ -52,8 +52,10 @@ io as io_v1beta1, job_service as job_service_v1beta1, job_state as job_state_v1beta1, + lineage_subgraph as lineage_subgraph_v1beta1, machine_resources as machine_resources_v1beta1, manual_batch_tuning_parameters as manual_batch_tuning_parameters_v1beta1, + metadata_schema as metadata_schema_v1beta1, metadata_service as metadata_service_v1beta1, metadata_store as metadata_store_v1beta1, model as model_v1beta1, @@ -113,9 +115,11 @@ io as io_v1, job_service as job_service_v1, job_state as job_state_v1, + lineage_subgraph as lineage_subgraph_v1, machine_resources as machine_resources_v1, manual_batch_tuning_parameters as manual_batch_tuning_parameters_v1, metadata_service as metadata_service_v1, + metadata_schema as metadata_schema_v1, metadata_store as metadata_store_v1, model as model_v1, model_evaluation as model_evaluation_v1, @@ -173,12 +177,14 @@ io_v1, job_service_v1, job_state_v1, + lineage_subgraph_v1, machine_resources_v1, manual_batch_tuning_parameters_v1, matching_engine_deployed_index_ref_v1, index_v1, index_endpoint_v1, metadata_service_v1, + metadata_schema_v1, metadata_store_v1, model_v1, model_evaluation_v1, @@ -233,12 +239,14 @@ io_v1beta1, job_service_v1beta1, job_state_v1beta1, + lineage_subgraph_v1beta1, machine_resources_v1beta1, manual_batch_tuning_parameters_v1beta1, matching_engine_deployed_index_ref_v1beta1, index_v1beta1, index_endpoint_v1beta1, metadata_service_v1beta1, + metadata_schema_v1beta1, metadata_store_v1beta1, model_v1beta1, model_evaluation_v1beta1, diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py index 3572cf222c..9f0afd9e70 100644 --- a/google/cloud/aiplatform/initializer.py +++ b/google/cloud/aiplatform/initializer.py @@ -33,6 +33,7 @@ from google.cloud.aiplatform import utils from google.cloud.aiplatform.metadata import metadata from google.cloud.aiplatform.utils import resource_manager_utils +from google.cloud.aiplatform.tensorboard import tensorboard_resource from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec_compat, @@ -58,6 +59,9 @@ def init( location: Optional[str] = None, experiment: Optional[str] = None, experiment_description: Optional[str] = None, + experiment_tensorboard: Optional[ + Union[str, tensorboard_resource.Tensorboard] + ] = None, staging_bucket: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, encryption_spec_key_name: Optional[str] = None, @@ -68,8 +72,15 @@ def init( project (str): The default project to use when making API calls. location (str): The default location to use when making API calls. If not set defaults to us-central-1. - experiment (str): The experiment name. - experiment_description (str): The description of the experiment. + experiment (str): Optional. The experiment name. + experiment_description (str): Optional. The description of the experiment. + experiment_tensorboard (Union[str, tensorboard_resource.Tensorboard]): + Optional. The Vertex AI TensorBoard instance, Tensorboard resource name, + or Tensorboard resource ID to use as a backing Tensorboard for the provided + experiment. + + Example tensorboard resource name format: + "projects/123/locations/us-central1/tensorboards/456" staging_bucket (str): The default staging bucket to use to stage artifacts when making API calls. In the form gs://... credentials (google.auth.credentials.Credentials): The default custom @@ -84,15 +95,29 @@ def init( resource is created. If set, this resource and all sub-resources will be secured by this key. + Raises: + ValueError: + If experiment_description is provided but experiment is not. + If experiment_tensorboard is provided but expeirment is not. """ + if experiment_description and experiment is None: + raise ValueError( + "Experiment needs to be set in `init` in order to add experiment descriptions." + ) + + if experiment_tensorboard and experiment is None: + raise ValueError( + "Experiment needs to be set in `init` in order to add experiment_tensorboard." + ) + # reset metadata_service config if project or location is updated. if (project and project != self._project) or ( location and location != self._location ): - if metadata.metadata_service.experiment_name: - logging.info("project/location updated, reset Metadata config.") - metadata.metadata_service.reset() + if metadata._experiment_tracker.experiment_name: + logging.info("project/location updated, reset Experiment config.") + metadata._experiment_tracker.reset() if project: self._project = project @@ -107,12 +132,10 @@ def init( self._encryption_spec_key_name = encryption_spec_key_name if experiment: - metadata.metadata_service.set_experiment( - experiment=experiment, description=experiment_description - ) - if experiment_description and experiment is None: - raise ValueError( - "Experiment name needs to be set in `init` in order to add experiment descriptions." + metadata._experiment_tracker.set_experiment( + experiment=experiment, + description=experiment_description, + backing_tensorboard=experiment_tensorboard, ) def get_encryption_spec( @@ -214,6 +237,11 @@ def encryption_spec_key_name(self) -> Optional[str]: """Default encryption spec key name, if provided.""" return self._encryption_spec_key_name + @property + def experiment_name(self) -> Optional[str]: + """Default experiment name, if provided.""" + return metadata._experiment_tracker.experiment_name + def get_client_options( self, location_override: Optional[str] = None, diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py index 0c5c2e2616..45e20731d5 100644 --- a/google/cloud/aiplatform/metadata/artifact.py +++ b/google/cloud/aiplatform/metadata/artifact.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,24 +15,75 @@ # limitations under the License. # -from typing import Optional, Dict +from typing import Optional, Dict, Union import proto +from google.auth import credentials as auth_credentials + +from google.cloud.aiplatform import base +from google.cloud.aiplatform import models from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.types import artifact as gca_artifact -from google.cloud.aiplatform.compat.types import metadata_service +from google.cloud.aiplatform.compat.types import ( + metadata_service as gca_metadata_service, +) +from google.cloud.aiplatform.metadata import metadata_store from google.cloud.aiplatform.metadata import resource +from google.cloud.aiplatform.metadata import utils as metadata_utils +from google.cloud.aiplatform.utils import rest_utils + +_LOGGER = base.Logger(__name__) -class _Artifact(resource._Resource): + +class Artifact(resource._Resource): """Metadata Artifact resource for Vertex AI""" + def __init__( + self, + artifact_name: str, + *, + metadata_store_id: str = "default", + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Retrieves an existing Metadata Artifact given a resource name or ID. + + Args: + artifact_name (str): + Required. A fully-qualified resource name or resource ID of the Artifact. + Example: "projects/123/locations/us-central1/metadataStores/default/artifacts/my-resource". + or "my-resource" when project and location are initialized or passed. + metadata_store_id (str): + Optional. MetadataStore to retrieve Artifact from. If not set, metadata_store_id is set to "default". + If artifact_name is a fully-qualified resource, its metadata_store_id overrides this one. + project (str): + Optional. Project to retrieve the artifact from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve the Artifact from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Artifact. Overrides + credentials set in aiplatform.init. + """ + + super().__init__( + resource_name=artifact_name, + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, + ) + _resource_noun = "artifacts" _getter_method = "get_artifact" _delete_method = "delete_artifact" _parse_resource_name_method = "parse_artifact_path" _format_resource_name_method = "artifact_path" + _list_method = "list_artifacts" @classmethod def _create_resource( @@ -41,17 +92,21 @@ def _create_resource( parent: str, resource_id: str, schema_title: str, + uri: Optional[str] = None, display_name: Optional[str] = None, schema_version: Optional[str] = None, description: Optional[str] = None, metadata: Optional[Dict] = None, - ) -> proto.Message: + state: gca_artifact.Artifact.State = gca_artifact.Artifact.State.STATE_UNSPECIFIED, + ) -> gca_artifact.Artifact: gapic_artifact = gca_artifact.Artifact( + uri=uri, schema_title=schema_title, schema_version=schema_version, display_name=display_name, description=description, metadata=metadata if metadata else {}, + state=state, ) return client.create_artifact( parent=parent, @@ -59,6 +114,98 @@ def _create_resource( artifact_id=resource_id, ) + @classmethod + def _create( + cls, + resource_id: str, + schema_title: str, + uri: Optional[str] = None, + display_name: Optional[str] = None, + schema_version: Optional[str] = None, + description: Optional[str] = None, + metadata: Optional[Dict] = None, + state: gca_artifact.Artifact.State = gca_artifact.Artifact.State.STATE_UNSPECIFIED, + metadata_store_id: Optional[str] = "default", + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> "Artifact": + """Creates a new Metadata resource. + + Args: + resource_id (str): + Required. The portion of the resource name with + the format: + projects/123/locations/us-central1/metadataStores///. + schema_title (str): + Required. schema_title identifies the schema title used by the resource. + display_name (str): + Optional. The user-defined name of the resource. + schema_version (str): + Optional. schema_version specifies the version used by the resource. + If not set, defaults to use the latest version. + description (str): + Optional. Describes the purpose of the resource to be created. + metadata (Dict): + Optional. Contains the metadata information that will be stored in the resource. + state (google.cloud.gapic.types.Artifact.State): + Optional. The state of this Artifact. This is a + property of the Artifact, and does not imply or + capture any ongoing process. This property is + managed by clients (such as Vertex AI + Pipelines), and the system does not prescribe or + check the validity of state transitions. + metadata_store_id (str): + The portion of the resource name with + the format: + projects/123/locations/us-central1/metadataStores/// + If not provided, the MetadataStore's ID will be set to "default". + project (str): + Project used to create this resource. Overrides project set in + aiplatform.init. + location (str): + Location used to create this resource. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Custom credentials used to create this resource. Overrides + credentials set in aiplatform.init. + + Returns: + resource (_Resource): + Instantiated representation of the managed Metadata resource. + + """ + api_client = cls._instantiate_client(location=location, credentials=credentials) + + parent = utils.full_resource_name( + resource_name=metadata_store_id, + resource_noun=metadata_store._MetadataStore._resource_noun, + parse_resource_name_method=metadata_store._MetadataStore._parse_resource_name, + format_resource_name_method=metadata_store._MetadataStore._format_resource_name, + project=project, + location=location, + ) + + resource = cls._create_resource( + client=api_client, + parent=parent, + resource_id=resource_id, + schema_title=schema_title, + uri=uri, + display_name=display_name, + schema_version=schema_version, + description=description, + metadata=metadata, + state=state, + ) + + self = cls._empty_constructor( + project=project, location=location, credentials=credentials + ) + self._gca_resource = resource + + return self + @classmethod def _update_resource( cls, @@ -93,8 +240,276 @@ def _list_resources( filter (str): Optional. filter string to restrict the list result """ - list_request = metadata_service.ListArtifactsRequest( + list_request = gca_metadata_service.ListArtifactsRequest( parent=parent, filter=filter, ) return client.list_artifacts(request=list_request) + + @classmethod + def create( + cls, + schema_title: str, + *, + resource_id: Optional[str] = None, + uri: Optional[str] = None, + display_name: Optional[str] = None, + schema_version: Optional[str] = None, + description: Optional[str] = None, + metadata: Optional[Dict] = None, + state: gca_artifact.Artifact.State = gca_artifact.Artifact.State.STATE_UNSPECIFIED, + metadata_store_id: Optional[str] = "default", + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> "Artifact": + """Creates a new Metadata Artifact. + + Args: + schema_title (str): + Required. schema_title identifies the schema title used by the Artifact. + + Please reference https://cloud.google.com/vertex-ai/docs/ml-metadata/system-schemas. + resource_id (str): + Optional. The portion of the Artifact name with + the format. This is globally unique in a metadataStore: + projects/123/locations/us-central1/metadataStores//artifacts/. + uri (str): + Optional. The uniform resource identifier of the artifact file. May be empty if there is no actual + artifact file. + display_name (str): + Optional. The user-defined name of the Artifact. + schema_version (str): + Optional. schema_version specifies the version used by the Artifact. + If not set, defaults to use the latest version. + description (str): + Optional. Describes the purpose of the Artifact to be created. + metadata (Dict): + Optional. Contains the metadata information that will be stored in the Artifact. + state (google.cloud.gapic.types.Artifact.State): + Optional. The state of this Artifact. This is a + property of the Artifact, and does not imply or + capture any ongoing process. This property is + managed by clients (such as Vertex AI + Pipelines), and the system does not prescribe or + check the validity of state transitions. + metadata_store_id (str): + Optional. The portion of the resource name with + the format: + projects/123/locations/us-central1/metadataStores//artifacts/ + If not provided, the MetadataStore's ID will be set to "default". + project (str): + Optional. Project used to create this Artifact. Overrides project set in + aiplatform.init. + location (str): + Optional. Location used to create this Artifact. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to create this Artifact. Overrides + credentials set in aiplatform.init. + + Returns: + Artifact: Instantiated representation of the managed Metadata Artifact. + """ + return cls._create( + resource_id=resource_id, + schema_title=schema_title, + uri=uri, + display_name=display_name, + schema_version=schema_version, + description=description, + metadata=metadata, + state=state, + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, + ) + + @property + def uri(self) -> Optional[str]: + "Uri for this Artifact." + return self.gca_resource.uri + + @classmethod + def get_with_uri( + cls, + uri: str, + *, + metadata_store_id: Optional[str] = "default", + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> "Artifact": + """Get an Artifact by it's uri. + + If more than one Artifact with this uri is in the metadata store then the Artifact with the latest + create_time is returned. + + Args: + uri(str): + Required. Uri of the Artifact to retrieve. + metadata_store_id (str): + Optional. MetadataStore to retrieve Artifact from. If not set, metadata_store_id is set to "default". + If artifact_name is a fully-qualified resource, its metadata_store_id overrides this one. + project (str): + Optional. Project to retrieve the artifact from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve the Artifact from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Artifact. Overrides + credentials set in aiplatform.init. + Returns: + Artifact: Artifact with given uri. + Raises: + ValueError: If no Artifact exists with the provided uri. + + """ + + matched_artifacts = cls.list( + filter=f'uri = "{uri}"', + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, + ) + + if not matched_artifacts: + raise ValueError( + f"No artifact with uri {uri} is in the `{metadata_store_id}` MetadataStore." + ) + + if len(matched_artifacts) > 1: + matched_artifacts.sort(key=lambda a: a.create_time, reverse=True) + resource_names = "\n".join(a.resource_name for a in matched_artifacts) + _LOGGER.warn( + f"Mutiple artifacts with uri {uri} were found: {resource_names}" + ) + _LOGGER.warn(f"Returning {matched_artifacts[0].resource_name}") + + return matched_artifacts[0] + + @property + def lineage_console_uri(self) -> str: + """Cloud console uri to view this Artifact Lineage.""" + metadata_store = self._parse_resource_name(self.resource_name)["metadata_store"] + return f"https://console.cloud.google.com/vertex-ai/locations/{self.location}/metadata-stores/{metadata_store}/artifacts/{self.name}?project={self.project}" + + def __repr__(self) -> str: + if self._gca_resource: + return f"{object.__repr__(self)} \nresource name: {self.resource_name}\nuri: {self.uri}\nschema_title:{self.gca_resource.schema_title}" + + return base.FutureManager.__repr__(self) + + +class _VertexResourceArtifactResolver: + + # TODO(b/235594717) Add support for managed datasets + _resource_to_artifact_type = {models.Model: "google.VertexModel"} + + @classmethod + def supports_metadata(cls, resource: base.VertexAiResourceNoun) -> bool: + """Returns True if Vertex resource is supported in Vertex Metadata otherwise False. + + Args: + resource (base.VertexAiResourceNoun): + Requried. Instance of Vertex AI Resource. + Returns: + True if Vertex resource is supported in Vertex Metadata otherwise False. + """ + return type(resource) in cls._resource_to_artifact_type + + @classmethod + def validate_resource_supports_metadata(cls, resource: base.VertexAiResourceNoun): + """Validates Vertex resource is supported in Vertex Metadata. + + Args: + resource (base.VertexAiResourceNoun): + Required. Instance of Vertex AI Resource. + Raises: + ValueError: If Vertex AI Resource is not support in Vertex Metadata. + """ + if not cls.supports_metadata(resource): + raise ValueError( + f"Vertex {type(resource)} is not yet supported in Vertex Metadata." + f"Only {list(cls._resource_to_artifact_type.keys())} are supported" + ) + + @classmethod + def resolve_vertex_resource( + cls, resource: Union[models.Model] + ) -> Optional[Artifact]: + """Resolves Vertex Metadata Artifact that represents this Vertex Resource. + + If there are multiple Artifacts in the metadata store that represent the provided resource. The one with the + latest create_time is returned. + + Args: + resource (base.VertexAiResourceNoun): + Required. Instance of Vertex AI Resource. + Returns: + Artifact: Artifact that represents this Vertex Resource. None if Resource not found in Metadata store. + """ + cls.validate_resource_supports_metadata(resource) + resource.wait() + metadata_type = cls._resource_to_artifact_type[type(resource)] + uri = rest_utils.make_gcp_resource_rest_url(resource=resource) + + artifacts = Artifact.list( + filter=metadata_utils._make_filter_string( + schema_title=metadata_type, + uri=uri, + ), + project=resource.project, + location=resource.location, + credentials=resource.credentials, + ) + + artifacts.sort(key=lambda a: a.create_time, reverse=True) + if artifacts: + # most recent + return artifacts[0] + + @classmethod + def create_vertex_resource_artifact(cls, resource: Union[models.Model]) -> Artifact: + """Creates Vertex Metadata Artifact that represents this Vertex Resource. + + Args: + resource (base.VertexAiResourceNoun): + Required. Instance of Vertex AI Resource. + Returns: + Artifact: Artifact that represents this Vertex Resource. + """ + cls.validate_resource_supports_metadata(resource) + resource.wait() + metadata_type = cls._resource_to_artifact_type[type(resource)] + uri = rest_utils.make_gcp_resource_rest_url(resource=resource) + + return Artifact.create( + schema_title=metadata_type, + display_name=getattr(resource.gca_resource, "display_name", None), + uri=uri, + metadata={"resourceName": resource.resource_name}, + project=resource.project, + location=resource.location, + credentials=resource.credentials, + ) + + @classmethod + def resolve_or_create_resource_artifact( + cls, resource: Union[models.Model] + ) -> Artifact: + """Create of gets Vertex Metadata Artifact that represents this Vertex Resource. + + Args: + resource (base.VertexAiResourceNoun): + Required. Instance of Vertex AI Resource. + Returns: + Artifact: Artifact that represents this Vertex Resource. + """ + artifact = cls.resolve_vertex_resource(resource=resource) + if artifact: + return artifact + return cls.create_vertex_resource_artifact(resource=resource) diff --git a/google/cloud/aiplatform/metadata/constants.py b/google/cloud/aiplatform/metadata/constants.py index 62e7d6e075..8776aef3b1 100644 --- a/google/cloud/aiplatform/metadata/constants.py +++ b/google/cloud/aiplatform/metadata/constants.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,21 +14,54 @@ # See the License for the specific language governing permissions and # limitations under the License. # +from google.cloud.aiplatform.compat.types import artifact SYSTEM_RUN = "system.Run" SYSTEM_EXPERIMENT = "system.Experiment" +SYSTEM_EXPERIMENT_RUN = "system.ExperimentRun" SYSTEM_PIPELINE = "system.Pipeline" +SYSTEM_PIPELINE_RUN = "system.PipelineRun" SYSTEM_METRICS = "system.Metrics" +_EXPERIMENTS_V2_TENSORBOARD_RUN = "google.VertexTensorboardRun" + _DEFAULT_SCHEMA_VERSION = "0.0.1" SCHEMA_VERSIONS = { SYSTEM_RUN: _DEFAULT_SCHEMA_VERSION, SYSTEM_EXPERIMENT: _DEFAULT_SCHEMA_VERSION, + SYSTEM_EXPERIMENT_RUN: _DEFAULT_SCHEMA_VERSION, SYSTEM_PIPELINE: _DEFAULT_SCHEMA_VERSION, SYSTEM_METRICS: _DEFAULT_SCHEMA_VERSION, } -# The EXPERIMENT_METADATA is needed until we support context deletion in backend service. -# TODO: delete EXPERIMENT_METADATA once backend supports context deletion. +_BACKING_TENSORBOARD_RESOURCE_KEY = "backing_tensorboard_resource" + + +_PARAM_KEY = "_params" +_METRIC_KEY = "_metrics" +_STATE_KEY = "_state" + +_PARAM_PREFIX = "param" +_METRIC_PREFIX = "metric" +_TIME_SERIES_METRIC_PREFIX = "time_series_metric" + +# This is currently used to filter in the Console. EXPERIMENT_METADATA = {"experiment_deleted": False} + +PIPELINE_PARAM_PREFIX = "input:" + +TENSORBOARD_CUSTOM_JOB_EXPERIMENT_FIELD = "tensorboard_link" + +GCP_ARTIFACT_RESOURCE_NAME_KEY = "resourceName" + +# constant to mark an Experiment context as originating from the SDK +# TODO(b/235593750) Remove this field +_VERTEX_EXPERIMENT_TRACKING_LABEL = "vertex_experiment_tracking" + +_TENSORBOARD_RUN_REFERENCE_ARTIFACT = artifact.Artifact( + name="google-vertex-tensorboard-run-v0-0-1", + schema_title=_EXPERIMENTS_V2_TENSORBOARD_RUN, + schema_version="0.0.1", + metadata={_VERTEX_EXPERIMENT_TRACKING_LABEL: True}, +) diff --git a/google/cloud/aiplatform/metadata/context.py b/google/cloud/aiplatform/metadata/context.py index 7687340a1c..d072a6e047 100644 --- a/google/cloud/aiplatform/metadata/context.py +++ b/google/cloud/aiplatform/metadata/context.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,13 +15,22 @@ # limitations under the License. # -from typing import Optional, Dict, Sequence +from typing import Optional, Dict, List, Sequence import proto +from google.cloud.aiplatform import base from google.cloud.aiplatform import utils -from google.cloud.aiplatform.compat.types import metadata_service +from google.cloud.aiplatform.metadata import utils as metadata_utils from google.cloud.aiplatform.compat.types import context as gca_context +from google.cloud.aiplatform.compat.types import ( + lineage_subgraph as gca_lineage_subgraph, +) +from google.cloud.aiplatform.compat.types import ( + metadata_service as gca_metadata_service, +) +from google.cloud.aiplatform.metadata import artifact +from google.cloud.aiplatform.metadata import execution from google.cloud.aiplatform.metadata import resource @@ -33,6 +42,12 @@ class _Context(resource._Resource): _delete_method = "delete_context" _parse_resource_name_method = "parse_context_path" _format_resource_name_method = "context_path" + _list_method = "list_contexts" + + @property + def parent_contexts(self) -> Sequence[str]: + """The parent context resource names of this context.""" + return self.gca_resource.parent_contexts def add_artifacts_and_executions( self, @@ -53,6 +68,19 @@ def add_artifacts_and_executions( executions=execution_resource_names, ) + def get_artifacts(self) -> List[artifact.Artifact]: + """Returns all Artifact attributed to this Context. + + Returns: + artifacts(List[Artifacts]): All Artifacts under this context. + """ + return artifact.Artifact.list( + filter=metadata_utils._make_filter_string(in_context=[self.resource_name]), + project=self.project, + location=self.location, + credentials=self.credentials, + ) + @classmethod def _create_resource( cls, @@ -113,8 +141,43 @@ def _list_resources( Optional. filter string to restrict the list result """ - list_request = metadata_service.ListContextsRequest( + list_request = gca_metadata_service.ListContextsRequest( parent=parent, filter=filter, ) return client.list_contexts(request=list_request) + + def add_context_children(self, contexts: List["_Context"]): + """Adds the provided contexts as children of this context. + + Args: + contexts (List[_Context]): Contexts to add as children. + """ + self.api_client.add_context_children( + context=self.resource_name, + child_contexts=[c.resource_name for c in contexts], + ) + + def query_lineage_subgraph(self) -> gca_lineage_subgraph.LineageSubgraph: + """Queries lineage subgraph of this context. + + Returns: + lineage subgraph(gca_lineage_subgraph.LineageSubgraph): Lineage subgraph of this Context. + """ + + return self.api_client.query_context_lineage_subgraph( + context=self.resource_name, retry=base._DEFAULT_RETRY + ) + + def get_executions(self) -> List[execution.Execution]: + """Returns Executions associated to this context. + + Returns: + executions (List[Executions]): Executions associated to this context. + """ + return execution.Execution.list( + filter=metadata_utils._make_filter_string(in_context=[self.resource_name]), + project=self.project, + location=self.location, + credentials=self.credentials, + ) diff --git a/google/cloud/aiplatform/metadata/execution.py b/google/cloud/aiplatform/metadata/execution.py index 2a21cce60d..9a85bce36f 100644 --- a/google/cloud/aiplatform/metadata/execution.py +++ b/google/cloud/aiplatform/metadata/execution.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,21 +14,26 @@ # See the License for the specific language governing permissions and # limitations under the License. # - -from typing import Optional, Dict, Sequence +from copy import deepcopy +from typing import Any, Dict, List, Optional, Union import proto -from google.api_core import exceptions +from google.auth import credentials as auth_credentials +from google.cloud.aiplatform import base +from google.cloud.aiplatform import models from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.types import event as gca_event from google.cloud.aiplatform.compat.types import execution as gca_execution -from google.cloud.aiplatform.compat.types import metadata_service +from google.cloud.aiplatform.compat.types import ( + metadata_service as gca_metadata_service, +) from google.cloud.aiplatform.metadata import artifact +from google.cloud.aiplatform.metadata import metadata_store from google.cloud.aiplatform.metadata import resource -class _Execution(resource._Resource): +class Execution(resource._Resource): """Metadata Execution resource for Vertex AI""" _resource_noun = "executions" @@ -36,73 +41,316 @@ class _Execution(resource._Resource): _delete_method = "delete_execution" _parse_resource_name_method = "parse_execution_path" _format_resource_name_method = "execution_path" + _list_method = "list_executions" + + def __init__( + self, + execution_name: str, + *, + metadata_store_id: str = "default", + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Retrieves an existing Metadata Execution given a resource name or ID. + + Args: + execution_name (str): + Required. A fully-qualified resource name or resource ID of the Execution. + Example: "projects/123/locations/us-central1/metadataStores/default/executions/my-resource". + or "my-resource" when project and location are initialized or passed. + metadata_store_id (str): + Optional. MetadataStore to retrieve Execution from. If not set, metadata_store_id is set to "default". + If execution_name is a fully-qualified resource, its metadata_store_id overrides this one. + project (str): + Optional. Project to retrieve the artifact from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve the Execution from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Execution. Overrides + credentials set in aiplatform.init. + """ + + super().__init__( + resource_name=execution_name, + metadata_store_id=metadata_store_id, + project=project, + location=location, + credentials=credentials, + ) + + @property + def state(self) -> gca_execution.Execution.State: + """State of this Execution.""" + return self._gca_resource.state + + @classmethod + def create( + cls, + schema_title: str, + *, + state: gca_execution.Execution.State = gca_execution.Execution.State.RUNNING, + resource_id: Optional[str] = None, + display_name: Optional[str] = None, + schema_version: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + description: Optional[str] = None, + metadata_store_id: str = "default", + project: Optional[str] = None, + location: Optional[str] = None, + credentials=Optional[auth_credentials.Credentials], + ) -> "Execution": + """ + Creates a new Metadata Execution. + + Args: + schema_title (str): + Required. schema_title identifies the schema title used by the Execution. + state (gca_execution.Execution.State.RUNNING): + Optional. State of this Execution. Defaults to RUNNING. + resource_id (str): + Optional. The portion of the Execution name with + the format. This is globally unique in a metadataStore: + projects/123/locations/us-central1/metadataStores//executions/. + display_name (str): + Optional. The user-defined name of the Execution. + schema_version (str): + Optional. schema_version specifies the version used by the Execution. + If not set, defaults to use the latest version. + metadata (Dict): + Optional. Contains the metadata information that will be stored in the Execution. + description (str): + Optional. Describes the purpose of the Execution to be created. + metadata_store_id (str): + Optional. The portion of the resource name with + the format: + projects/123/locations/us-central1/metadataStores//artifacts/ + If not provided, the MetadataStore's ID will be set to "default". + project (str): + Optional. Project used to create this Execution. Overrides project set in + aiplatform.init. + location (str): + Optional. Location used to create this Execution. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to create this Execution. Overrides + credentials set in aiplatform.init. + + Returns: + Execution: Instantiated representation of the managed Metadata Execution. + + """ + self = cls._empty_constructor( + project=project, location=location, credentials=credentials + ) + super(base.VertexAiResourceNounWithFutureManager, self).__init__() + + resource = Execution._create_resource( + client=self.api_client, + parent=metadata_store._MetadataStore._format_resource_name( + project=self.project, + location=self.location, + metadata_store=metadata_store_id, + ), + schema_title=schema_title, + resource_id=resource_id, + metadata=metadata, + description=description, + display_name=display_name, + schema_version=schema_version, + state=state, + ) + self._gca_resource = resource + + return self + + def __enter__(self): + if self.state is not gca_execution.Execution.State.RUNNING: + self.update(state=gca_execution.Execution.State.RUNNING) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + state = ( + gca_execution.Execution.State.FAILED + if exc_type + else gca_execution.Execution.State.COMPLETE + ) + self.update(state=state) + + def assign_input_artifacts( + self, artifacts: List[Union[artifact.Artifact, models.Model]] + ): + """Assigns Artifacts as inputs to this Executions. + + Args: + artifacts (List[Union[artifact.Artifact, models.Model]]): + Required. Artifacts to assign as input. + """ + self._add_artifact(artifacts=artifacts, input=True) - def add_artifact( + def assign_output_artifacts( + self, artifacts: List[Union[artifact.Artifact, models.Model]] + ): + """Assigns Artifacts as outputs to this Executions. + + Args: + artifacts (List[Union[artifact.Artifact, models.Model]]): + Required. Artifacts to assign as input. + """ + self._add_artifact(artifacts=artifacts, input=False) + + def _add_artifact( self, - artifact_resource_name: str, + artifacts: List[Union[artifact.Artifact, models.Model]], input: bool, ): """Connect Artifact to a given Execution. Args: - artifact_resource_name (str): + artifact_resource_names (List[str]): Required. The full resource name of the Artifact to connect to the Execution through an Event. input (bool) Required. Whether Artifact is an input event to the Execution or not. """ - event = gca_event.Event( - artifact=artifact_resource_name, - type_=gca_event.Event.Type.INPUT if input else gca_event.Event.Type.OUTPUT, - ) + artifact_resource_names = [] + for a in artifacts: + if isinstance(a, artifact.Artifact): + artifact_resource_names.append(a.resource_name) + else: + artifact_resource_names.append( + artifact._VertexResourceArtifactResolver.resolve_or_create_resource_artifact( + a + ).resource_name + ) + + events = [ + gca_event.Event( + artifact=artifact_resource_name, + type_=gca_event.Event.Type.INPUT + if input + else gca_event.Event.Type.OUTPUT, + ) + for artifact_resource_name in artifact_resource_names + ] self.api_client.add_execution_events( execution=self.resource_name, - events=[event], + events=events, ) - def query_input_and_output_artifacts(self) -> Sequence[artifact._Artifact]: - """query the input and output artifacts connected to the execution. + def _get_artifacts( + self, event_type: gca_event.Event.Type + ) -> List[artifact.Artifact]: + """Get Executions input or output Artifacts. + Args: + event_type (gca_event.Event.Type): + Required. The Event type, input or output. Returns: - A Sequence of _Artifacts + List of Artifacts. """ + subgraph = self.api_client.query_execution_inputs_and_outputs( + execution=self.resource_name + ) - try: - artifacts = self.api_client.query_execution_inputs_and_outputs( - execution=self.resource_name - ).artifacts - except exceptions.NotFound: - return [] + artifact_map = { + artifact_metadata.name: artifact_metadata + for artifact_metadata in subgraph.artifacts + } - return [ - artifact._Artifact( - resource=metadata_artifact, + gca_artifacts = [ + artifact_map[event.artifact] + for event in subgraph.events + if event.type_ == event_type + ] + + artifacts = [] + for gca_artifact in gca_artifacts: + this_artifact = artifact.Artifact._empty_constructor( project=self.project, location=self.location, credentials=self.credentials, ) - for metadata_artifact in artifacts - ] + this_artifact._gca_resource = gca_artifact + artifacts.append(this_artifact) + + return artifacts + + def get_input_artifacts(self) -> List[artifact.Artifact]: + """Get the input Artifacts of this Execution. + + Returns: + List of input Artifacts. + """ + return self._get_artifacts(event_type=gca_event.Event.Type.INPUT) + + def get_output_artifacts(self) -> List[artifact.Artifact]: + """Get the output Artifacts of this Execution. + + Returns: + List of output Artifacts. + """ + return self._get_artifacts(event_type=gca_event.Event.Type.OUTPUT) @classmethod def _create_resource( cls, client: utils.MetadataClientWithOverride, parent: str, - resource_id: str, schema_title: str, + state: gca_execution.Execution.State = gca_execution.Execution.State.RUNNING, + resource_id: Optional[str] = None, display_name: Optional[str] = None, schema_version: Optional[str] = None, description: Optional[str] = None, metadata: Optional[Dict] = None, - ) -> proto.Message: + ) -> gca_execution.Execution: + """ + Creates a new Metadata Execution. + + Args: + client (utils.MetadataClientWithOverride): + Required. Instantiated Metadata Service Client. + parent (str): + Required: MetadataStore parent in which to create this Execution. + schema_title (str): + Required. schema_title identifies the schema title used by the Execution. + state (gca_execution.Execution.State): + Optional. State of this Execution. Defaults to RUNNING. + resource_id (str): + Optional. The {execution} portion of the resource name with the + format: + ``projects/{project}/locations/{location}/metadataStores/{metadatastore}/executions/{execution}`` + If not provided, the Execution's ID will be a UUID generated + by the service. Must be 4-128 characters in length. Valid + characters are ``/[a-z][0-9]-/``. Must be unique across all + Executions in the parent MetadataStore. (Otherwise the + request will fail with ALREADY_EXISTS, or PERMISSION_DENIED + if the caller can't view the preexisting Execution.) + display_name (str): + Optional. The user-defined name of the Execution. + schema_version (str): + Optional. schema_version specifies the version used by the Execution. + If not set, defaults to use the latest version. + description (str): + Optional. Describes the purpose of the Execution to be created. + metadata (Dict): + Optional. Contains the metadata information that will be stored in the Execution. + + Returns: + Execution: Instantiated representation of the managed Metadata Execution. + + """ gapic_execution = gca_execution.Execution( schema_title=schema_title, schema_version=schema_version, display_name=display_name, description=description, metadata=metadata if metadata else {}, + state=state, ) return client.create_execution( parent=parent, @@ -128,7 +376,7 @@ def _list_resources( Optional. filter string to restrict the list result """ - list_request = metadata_service.ListExecutionsRequest( + list_request = gca_metadata_service.ListExecutionsRequest( parent=parent, filter=filter, ) @@ -150,3 +398,30 @@ def _update_resource( """ return client.update_execution(execution=resource) + + def update( + self, + state: Optional[gca_execution.Execution.State] = None, + description: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ): + """Update this Execution. + + Args: + state (gca_execution.Execution.State): + Optional. State of this Execution. + description (str): + Optional. Describes the purpose of the Execution to be created. + metadata (Dict[str, Any): + Optional. Contains the metadata information that will be stored in the Execution. + """ + + gca_resource = deepcopy(self._gca_resource) + if state: + gca_resource.state = state + if description: + gca_resource.description = description + self._nested_update_metadata(gca_resource=gca_resource, metadata=metadata) + self._gca_resource = self._update_resource( + self.api_client, resource=gca_resource + ) diff --git a/google/cloud/aiplatform/metadata/experiment_resources.py b/google/cloud/aiplatform/metadata/experiment_resources.py new file mode 100644 index 0000000000..353f05aa70 --- /dev/null +++ b/google/cloud/aiplatform/metadata/experiment_resources.py @@ -0,0 +1,721 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 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. +# + +import abc +from dataclasses import dataclass +import logging +from typing import Dict, List, NamedTuple, Optional, Union, Tuple, Type + +from google.auth import credentials as auth_credentials + +from google.cloud.aiplatform import base +from google.cloud.aiplatform.metadata import artifact +from google.cloud.aiplatform.metadata import constants +from google.cloud.aiplatform.metadata import context +from google.cloud.aiplatform.metadata import execution +from google.cloud.aiplatform.metadata import metadata +from google.cloud.aiplatform.metadata import metadata_store +from google.cloud.aiplatform.metadata import resource +from google.cloud.aiplatform.metadata import utils as metadata_utils +from google.cloud.aiplatform.tensorboard import tensorboard_resource + +_LOGGER = base.Logger(__name__) + + +@dataclass +class _ExperimentRow: + """Class for representing a run row in an Experiments Dataframe. + + Attributes: + params (Dict[str, Union[float, int, str]]): Optional. The parameters of this run. + metrics (Dict[str, Union[float, int, str]]): Optional. The metrics of this run. + time_series_metrics (Dict[str, float]): Optional. The latest time series metrics of this run. + experiment_run_type (Optional[str]): Optional. The type of this run. + name (str): Optional. The name of this run. + state (str): Optional. The state of this run. + """ + + params: Optional[Dict[str, Union[float, int, str]]] = None + metrics: Optional[Dict[str, Union[float, int, str]]] = None + time_series_metrics: Optional[Dict[str, float]] = None + experiment_run_type: Optional[str] = None + name: Optional[str] = None + state: Optional[str] = None + + def to_dict(self) -> Dict[str, Union[float, int, str]]: + """Converts this experiment row into a dictionary. + + Returns: + Row as a dictionary. + """ + result = { + "run_type": self.experiment_run_type, + "run_name": self.name, + "state": self.state, + } + for prefix, field in [ + (constants._PARAM_PREFIX, self.params), + (constants._METRIC_PREFIX, self.metrics), + (constants._TIME_SERIES_METRIC_PREFIX, self.time_series_metrics), + ]: + if field: + result.update( + {f"{prefix}.{key}": value for key, value in field.items()} + ) + return result + + +class Experiment: + """Represents a Vertex AI Experiment resource.""" + + def __init__( + self, + experiment_name: str, + *, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """ + + ``` + my_experiment = aiplatform.Experiment('my-experiment') + ``` + + Args: + experiment_name (str): Required. The name or resource name of this experiment. + + Resource name is of the format: projects/123/locations/us-central1/experiments/my-experiment + project (str): + Optional. Project where this experiment is located. Overrides project set in + aiplatform.init. + location (str): + Optional. Location where this experiment is located. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to retrieve this experiment. Overrides + credentials set in aiplatform.init. + """ + + metadata_args = dict( + resource_name=experiment_name, + project=project, + location=location, + credentials=credentials, + ) + + with _SetLoggerLevel(resource): + experiment_context = context._Context(**metadata_args) + self._validate_experiment_context(experiment_context) + + self._metadata_context = experiment_context + + @staticmethod + def _validate_experiment_context(experiment_context: context._Context): + """Validates this context is an experiment context. + + Args: + experiment_context (context._Context): Metadata context. + Raises: + ValueError: If Metadata context is not an experiment context or a TensorboardExperiment. + """ + if experiment_context.schema_title != constants.SYSTEM_EXPERIMENT: + raise ValueError( + f"Experiment name {experiment_context.name} is of type " + f"({experiment_context.schema_title}) in this MetadataStore. " + f"It must of type {constants.SYSTEM_EXPERIMENT}." + ) + if Experiment._is_tensorboard_experiment(experiment_context): + raise ValueError( + f"Experiment name {experiment_context.name} is a TensorboardExperiment context " + f"and cannot be used as a Vertex AI Experiment." + ) + + @staticmethod + def _is_tensorboard_experiment(context: context._Context) -> bool: + """Returns True if Experiment is a Tensorboard Experiment created by CustomJob.""" + return constants.TENSORBOARD_CUSTOM_JOB_EXPERIMENT_FIELD in context.metadata + + @property + def name(self) -> str: + """The name of this experiment.""" + return self._metadata_context.name + + @classmethod + def create( + cls, + experiment_name: str, + *, + description: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> "Experiment": + """Creates a new experiment in Vertex AI Experiments. + + ``` + my_experiment = aiplatform.Experiment.create('my-experiment', description='my description') + ``` + + Args: + experiment_name (str): Required. The name of this experiment. + description (str): Optional. Describes this experiment's purpose. + project (str): + Optional. Project where this experiment will be created. Overrides project set in + aiplatform.init. + location (str): + Optional. Location where this experiment will be created. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to create this experiment. Overrides + credentials set in aiplatform.init. + Returns: + The newly created experiment. + """ + + metadata_store._MetadataStore.ensure_default_metadata_store_exists( + project=project, location=location, credentials=credentials + ) + + with _SetLoggerLevel(resource): + experiment_context = context._Context._create( + resource_id=experiment_name, + display_name=experiment_name, + description=description, + schema_title=constants.SYSTEM_EXPERIMENT, + schema_version=metadata._get_experiment_schema_version(), + metadata=constants.EXPERIMENT_METADATA, + project=project, + location=location, + credentials=credentials, + ) + + self = cls.__new__() + self._metadata_context = experiment_context + + return self + + @classmethod + def get_or_create( + cls, + experiment_name: str, + *, + description: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> "Experiment": + """Gets experiment if one exists with this experiment_name in Vertex AI Experiments. + + Otherwise creates this experiment. + + ``` + my_experiment = aiplatform.Experiment.get_or_create('my-experiment', description='my description') + ``` + + Args: + experiment_name (str): Required. The name of this experiment. + description (str): Optional. Describes this experiment's purpose. + project (str): + Optional. Project where this experiment will be retrieved from or created. Overrides project set in + aiplatform.init. + location (str): + Optional. Location where this experiment will be retrieved from or created. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to retrieve or create this experiment. Overrides + credentials set in aiplatform.init. + Returns: + Vertex AI experiment. + """ + + metadata_store._MetadataStore.ensure_default_metadata_store_exists( + project=project, location=location, credentials=credentials + ) + + with _SetLoggerLevel(resource): + experiment_context = context._Context.get_or_create( + resource_id=experiment_name, + display_name=experiment_name, + description=description, + schema_title=constants.SYSTEM_EXPERIMENT, + schema_version=metadata._get_experiment_schema_version(), + metadata=constants.EXPERIMENT_METADATA, + project=project, + location=location, + credentials=credentials, + ) + + cls._validate_experiment_context(experiment_context) + + if description and description != experiment_context.description: + experiment_context.update(description=description) + + self = cls.__new__(cls) + self._metadata_context = experiment_context + + return self + + @classmethod + def list( + cls, + *, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> List["Experiment"]: + """List all Vertex AI Experiments in the given project. + + ``` + my_experiments = aiplatform.Experiment.list() + ``` + + Args: + project (str): + Optional. Project to list these experiments from. Overrides project set in + aiplatform.init. + location (str): + Optional. Location to list these experiments from. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to list these experiments. Overrides + credentials set in aiplatform.init. + Returns: + List of Vertex AI experiments. + """ + + filter_str = metadata_utils._make_filter_string( + schema_title=constants.SYSTEM_EXPERIMENT + ) + + with _SetLoggerLevel(resource): + experiment_contexts = context._Context.list( + filter=filter_str, + project=project, + location=location, + credentials=credentials, + ) + + experiments = [] + for experiment_context in experiment_contexts: + # Filters Tensorboard Experiments + if not cls._is_tensorboard_experiment(experiment_context): + experiment = cls.__new__(cls) + experiment._metadata_context = experiment_context + experiments.append(experiment) + return experiments + + @property + def resource_name(self) -> str: + """The Metadata context resource name of this experiment.""" + return self._metadata_context.resource_name + + def delete(self, *, delete_backing_tensorboard_runs: bool = False): + """Deletes this experiment all the experiment runs under this experiment + + Does not delete Pipeline runs, Artifacts, or Executions associated to this experiment + or experiment runs in this experiment. + + ``` + my_experiment = aiplatform.Experiment('my-experiment') + my_experiment.delete(delete_backing_tensorboard_runs=True) + ``` + + Args: + delete_backing_tensorboard_runs (bool): + Optional. If True will also delete the Tensorboard Runs associated to the experiment + runs under this experiment that we used to store time series metrics. + """ + + experiment_runs = _SUPPORTED_LOGGABLE_RESOURCES[context._Context][ + constants.SYSTEM_EXPERIMENT_RUN + ].list(experiment=self) + for experiment_run in experiment_runs: + experiment_run.delete( + delete_backing_tensorboard_run=delete_backing_tensorboard_runs + ) + self._metadata_context.delete() + + def get_data_frame(self) -> "pd.DataFrame": # noqa: F821 + """Get parameters, metrics, and time series metrics of all runs in this experiment as Dataframe. + + ``` + my_experiment = aiplatform.Experiment('my-experiment') + df = my_experiment.get_data_frame() + ``` + + Returns: + pd.DataFrame: Pandas Dataframe of Experiment Runs. + + Raises: + ImportError: If pandas is not installed. + """ + try: + import pandas as pd + except ImportError: + raise ImportError( + "Pandas is not installed and is required to get dataframe as the return format. " + 'Please install the SDK using "pip install python-aiplatform[metadata]"' + ) + + service_request_args = dict( + project=self._metadata_context.project, + location=self._metadata_context.location, + credentials=self._metadata_context.credentials, + ) + + filter_str = metadata_utils._make_filter_string( + schema_title=sorted( + list(_SUPPORTED_LOGGABLE_RESOURCES[context._Context].keys()) + ), + parent_contexts=[self._metadata_context.resource_name], + ) + contexts = context._Context.list(filter_str, **service_request_args) + + filter_str = metadata_utils._make_filter_string( + schema_title=list( + _SUPPORTED_LOGGABLE_RESOURCES[execution.Execution].keys() + ), + in_context=[self._metadata_context.resource_name], + ) + + executions = execution.Execution.list(filter_str, **service_request_args) + + rows = [] + for metadata_context in contexts: + row_dict = ( + _SUPPORTED_LOGGABLE_RESOURCES[context._Context][ + metadata_context.schema_title + ] + ._query_experiment_row(metadata_context) + .to_dict() + ) + row_dict.update({"experiment_name": self.name}) + rows.append(row_dict) + + # backward compatibility + for metadata_execution in executions: + row_dict = ( + _SUPPORTED_LOGGABLE_RESOURCES[execution.Execution][ + metadata_execution.schema_title + ] + ._query_experiment_row(metadata_execution) + .to_dict() + ) + row_dict.update({"experiment_name": self.name}) + rows.append(row_dict) + + df = pd.DataFrame(rows) + + column_name_sort_map = { + "experiment_name": -1, + "run_name": 1, + "run_type": 2, + "state": 3, + } + + def column_sort_key(key: str) -> int: + """Helper method to reorder columns.""" + order = column_name_sort_map.get(key) + if order: + return order + elif key.startswith("param"): + return 5 + elif key.startswith("metric"): + return 6 + else: + return 7 + + columns = df.columns + columns = sorted(columns, key=column_sort_key) + df = df.reindex(columns, axis=1) + + return df + + def _lookup_backing_tensorboard(self) -> Optional[tensorboard_resource.Tensorboard]: + """Returns backing tensorboard if one is set. + + Returns: + Tensorboard resource if one exists. + """ + tensorboard_resource_name = self._metadata_context.metadata.get( + constants._BACKING_TENSORBOARD_RESOURCE_KEY + ) + + if not tensorboard_resource_name: + with _SetLoggerLevel(resource): + self._metadata_context.sync_resource() + tensorboard_resource_name = self._metadata_context.metadata.get( + constants._BACKING_TENSORBOARD_RESOURCE_KEY + ) + + if tensorboard_resource_name: + return tensorboard_resource.Tensorboard( + tensorboard_resource_name, + credentials=self._metadata_context.credentials, + ) + + def get_backing_tensorboard_resource( + self, + ) -> Optional[tensorboard_resource.Tensorboard]: + """Get the backing tensorboard for this experiment in one exists. + + ``` + my_experiment = aiplatform.Experiment('my-experiment') + tb = my_experiment.get_backing_tensorboard_resource() + ``` + + Returns: + Backing Tensorboard resource for this experiment if one exists. + """ + return self._lookup_backing_tensorboard() + + def assign_backing_tensorboard( + self, tensorboard: Union[tensorboard_resource.Tensorboard, str] + ): + """Assigns tensorboard as backing tensorboard to support time series metrics logging. + + ``` + tb = aiplatform.Tensorboard('tensorboard-resource-id') + my_experiment = aiplatform.Experiment('my-experiment') + my_experiment.assign_backing_tensorboard(tb) + ``` + + Args: + tensorboard (Union[aiplatform.Tensorboard, str]): + Required. Tensorboard resource or resource name to associate to this experiment. + + Raises: + ValueError: If this experiment already has a previously set backing tensorboard resource. + ValueError: If Tensorboard is not in same project and location as this experiment. + """ + + backing_tensorboard = self._lookup_backing_tensorboard() + if backing_tensorboard: + tensorboard_resource_name = ( + tensorboard + if isinstance(tensorboard, str) + else tensorboard.resource_name + ) + if tensorboard_resource_name != backing_tensorboard.resource_name: + raise ValueError( + f"Experiment {self._metadata_context.name} already associated '" + f"to tensorboard resource {backing_tensorboard.resource_name}" + ) + + if isinstance(tensorboard, str): + tensorboard = tensorboard_resource.Tensorboard( + tensorboard, + project=self._metadata_context.project, + location=self._metadata_context.location, + credentials=self._metadata_context.credentials, + ) + + if tensorboard.project not in self._metadata_context._project_tuple: + raise ValueError( + f"Tensorboard is in project {tensorboard.project} but must be in project {self._metadata_context.project}" + ) + if tensorboard.location != self._metadata_context.location: + raise ValueError( + f"Tensorboard is in location {tensorboard.location} but must be in location {self._metadata_context.location}" + ) + + self._metadata_context.update( + metadata={ + constants._BACKING_TENSORBOARD_RESOURCE_KEY: tensorboard.resource_name + } + ) + + def _log_experiment_loggable(self, experiment_loggable: "_ExperimentLoggable"): + """Associates a Vertex resource that can be logged to an Experiment as run of this experiment. + + Args: + experiment_loggable (_ExperimentLoggable): + A Vertex Resource that can be logged to an Experiment directly. + """ + context = experiment_loggable._get_context() + self._metadata_context.add_context_children([context]) + + +class _SetLoggerLevel: + """Helper method to suppress logging.""" + + def __init__(self, module): + self._module = module + + def __enter__(self): + logging.getLogger(self._module.__name__).setLevel(logging.WARNING) + + def __exit__(self, exc_type, exc_value, traceback): + logging.getLogger(self._module.__name__).setLevel(logging.INFO) + + +class _VertexResourceWithMetadata(NamedTuple): + """Represents a resource coupled with it's metadata representation""" + + resource: base.VertexAiResourceNoun + metadata: Union[artifact.Artifact, execution.Execution, context._Context] + + +class _ExperimentLoggableSchema(NamedTuple): + """Used with _ExperimentLoggable to capture Metadata representation information about resoure. + + For example: + _ExperimentLoggableSchema(title='system.PipelineRun', type=context._Context) + + Defines the schema and metadata type to lookup PipelineJobs. + """ + + title: str + type: Union[Type[context._Context], Type[execution.Execution]] = context._Context + + +class _ExperimentLoggable(abc.ABC): + """Abstract base class to define a Vertex Resource as loggable against an Experiment. + + For example: + class PipelineJob(..., experiment_loggable_schemas= + (_ExperimentLoggableSchema(title='system.PipelineRun'), ) + + """ + + def __init_subclass__( + cls, *, experiment_loggable_schemas: Tuple[_ExperimentLoggableSchema], **kwargs + ): + """Register the metadata_schema for the subclass so Experiment can use it to retrieve the associated types. + + usage: + + class PipelineJob(..., experiment_loggable_schemas= + (_ExperimentLoggableSchema(title='system.PipelineRun'), ) + + Args: + experiment_loggable_schemas: + Tuple of the schema_title and type pairs that represent this resource. Note that a single item in the + tuple will be most common. Currently only experiment run has multiple representation for backwards + compatibility. Almost all schemas should be Contexts and Execution is currently only supported + for backwards compatibility of experiment runs. + + """ + super().__init_subclass__(**kwargs) + + # register the type when module is loaded + for schema in experiment_loggable_schemas: + _SUPPORTED_LOGGABLE_RESOURCES[schema.type][schema.title] = cls + + @abc.abstractmethod + def _get_context(self) -> context._Context: + """Should return the metadata context that represents this resource. + + The subclass should enforce this context exists. + + Returns: + Context that represents this resource. + """ + pass + + @classmethod + @abc.abstractmethod + def _query_experiment_row( + cls, node: Union[context._Context, execution.Execution] + ) -> _ExperimentRow: + """Should return parameters and metrics for this resource as a run row. + + Args: + node: The metadata node that represents this resource. + Returns: + A populated run row for this resource. + """ + pass + + def _validate_experiment(self, experiment: Union[str, Experiment]): + """Validates experiment is accessible. Can be used by subclass to throw before creating the intended resource. + + Args: + experiment (Union[str, Experiment]): The experiment that this resource will be associated to. + + Raises: + RuntimeError: If service raises any exception when trying to access this experiment. + ValueError: If resource project or location do not match experiment project or location. + """ + + if isinstance(experiment, str): + try: + experiment = Experiment.get_or_create( + experiment, + project=self.project, + location=self.location, + credentials=self.credentials, + ) + except Exception as e: + raise RuntimeError( + f"Experiment {experiment} could not be found or created. {self.__class__.__name__} not created" + ) from e + + if self.project not in experiment._metadata_context._project_tuple: + raise ValueError( + f"{self.__class__.__name__} project {self.project} does not match experiment " + f"{experiment.name} project {experiment.project}" + ) + + if experiment._metadata_context.location != self.location: + raise ValueError( + f"{self.__class__.__name__} location {self.location} does not match experiment " + f"{experiment.name} location {experiment.location}" + ) + + def _associate_to_experiment(self, experiment: Union[str, Experiment]): + """Associates this resource to the provided Experiment. + + Args: + experiment (Union[str, Experiment]): Required. Experiment name or experiment instance. + + Raises: + RuntimeError: If Metadata service cannot associate resource to Experiment. + """ + experiment_name = experiment if isinstance(experiment, str) else experiment.name + _LOGGER.info( + "Associating %s to Experiment: %s" % (self.resource_name, experiment_name) + ) + + try: + if isinstance(experiment, str): + experiment = Experiment.get_or_create( + experiment, + project=self.project, + location=self.location, + credentials=self.credentials, + ) + experiment._log_experiment_loggable(self) + except Exception as e: + raise RuntimeError( + f"{self.resource_name} could not be associated with Experiment {experiment.name}" + ) from e + + +# maps context names to their resources classes +# used by the Experiment implementation to filter for representations in the metadata store +# populated at module import time from class that inherit _ExperimentLoggable +# example mapping: +# {Metadata Type} -> {schema title} -> {vertex sdk class} +# Context -> 'system.PipelineRun' -> aiplatform.PipelineJob +# Context -> 'system.ExperimentRun' -> aiplatform.ExperimentRun +# Execution -> 'system.Run' -> aiplatform.ExperimentRun +_SUPPORTED_LOGGABLE_RESOURCES: Dict[ + Union[Type[context._Context], Type[execution.Execution]], + Dict[str, _ExperimentLoggable], +] = {execution.Execution: dict(), context._Context: dict()} diff --git a/google/cloud/aiplatform/metadata/experiment_run_resource.py b/google/cloud/aiplatform/metadata/experiment_run_resource.py new file mode 100644 index 0000000000..3a7e6f2271 --- /dev/null +++ b/google/cloud/aiplatform/metadata/experiment_run_resource.py @@ -0,0 +1,1192 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 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. + +import collections +import concurrent.futures +import functools +from typing import Callable, Dict, List, Optional, Set, Union, Any + +from google.api_core import exceptions +from google.auth import credentials as auth_credentials +from google.protobuf import timestamp_pb2 + +from google.cloud.aiplatform import base +from google.cloud.aiplatform import initializer +from google.cloud.aiplatform import pipeline_jobs +from google.cloud.aiplatform.compat.types import artifact as gca_artifact +from google.cloud.aiplatform.compat.types import execution as gca_execution +from google.cloud.aiplatform.compat.types import ( + tensorboard_time_series as gca_tensorboard_time_series, +) +from google.cloud.aiplatform.metadata import artifact +from google.cloud.aiplatform.metadata import constants +from google.cloud.aiplatform.metadata import context +from google.cloud.aiplatform.metadata import execution +from google.cloud.aiplatform.metadata import experiment_resources +from google.cloud.aiplatform.metadata import metadata +from google.cloud.aiplatform.metadata import resource +from google.cloud.aiplatform.metadata import utils as metadata_utils +from google.cloud.aiplatform.tensorboard import tensorboard_resource +from google.cloud.aiplatform.utils import rest_utils + + +_LOGGER = base.Logger(__name__) + + +def _format_experiment_run_resource_id(experiment_name: str, run_name: str) -> str: + """Formats the the experiment run resource id as a concatenation of experiment name and run name. + + Args: + experiment_name (str): Name of the experiment which is it's resource id. + run_name (str): Name of the run. + Returns: + The resource id to be used with this run. + """ + return f"{experiment_name}-{run_name}" + + +def _v1_not_supported(method: Callable) -> Callable: + """Helpers wrapper for backward compatibility. Raises when using an API not support for legacy runs.""" + + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + if isinstance(self._metadata_node, execution.Execution): + raise NotImplementedError( + f"{self._run_name} is an Execution run created during Vertex Experiment Preview and does not support" + f" {method.__name__}. Please create a new Experiment run to use this method." + ) + else: + return method(self, *args, **kwargs) + + return wrapper + + +class ExperimentRun( + experiment_resources._ExperimentLoggable, + experiment_loggable_schemas=( + experiment_resources._ExperimentLoggableSchema( + title=constants.SYSTEM_EXPERIMENT_RUN, type=context._Context + ), + # backwards compatibility with Preview Experiment runs + experiment_resources._ExperimentLoggableSchema( + title=constants.SYSTEM_RUN, type=execution.Execution + ), + ), +): + """A Vertex AI Experiment run""" + + def __init__( + self, + run_name: str, + experiment: Union[experiment_resources.Experiment, str], + *, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """ + + ``` + my_run = aiplatform.ExperimentRun('my-run, experiment='my-experiment') + ``` + + Args: + run (str): Required. The name of this run. + experiment (Union[experiment_resources.Experiment, str]): + Required. The name or instance of this experiment. + project (str): + Optional. Project where this experiment run is located. Overrides project set in + aiplatform.init. + location (str): + Optional. Location where this experiment run is located. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to retrieve this experiment run. Overrides + credentials set in aiplatform.init. + """ + + self._experiment = self._get_experiment(experiment=experiment) + self._run_name = run_name + + run_id = _format_experiment_run_resource_id( + experiment_name=self._experiment.name, run_name=run_name + ) + + metadata_args = dict( + project=project, + location=location, + credentials=credentials, + ) + + def _get_context() -> context._Context: + with experiment_resources._SetLoggerLevel(resource): + run_context = context._Context( + **{**metadata_args, "resource_name": run_id} + ) + if run_context.schema_title != constants.SYSTEM_EXPERIMENT_RUN: + raise ValueError( + f"Run {run_name} must be of type {constants.SYSTEM_EXPERIMENT_RUN}" + f" but is of type {run_context.schema_title}" + ) + return run_context + + try: + self._metadata_node = _get_context() + except exceptions.NotFound as context_not_found: + try: + # backward compatibility + self._v1_resolve_experiment_run( + { + **metadata_args, + "execution_name": run_id, + } + ) + except exceptions.NotFound: + raise context_not_found + else: + self._backing_tensorboard_run = self._lookup_tensorboard_run_artifact() + + # initially set to None. Will initially update from resource then track locally. + self._largest_step: Optional[int] = None + + def _v1_resolve_experiment_run(self, metadata_args: Dict[str, Any]): + """Resolves preview Experiment. + + Args: + metadata_args (Dict[str, Any): Arguments to pass to Execution constructor. + """ + + def _get_execution(): + with experiment_resources._SetLoggerLevel(resource): + run_execution = execution.Execution(**metadata_args) + if run_execution.schema_title != constants.SYSTEM_RUN: + # note this will raise the context not found exception in the constructor + raise exceptions.NotFound("Experiment run not found.") + return run_execution + + self._metadata_node = _get_execution() + self._metadata_metric_artifact = self._v1_get_metric_artifact() + + def _v1_get_metric_artifact(self) -> artifact.Artifact: + """Resolves metric artifact for backward compatibility. + + Returns: + Instance of Artifact that represents this run's metric artifact. + """ + metadata_args = dict( + artifact_name=self._v1_format_artifact_name(self._metadata_node.name), + project=self.project, + location=self.location, + credentials=self.credentials, + ) + + with experiment_resources._SetLoggerLevel(resource): + metric_artifact = artifact.Artifact(**metadata_args) + + if metric_artifact.schema_title != constants.SYSTEM_METRICS: + # note this will raise the context not found exception in the constructor + raise exceptions.NotFound("Experiment run not found.") + + return metric_artifact + + @staticmethod + def _v1_format_artifact_name(run_id: str) -> str: + """Formats resource id of legacy metric artifact for this run.""" + return f"{run_id}-metrics" + + def _get_context(self) -> context._Context: + """Returns this metadata context that represents this run. + + Returns: + Context instance of this run. + """ + return self._metadata_node + + @property + def resource_id(self) -> str: + """The resource ID of this experiment run's Metadata context. + + The resource ID is the final part of the resource name: + ``projects/{project}/locations/{location}/metadataStores/{metadatastore}/contexts/{resource ID}`` + """ + return self._metadata_node.name + + @property + def name(self) -> str: + """This run's name used to identify this run within it's Experiment.""" + return self._run_name + + @property + def resource_name(self) -> str: + """This run's Metadata context resource name. + + In the format: ``projects/{project}/locations/{location}/metadataStores/{metadatastore}/contexts/{context}`` + """ + return self._metadata_node.resource_name + + @property + def project(self) -> str: + """The project that this experiment run is located in.""" + return self._metadata_node.project + + @property + def location(self) -> str: + """The location that this experiment is located in.""" + return self._metadata_node.location + + @property + def credentials(self) -> auth_credentials.Credentials: + """The credentials used to access this experiment run.""" + return self._metadata_node.credentials + + @property + def state(self) -> gca_execution.Execution.State: + """The state of this run.""" + if self._is_legacy_experiment_run(): + return self._metadata_node.state + else: + return getattr( + gca_execution.Execution.State, + self._metadata_node.metadata[constants._STATE_KEY], + ) + + @staticmethod + def _get_experiment( + experiment: Optional[Union[experiment_resources.Experiment, str]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> experiment_resources.Experiment: + """Helper method ot get the experiment by name(str) or instance. + + Args: + experiment(str): + Optional. The name of this experiment. Defaults to experiment set in aiplatform.init if not provided. + project (str): + Optional. Project where this experiment is located. Overrides project set in + aiplatform.init. + location (str): + Optional. Location where this experiment is located. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to retrieve this experiment. Overrides + credentials set in aiplatform.init. + Raises: + ValueError if experiment is None and experiment has not been set using aiplatform.init. + """ + + experiment = experiment or initializer.global_config.experiment + + if not experiment: + raise ValueError( + "experiment must be provided or experiment should be set using aiplatform.init" + ) + + if not isinstance(experiment, experiment_resources.Experiment): + experiment = experiment_resources.Experiment( + experiment_name=experiment, + project=project, + location=location, + credentials=credentials, + ) + return experiment + + def _is_backing_tensorboard_run_artifact(self, artifact: artifact.Artifact) -> bool: + """Helper method to confirm tensorboard run metadata artifact is this run's tensorboard artifact. + + Args: + artifact (artifact.Artifact): Required. Instance of metadata Artifact. + Returns: + bool whether the provided artifact is this run's TensorboardRun's artifact. + """ + return all( + [ + artifact.metadata.get(constants._VERTEX_EXPERIMENT_TRACKING_LABEL), + artifact.name == self._tensorboard_run_id(self._metadata_node.name), + artifact.schema_title + == constants._TENSORBOARD_RUN_REFERENCE_ARTIFACT.schema_title, + ] + ) + + def _is_legacy_experiment_run(self) -> bool: + """Helper method that return True if this is a legacy experiment run.""" + return isinstance(self._metadata_node, execution.Execution) + + def update_state(self, state: gca_execution.Execution.State): + """Update the state of this experiment run. + + ``` + my_run = aiplatform.ExperimentRun('my-run', experiment='my-experiment') + my_run.update_state(state=aiplatform.gapic.Execution.State.COMPLETE) + ``` + + Args: + state (aiplatform.gapic.Execution.State): State of this run. + """ + if self._is_legacy_experiment_run(): + self._metadata_node.update(state=state) + else: + self._metadata_node.update(metadata={constants._STATE_KEY: state.name}) + + def _lookup_tensorboard_run_artifact( + self, + ) -> Optional[experiment_resources._VertexResourceWithMetadata]: + """Helpers method to resolve this run's TensorboardRun Artifact if it exists. + + Returns: + Tuple of Tensorboard Run Artifact and TensorboardRun is it exists. + """ + with experiment_resources._SetLoggerLevel(resource): + try: + tensorboard_run_artifact = artifact.Artifact( + artifact_name=self._tensorboard_run_id(self._metadata_node.name), + project=self._metadata_node.project, + location=self._metadata_node.location, + credentials=self._metadata_node.credentials, + ) + except exceptions.NotFound: + tensorboard_run_artifact = None + + if tensorboard_run_artifact and self._is_backing_tensorboard_run_artifact( + tensorboard_run_artifact + ): + return experiment_resources._VertexResourceWithMetadata( + resource=tensorboard_resource.TensorboardRun( + tensorboard_run_artifact.metadata[ + constants.GCP_ARTIFACT_RESOURCE_NAME_KEY + ] + ), + metadata=tensorboard_run_artifact, + ) + + @classmethod + def list( + cls, + *, + experiment: Optional[Union[experiment_resources.Experiment, str]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> List["ExperimentRun"]: + """List the experiment runs for a given aiplatform.Experiment. + + ``` + my_runs = aiplatform.ExperimentRun.list(experiment='my-experiment') + ``` + + Args: + experiment (Union[aiplatform.Experiment, str]): + Optional. The experiment name or instance to list the experiment run from. If not provided, + will use the experiment set in aiplatform.init. + project (str): + Optional. Project where this experiment is located. Overrides project set in + aiplatform.init. + location (str): + Optional. Location where this experiment is located. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to retrieve this experiment. Overrides + credentials set in aiplatform.init. + Returns: + List of experiment runs. + """ + + experiment = cls._get_experiment( + experiment=experiment, + project=project, + location=location, + credentials=credentials, + ) + + metadata_args = dict( + project=experiment._metadata_context.project, + location=experiment._metadata_context.location, + credentials=experiment._metadata_context.credentials, + ) + + filter_str = metadata_utils._make_filter_string( + schema_title=constants.SYSTEM_EXPERIMENT_RUN, + parent_contexts=[experiment.resource_name], + ) + + run_contexts = context._Context.list(filter=filter_str, **metadata_args) + + filter_str = metadata_utils._make_filter_string( + schema_title=constants.SYSTEM_RUN, in_context=[experiment.resource_name] + ) + + run_executions = execution.Execution.list(filter=filter_str, **metadata_args) + + def _initialize_experiment_run(context: context._Context) -> ExperimentRun: + this_experiment_run = cls.__new__(cls) + this_experiment_run._experiment = experiment + this_experiment_run._run_name = context.display_name + this_experiment_run._metadata_node = context + + with experiment_resources._SetLoggerLevel(resource): + tb_run = this_experiment_run._lookup_tensorboard_run_artifact() + if tb_run: + this_experiment_run._backing_tensorboard_run = tb_run + else: + this_experiment_run._backing_tensorboard_run = None + + this_experiment_run._largest_step = None + + return this_experiment_run + + def _initialize_v1_experiment_run( + execution: execution.Execution, + ) -> ExperimentRun: + this_experiment_run = cls.__new__(cls) + this_experiment_run._experiment = experiment + this_experiment_run._run_name = execution.display_name + this_experiment_run._metadata_node = execution + this_experiment_run._metadata_metric_artifact = ( + this_experiment_run._v1_get_metric_artifact() + ) + + return this_experiment_run + + if run_contexts or run_executions: + with concurrent.futures.ThreadPoolExecutor( + max_workers=max([len(run_contexts), len(run_executions)]) + ) as executor: + submissions = [ + executor.submit(_initialize_experiment_run, context) + for context in run_contexts + ] + experiment_runs = [submission.result() for submission in submissions] + + submissions = [ + executor.submit(_initialize_v1_experiment_run, execution) + for execution in run_executions + ] + + for submission in submissions: + experiment_runs.append(submission.result()) + + return experiment_runs + else: + return [] + + @classmethod + def _query_experiment_row( + cls, node: Union[context._Context, execution.Execution] + ) -> experiment_resources._ExperimentRow: + """Retrieves the runs metric and parameters into an experiment run row. + + Args: + node (Union[context._Context, execution.Execution]): + Required. Metadata node instance that represents this run. + Returns: + Experiment run row that represents this run. + """ + this_experiment_run = cls.__new__(cls) + this_experiment_run._metadata_node = node + + row = experiment_resources._ExperimentRow( + experiment_run_type=node.schema_title, + name=node.display_name, + ) + + if isinstance(node, context._Context): + this_experiment_run._backing_tensorboard_run = ( + this_experiment_run._lookup_tensorboard_run_artifact() + ) + row.params = node.metadata[constants._PARAM_KEY] + row.metrics = node.metadata[constants._METRIC_KEY] + row.time_series_metrics = ( + this_experiment_run._get_latest_time_series_metric_columns() + ) + row.state = node.metadata[constants._STATE_KEY] + else: + this_experiment_run._metadata_metric_artifact = ( + this_experiment_run._v1_get_metric_artifact() + ) + row.params = node.metadata + row.metrics = this_experiment_run._metadata_metric_artifact.metadata + row.state = node.state.name + return row + + def _get_logged_pipeline_runs(self) -> List[context._Context]: + """Returns Pipeline Run contexts logged to this Experiment Run. + + Returns: + List of Pipeline system.PipelineRun contexts. + """ + + service_request_args = dict( + project=self._metadata_node.project, + location=self._metadata_node.location, + credentials=self._metadata_node.credentials, + ) + + filter_str = metadata_utils._make_filter_string( + schema_title=constants.SYSTEM_PIPELINE_RUN, + parent_contexts=[self._metadata_node.resource_name], + ) + + return context._Context.list(filter=filter_str, **service_request_args) + + def _get_latest_time_series_metric_columns(self) -> Dict[str, Union[float, int]]: + """Determines the latest step for each time series metric. + + Returns: + Dictionary mapping time series metric key to the latest step of that metric. + """ + if self._backing_tensorboard_run: + time_series_metrics = ( + self._backing_tensorboard_run.resource.read_time_series_data() + ) + + return { + display_name: data.values[-1].scalar.value + for display_name, data in time_series_metrics.items() + if data.value_type + == gca_tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR + } + return {} + + def _log_pipeline_job(self, pipeline_job: pipeline_jobs.PipelineJob): + """Associate this PipelineJob's Context to the current ExperimentRun Context as a child context. + + Args: + pipeline_job (pipeline_jobs.PipelineJob): + Required. The PipelineJob to associate. + """ + + pipeline_job_context = pipeline_job._get_context() + self._metadata_node.add_context_children([pipeline_job_context]) + + @_v1_not_supported + def log( + self, + *, + pipeline_job: Optional[pipeline_jobs.PipelineJob] = None, + ): + """Log a Vertex Resource to this experiment run. + + ``` + my_run = aiplatform.ExperimentRun('my-run', experiment='my-experiment') + my_job = aiplatform.PipelineJob(...) + my_job.submit() + my_run.log(my_job) + ``` + + Args: + pipeline_job (aiplatform.PipelineJob): Optional. A Vertex PipelineJob. + """ + if pipeline_job: + self._log_pipeline_job(pipeline_job=pipeline_job) + + @staticmethod + def _validate_run_id(run_id: str): + """Validates the run id + + Args: + run_id(str): Required. The run id to validate. + Raises: + ValueError if run id is too long. + """ + + if len(run_id) > 128: + raise ValueError( + f"Length of Experiment ID and Run ID cannot be greater than 128. " + f"{run_id} is of length {len(run_id)}" + ) + + @classmethod + def create( + cls, + run_name: str, + *, + experiment: Optional[Union[experiment_resources.Experiment, str]] = None, + tensorboard: Optional[Union[tensorboard_resource.Tensorboard, str]] = None, + state: gca_execution.Execution.State = gca_execution.Execution.State.RUNNING, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> "ExperimentRun": + """Creates a new experiment run in Vertex AI Experiments. + + ``` + my_run = aiplatform.ExperimentRun.create('my-run', experiment='my-experiment') + ``` + + Args: + run_name (str): Required. The name of this run. + experiment (Union[aiplatform.Experiment, str]): + Optional. The name or instance of the experiment to create this run under. + If not provided, will default to the experiment set in `aiplatform.init`. + tensorboard (Union[aiplatform.Tensorboard, str]): + Optional. The resource name or instance of Vertex Tensorbaord to use as the backing + Tensorboard for time series metric logging. If not provided, will default to the + the backing tensorboard of parent experiment if set. Must be in same project and location + as this experiment run. + state (aiplatform.gapic.Execution.State): + Optional. The state of this run. Defaults to RUNNING. + project (str): + Optional. Project where this experiment will be created. Overrides project set in + aiplatform.init. + location (str): + Optional. Location where this experiment will be created. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to create this experiment. Overrides + credentials set in aiplatform.init. + Returns: + The newly created experiment run. + """ + + experiment = cls._get_experiment(experiment) + + run_id = _format_experiment_run_resource_id( + experiment_name=experiment.name, run_name=run_name + ) + + cls._validate_run_id(run_id) + + def _create_context(): + with experiment_resources._SetLoggerLevel(resource): + return context._Context._create( + resource_id=run_id, + display_name=run_name, + schema_title=constants.SYSTEM_EXPERIMENT_RUN, + schema_version=constants.SCHEMA_VERSIONS[ + constants.SYSTEM_EXPERIMENT_RUN + ], + metadata={ + constants._PARAM_KEY: {}, + constants._METRIC_KEY: {}, + constants._STATE_KEY: state.name, + }, + project=project, + location=location, + credentials=credentials, + ) + + metadata_context = _create_context() + + if metadata_context is None: + raise RuntimeError( + f"Experiment Run with name {run_name} in {experiment.name} already exists." + ) + + experiment_run = cls.__new__(cls) + experiment_run._experiment = experiment + experiment_run._run_name = metadata_context.display_name + experiment_run._metadata_node = metadata_context + experiment_run._backing_tensorboard_run = None + experiment_run._largest_step = None + + if tensorboard: + cls._assign_backing_tensorboard( + self=experiment_run, tensorboard=tensorboard + ) + else: + cls._assign_to_experiment_backing_tensorboard(self=experiment_run) + + experiment_run._associate_to_experiment(experiment) + return experiment_run + + def _assign_to_experiment_backing_tensorboard(self): + """Assigns parent Experiment backing tensorboard resource to this Experiment Run.""" + backing_tensorboard_resource = ( + self._experiment.get_backing_tensorboard_resource() + ) + + if backing_tensorboard_resource: + self.assign_backing_tensorboard(tensorboard=backing_tensorboard_resource) + + @staticmethod + def _format_tensorboard_experiment_display_name(experiment_name: str) -> str: + """Formats Tensorboard experiment name that backs this run. + Args: + experiment_name (str): Required. The name of the experiment. + Returns: + Formatted Tensorboard Experiment name + """ + # post fix helps distinguish from the Vertex Experiment in console + return f"{experiment_name} Backing Tensorboard Experiment" + + def _assign_backing_tensorboard( + self, tensorboard: Union[tensorboard_resource.Tensorboard, str] + ): + """Assign tensorboard as the backing tensorboard to this run. + + Args: + tensorboard (Union[tensorboard_resource.Tensorboard, str]): + Required. Tensorboard instance or resource name. + """ + if isinstance(tensorboard, str): + tensorboard = tensorboard_resource.Tensorboard( + tensorboard, credentials=self._metadata_node.credentials + ) + + tensorboard_resource_name_parts = tensorboard._parse_resource_name( + tensorboard.resource_name + ) + tensorboard_experiment_resource_name = ( + tensorboard_resource.TensorboardExperiment._format_resource_name( + experiment=self._experiment.name, **tensorboard_resource_name_parts + ) + ) + try: + tensorboard_experiment = tensorboard_resource.TensorboardExperiment( + tensorboard_experiment_resource_name, + credentials=tensorboard.credentials, + ) + except exceptions.NotFound: + with experiment_resources._SetLoggerLevel(tensorboard_resource): + tensorboard_experiment = ( + tensorboard_resource.TensorboardExperiment.create( + tensorboard_experiment_id=self._experiment.name, + display_name=self._format_tensorboard_experiment_display_name( + self._experiment.name + ), + tensorboard_name=tensorboard.resource_name, + credentials=tensorboard.credentials, + ) + ) + + tensorboard_experiment_name_parts = tensorboard_experiment._parse_resource_name( + tensorboard_experiment.resource_name + ) + tensorboard_run_resource_name = ( + tensorboard_resource.TensorboardRun._format_resource_name( + run=self._run_name, **tensorboard_experiment_name_parts + ) + ) + try: + tensorboard_run = tensorboard_resource.TensorboardRun( + tensorboard_run_resource_name + ) + except exceptions.NotFound: + with experiment_resources._SetLoggerLevel(tensorboard_resource): + tensorboard_run = tensorboard_resource.TensorboardRun.create( + tensorboard_run_id=self._run_name, + tensorboard_experiment_name=tensorboard_experiment.resource_name, + credentials=tensorboard.credentials, + ) + + gcp_resource_url = rest_utils.make_gcp_resource_rest_url(tensorboard_run) + + with experiment_resources._SetLoggerLevel(resource): + tensorboard_run_metadata_artifact = artifact.Artifact._create( + uri=gcp_resource_url, + resource_id=self._tensorboard_run_id(self._metadata_node.name), + metadata={ + "resourceName": tensorboard_run.resource_name, + constants._VERTEX_EXPERIMENT_TRACKING_LABEL: True, + }, + schema_title=constants._TENSORBOARD_RUN_REFERENCE_ARTIFACT.schema_title, + schema_version=constants._TENSORBOARD_RUN_REFERENCE_ARTIFACT.schema_version, + state=gca_artifact.Artifact.State.LIVE, + ) + + self._metadata_node.add_artifacts_and_executions( + artifact_resource_names=[tensorboard_run_metadata_artifact.resource_name] + ) + + self._backing_tensorboard_run = ( + experiment_resources._VertexResourceWithMetadata( + resource=tensorboard_run, metadata=tensorboard_run_metadata_artifact + ) + ) + + @staticmethod + def _tensorboard_run_id(run_id: str) -> str: + """Helper method to format the tensorboard run artifact resource id for a run. + + Args: + run_id: The resource id of the experiment run. + + Returns: + Resource id for the associated tensorboard run artifact. + """ + return f"{run_id}-tb-run" + + @_v1_not_supported + def assign_backing_tensorboard( + self, tensorboard: Union[tensorboard_resource.Tensorboard, str] + ): + """Assigns tensorboard as backing tensorboard to support timeseries metrics logging for this run. + + Args: + tensorboard (Union[aiplatform.Tensorboard, str]): + Required. Tensorboard instance or resource name. + """ + + backing_tensorboard = self._lookup_tensorboard_run_artifact() + if backing_tensorboard: + raise ValueError( + f"Experiment run {self._run_name} already associated to tensorboard resource {backing_tensorboard.resource.resource_name}" + ) + + self._assign_backing_tensorboard(tensorboard=tensorboard) + + def _get_latest_time_series_step(self) -> int: + """Gets latest time series step of all time series from Tensorboard resource. + + Returns: + Latest step of all time series metrics. + """ + data = self._backing_tensorboard_run.resource.read_time_series_data() + return max(ts.values[-1].step if ts.values else 0 for ts in data.values()) + + @_v1_not_supported + def log_time_series_metrics( + self, + metrics: Dict[str, float], + step: Optional[int] = None, + wall_time: Optional[timestamp_pb2.Timestamp] = None, + ): + """Logs time series metrics to backing TensorboardRun of this Experiment Run. + + ``` + run.log_time_series_metrics({'accuracy': 0.9}, step=10) + ``` + + Args: + metrics (Dict[str, Union[str, float]]): + Required. Dictionary of where keys are metric names and values are metric values. + step (int): + Optional. Step index of this data point within the run. + + If not provided, the latest + step amongst all time series metrics already logged will be used. + wall_time (timestamp_pb2.Timestamp): + Optional. Wall clock timestamp when this data point is + generated by the end user. + + If not provided, this will be generated based on the value from time.time() + Raises: + RuntimeError: If current experiment run doesn't have a backing Tensorboard resource. + """ + + if not self._backing_tensorboard_run: + self._assign_to_experiment_backing_tensorboard() + if not self._backing_tensorboard_run: + raise RuntimeError( + "Please set this experiment run with backing tensorboard resource to use log_time_series_metrics." + ) + + self._soft_create_time_series(metric_keys=set(metrics.keys())) + + if not step: + step = self._largest_step or self._get_latest_time_series_step() + step += 1 + self._largest_step = step + + self._backing_tensorboard_run.resource.write_tensorboard_scalar_data( + time_series_data=metrics, step=step, wall_time=wall_time + ) + + def _soft_create_time_series(self, metric_keys: Set[str]): + """Creates TensorboardTimeSeries for the metric keys if one currently does not exist. + + Args: + metric_keys (Set[str]): Keys of the metrics. + """ + + if any( + key + not in self._backing_tensorboard_run.resource._time_series_display_name_to_id_mapping + for key in metric_keys + ): + self._backing_tensorboard_run.resource._sync_time_series_display_name_to_id_mapping() + + for key in metric_keys: + if ( + key + not in self._backing_tensorboard_run.resource._time_series_display_name_to_id_mapping + ): + with experiment_resources._SetLoggerLevel(tensorboard_resource): + self._backing_tensorboard_run.resource.create_tensorboard_time_series( + display_name=key + ) + + def log_params(self, params: Dict[str, Union[float, int, str]]): + """Log single or multiple parameters with specified key value pairs. + + Parameters with the same key will be overwritten. + + ``` + my_run = aiplatform.ExperimentRun('my-run', experiment='my-experiment') + my_run.log_params({'learning_rate': 0.1, 'dropout_rate': 0.2}) + ``` + + Args: + params (Dict[str, Union[float, int, str]]): + Required. Parameter key/value pairs. + + Raises: + ValueError: If key is not str or value is not float, int, str. + """ + # query the latest run execution resource before logging. + for key, value in params.items(): + if not isinstance(key, str): + raise TypeError( + f"{key} is of type {type(key).__name__} must of type str" + ) + if not isinstance(value, (float, int, str)): + raise TypeError( + f"Value for key {key} is of type {type(value).__name__} but must be one of float, int, str" + ) + + if self._is_legacy_experiment_run(): + self._metadata_node.update(metadata=params) + else: + self._metadata_node.update(metadata={constants._PARAM_KEY: params}) + + def log_metrics(self, metrics: Dict[str, Union[float, int, str]]): + """Log single or multiple Metrics with specified key and value pairs. + + Metrics with the same key will be overwritten. + + ``` + my_run = aiplatform.ExperimentRun('my-run', experiment='my-experiment') + my_run.log_metrics({'accuracy': 0.9, 'recall': 0.8}) + ``` + + Args: + metrics (Dict[str, Union[float, int]]): + Required. Metrics key/value pairs. + Raises: + TypeError: If keys are not str or values are not float, int, or str. + """ + for key, value in metrics.items(): + if not isinstance(key, str): + raise TypeError( + f"{key} is of type {type(key).__name__} must of type str" + ) + if not isinstance(value, (float, int, str)): + raise TypeError( + f"Value for key {key} is of type {type(value).__name__} but must be one of float, int, str" + ) + + if self._is_legacy_experiment_run(): + self._metadata_metric_artifact.update(metadata=metrics) + else: + # TODO: query the latest metrics artifact resource before logging. + self._metadata_node.update(metadata={constants._METRIC_KEY: metrics}) + + @_v1_not_supported + def get_time_series_data_frame(self) -> "pd.DataFrame": # noqa: F821 + """Returns all time series in this Run as a DataFrame. + + Returns: + pd.DataFrame: Time series metrics in this Run as a Dataframe. + """ + try: + import pandas as pd + except ImportError: + raise ImportError( + "Pandas is not installed and is required to get dataframe as the return format. " + 'Please install the SDK using "pip install python-aiplatform[metadata]"' + ) + + if not self._backing_tensorboard_run: + return pd.DataFrame({}) + data = self._backing_tensorboard_run.resource.read_time_series_data() + + if not data: + return pd.DataFrame({}) + + return ( + pd.DataFrame( + { + name: entry.scalar.value, + "step": entry.step, + "wall_time": entry.wall_time, + } + for name, ts in data.items() + for entry in ts.values + ) + .groupby(["step", "wall_time"]) + .first() + .reset_index() + ) + + @_v1_not_supported + def get_logged_pipeline_jobs(self) -> List[pipeline_jobs.PipelineJob]: + """Get all PipelineJobs associated to this experiment run. + + Returns: + List of PipelineJobs associated this run. + """ + + pipeline_job_contexts = self._get_logged_pipeline_runs() + + return [ + pipeline_jobs.PipelineJob.get( + c.display_name, + project=c.project, + location=c.location, + credentials=c.credentials, + ) + for c in pipeline_job_contexts + ] + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + state = ( + gca_execution.Execution.State.FAILED + if exc_type + else gca_execution.Execution.State.COMPLETE + ) + + if metadata._experiment_tracker.experiment_run is self: + metadata._experiment_tracker.end_run(state=state) + else: + self.end_run(state) + + def end_run( + self, + *, + state: gca_execution.Execution.State = gca_execution.Execution.State.COMPLETE, + ): + """Ends this experiment run and sets state to COMPLETE. + + Args: + state (aiplatform.gapic.Execution.State): + Optional. Override the state at the end of run. Defaults to COMPLETE. + """ + self.update_state(state) + + def delete(self, *, delete_backing_tensorboard_run: bool = False): + """Deletes this experiment run. + + Does not delete the executions, artifacts, or resources logged to this run. + + Args: + delete_backing_tensorboard_run (bool): + Optional. Whether to delete the backing tensorboard run that stores time series metrics for this run. + """ + if delete_backing_tensorboard_run: + if not self._is_legacy_experiment_run(): + if not self._backing_tensorboard_run: + self._backing_tensorboard_run = ( + self._lookup_tensorboard_run_artifact() + ) + if self._backing_tensorboard_run: + self._backing_tensorboard_run.resource.delete() + self._backing_tensorboard_run.metadata.delete() + else: + _LOGGER.warn( + f"Experiment run {self.name} does not have a backing tensorboard run." + " Skipping deletion." + ) + else: + _LOGGER.warn( + f"Experiment run {self.name} does not have a backing tensorboard run." + " Skipping deletion." + ) + + self._metadata_node.delete() + + if self._is_legacy_experiment_run(): + self._metadata_metric_artifact.delete() + + @_v1_not_supported + def get_artifacts(self) -> List[artifact.Artifact]: + """Get the list of artifacts associated to this run. + + Returns: + List of artifacts associated to this run. + """ + return self._metadata_node.get_artifacts() + + @_v1_not_supported + def get_executions(self) -> List[execution.Execution]: + """Get the List of Executions associated to this run + + Returns: + List of executions associated to this run. + """ + return self._metadata_node.get_executions() + + def get_params(self) -> Dict[str, Union[int, float, str]]: + """Get the parameters logged to this run. + + Returns: + Parameters logged to this experiment run. + """ + if self._is_legacy_experiment_run(): + return self._metadata_node.metadata + else: + return self._metadata_node.metadata[constants._PARAM_KEY] + + def get_metrics(self) -> Dict[str, Union[float, int, str]]: + """Get the summary metrics logged to this run. + + Returns: + Summary metrics logged to this experiment run. + """ + if self._is_legacy_experiment_run(): + return self._metadata_metric_artifact.metadata + else: + return self._metadata_node.metadata[constants._METRIC_KEY] + + @_v1_not_supported + def associate_execution(self, execution: execution.Execution): + """Associate an execution to this experiment run. + + Args: + execution (aiplatform.Execution): Execution to associate to this run. + """ + self._metadata_node.add_artifacts_and_executions( + execution_resource_names=[execution.resource_name] + ) + + def _association_wrapper(self, f: Callable[..., Any]) -> Callable[..., Any]: + """Wraps methods and automatically associates all passed in Artifacts or Executions to this ExperimentRun. + + This is used to wrap artifact passing methods of Executions so they get associated to this run. + """ + + @functools.wraps(f) + def wrapper(*args, **kwargs): + artifacts = [] + executions = [] + for value in [*args, *kwargs.values()]: + value = value if isinstance(value, collections.Iterable) else [value] + for item in value: + if isinstance(item, execution.Execution): + executions.append(item) + elif isinstance(item, artifact.Artifact): + artifacts.append(item) + elif artifact._VertexResourceArtifactResolver.supports_metadata( + item + ): + artifacts.append( + artifact._VertexResourceArtifactResolver.resolve_or_create_resource_artifact( + item + ) + ) + + if artifacts or executions: + self._metadata_node.add_artifacts_and_executions( + artifact_resource_names=[a.resource_name for a in artifacts], + execution_resource_names=[e.resource_name for e in executions], + ) + + result = f(*args, **kwargs) + return result + + return wrapper diff --git a/google/cloud/aiplatform/metadata/metadata.py b/google/cloud/aiplatform/metadata/metadata.py index 0012c8f975..1b533d5176 100644 --- a/google/cloud/aiplatform/metadata/metadata.py +++ b/google/cloud/aiplatform/metadata/metadata.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,161 +15,381 @@ # limitations under the License. # -from typing import Dict, Union, Optional +from typing import Dict, Union, Optional, Any + +from google.api_core import exceptions +from google.auth import credentials as auth_credentials +from google.protobuf import timestamp_pb2 + +from google.cloud.aiplatform import base +from google.cloud.aiplatform import gapic +from google.cloud.aiplatform import pipeline_jobs +from google.cloud.aiplatform.compat.types import execution as gca_execution from google.cloud.aiplatform.metadata import constants -from google.cloud.aiplatform.metadata.artifact import _Artifact -from google.cloud.aiplatform.metadata.context import _Context -from google.cloud.aiplatform.metadata.execution import _Execution -from google.cloud.aiplatform.metadata.metadata_store import _MetadataStore +from google.cloud.aiplatform.metadata import context +from google.cloud.aiplatform.metadata import execution +from google.cloud.aiplatform.metadata import experiment_resources +from google.cloud.aiplatform.metadata import experiment_run_resource +from google.cloud.aiplatform.tensorboard import tensorboard_resource + +_LOGGER = base.Logger(__name__) + + +def _get_experiment_schema_version() -> str: + """Helper method to get experiment schema version + Returns: + str: schema version of the currently set experiment tracking version + """ + return constants.SCHEMA_VERSIONS[constants.SYSTEM_EXPERIMENT] -class _MetadataService: + +# Legacy Experiment tracking +# Maintaining creation APIs for backwards compatibility testing +class _LegacyExperimentService: """Contains the exposed APIs to interact with the Managed Metadata Service.""" + @staticmethod + def get_pipeline_df(pipeline: str) -> "pd.DataFrame": # noqa: F821 + """Returns a Pandas DataFrame of the parameters and metrics associated with one pipeline. + + Args: + pipeline: Name of the Pipeline to filter results. + + Returns: + Pandas Dataframe of Pipeline with metrics and parameters. + """ + + source = "pipeline" + pipeline_resource_name = ( + _LegacyExperimentService._get_experiment_or_pipeline_resource_name( + name=pipeline, source=source, expected_schema=constants.SYSTEM_PIPELINE + ) + ) + + return _LegacyExperimentService._query_runs_to_data_frame( + context_id=pipeline, + context_resource_name=pipeline_resource_name, + source=source, + ) + + @staticmethod + def _get_experiment_or_pipeline_resource_name( + name: str, source: str, expected_schema: str + ) -> str: + """Get the full resource name of the Context representing an Experiment or Pipeline. + + Args: + name (str): + Name of the Experiment or Pipeline. + source (str): + Identify whether the this is an Experiment or a Pipeline. + expected_schema (str): + expected_schema identifies the expected schema used for Experiment or Pipeline. + + Returns: + The full resource name of the Experiment or Pipeline Context. + + Raise: + NotFound exception if experiment or pipeline does not exist. + """ + + this_context = context._Context(resource_name=name) + + if this_context.schema_title != expected_schema: + raise ValueError( + f"Please provide a valid {source} name. {name} is not a {source}." + ) + return this_context.resource_name + + @staticmethod + def _query_runs_to_data_frame( + context_id: str, context_resource_name: str, source: str + ) -> "pd.DataFrame": # noqa: F821 + """Get metrics and parameters associated with a given Context into a Dataframe. + + Args: + context_id (str): + Name of the Experiment or Pipeline. + context_resource_name (str): + Full resource name of the Context associated with an Experiment or Pipeline. + source (str): + Identify whether the this is an Experiment or a Pipeline. + + Returns: + The full resource name of the Experiment or Pipeline Context. + """ + + try: + import pandas as pd + except ImportError: + raise ImportError( + "Pandas is not installed and is required to get dataframe as the return format. " + 'Please install the SDK using "pip install python-aiplatform[metadata]"' + ) + + filter = f'schema_title="{constants.SYSTEM_RUN}" AND in_context("{context_resource_name}")' + run_executions = execution.Execution.list(filter=filter) + + context_summary = [] + for run_execution in run_executions: + run_dict = { + f"{source}_name": context_id, + "run_name": run_execution.display_name, + } + run_dict.update( + _LegacyExperimentService._execution_to_column_named_metadata( + "param", run_execution.metadata + ) + ) + + for metric_artifact in run_execution.get_output_artifacts(): + run_dict.update( + _LegacyExperimentService._execution_to_column_named_metadata( + "metric", metric_artifact.metadata + ) + ) + + context_summary.append(run_dict) + + return pd.DataFrame(context_summary) + + @staticmethod + def _execution_to_column_named_metadata( + metadata_type: str, metadata: Dict, filter_prefix: Optional[str] = None + ) -> Dict[str, Union[int, float, str]]: + """Returns a dict of the Execution/Artifact metadata with column names. + + Args: + metadata_type: The type of this execution properties (param, metric). + metadata: Either an Execution or Artifact metadata field. + filter_prefix: + Remove this prefix from the key of metadata field. Mainly used for removing + "input:" from PipelineJob parameter keys + + Returns: + Dict of custom properties with keys mapped to column names + """ + column_key_to_value = {} + for key, value in metadata.items(): + if filter_prefix and key.startswith(filter_prefix): + key = key[len(filter_prefix) :] + column_key_to_value[".".join([metadata_type, key])] = value + + return column_key_to_value + + +class _ExperimentTracker: + """Tracks Experiments and Experiment Runs wil high level APIs""" + def __init__(self): - self._experiment = None - self._run = None - self._metrics = None + self._experiment: Optional[experiment_resources.Experiment] = None + self._experiment_run: Optional[experiment_run_resource.ExperimentRun] = None def reset(self): - """Reset all _MetadataService fields to None""" + """Resets this experiment tracker, clearing the current experiment and run.""" self._experiment = None - self._run = None - self._metrics = None + self._experiment_run = None @property def experiment_name(self) -> Optional[str]: - """Return the experiment name of the _MetadataService, if experiment is not set, return None""" + """Return the currently set experiment name, if experiment is not set, return None""" if self._experiment: - return self._experiment.display_name + return self._experiment.name return None @property - def run_name(self) -> Optional[str]: - """Return the run name of the _MetadataService, if run is not set, return None""" - if self._run: - return self._run.display_name - return None + def experiment(self) -> Optional[experiment_resources.Experiment]: + "Returns the currently set Experiment." + return self._experiment - def set_experiment(self, experiment: str, description: Optional[str] = None): - """Setup a experiment to current session. + @property + def experiment_run(self) -> Optional[experiment_run_resource.ExperimentRun]: + """Returns the currently set experiment run.""" + return self._experiment_run + + def set_experiment( + self, + experiment: str, + *, + description: Optional[str] = None, + backing_tensorboard: Optional[ + Union[str, tensorboard_resource.Tensorboard] + ] = None, + ): + """Set the experiment. Will retrieve the Experiment if it exists or create one with the provided name. Args: experiment (str): - Required. Name of the experiment to assign current session with. + Required. Name of the experiment to set. description (str): Optional. Description of an experiment. + backing_tensorboard Union[str, aiplatform.Tensorboard]: + Optional. If provided, assigns tensorboard as backing tensorboard to support time series metrics + logging. """ + self.reset() - _MetadataStore.get_or_create() - context = _Context.get_or_create( - resource_id=experiment, - display_name=experiment, - description=description, - schema_title=constants.SYSTEM_EXPERIMENT, - schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_EXPERIMENT], - metadata=constants.EXPERIMENT_METADATA, + experiment = experiment_resources.Experiment.get_or_create( + experiment_name=experiment, description=description ) - if context.schema_title != constants.SYSTEM_EXPERIMENT: - raise ValueError( - f"Experiment name {experiment} has been used to create other type of resources " - f"({context.schema_title}) in this MetadataStore, please choose a different experiment name." - ) - - if description and context.description != description: - context.update(metadata=context.metadata, description=description) - self._experiment = context + if backing_tensorboard: + experiment.assign_backing_tensorboard(tensorboard=backing_tensorboard) + + self._experiment = experiment + + def start_run( + self, + run: str, + *, + tensorboard: Union[tensorboard_resource.Tensorboard, str, None] = None, + resume=False, + ) -> experiment_run_resource.ExperimentRun: + """Start a run to current session. + + ``` + aiplatform.init(experiment='my-experiment') + aiplatform.start_run('my-run') + aiplatform.log_params({'learning_rate':0.1}) + ``` + + Use as context manager. Run will be ended on context exit: + ``` + aiplatform.init(experiment='my-experiment') + with aiplatform.start_run('my-run') as my_run: + my_run.log_params({'learning_rate':0.1}) + ``` + + Resume a previously started run: + ``` + aiplatform.init(experiment='my-experiment') + with aiplatform.start_run('my-run') as my_run: + my_run.log_params({'learning_rate':0.1}) + ``` - def start_run(self, run: str): - """Setup a run to current session. Args: - run (str): + run(str): Required. Name of the run to assign current session with. - Raise: - ValueError if experiment is not set. Or if run execution or metrics artifact - is already created but with a different schema. + tensorboard Union[str, tensorboard_resource.Tensorboard]: + Optional. Backing Tensorboard Resource to enable and store time series metrics + logged to this Experiment Run using `log_time_series_metrics`. + + If not provided will the the default backing tensorboard of the currently + set experiment. + resume (bool): + Whether to resume this run. If False a new run will be created. + Raises: + ValueError: + if experiment is not set. Or if run execution or metrics artifact is already created + but with a different schema. """ if not self._experiment: raise ValueError( "No experiment set for this run. Make sure to call aiplatform.init(experiment='my-experiment') " - "before trying to start_run. " + "before invoking start_run. " ) - run_execution_id = f"{self._experiment.name}-{run}" - run_execution = _Execution.get_or_create( - resource_id=run_execution_id, - display_name=run, - schema_title=constants.SYSTEM_RUN, - schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_RUN], - ) - if run_execution.schema_title != constants.SYSTEM_RUN: - raise ValueError( - f"Run name {run} has been used to create other type of resources ({run_execution.schema_title}) " - "in this MetadataStore, please choose a different run name." + + if self._experiment_run: + self.end_run() + + if resume: + self._experiment_run = experiment_run_resource.ExperimentRun( + run_name=run, experiment=self._experiment ) - self._experiment.add_artifacts_and_executions( - execution_resource_names=[run_execution.resource_name] - ) + if tensorboard: + self._experiment_run.assign_backing_tensorboard(tensorboard=tensorboard) - metrics_artifact_id = f"{self._experiment.name}-{run}-metrics" - metrics_artifact = _Artifact.get_or_create( - resource_id=metrics_artifact_id, - display_name=metrics_artifact_id, - schema_title=constants.SYSTEM_METRICS, - schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_METRICS], - ) - if metrics_artifact.schema_title != constants.SYSTEM_METRICS: - raise ValueError( - f"Run name {run} has been used to create other type of resources ({metrics_artifact.schema_title}) " - "in this MetadataStore, please choose a different run name." + self._experiment_run.update_state(state=gapic.Execution.State.RUNNING) + + else: + self._experiment_run = experiment_run_resource.ExperimentRun.create( + run_name=run, experiment=self._experiment, tensorboard=tensorboard ) - run_execution.add_artifact( - artifact_resource_name=metrics_artifact.resource_name, input=False - ) - self._run = run_execution - self._metrics = metrics_artifact + return self._experiment_run + + def end_run(self, state: gapic.Execution.State = gapic.Execution.State.COMPLETE): + """Ends the the current experiment run. + + ``` + aiplatform.start_run('my-run') + ... + aiplatform.end_run() + ``` + + """ + self._validate_experiment_and_run(method_name="end_run") + try: + self._experiment_run.end_run(state=state) + except exceptions.NotFound: + _LOGGER.warn( + f"Experiment run {self._experiment_run.name} was not found." + "It may have been deleted" + ) + finally: + self._experiment_run = None def log_params(self, params: Dict[str, Union[float, int, str]]): """Log single or multiple parameters with specified key and value pairs. + Parameters with the same key will be overwritten. + + ``` + aiplatform.start_run('my-run') + aiplatform.log_params({'learning_rate': 0.1, 'dropout_rate': 0.2}) + ``` + Args: - params (Dict): + params (Dict[str, Union[float, int, str]]): Required. Parameter key/value pairs. """ self._validate_experiment_and_run(method_name="log_params") # query the latest run execution resource before logging. - execution = _Execution.get_or_create( - resource_id=self._run.name, - schema_title=constants.SYSTEM_RUN, - schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_RUN], - ) - execution.update(metadata=params) + self._experiment_run.log_params(params=params) - def log_metrics(self, metrics: Dict[str, Union[float, int]]): + def log_metrics(self, metrics: Dict[str, Union[float, int, str]]): """Log single or multiple Metrics with specified key and value pairs. + Metrics with the same key will be overwritten. + + ``` + aiplatform.start_run('my-run', experiment='my-experiment') + aiplatform.log_metrics({'accuracy': 0.9, 'recall': 0.8}) + ``` + Args: - metrics (Dict): - Required. Metrics key/value pairs. Only float and int are supported format for value. - Raises: - TypeError: If value contains unsupported types. - ValueError: If Experiment or Run is not set. + metrics (Dict[str, Union[float, int, str]]): + Required. Metrics key/value pairs. """ self._validate_experiment_and_run(method_name="log_metrics") - self._validate_metrics_value_type(metrics) # query the latest metrics artifact resource before logging. - artifact = _Artifact.get_or_create( - resource_id=self._metrics.name, - schema_title=constants.SYSTEM_METRICS, - schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_METRICS], - ) - artifact.update(metadata=metrics) + self._experiment_run.log_metrics(metrics=metrics) + + def _validate_experiment_and_run(self, method_name: str): + """Validates Experiment and Run are set and raises informative error message. + + Args: + method_name: The name of th method to raise from. + + Raises: + ValueError: If Experiment or Run are not set. + """ + + if not self._experiment: + raise ValueError( + f"No experiment set. Make sure to call aiplatform.init(experiment='my-experiment') " + f"before trying to {method_name}. " + ) + if not self._experiment_run: + raise ValueError( + f"No run set. Make sure to call aiplatform.start_run('my-run') before trying to {method_name}. " + ) def get_experiment_df( self, experiment: Optional[str] = None @@ -187,6 +407,8 @@ def get_experiment_df( aiplatform.log_params({'learning_rate': 0.2}) aiplatform.log_metrics({'accuracy': 0.95}) + aiplatform.get_experiments_df() + Will result in the following DataFrame ___________________________________________________________________________ | experiment_name | run_name | param.learning_rate | metric.accuracy | @@ -208,173 +430,219 @@ def get_experiment_df( """ if not experiment: - experiment = self._experiment.name - - source = "experiment" - experiment_resource_name = self._get_experiment_or_pipeline_resource_name( - name=experiment, - source=source, - expected_schema=constants.SYSTEM_EXPERIMENT, - ) - - return self._query_runs_to_data_frame( - context_id=experiment, - context_resource_name=experiment_resource_name, - source=source, - ) - - def get_pipeline_df(self, pipeline: str) -> "pd.DataFrame": # noqa: F821 - """Returns a Pandas DataFrame of the parameters and metrics associated with one pipeline. + experiment = self._experiment + else: + experiment = experiment_resources.Experiment(experiment) + + return experiment.get_data_frame() + + def log( + self, + *, + pipeline_job: Optional[pipeline_jobs.PipelineJob] = None, + ): + """Log Vertex AI Resources to the current experiment run. + + ``` + aiplatform.start_run('my-run') + my_job = aiplatform.PipelineJob(...) + my_job.submit() + aiplatform.log(my_job) + ``` Args: - pipeline: Name of the Pipeline to filter results. - - Returns: - Pandas Dataframe of Pipeline with metrics and parameters. - - Raise: - NotFound exception if experiment does not exist. - ValueError if given experiment is not associated with a wrong schema. + pipeline_job (pipeline_jobs.PipelineJob): + Optional. Vertex PipelineJob to associate to this Experiment Run. """ + self._validate_experiment_and_run(method_name="log") + self._experiment_run.log(pipeline_job=pipeline_job) - source = "pipeline" - pipeline_resource_name = self._get_experiment_or_pipeline_resource_name( - name=pipeline, source=source, expected_schema=constants.SYSTEM_PIPELINE - ) - - return self._query_runs_to_data_frame( - context_id=pipeline, - context_resource_name=pipeline_resource_name, - source=source, - ) - - def _validate_experiment_and_run(self, method_name: str): - if not self._experiment: - raise ValueError( - f"No experiment set. Make sure to call aiplatform.init(experiment='my-experiment') " - f"before trying to {method_name}. " - ) - if not self._run: - raise ValueError( - f"No run set. Make sure to call aiplatform.start_run('my-run') before trying to {method_name}. " - ) + def log_time_series_metrics( + self, + metrics: Dict[str, Union[float]], + step: Optional[int] = None, + wall_time: Optional[timestamp_pb2.Timestamp] = None, + ): + """Logs time series metrics to to this Experiment Run. - @staticmethod - def _validate_metrics_value_type(metrics: Dict[str, Union[float, int]]): - """Verify that metrics value are with supported types. + Requires the experiment or experiment run has a backing Vertex Tensorboard resource. - Args: - metrics (Dict): - Required. Metrics key/value pairs. Only float and int are supported format for value. - Raises: - TypeError: If value contains unsupported types. - """ + ``` + my_tensorboard = aiplatform.Tensorboard(...) + aiplatform.init(experiment='my-experiment', experiment_tensorboard=my_tensorboard) + aiplatform.start_run('my-run') - for key, value in metrics.items(): - if isinstance(value, int) or isinstance(value, float): - continue - raise TypeError( - f"metrics contain unsupported value types. key: {key}; value: {value}; type: {type(value)}" - ) + # increments steps as logged + for i in range(10): + aiplatform.log_time_series_metrics({'loss': loss}) - @staticmethod - def _get_experiment_or_pipeline_resource_name( - name: str, source: str, expected_schema: str - ) -> str: - """Get the full resource name of the Context representing an Experiment or Pipeline. + # explicitly log steps + for i in range(10): + aiplatform.log_time_series_metrics({'loss': loss}, step=i) + ``` Args: - name (str): - Name of the Experiment or Pipeline. - source (str): - Identify whether the this is an Experiment or a Pipeline. - expected_schema (str): - expected_schema identifies the expected schema used for Experiment or Pipeline. + metrics (Dict[str, Union[str, float]]): + Required. Dictionary of where keys are metric names and values are metric values. + step (int): + Optional. Step index of this data point within the run. - Returns: - The full resource name of the Experiment or Pipeline Context. + If not provided, the latest + step amongst all time series metrics already logged will be used. + wall_time (timestamp_pb2.Timestamp): + Optional. Wall clock timestamp when this data point is + generated by the end user. - Raise: - NotFound exception if experiment or pipeline does not exist. - """ + If not provided, this will be generated based on the value from time.time() - context = _Context(resource_name=name) - - if context.schema_title != expected_schema: - raise ValueError( - f"Please provide a valid {source} name. {name} is not a {source}." - ) - return context.resource_name - - def _query_runs_to_data_frame( - self, context_id: str, context_resource_name: str, source: str - ) -> "pd.DataFrame": # noqa: F821 - """Get metrics and parameters associated with a given Context into a Dataframe. + Raises: + RuntimeError: If current experiment run doesn't have a backing Tensorboard resource. + """ + self._validate_experiment_and_run(method_name="log_time_series_metrics") + self._experiment_run.log_time_series_metrics( + metrics=metrics, step=step, wall_time=wall_time + ) + def start_execution( + self, + *, + schema_title: Optional[str] = None, + display_name: Optional[str] = None, + resource_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + schema_version: Optional[str] = None, + description: Optional[str] = None, + resume: bool = False, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> execution.Execution: + """ + Create and starts a new Metadata Execution or resumes a previously created Execution. + + To start a new execution: + + ``` + with aiplatform.start_execution(schema_title='system.ContainerExecution', display_name='trainer) as exc: + exc.assign_input_artifacts([my_artifact]) + model = aiplatform.Artifact.create(uri='gs://my-uri', schema_title='system.Model') + exc.assign_output_artifacts([model]) + ``` + + To continue a previously created execution: + ``` + with aiplatform.start_execution(resource_id='my-exc', resume=True) as exc: + ... + ``` Args: - context_id (str): - Name of the Experiment or Pipeline. - context_resource_name (str): - Full resource name of the Context associated with an Experiment or Pipeline. - source (str): - Identify whether the this is an Experiment or a Pipeline. + schema_title (str): + Optional. schema_title identifies the schema title used by the Execution. Required if starting + a new Execution. + resource_id (str): + Optional. The portion of the Execution name with + the format. This is globally unique in a metadataStore: + projects/123/locations/us-central1/metadataStores//executions/. + display_name (str): + Optional. The user-defined name of the Execution. + schema_version (str): + Optional. schema_version specifies the version used by the Execution. + If not set, defaults to use the latest version. + metadata (Dict): + Optional. Contains the metadata information that will be stored in the Execution. + description (str): + Optional. Describes the purpose of the Execution to be created. + metadata_store_id (str): + Optional. The portion of the resource name with + the format: + projects/123/locations/us-central1/metadataStores//artifacts/ + If not provided, the MetadataStore's ID will be set to "default". + project (str): + Optional. Project used to create this Execution. Overrides project set in + aiplatform.init. + location (str): + Optional. Location used to create this Execution. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials used to create this Execution. Overrides + credentials set in aiplatform.init. Returns: - The full resource name of the Experiment or Pipeline Context. - """ + Execution: Instantiated representation of the managed Metadata Execution. - filter = f'schema_title="{constants.SYSTEM_RUN}" AND in_context("{context_resource_name}")' - run_executions = _Execution.list(filter=filter) + Raises: + ValueError: If experiment run is set and project or location do not match experiment run. + ValueError: If resume set to `True` and resource_id is not provided. + ValueError: If creating a new executin and schema_title is not provided. + """ - context_summary = [] - for run_execution in run_executions: - run_dict = { - f"{source}_name": context_id, - "run_name": run_execution.display_name, - } - run_dict.update( - self._execution_to_column_named_metadata( - "param", run_execution.metadata + if ( + self._experiment_run + and not self._experiment_run._is_legacy_experiment_run() + ): + if project and project != self._experiment_run.project: + raise ValueError( + f"Currently set Experiment run project {self._experiment_run.project} must" + f"match provided project {project}" ) - ) - - for metric_artifact in run_execution.query_input_and_output_artifacts(): - run_dict.update( - self._execution_to_column_named_metadata( - "metric", metric_artifact.metadata - ) + if location and location != self._experiment_run.location: + raise ValueError( + f"Currently set Experiment run location {self._experiment_run.location} must" + f"match provided location {project}" ) - context_summary.append(run_dict) + if resume: + if not resource_id: + raise ValueError("resource_id is required when resume=True") - try: - import pandas as pd - except ImportError: - raise ImportError( - "Pandas is not installed and is required to get dataframe as the return format. " - 'Please install the SDK using "pip install python-aiplatform[full]"' + run_execution = execution.Execution( + execution_name=resource_id, + project=project, + location=location, + credentials=credentials, ) - return pd.DataFrame(context_summary) + # TODO(handle updates if resuming) - @staticmethod - def _execution_to_column_named_metadata( - metadata_type: str, - metadata: Dict, - ) -> Dict[str, Union[int, float, str]]: - """Returns a dict of the Execution/Artifact metadata with column names. + run_execution.update(state=gca_execution.Execution.State.RUNNING) + else: + if not schema_title: + raise ValueError( + "schema_title must be provided when starting a new Execution" + ) - Args: - metadata_type: The type of this execution properties (param, metric). - metadata: Either an Execution or Artifact metadata field. + run_execution = execution.Execution.create( + display_name=display_name, + schema_title=schema_title, + schema_version=schema_version, + metadata=metadata, + description=description, + resource_id=resource_id, + project=project, + location=location, + credentials=credentials, + ) - Returns: - Dict of custom properties with keys mapped to column names - """ + if self.experiment_run: + if self.experiment_run._is_legacy_experiment_run(): + _LOGGER.warn( + f"{self.experiment_run._run_name} is an Experiment run created in Vertex Experiment Preview", + " and does not support tracking Executions." + " Please create a new Experiment run to track executions against an Experiment run.", + ) + else: + self.experiment_run.associate_execution(run_execution) + run_execution.assign_input_artifacts = ( + self.experiment_run._association_wrapper( + run_execution.assign_input_artifacts + ) + ) + run_execution.assign_output_artifacts = ( + self.experiment_run._association_wrapper( + run_execution.assign_output_artifacts + ) + ) - return { - ".".join([metadata_type, key]): value for key, value in metadata.items() - } + return run_execution -metadata_service = _MetadataService() +_experiment_tracker = _ExperimentTracker() diff --git a/google/cloud/aiplatform/metadata/metadata_store.py b/google/cloud/aiplatform/metadata/metadata_store.py index 6b23ee880c..2f0c8e2955 100644 --- a/google/cloud/aiplatform/metadata/metadata_store.py +++ b/google/cloud/aiplatform/metadata/metadata_store.py @@ -242,3 +242,43 @@ def _get( ) except exceptions.NotFound: logging.info(f"MetadataStore {metadata_store_name} not found.") + + @classmethod + def ensure_default_metadata_store_exists( + cls, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + encryption_key_spec_name: Optional[str] = None, + ): + """Helpers method to ensure the `default` MetadataStore exists in this project and location. + + Args: + project (str): + Optional. Project to retrieve resource from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve resource from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to upload this model. Overrides + credentials set in aiplatform.init. + encryption_spec_key_name (str): + Optional. The Cloud KMS resource identifier of the customer + managed encryption key used to protect the metadata store. Has the + form: + ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``. + The key needs to be in the same region as where the compute + resource is created. + + If set, this MetadataStore and all sub-resources of this MetadataStore will be secured by this key. + + Overrides encryption_spec_key_name set in aiplatform.init. + """ + + cls.get_or_create( + project=project, + location=location, + credentials=credentials, + encryption_spec_key_name=encryption_key_spec_name, + ) diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py index 8816fe3455..060d138c3a 100644 --- a/google/cloud/aiplatform/metadata/resource.py +++ b/google/cloud/aiplatform/metadata/resource.py @@ -16,10 +16,10 @@ # import abc -import logging +import collections import re from copy import deepcopy -from typing import Optional, Dict, Union, Sequence +from typing import Dict, Optional, Union, Any, List import proto from google.api_core import exceptions @@ -32,6 +32,8 @@ from google.cloud.aiplatform.compat.types import context as gca_context from google.cloud.aiplatform.compat.types import execution as gca_execution +_LOGGER = base.Logger(__name__) + class _Resource(base.VertexAiResourceNounWithFutureManager, abc.ABC): """Metadata Resource for Vertex AI""" @@ -71,7 +73,7 @@ def __init__( Optional location to retrieve the resource from. If not set, location set in aiplatform.init will be used. credentials (auth_credentials.Credentials): - Custom credentials to use to upload this model. Overrides + Custom credentials to use to retrieve this resource. Overrides credentials set in aiplatform.init. """ @@ -172,7 +174,7 @@ def get_or_create( credentials=credentials, ) if not resource: - logging.info(f"Creating Resource {resource_id}") + _LOGGER.info(f"Creating Resource {resource_id}") resource = cls._create( resource_id=resource_id, schema_title=schema_title, @@ -187,9 +189,45 @@ def get_or_create( ) return resource + def sync_resource(self): + """Syncs local resource with the resource in metadata store.""" + self._gca_resource = getattr(self.api_client, self._getter_method)( + name=self.resource_name, retry=base._DEFAULT_RETRY + ) + + @staticmethod + def _nested_update_metadata( + gca_resource: Union[ + gca_context.Context, gca_execution.Execution, gca_artifact.Artifact + ], + metadata: Optional[Dict[str, Any]] = None, + ): + """Helper method to update gca_resource in place. + + Performs a one-level deep nested update on the metadata field. + + Args: + gca_resource (Union[gca_context.Context, gca_execution.Execution, gca_artifact.Artifact]): + Required. Metadata Protobuf resource. This proto's metadata will be + updated in place. + metadata (Dict[str, Any]): + Optional. Metadata dictionary to merge into gca_resource.metadata. + """ + + if metadata: + if gca_resource.metadata: + for key, value in metadata.items(): + # Note: This only support nested dictionaries one level deep + if isinstance(value, collections.Mapping): + gca_resource.metadata[key].update(value) + else: + gca_resource.metadata[key] = value + else: + gca_resource.metadata = metadata + def update( self, - metadata: Dict, + metadata: Optional[Dict] = None, description: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, ): @@ -197,25 +235,23 @@ def update( Args: metadata (Dict): - Required. metadata contains the updated metadata information. + Optional. metadata contains the updated metadata information. description (str): Optional. Description describes the resource to be updated. credentials (auth_credentials.Credentials): Custom credentials to use to update this resource. Overrides credentials set in aiplatform.init. - """ gca_resource = deepcopy(self._gca_resource) - if gca_resource.metadata: - gca_resource.metadata.update(metadata) - else: - gca_resource.metadata = metadata + if metadata: + self._nested_update_metadata(gca_resource=gca_resource, metadata=metadata) if description: gca_resource.description = description api_client = self._instantiate_client(credentials=credentials) + # TODO: if etag is not valid sync and retry update_gca_resource = self._update_resource( client=api_client, resource=gca_resource, @@ -230,7 +266,7 @@ def list( project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, - ) -> Sequence["_Resource"]: + ) -> List["_Resource"]: """List Metadata resources that match the list filter in target metadataStore. Args: @@ -257,8 +293,6 @@ def list( a list of managed Metadata resource. """ - api_client = cls._instantiate_client(location=location, credentials=credentials) - parent = ( initializer.global_config.common_location_path( project=project, location=location @@ -266,27 +300,13 @@ def list( + f"/metadataStores/{metadata_store_id}" ) - try: - resources = cls._list_resources( - client=api_client, - parent=parent, - filter=filter, - ) - except exceptions.NotFound: - logging.info( - f"No matching resources in metadataStore: {metadata_store_id} with filter: {filter}" - ) - return [] - - return [ - cls( - resource=resource, - project=project, - location=location, - credentials=credentials, - ) - for resource in resources - ] + return super().list( + filter=filter, + project=project, + location=location, + credentials=credentials, + parent=parent, + ) @classmethod def _create( @@ -301,7 +321,7 @@ def _create( project: Optional[str] = None, location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, - ): + ) -> Optional["_Resource"]: """Creates a new Metadata resource. Args: @@ -361,16 +381,19 @@ def _create( metadata=metadata, ) except exceptions.AlreadyExists: - logging.info(f"Resource '{resource_id}' already exist") + _LOGGER.info(f"Resource '{resource_id}' already exist") return - return cls( - resource=resource, + self = cls._empty_constructor( project=project, location=location, credentials=credentials, ) + self._gca_resource = resource + + return self + @classmethod def _get( cls, @@ -410,14 +433,14 @@ def _get( try: return cls( - resource_name=resource_name, + resource_name, metadata_store_id=metadata_store_id, project=project, location=location, credentials=credentials, ) except exceptions.NotFound: - logging.info(f"Resource {resource_name} not found.") + _LOGGER.info(f"Resource {resource_name} not found.") @classmethod @abc.abstractmethod diff --git a/google/cloud/aiplatform/metadata/utils.py b/google/cloud/aiplatform/metadata/utils.py new file mode 100644 index 0000000000..57eb4c0691 --- /dev/null +++ b/google/cloud/aiplatform/metadata/utils.py @@ -0,0 +1,54 @@ +# -*- coding: utf-8 -*- +# Copyright 2022 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 typing import List, Optional, Union + + +def _make_filter_string( + schema_title: Optional[Union[str, List[str]]] = None, + in_context: Optional[List[str]] = None, + parent_contexts: Optional[List[str]] = None, + uri: Optional[str] = None, +) -> str: + """Helper method to format filter strings for Metadata querying. + + No enforcement of correctness. + + Args: + schema_title (Union[str, List[str]]): Optional. schema_titles to filter for. + in_context (List[str]): + Optional. Context resource names that the node should be in. Only for Artifacts/Executions. + parent_contexts (List[str]): Optional. Parent contexts the context should be in. Only for Contexts. + uri (str): Optional. uri to match for. Only for Artifacts. + Returns: + String that can be used for Metadata service filtering. + """ + parts = [] + if schema_title: + if isinstance(schema_title, str): + parts.append(f'schema_title="{schema_title}"') + else: + substring = " OR ".join(f'schema_title="{s}"' for s in schema_title) + parts.append(f"({substring})") + if in_context: + for context in in_context: + parts.append(f'in_context("{context}")') + if parent_contexts: + parent_context_str = ",".join([f'"{c}"' for c in parent_contexts]) + parts.append(f"parent_contexts:{parent_context_str}") + if uri: + parts.append(f'uri="{uri}"') + return " AND ".join(parts) diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index bc50a47aa2..a1ea72e8fd 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/google/cloud/aiplatform/pipeline_jobs.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2021 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,12 +19,18 @@ import logging import time import re -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Union from google.auth import credentials as auth_credentials from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform import utils +from google.cloud.aiplatform.metadata import artifact +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 yaml_utils from google.cloud.aiplatform.utils import pipeline_utils from google.protobuf import json_format @@ -75,7 +81,15 @@ def _set_enable_caching_value( task["cachingOptions"] = {"enableCache": enable_caching} -class PipelineJob(base.VertexAiStatefulResource): +class PipelineJob( + base.VertexAiStatefulResource, + experiment_resources._ExperimentLoggable, + experiment_loggable_schemas=( + experiment_resources._ExperimentLoggableSchema( + title=metadata_constants.SYSTEM_PIPELINE_RUN + ), + ), +): client_class = utils.PipelineJobClientWithOverride _resource_noun = "pipelineJobs" @@ -271,6 +285,8 @@ def submit( service_account: Optional[str] = None, network: Optional[str] = None, create_request_timeout: Optional[float] = None, + *, + experiment: Optional[Union[str, experiment_resources.Experiment]] = None, ) -> None: """Run this configured PipelineJob. @@ -286,6 +302,14 @@ def submit( If left unspecified, the job is not peered with any network. create_request_timeout (float): Optional. The timeout for the create request in seconds. + experiment (Union[str, experiments_resource.Experiment]): + Optional. The Vertex AI experiment name or instance to associate to this PipelineJob. + + Metrics produced by the PipelineJob as system.Metric Artifacts + will be associated as metrics to the current Experiment Run. + + Pipeline parameters will be associated as parameters to the + current Experiment Run. """ if service_account: self._gca_resource.service_account = service_account @@ -297,6 +321,9 @@ def submit( if self._gca_resource.pipeline_spec.get("sdkVersion", "").startswith("tfx"): _LOGGER.setLevel(logging.INFO) + if experiment: + self._validate_experiment(experiment) + _LOGGER.log_create_with_lro(self.__class__) self._gca_resource = self.api_client.create_pipeline_job( @@ -312,8 +339,11 @@ def submit( _LOGGER.info("View Pipeline Job:\n%s" % self._dashboard_uri()) + if experiment: + self._associate_to_experiment(experiment) + def wait(self): - """Wait for thie PipelineJob to complete.""" + """Wait for this PipelineJob to complete.""" if self._latest_future is None: self._block_until_complete() else: @@ -329,6 +359,11 @@ def state(self) -> Optional[gca_pipeline_state.PipelineState]: self._sync_gca_resource() return self._gca_resource.state + @property + def task_details(self) -> List[gca_pipeline_job.PipelineTaskDetail]: + self._sync_gca_resource() + return list(self._gca_resource.job_detail.task_details) + @property def has_failed(self) -> bool: """Returns True if pipeline has failed. @@ -472,6 +507,112 @@ def wait_for_resource_creation(self) -> None: """Waits until resource has been created.""" self._wait_for_resource_creation() + def done(self) -> bool: + """Helper method that return True is PipelineJob is done. False otherwise.""" + if not self._gca_resource: + return False + + return self.state in _PIPELINE_COMPLETE_STATES + + def _has_failed(self) -> bool: + """Return True if PipelineJob has Failed.""" + if not self._gca_resource: + return False + + return self.state in _PIPELINE_ERROR_STATES + + def _get_context(self) -> context._Context: + """Returns the PipelineRun Context for this PipelineJob in the MetadataStore. + + Returns: + System.PipelineRUn Context instance that represents this PipelineJob. + + Raises: + RuntimeError if Pipeline has failed or system.PipelineRun context is not found. + """ + self.wait_for_resource_creation() + pipeline_run_context = self._gca_resource.job_detail.pipeline_run_context + + # PipelineJob context is created asynchronously so we need to poll until it exists. + while not self.done(): + pipeline_run_context = self._gca_resource.job_detail.pipeline_run_context + if pipeline_run_context: + break + time.sleep(1) + + if not pipeline_run_context: + if self._has_failed: + raise RuntimeError( + f"Cannot associate PipelineJob to Experiment: {self.gca_resource.error}" + ) + else: + raise RuntimeError( + "Cannot associate PipelineJob to Experiment because PipelineJob context could not be found." + ) + + return context._Context( + resource=pipeline_run_context, + project=self.project, + location=self.location, + credentials=self.credentials, + ) + + @classmethod + def _query_experiment_row( + cls, node: context._Context + ) -> experiment_resources._ExperimentRow: + """Queries the PipelineJob metadata as an experiment run parameter and metric row. + + Parameters are retrieved from the system.Run Execution.metadata of the PipelineJob. + + Metrics are retrieved from the system.Metric Artifacts.metadata produced by this PipelineJob. + + Args: + node (context._Context): + Required. System.PipelineRun context that represents a PipelineJob Run. + Returns: + Experiment run row representing this PipelineJob. + """ + + system_run_executions = execution.Execution.list( + project=node.project, + location=node.location, + credentials=node.credentials, + filter=metadata_utils._make_filter_string( + in_context=[node.resource_name], + schema_title=metadata_constants.SYSTEM_RUN, + ), + ) + + metric_artifacts = artifact.Artifact.list( + project=node.project, + location=node.location, + credentials=node.credentials, + filter=metadata_utils._make_filter_string( + in_context=[node.resource_name], + schema_title=metadata_constants.SYSTEM_METRICS, + ), + ) + + row = experiment_resources._ExperimentRow( + experiment_run_type=node.schema_title, name=node.display_name + ) + + if system_run_executions: + row.params = { + key[len(metadata_constants.PIPELINE_PARAM_PREFIX) :]: value + for key, value in system_run_executions[0].metadata.items() + } + row.state = system_run_executions[0].state.name + + for metric_artifact in metric_artifacts: + if row.metrics: + row.metrics.update(metric_artifact.metadata) + else: + row.metrics = metric_artifact.metadata + + return row + def clone( self, display_name: Optional[str] = None, diff --git a/google/cloud/aiplatform/tensorboard/__init__.py b/google/cloud/aiplatform/tensorboard/__init__.py index 63281fe972..58eb7c3640 100644 --- a/google/cloud/aiplatform/tensorboard/__init__.py +++ b/google/cloud/aiplatform/tensorboard/__init__.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,13 @@ Tensorboard, TensorboardExperiment, TensorboardRun, + TensorboardTimeSeries, ) -__all__ = ("Tensorboard", "TensorboardExperiment", "TensorboardRun") +__all__ = ( + "Tensorboard", + "TensorboardExperiment", + "TensorboardRun", + "TensorboardTimeSeries", +) diff --git a/google/cloud/aiplatform/tensorboard/tensorboard_resource.py b/google/cloud/aiplatform/tensorboard/tensorboard_resource.py index b838ad039a..76e89ca6bd 100644 --- a/google/cloud/aiplatform/tensorboard/tensorboard_resource.py +++ b/google/cloud/aiplatform/tensorboard/tensorboard_resource.py @@ -15,19 +15,29 @@ # limitations under the License. # -from typing import Dict, List, Optional, Sequence, Tuple +from typing import Dict, List, Optional, Sequence, Tuple, Union from google.auth import credentials as auth_credentials from google.protobuf import field_mask_pb2 +from google.protobuf import timestamp_pb2 from google.cloud.aiplatform import base +from google.cloud.aiplatform import initializer +from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.types import tensorboard as gca_tensorboard +from google.cloud.aiplatform.compat.types import ( + tensorboard_data as gca_tensorboard_data, +) from google.cloud.aiplatform.compat.types import ( tensorboard_experiment as gca_tensorboard_experiment, - tensorboard_run as gca_tensorboard_run, ) -from google.cloud.aiplatform import initializer -from google.cloud.aiplatform import utils +from google.cloud.aiplatform.compat.types import tensorboard_run as gca_tensorboard_run +from google.cloud.aiplatform.compat.types import ( + tensorboard_service as gca_tensorboard_service, +) +from google.cloud.aiplatform.compat.types import ( + tensorboard_time_series as gca_tensorboard_time_series, +) _LOGGER = base.Logger(__name__) @@ -378,7 +388,7 @@ def create( Example Usage: - tb = aiplatform.TensorboardExperiment.create( + tb_exp = aiplatform.TensorboardExperiment.create( tensorboard_experiment_id='my-experiment' tensorboard_id='456' display_name='my display name', @@ -560,16 +570,16 @@ def __init__( location: Optional[str] = None, credentials: Optional[auth_credentials.Credentials] = None, ): - """Retrieves an existing tensorboard experiment given a tensorboard experiment name or ID. + """Retrieves an existing tensorboard run given a tensorboard run name or ID. Example Usage: - tb_exp = aiplatform.TensorboardRun( + tb_run = aiplatform.TensorboardRun( tensorboard_run_name= "projects/123/locations/us-central1/tensorboards/456/experiments/678/run/8910" ) - tb_exp = aiplatform.TensorboardExperiment( - tensorboard_experiment_name= "8910", + tb_run = aiplatform.TensorboardRun( + tensorboard_run_name= "8910", tensorboard_id = "456", tensorboard_experiment_id = "678" ) @@ -617,6 +627,10 @@ def __init__( else tensorboard_id, ) + self._time_series_display_name_to_id_mapping = ( + self._get_time_series_display_name_to_id_mapping() + ) + @classmethod def create( cls, @@ -632,12 +646,13 @@ def create( request_metadata: Sequence[Tuple[str, str]] = (), create_request_timeout: Optional[float] = None, ) -> "TensorboardRun": - """Creates a new tensorboard. + """Creates a new tensorboard run. Example Usage: - tb = aiplatform.TensorboardExperiment.create( - tensorboard_experiment_id='my-experiment' + tb_run = aiplatform.TensorboardRun.create( + tensorboard_run_id='my-run' + tensorboard_experiment_name='my-experiment' tensorboard_id='456' display_name='my display name', description='my description', @@ -662,8 +677,7 @@ def create( If resource ID is provided then tensorboard_id must be provided. tensorboard_id (str): - Optional. The resource ID of the Tensorboard to create - the TensorboardRun in. Format of resource name. + Optional. The resource ID of the Tensorboard to create the TensorboardRun in. display_name (str): Optional. The user-defined name of the Tensorboard Run. This value must be unique among all TensorboardRuns belonging to the @@ -696,9 +710,8 @@ def create( create_request_timeout (float): Optional. The timeout for the create request in seconds. Returns: - TensorboardExperiment: The TensorboardExperiment resource. + TensorboardRun: The TensorboardRun resource. """ - if display_name: utils.validate_display_name(display_name) @@ -758,7 +771,7 @@ def list( Example Usage: aiplatform.TensorboardRun.list( - tensorboard_name='projects/my-project/locations/us-central1/tensorboards/123/experiments/456' + tensorboard_experiment_name='projects/my-project/locations/us-central1/tensorboards/123/experiments/456' ) Args: @@ -802,6 +815,470 @@ def list( location=location, ) + tensorboard_runs = super()._list( + filter=filter, + order_by=order_by, + project=project, + location=location, + credentials=credentials, + parent=parent, + ) + + for tensorboard_run in tensorboard_runs: + tensorboard_run._sync_time_series_display_name_to_id_mapping() + + return tensorboard_runs + + def write_tensorboard_scalar_data( + self, + time_series_data: Dict[str, float], + step: int, + wall_time: Optional[timestamp_pb2.Timestamp] = None, + ): + """Writes tensorboard scalar data to this run. + + Args: + time_series_data (Dict[str, float]): + Required. Dictionary of where keys are TensorboardTimeSeries display name and values are the scalar value.. + step (int): + Required. Step index of this data point within the run. + wall_time (timestamp_pb2.Timestamp): + Optional. Wall clock timestamp when this data point is + generated by the end user. + + If not provided, this will be generated based on the value from time.time() + """ + + if not wall_time: + wall_time = utils.get_timestamp_proto() + + ts_data = [] + + if any( + key not in self._time_series_display_name_to_id_mapping + for key in time_series_data.keys() + ): + self._sync_time_series_display_name_to_id_mapping() + + for display_name, value in time_series_data.items(): + time_series_id = self._time_series_display_name_to_id_mapping.get( + display_name + ) + + if not time_series_id: + raise RuntimeError( + f"TensorboardTimeSeries with display name {display_name} has not been created in TensorboardRun {self.resource_name}." + ) + + ts_data.append( + gca_tensorboard_data.TimeSeriesData( + tensorboard_time_series_id=time_series_id, + value_type=gca_tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR, + values=[ + gca_tensorboard_data.TimeSeriesDataPoint( + scalar=gca_tensorboard_data.Scalar(value=value), + wall_time=wall_time, + step=step, + ) + ], + ) + ) + + self.api_client.write_tensorboard_run_data( + tensorboard_run=self.resource_name, time_series_data=ts_data + ) + + def _get_time_series_display_name_to_id_mapping(self) -> Dict[str, str]: + """Returns a mapping of the TimeSeries display names to resource IDs for this Run. + + Returns: + Dict[str, str] - Dictionary mapping TensorboardTimeSeries display names to + resource IDs of TensorboardTimeSeries in this TensorboardRun.""" + time_series = TensorboardTimeSeries.list( + tensorboard_run_name=self.resource_name, credentials=self.credentials + ) + + return {ts.display_name: ts.name for ts in time_series} + + def _sync_time_series_display_name_to_id_mapping(self): + """Updates the local map of TimeSeries diplay name to resource ID.""" + self._time_series_display_name_to_id_mapping = ( + self._get_time_series_display_name_to_id_mapping() + ) + + def create_tensorboard_time_series( + self, + display_name: str, + value_type: Union[ + gca_tensorboard_time_series.TensorboardTimeSeries.ValueType, str + ] = "SCALAR", + plugin_name: str = "scalars", + plugin_data: Optional[bytes] = None, + description: Optional[str] = None, + ) -> "TensorboardTimeSeries": + """Creates a new tensorboard time series. + + Example Usage: + + tb_ts = tensorboard_run.create_tensorboard_time_series( + display_name='my display name', + tensorboard_run_name='my-run' + tensorboard_id='456' + tensorboard_experiment_id='my-experiment' + description='my description', + labels={ + 'key1': 'value1', + 'key2': 'value2' + } + ) + + Args: + display_name (str): + Optional. User provided name of this + TensorboardTimeSeries. This value should be + unique among all TensorboardTimeSeries resources + belonging to the same TensorboardRun resource + (parent resource). + value_type (Union[gca_tensorboard_time_series.TensorboardTimeSeries.ValueType, str]): + Optional. Type of TensorboardTimeSeries value. One of 'SCALAR', 'TENSOR', 'BLOB_SEQUENCE'. + plugin_name (str): + Optional. Name of the plugin this time series pertain to. Such as Scalar, Tensor, Blob. + plugin_data (bytes): + Optional. Data of the current plugin, with the size limited to 65KB. + description (str): + Optional. Description of this TensorboardTimeseries. + Returns: + TensorboardTimeSeries: The TensorboardTimeSeries resource. + """ + + tb_time_series = TensorboardTimeSeries.create( + display_name=display_name, + tensorboard_run_name=self.resource_name, + value_type=value_type, + plugin_name=plugin_name, + plugin_data=plugin_data, + description=description, + credentials=self.credentials, + ) + + self._time_series_display_name_to_id_mapping[ + tb_time_series.display_name + ] = tb_time_series.name + + return tb_time_series + + def read_time_series_data(self) -> Dict[str, gca_tensorboard_data.TimeSeriesData]: + """Read the time series data of this run. + + ``` + time_series_data = tensorboard_run.read_time_series_data() + + print(time_series_data['loss'].values[-1].scalar.value) + ``` + + Returns: + Dictionary of time series metric id to TimeSeriesData. + """ + self._sync_time_series_display_name_to_id_mapping() + + resource_name_parts = self._parse_resource_name(self.resource_name) + inverted_mapping = { + resource_id: display_name + for display_name, resource_id in self._time_series_display_name_to_id_mapping.items() + } + + time_series_resource_names = [ + TensorboardTimeSeries._format_resource_name( + time_series=resource_id, **resource_name_parts + ) + for resource_id in inverted_mapping.keys() + ] + + resource_name_parts.pop("experiment") + resource_name_parts.pop("run") + + tensorboard_resource_name = Tensorboard._format_resource_name( + **resource_name_parts + ) + + read_response = self.api_client.batch_read_tensorboard_time_series_data( + request=gca_tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest( + tensorboard=tensorboard_resource_name, + time_series=time_series_resource_names, + ) + ) + + return { + inverted_mapping[data.tensorboard_time_series_id]: data + for data in read_response.time_series_data + } + + +class TensorboardTimeSeries(_TensorboardServiceResource): + """Managed tensorboard resource for Vertex AI.""" + + _resource_noun = "timeSeries" + _getter_method = "get_tensorboard_time_series" + _list_method = "list_tensorboard_time_series" + _delete_method = "delete_tensorboard_time_series" + _parse_resource_name_method = "parse_tensorboard_time_series_path" + _format_resource_name_method = "tensorboard_time_series_path" + + def __init__( + self, + tensorboard_time_series_name: str, + tensorboard_id: Optional[str] = None, + tensorboard_experiment_id: Optional[str] = None, + tensorboard_run_id: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ): + """Retrieves an existing tensorboard time series given a tensorboard time series name or ID. + + Example Usage: + + tb_ts = aiplatform.TensorboardTimeSeries( + tensorboard_time_series_name="projects/123/locations/us-central1/tensorboards/456/experiments/789/run/1011/timeSeries/mse" + ) + + tb_ts = aiplatform.TensorboardTimeSeries( + tensorboard_time_series_name= "mse", + tensorboard_id = "456", + tensorboard_experiment_id = "789" + tensorboard_run_id = "1011" + ) + + Args: + tensorboard_time_series_name (str): + Required. A fully-qualified tensorboard time series resource name or resource ID. + Example: "projects/123/locations/us-central1/tensorboards/456/experiments/789/run/1011/timeSeries/mse" or + "mse" when tensorboard_id, tensorboard_experiment_id, tensorboard_run_id are passed + and project and location are initialized or passed. + tensorboard_id (str): + Optional. A tensorboard resource ID. + tensorboard_experiment_id (str): + Optional. A tensorboard experiment resource ID. + tensorboard_run_id (str): + Optional. A tensorboard run resource ID. + project (str): + Optional. Project to retrieve tensorboard from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve tensorboard from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve this Tensorboard. Overrides + credentials set in aiplatform.init. + Raises: + ValueError: if only one of tensorboard_id or tensorboard_experiment_id is provided. + """ + if not ( + bool(tensorboard_id) + == bool(tensorboard_experiment_id) + == bool(tensorboard_run_id) + ): + raise ValueError( + "tensorboard_id, tensorboard_experiment_id, tensorboard_run_id must all be provided or none should be provided." + ) + + super().__init__( + project=project, + location=location, + credentials=credentials, + resource_name=tensorboard_time_series_name, + ) + self._gca_resource = self._get_gca_resource( + resource_name=tensorboard_time_series_name, + parent_resource_name_fields={ + Tensorboard._resource_noun: tensorboard_id, + TensorboardExperiment._resource_noun: tensorboard_experiment_id, + TensorboardRun._resource_noun: tensorboard_run_id, + } + if tensorboard_id + else tensorboard_id, + ) + + @classmethod + def create( + cls, + display_name: str, + tensorboard_run_name: str, + tensorboard_id: Optional[str] = None, + tensorboard_experiment_id: Optional[str] = None, + value_type: Union[ + gca_tensorboard_time_series.TensorboardTimeSeries.ValueType, str + ] = "SCALAR", + plugin_name: str = "scalars", + plugin_data: Optional[bytes] = None, + description: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> "TensorboardTimeSeries": + """Creates a new tensorboard time series. + + Example Usage: + + tb_ts = aiplatform.TensorboardTimeSeries.create( + display_name='my display name', + tensorboard_run_name='my-run' + tensorboard_id='456' + tensorboard_experiment_id='my-experiment' + description='my description', + labels={ + 'key1': 'value1', + 'key2': 'value2' + } + ) + + Args: + display_name (str): + Optional. User provided name of this + TensorboardTimeSeries. This value should be + unique among all TensorboardTimeSeries resources + belonging to the same TensorboardRun resource + (parent resource). + tensorboard_run_name (str): + Required. The resource name or ID of the TensorboardRun + to create the TensorboardTimeseries in. Resource name format: + ``projects/{project}/locations/{location}/tensorboards/{tensorboard}/experiments/{experiment}/runs/{run}`` + + If resource ID is provided then tensorboard_id and tensorboard_experiment_id must be provided. + tensorboard_id (str): + Optional. The resource ID of the Tensorboard to create the TensorboardTimeSeries in. + tensorboard_experiment_id (str): + Optional. The ID of the TensorboardExperiment to create the TensorboardTimeSeries in. + value_type (Union[gca_tensorboard_time_series.TensorboardTimeSeries.ValueType, str]): + Optional. Type of TensorboardTimeSeries value. One of 'SCALAR', 'TENSOR', 'BLOB_SEQUENCE'. + plugin_name (str): + Optional. Name of the plugin this time series pertain to. + plugin_data (bytes): + Optional. Data of the current plugin, with the size limited to 65KB. + description (str): + Optional. Description of this TensorboardTimeseries. + project (str): + Optional. Project to upload this model to. Overrides project set in + aiplatform.init. + location (str): + Optional. Location to upload this model to. Overrides location set in + aiplatform.init. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to upload this model. Overrides + credentials set in aiplatform.init. + Returns: + TensorboardTimeSeries: The TensorboardTimeSeries resource. + """ + + if isinstance(value_type, str): + value_type = getattr( + gca_tensorboard_time_series.TensorboardTimeSeries.ValueType, value_type + ) + + api_client = cls._instantiate_client(location=location, credentials=credentials) + + parent = utils.full_resource_name( + resource_name=tensorboard_run_name, + resource_noun=TensorboardRun._resource_noun, + parse_resource_name_method=TensorboardRun._parse_resource_name, + format_resource_name_method=TensorboardRun._format_resource_name, + parent_resource_name_fields={ + Tensorboard._resource_noun: tensorboard_id, + TensorboardExperiment._resource_noun: tensorboard_experiment_id, + }, + project=project, + location=location, + ) + + gapic_tensorboard_time_series = ( + gca_tensorboard_time_series.TensorboardTimeSeries( + display_name=display_name, + description=description, + value_type=value_type, + plugin_name=plugin_name, + plugin_data=plugin_data, + ) + ) + + _LOGGER.log_create_with_lro(cls) + + tensorboard_time_series = api_client.create_tensorboard_time_series( + parent=parent, tensorboard_time_series=gapic_tensorboard_time_series + ) + + _LOGGER.log_create_complete(cls, tensorboard_time_series, "tb_time_series") + + self = cls._empty_constructor( + project=project, location=location, credentials=credentials + ) + self._gca_resource = tensorboard_time_series + + return self + + @classmethod + def list( + cls, + tensorboard_run_name: str, + tensorboard_id: Optional[str] = None, + tensorboard_experiment_id: Optional[str] = None, + filter: Optional[str] = None, + order_by: Optional[str] = None, + project: Optional[str] = None, + location: Optional[str] = None, + credentials: Optional[auth_credentials.Credentials] = None, + ) -> List["TensorboardTimeSeries"]: + """List all instances of TensorboardTimeSeries in TensorboardRun. + + Example Usage: + + aiplatform.TensorboardTimeSeries.list( + tensorboard_run_name='projects/my-project/locations/us-central1/tensorboards/123/experiments/my-experiment/runs/my-run' + ) + + Args: + tensorboard_run_name (str): + Required. The resource name or ID of the TensorboardRun + to list the TensorboardTimeseries from. Resource name format: + ``projects/{project}/locations/{location}/tensorboards/{tensorboard}/experiments/{experiment}/runs/{run}`` + + If resource ID is provided then tensorboard_id and tensorboard_experiment_id must be provided. + tensorboard_id (str): + Optional. The resource ID of the Tensorboard to list the TensorboardTimeSeries from. + tensorboard_experiment_id (str): + Optional. The ID of the TensorboardExperiment to list the TensorboardTimeSeries from. + filter (str): + Optional. An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported. + order_by (str): + 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` + project (str): + Optional. Project to retrieve list from. If not set, project + set in aiplatform.init will be used. + location (str): + Optional. Location to retrieve list from. If not set, location + set in aiplatform.init will be used. + credentials (auth_credentials.Credentials): + Optional. Custom credentials to use to retrieve list. Overrides + credentials set in aiplatform.init. + Returns: + List[TensorboardTimeSeries] - A list of TensorboardTimeSeries + """ + + parent = utils.full_resource_name( + resource_name=tensorboard_run_name, + resource_noun=TensorboardRun._resource_noun, + parse_resource_name_method=TensorboardRun._parse_resource_name, + format_resource_name_method=TensorboardRun._format_resource_name, + parent_resource_name_fields={ + Tensorboard._resource_noun: tensorboard_id, + TensorboardExperiment._resource_noun: tensorboard_experiment_id, + }, + project=project, + location=location, + ) + return super()._list( filter=filter, order_by=order_by, diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py index 0ea641eee3..9ec2f27779 100644 --- a/google/cloud/aiplatform/utils/__init__.py +++ b/google/cloud/aiplatform/utils/__init__.py @@ -377,6 +377,16 @@ def _default_version(self) -> str: def _version_map(self) -> Tuple: pass + @property + def api_endpoint(self) -> str: + """Default API endpoint used by this client.""" + client = self._clients[self._default_version] + + if self._is_temporary: + return client._client_options.api_endpoint + else: + return client._transport._host.split(":")[0] + def __init__( self, client_options: client_options.ClientOptions, diff --git a/google/cloud/aiplatform/utils/rest_utils.py b/google/cloud/aiplatform/utils/rest_utils.py new file mode 100644 index 0000000000..4d8db45c47 --- /dev/null +++ b/google/cloud/aiplatform/utils/rest_utils.py @@ -0,0 +1,32 @@ +# -*- coding: utf-8 -*- +# Copyright 2022 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 import base + + +def make_gcp_resource_rest_url(resource: base.VertexAiResourceNoun) -> str: + """Helper function to format the GCP resource url for google.X metadata schemas. + + Args: + resource (base.VertexAiResourceNoun): Required. A Vertex resource instance. + Returns: + The formatted url of resource. + """ + resource_name = resource.resource_name + version = resource.api_client._default_version + api_uri = resource.api_client.api_endpoint + + return f"https://{api_uri}/{version}/{resource_name}" diff --git a/samples/model-builder/automl_image_classification_training_job_sample.py b/samples/model-builder/automl_image_classification_training_job_sample.py index 502caf008d..352d057bcd 100644 --- a/samples/model-builder/automl_image_classification_training_job_sample.py +++ b/samples/model-builder/automl_image_classification_training_job_sample.py @@ -18,7 +18,10 @@ # [START aiplatform_sdk_automl_image_classification_training_job_sample] def automl_image_classification_training_job_sample( - project: str, location: str, dataset_id: str, display_name: str, + project: str, + location: str, + dataset_id: str, + display_name: str, ): aiplatform.init(project=project, location=location) diff --git a/samples/model-builder/batch_create_features_sample.py b/samples/model-builder/batch_create_features_sample.py index 7f91afbe92..2dd1c11dab 100644 --- a/samples/model-builder/batch_create_features_sample.py +++ b/samples/model-builder/batch_create_features_sample.py @@ -32,14 +32,8 @@ def batch_create_features_sample( ) FEATURE_CONFIGS = { - "age": { - "value_type": "INT64", - "description": "User age" - }, - "gender": { - "value_type": "STRING", - "description": "User gender" - }, + "age": {"value_type": "INT64", "description": "User age"}, + "gender": {"value_type": "STRING", "description": "User gender"}, "liked_genres": { "value_type": "STRING_ARRAY", "description": "An array of genres this user liked", diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py index 73a4ee5f2d..d48547c9e3 100644 --- a/samples/model-builder/conftest.py +++ b/samples/model-builder/conftest.py @@ -503,3 +503,204 @@ def mock_import_feature_values(mock_entity_type): mock_entity_type, "ingest_from_gcs" ) as mock_import_feature_values: yield mock_import_feature_values + + +""" +---------------------------------------------------------------------------- +Experiment Tracking Fixtures +---------------------------------------------------------------------------- +""" + + +@pytest.fixture +def mock_execution(): + mock = MagicMock(aiplatform.Execution) + mock.assign_input_artifacts.return_value = None + mock.assign_output_artifacts.return_value = None + mock.__enter__.return_value = mock + yield mock + + +@pytest.fixture +def mock_artifact(): + mock = MagicMock(aiplatform.Artifact) + yield mock + + +@pytest.fixture +def mock_experiment(): + mock = MagicMock(aiplatform.Experiment) + yield mock + + +@pytest.fixture +def mock_experiment_run(): + mock = MagicMock(aiplatform.ExperimentRun) + yield mock + + +@pytest.fixture +def mock_pipeline_job(): + mock = MagicMock(aiplatform.PipelineJob) + yield mock + + +@pytest.fixture +def mock_df(): + mock = MagicMock() + yield mock + + +@pytest.fixture +def mock_metrics(): + mock = MagicMock() + yield mock + + +@pytest.fixture +def mock_params(): + mock = MagicMock() + yield mock + + +@pytest.fixture +def mock_time_series_metrics(): + mock = MagicMock() + yield mock + + +@pytest.fixture +def mock_get_execution(mock_execution): + with patch.object(aiplatform, "Execution") as mock_get_execution: + mock_get_execution.return_value = mock_execution + yield mock_get_execution + + +@pytest.fixture +def mock_get_artifact(mock_artifact): + with patch.object(aiplatform, "Artifact") as mock_get_artifact: + mock_get_artifact.return_value = mock_artifact + yield mock_get_artifact + + +@pytest.fixture +def mock_pipeline_job_create(mock_pipeline_job): + with patch.object(aiplatform, "PipelineJob") as mock_pipeline_job_create: + mock_pipeline_job_create.return_value = mock_pipeline_job + yield mock_pipeline_job_create + + +@pytest.fixture +def mock_pipeline_job_submit(mock_pipeline_job): + with patch.object(mock_pipeline_job, "submit") as mock_pipeline_job_submit: + mock_pipeline_job_submit.return_value = None + yield mock_pipeline_job_submit + + +@pytest.fixture +def mock_create_artifact(mock_artifact): + with patch.object(aiplatform.Artifact, "create") as mock_create_artifact: + mock_create_artifact.return_value = mock_artifact + yield mock_create_artifact + + +@pytest.fixture +def mock_start_run(mock_experiment_run): + with patch.object(aiplatform, "start_run") as mock_start_run: + mock_start_run.return_value = mock_experiment_run + yield mock_start_run + + +@pytest.fixture +def mock_start_execution(mock_execution): + with patch.object(aiplatform, "start_execution") as mock_start_execution: + mock_start_execution.return_value = mock_execution + yield mock_start_execution + + +@pytest.fixture +def mock_end_run(): + with patch.object(aiplatform, "end_run") as mock_end_run: + mock_end_run.return_value = None + yield mock_end_run + + +@pytest.fixture +def mock_log_metrics(): + with patch.object(aiplatform, "log_metrics") as mock_log_metrics: + mock_log_metrics.return_value = None + yield mock_log_metrics + + +@pytest.fixture +def mock_log_time_series_metrics(): + with patch.object( + aiplatform, "log_time_series_metrics" + ) as mock_log_time_series_metrics: + mock_log_time_series_metrics.return_value = None + yield mock_log_time_series_metrics + + +@pytest.fixture +def mock_log_params(): + with patch.object(aiplatform, "log_params") as mock_log_params: + mock_log_params.return_value = None + yield mock_log_params + + +@pytest.fixture +def mock_log_pipeline_job(): + with patch.object(aiplatform, "log") as mock_log_pipeline_job: + mock_log_pipeline_job.return_value = None + yield mock_log_pipeline_job + + +@pytest.fixture +def mock_get_run(mock_experiment_run): + with patch.object(aiplatform, "ExperimentRun") as mock_get_run: + mock_get_run.return_value = mock_experiment_run + yield mock_get_run + + +@pytest.fixture +def mock_get_experiment(mock_experiment): + with patch.object(aiplatform, "Experiment") as mock_get_experiment: + mock_get_experiment.return_value = mock_experiment + yield mock_get_experiment + + +@pytest.fixture +def mock_get_with_uri(mock_artifact): + with patch.object(aiplatform.Artifact, "get_with_uri") as mock_get_with_uri: + mock_get_with_uri.return_value = mock_artifact + yield mock_get_with_uri + + +@pytest.fixture +def mock_get_experiment_df(mock_df): + with patch.object(aiplatform, "get_experiment_df") as mock_get_experiment_df: + mock_get_experiment_df.return_value = mock_df + yield mock_get_experiment_df + + +@pytest.fixture +def mock_get_metrics(mock_metrics, mock_experiment_run): + with patch.object(mock_experiment_run, "get_metrics") as mock_get_metrics: + mock_get_metrics.return_value = mock_metrics + yield mock_get_metrics + + +@pytest.fixture +def mock_get_params(mock_params, mock_experiment_run): + with patch.object(mock_experiment_run, "get_params") as mock_get_params: + mock_get_params.return_value = mock_params + yield mock_get_params + + +@pytest.fixture +def mock_get_time_series_metrics(mock_time_series_metrics, mock_experiment_run): + with patch.object( + mock_experiment_run, "get_time_series_data_frame" + ) as mock_get_time_series_metrics: + mock_get_time_series_metrics.return_value = mock_time_series_metrics + yield mock_get_time_series_metrics diff --git a/samples/model-builder/create_and_import_dataset_tabular_gcs_sample.py b/samples/model-builder/create_and_import_dataset_tabular_gcs_sample.py index cac7a64d89..ba7f6f3b49 100644 --- a/samples/model-builder/create_and_import_dataset_tabular_gcs_sample.py +++ b/samples/model-builder/create_and_import_dataset_tabular_gcs_sample.py @@ -19,13 +19,17 @@ # [START aiplatform_sdk_create_and_import_dataset_tabular_gcs_sample] def create_and_import_dataset_tabular_gcs_sample( - display_name: str, project: str, location: str, gcs_source: Union[str, List[str]], + display_name: str, + project: str, + location: str, + gcs_source: Union[str, List[str]], ): aiplatform.init(project=project, location=location) dataset = aiplatform.TabularDataset.create( - display_name=display_name, gcs_source=gcs_source, + display_name=display_name, + gcs_source=gcs_source, ) dataset.wait() diff --git a/samples/model-builder/create_and_import_dataset_tabular_gcs_sample_test.py b/samples/model-builder/create_and_import_dataset_tabular_gcs_sample_test.py index ca8679be01..3f486a452b 100644 --- a/samples/model-builder/create_and_import_dataset_tabular_gcs_sample_test.py +++ b/samples/model-builder/create_and_import_dataset_tabular_gcs_sample_test.py @@ -32,5 +32,6 @@ def test_create_and_import_dataset_tabular_gcs_sample( project=constants.PROJECT, location=constants.LOCATION ) mock_create_tabular_dataset.assert_called_once_with( - display_name=constants.DISPLAY_NAME, gcs_source=constants.GCS_SOURCES, + display_name=constants.DISPLAY_NAME, + gcs_source=constants.GCS_SOURCES, ) diff --git a/samples/model-builder/create_endpoint_sample.py b/samples/model-builder/create_endpoint_sample.py index 7867aefea4..38536ca598 100644 --- a/samples/model-builder/create_endpoint_sample.py +++ b/samples/model-builder/create_endpoint_sample.py @@ -17,12 +17,16 @@ # [START aiplatform_sdk_create_endpoint_sample] def create_endpoint_sample( - project: str, display_name: str, location: str, + project: str, + display_name: str, + location: str, ): aiplatform.init(project=project, location=location) endpoint = aiplatform.Endpoint.create( - display_name=display_name, project=project, location=location, + display_name=display_name, + project=project, + location=location, ) print(endpoint.display_name) diff --git a/samples/model-builder/create_entity_type_sample.py b/samples/model-builder/create_entity_type_sample.py index eee0ddb8f5..bdd3678dcd 100644 --- a/samples/model-builder/create_entity_type_sample.py +++ b/samples/model-builder/create_entity_type_sample.py @@ -18,7 +18,10 @@ def create_entity_type_sample( - project: str, location: str, entity_type_id: str, featurestore_name: str, + project: str, + location: str, + entity_type_id: str, + featurestore_name: str, ): aiplatform.init(project=project, location=location) diff --git a/samples/model-builder/create_training_pipeline_image_classification_sample_test.py b/samples/model-builder/create_training_pipeline_image_classification_sample_test.py index 4f02441b4d..83d49b4d91 100644 --- a/samples/model-builder/create_training_pipeline_image_classification_sample_test.py +++ b/samples/model-builder/create_training_pipeline_image_classification_sample_test.py @@ -47,7 +47,7 @@ def test_create_training_pipeline_image_classification_sample( display_name=constants.DISPLAY_NAME, model_type=constants.MODEL_TYPE, multi_label=False, - prediction_type='classification' + prediction_type="classification", ) mock_run_automl_image_training_job.assert_called_once_with( dataset=mock_image_dataset, diff --git a/samples/model-builder/deploy_model_with_automatic_resources_sample.py b/samples/model-builder/deploy_model_with_automatic_resources_sample.py index 9efe3967f9..22e0e779c5 100644 --- a/samples/model-builder/deploy_model_with_automatic_resources_sample.py +++ b/samples/model-builder/deploy_model_with_automatic_resources_sample.py @@ -32,9 +32,9 @@ def deploy_model_with_automatic_resources_sample( sync: bool = True, ): """ - model_name: A fully-qualified model resource name or model ID. - Example: "projects/123/locations/us-central1/models/456" or - "456" when project and location are initialized or passed. + model_name: A fully-qualified model resource name or model ID. + Example: "projects/123/locations/us-central1/models/456" or + "456" when project and location are initialized or passed. """ aiplatform.init(project=project, location=location) diff --git a/samples/model-builder/deploy_model_with_automatic_resources_test.py b/samples/model-builder/deploy_model_with_automatic_resources_test.py index fff08b6e7e..b4485961ff 100644 --- a/samples/model-builder/deploy_model_with_automatic_resources_test.py +++ b/samples/model-builder/deploy_model_with_automatic_resources_test.py @@ -18,7 +18,10 @@ def test_deploy_model_with_automatic_resources_sample( - mock_sdk_init, mock_model, mock_init_model, mock_deploy_model, + mock_sdk_init, + mock_model, + mock_init_model, + mock_deploy_model, ): deploy_model_with_automatic_resources_sample.deploy_model_with_automatic_resources_sample( diff --git a/samples/model-builder/deploy_model_with_dedicated_resources_sample.py b/samples/model-builder/deploy_model_with_dedicated_resources_sample.py index a0a9e0ffa3..cdb021a77b 100644 --- a/samples/model-builder/deploy_model_with_dedicated_resources_sample.py +++ b/samples/model-builder/deploy_model_with_dedicated_resources_sample.py @@ -38,9 +38,9 @@ def deploy_model_with_dedicated_resources_sample( sync: bool = True, ): """ - model_name: A fully-qualified model resource name or model ID. - Example: "projects/123/locations/us-central1/models/456" or - "456" when project and location are initialized or passed. + model_name: A fully-qualified model resource name or model ID. + Example: "projects/123/locations/us-central1/models/456" or + "456" when project and location are initialized or passed. """ aiplatform.init(project=project, location=location) diff --git a/samples/model-builder/experiment_tracking/assign_artifact_as_execution_input_sample.py b/samples/model-builder/experiment_tracking/assign_artifact_as_execution_input_sample.py new file mode 100644 index 0000000000..e51e8cfd6f --- /dev/null +++ b/samples/model-builder/experiment_tracking/assign_artifact_as_execution_input_sample.py @@ -0,0 +1,26 @@ +# Copyright 2022 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 +# +# https://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 import aiplatform + + +# [START aiplatform_sdk_assign_artifact_as_execution_input_sample] +def assign_artifact_as_execution_input_sample( + execution: aiplatform.Execution, + artifact: aiplatform.Artifact, +): + execution.assign_input_artifacts([artifact]) + + +# [END aiplatform_sdk_assign_artifact_as_execution_input_sample] diff --git a/samples/model-builder/experiment_tracking/assign_artifact_as_execution_input_sample_test.py b/samples/model-builder/experiment_tracking/assign_artifact_as_execution_input_sample_test.py new file mode 100644 index 0000000000..316caf5393 --- /dev/null +++ b/samples/model-builder/experiment_tracking/assign_artifact_as_execution_input_sample_test.py @@ -0,0 +1,29 @@ +# Copyright 2022 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 +# +# https://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. + +import assign_artifact_as_execution_input_sample +from google.cloud import aiplatform + + +def test_assign_artifact_as_execution_input_sample( + mock_get_execution, + mock_get_artifact, +): + exc = aiplatform.Execution() + art = aiplatform.Artifact() + assign_artifact_as_execution_input_sample.assign_artifact_as_execution_input_sample( + execution=exc, artifact=art + ) + + exc.assign_input_artifacts.assert_called_with([art]) diff --git a/samples/model-builder/experiment_tracking/assign_artifact_as_execution_output_sample.py b/samples/model-builder/experiment_tracking/assign_artifact_as_execution_output_sample.py new file mode 100644 index 0000000000..dabf50985e --- /dev/null +++ b/samples/model-builder/experiment_tracking/assign_artifact_as_execution_output_sample.py @@ -0,0 +1,26 @@ +# Copyright 2022 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 +# +# https://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 import aiplatform + + +# [START aiplatform_sdk_assign_artifact_as_execution_output_sample] +def assign_artifact_as_execution_output_sample( + execution: aiplatform.Execution, + artifact: aiplatform.Artifact, +): + execution.assign_output_artifacts([artifact]) + + +# [END aiplatform_sdk_assign_artifact_as_execution_output_sample] diff --git a/samples/model-builder/experiment_tracking/assign_artifact_as_execution_output_sample_test.py b/samples/model-builder/experiment_tracking/assign_artifact_as_execution_output_sample_test.py new file mode 100644 index 0000000000..049acdf1e3 --- /dev/null +++ b/samples/model-builder/experiment_tracking/assign_artifact_as_execution_output_sample_test.py @@ -0,0 +1,29 @@ +# Copyright 2022 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 +# +# https://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. + +import assign_artifact_as_execution_output_sample +from google.cloud import aiplatform + + +def test_assign_artifact_as_execution_output_sample( + mock_get_execution, + mock_get_artifact, +): + exc = aiplatform.Execution() + art = aiplatform.Artifact() + assign_artifact_as_execution_output_sample.assign_artifact_as_execution_output_sample( + execution=exc, artifact=art + ) + + exc.assign_output_artifacts.assert_called_with([art]) diff --git a/samples/model-builder/experiment_tracking/create_artifact_sample.py b/samples/model-builder/experiment_tracking/create_artifact_sample.py new file mode 100644 index 0000000000..90a7432bea --- /dev/null +++ b/samples/model-builder/experiment_tracking/create_artifact_sample.py @@ -0,0 +1,47 @@ +# Copyright 2022 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 +# +# https://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 typing import Dict, Optional + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_create_artifact_sample] +def create_artifact_sample( + schema_title: str, + project: str, + location: str, + uri: Optional[str] = None, + resource_id: Optional[str] = None, + display_name: Optional[str] = None, + schema_version: Optional[str] = None, + description: Optional[str] = None, + metadata: Optional[Dict] = None, +): + artifact = aiplatform.Artifact.create( + schema_title=schema_title, + uri=uri, + resource_id=resource_id, + display_name=display_name, + schema_version=schema_version, + description=description, + metadata=metadata, + project=project, + location=location, + ) + + return artifact + + +# [END aiplatform_sdk_create_artifact_sample] diff --git a/samples/model-builder/experiment_tracking/create_artifact_sample_test.py b/samples/model-builder/experiment_tracking/create_artifact_sample_test.py new file mode 100644 index 0000000000..dea7747e0e --- /dev/null +++ b/samples/model-builder/experiment_tracking/create_artifact_sample_test.py @@ -0,0 +1,45 @@ +# Copyright 2022 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 +# +# https://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. + +import create_artifact_sample + +import test_constants as constants + + +def test_create_artifact_sample(mock_artifact, mock_create_artifact): + artifact = create_artifact_sample.create_artifact_sample( + schema_title=constants.SCHEMA_TITLE, + uri=constants.MODEL_ARTIFACT_URI, + resource_id=constants.RESOURCE_ID, + display_name=constants.DISPLAY_NAME, + schema_version=constants.SCHEMA_VERSION, + description=constants.DESCRIPTION, + metadata=constants.METADATA, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_create_artifact.assert_called_with( + schema_title=constants.SCHEMA_TITLE, + uri=constants.MODEL_ARTIFACT_URI, + resource_id=constants.RESOURCE_ID, + display_name=constants.DISPLAY_NAME, + schema_version=constants.SCHEMA_VERSION, + description=constants.DESCRIPTION, + metadata=constants.METADATA, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + assert artifact is mock_artifact diff --git a/samples/model-builder/experiment_tracking/create_experiment_run_sample.py b/samples/model-builder/experiment_tracking/create_experiment_run_sample.py new file mode 100644 index 0000000000..0767139eba --- /dev/null +++ b/samples/model-builder/experiment_tracking/create_experiment_run_sample.py @@ -0,0 +1,33 @@ +# Copyright 2022 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 +# +# https://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 typing import Optional, Union + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_create_experiment_run_sample] +def create_experiment_run_sample( + experiment_name: str, + run_name: str, + experiment_run_tensorboard: Optional[Union[str, aiplatform.Tensorboard]], + project: str, + location: str, +): + aiplatform.init(experiment_name=experiment_name, project=project, location=location) + + aiplatform.start_run(run=run_name, tensorboard=experiment_run_tensorboard) + + +# [END aiplatform_sdk_create_experiment_run_sample] diff --git a/samples/model-builder/experiment_tracking/create_experiment_run_sample_test.py b/samples/model-builder/experiment_tracking/create_experiment_run_sample_test.py new file mode 100644 index 0000000000..c0affc0da1 --- /dev/null +++ b/samples/model-builder/experiment_tracking/create_experiment_run_sample_test.py @@ -0,0 +1,39 @@ +# Copyright 2022 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 +# +# https://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. + +import create_experiment_run_sample + +import test_constants as constants + + +def test_create_experiment_run_sample(mock_sdk_init, mock_start_run): + + create_experiment_run_sample.create_experiment_run_sample( + experiment_name=constants.EXPERIMENT_NAME, + run_name=constants.EXPERIMENT_RUN_NAME, + experiment_run_tensorboard=constants.TENSORBOARD_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_sdk_init.assert_called_with( + experiment_name=constants.EXPERIMENT_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_start_run.assert_called_with( + run=constants.EXPERIMENT_RUN_NAME, + tensorboard=constants.TENSORBOARD_NAME, + ) diff --git a/samples/model-builder/experiment_tracking/create_experiment_sample.py b/samples/model-builder/experiment_tracking/create_experiment_sample.py new file mode 100644 index 0000000000..f73954cc95 --- /dev/null +++ b/samples/model-builder/experiment_tracking/create_experiment_sample.py @@ -0,0 +1,37 @@ +# Copyright 2022 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 +# +# https://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 typing import Optional, Union + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_create_experiment_sample] +def create_experiment_sample( + experiment_name: str, + experiment_description: str, + experiment_tensorboard: Optional[Union[str, aiplatform.Tensorboard]], + project: str, + location: str, +): + aiplatform.init( + experiment_name=experiment_name, + experiment_description=experiment_description, + experiment_tensorboard=experiment_tensorboard, + project=project, + location=location, + ) + + +# [END aiplatform_sdk_create_experiment_sample] diff --git a/samples/model-builder/experiment_tracking/create_experiment_sample_test.py b/samples/model-builder/experiment_tracking/create_experiment_sample_test.py new file mode 100644 index 0000000000..c4c3ae3843 --- /dev/null +++ b/samples/model-builder/experiment_tracking/create_experiment_sample_test.py @@ -0,0 +1,37 @@ +# Copyright 2022 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 +# +# https://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. + + +import create_experiment_sample + +import test_constants as constants + + +def test_create_experiment_sample(mock_sdk_init): + + create_experiment_sample.create_experiment_sample( + experiment_name=constants.EXPERIMENT_NAME, + experiment_description=constants.DESCRIPTION, + experiment_tensorboard=constants.TENSORBOARD_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_sdk_init.assert_called_with( + experiment_name=constants.EXPERIMENT_NAME, + experiment_description=constants.DESCRIPTION, + experiment_tensorboard=constants.TENSORBOARD_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) diff --git a/samples/model-builder/experiment_tracking/delete_experiment_run_sample.py b/samples/model-builder/experiment_tracking/delete_experiment_run_sample.py new file mode 100644 index 0000000000..7f8704bb00 --- /dev/null +++ b/samples/model-builder/experiment_tracking/delete_experiment_run_sample.py @@ -0,0 +1,35 @@ +# Copyright 2022 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 +# +# https://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 typing import Union + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_delete_experiment_run_sample] +def delete_experiment_run_sample( + run_name: str, + experiment: Union[str, aiplatform.Experiment], + project: str, + location: str, + delete_backing_tensorboard_run: bool = False, +): + experiment_run = aiplatform.ExperimentRun( + run_name=run_name, experiment=experiment, project=project, location=location + ) + + experiment_run.delete(delete_backing_tensorboard_run=delete_backing_tensorboard_run) + + +# [END aiplatform_sdk_delete_experiment_run_sample] diff --git a/samples/model-builder/experiment_tracking/delete_experiment_run_sample_test.py b/samples/model-builder/experiment_tracking/delete_experiment_run_sample_test.py new file mode 100644 index 0000000000..2da57e3f47 --- /dev/null +++ b/samples/model-builder/experiment_tracking/delete_experiment_run_sample_test.py @@ -0,0 +1,39 @@ +# Copyright 2022 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 +# +# https://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. + +import delete_experiment_run_sample + +import test_constants + + +def test_delete_experiment_run_sample( + mock_experiment_run, + mock_get_run, +): + delete_experiment_run_sample.delete_experiment_run_sample( + run_name=test_constants.EXPERIMENT_RUN_NAME, + experiment=test_constants.EXPERIMENT_NAME, + project=test_constants.PROJECT, + location=test_constants.LOCATION, + delete_backing_tensorboard_run=True, + ) + + mock_get_run.assert_called_with( + run_name=test_constants.EXPERIMENT_RUN_NAME, + experiment=test_constants.EXPERIMENT_NAME, + project=test_constants.PROJECT, + location=test_constants.LOCATION, + ) + + mock_experiment_run.delete.assert_called_with(delete_backing_tensorboard_run=True) diff --git a/samples/model-builder/experiment_tracking/delete_experiment_sample.py b/samples/model-builder/experiment_tracking/delete_experiment_sample.py new file mode 100644 index 0000000000..0b95e438b9 --- /dev/null +++ b/samples/model-builder/experiment_tracking/delete_experiment_sample.py @@ -0,0 +1,33 @@ +# Copyright 2022 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 +# +# https://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 import aiplatform + + +# [START aiplatform_sdk_delete_experiment_sample] +def delete_experiment_sample( + experiment_name: str, + project: str, + location: str, + delete_backing_tensorboard_runs: bool = False, +): + experiment = aiplatform.Experiment( + experiment_name=experiment_name, project=project, location=location + ) + + experiment.delete(delete_backing_tensorboard_runs=delete_backing_tensorboard_runs) + + +# [END aiplatform_sdk_delete_experiment_sample] diff --git a/samples/model-builder/experiment_tracking/delete_experiment_sample_test.py b/samples/model-builder/experiment_tracking/delete_experiment_sample_test.py new file mode 100644 index 0000000000..a397b185aa --- /dev/null +++ b/samples/model-builder/experiment_tracking/delete_experiment_sample_test.py @@ -0,0 +1,35 @@ +# Copyright 2022 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 +# +# https://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. + + +import delete_experiment_sample + +import test_constants + + +def test_delete_experiment_sample(mock_experiment, mock_get_experiment): + delete_experiment_sample.delete_experiment_sample( + experiment_name=test_constants.EXPERIMENT_NAME, + project=test_constants.PROJECT, + location=test_constants.LOCATION, + delete_backing_tensorboard_runs=True, + ) + + mock_get_experiment.assert_called_with( + experiment_name=test_constants.EXPERIMENT_NAME, + project=test_constants.PROJECT, + location=test_constants.LOCATION, + ) + + mock_experiment.delete.assert_called_with(delete_backing_tensorboard_runs=True) diff --git a/samples/model-builder/experiment_tracking/end_experiment_run_sample.py b/samples/model-builder/experiment_tracking/end_experiment_run_sample.py new file mode 100644 index 0000000000..5161c15937 --- /dev/null +++ b/samples/model-builder/experiment_tracking/end_experiment_run_sample.py @@ -0,0 +1,33 @@ +# Copyright 2022 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 +# +# https://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 import aiplatform + + +# [START aiplatform_sdk_end_experiment_run_sample] +def end_experiment_run_sample( + experiment_name: str, + run_name: str, + project: str, + location: str, +): + aiplatform.init(experiment_name=experiment_name, project=project, location=location) + + aiplatform.start_run(run=run_name, resume=True) + + aiplatform.end_run() + + +# [END aiplatform_sdk_end_experiment_run_sample] diff --git a/samples/model-builder/experiment_tracking/end_experiment_run_sample_test.py b/samples/model-builder/experiment_tracking/end_experiment_run_sample_test.py new file mode 100644 index 0000000000..551f57d3cd --- /dev/null +++ b/samples/model-builder/experiment_tracking/end_experiment_run_sample_test.py @@ -0,0 +1,37 @@ +# Copyright 2022 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 +# +# https://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. + +import end_experiment_run_sample + +import test_constants as constants + + +def test_end_experiment_run_sample(mock_sdk_init, mock_start_run, mock_end_run): + + end_experiment_run_sample.end_experiment_run_sample( + experiment_name=constants.EXPERIMENT_NAME, + run_name=constants.EXPERIMENT_RUN_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_sdk_init.assert_called_with( + experiment_name=constants.EXPERIMENT_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_start_run.assert_called_with(run=constants.EXPERIMENT_RUN_NAME, resume=True) + + mock_end_run.assert_called_with() diff --git a/samples/model-builder/experiment_tracking/get_artifact_sample.py b/samples/model-builder/experiment_tracking/get_artifact_sample.py new file mode 100644 index 0000000000..e0b9fc5500 --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_artifact_sample.py @@ -0,0 +1,31 @@ +# Copyright 2022 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 +# +# https://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 import aiplatform + + +# [START aiplatform_sdk_get_artifact_sample] +def get_artifact_sample( + uri: str, + project: str, + location: str, +): + artifact = aiplatform.Artifact.get_with_uri( + uri=uri, project=project, location=location + ) + + return artifact + + +# [END aiplatform_sdk_get_artifact_sample] diff --git a/samples/model-builder/experiment_tracking/get_artifact_sample_test.py b/samples/model-builder/experiment_tracking/get_artifact_sample_test.py new file mode 100644 index 0000000000..7387927ef3 --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_artifact_sample_test.py @@ -0,0 +1,33 @@ +# Copyright 2022 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 +# +# https://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. + +import get_artifact_sample + +import test_constants + + +def test_get_artifact_sample(mock_artifact, mock_get_with_uri): + artifact = get_artifact_sample.get_artifact_sample( + uri=test_constants.MODEL_ARTIFACT_URI, + project=test_constants.PROJECT, + location=test_constants.LOCATION, + ) + + mock_get_with_uri.assert_called_with( + uri=test_constants.MODEL_ARTIFACT_URI, + project=test_constants.PROJECT, + location=test_constants.LOCATION, + ) + + assert artifact is mock_artifact diff --git a/samples/model-builder/experiment_tracking/get_experiment_data_frame_sample.py b/samples/model-builder/experiment_tracking/get_experiment_data_frame_sample.py new file mode 100644 index 0000000000..0a65a82f7b --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_experiment_data_frame_sample.py @@ -0,0 +1,32 @@ +# Copyright 2022 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 +# +# https://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 import aiplatform + + +# [START aiplatform_sdk_get_experiments_data_frame_sample] +def get_experiments_data_frame_sample( + experiment_name: str, + project: str, + location: str, +): + aiplatform.init(experiment_name=experiment_name, project=project, location=location) + + experiments_df = aiplatform.get_experiment_df() + + return experiments_df + + +# [END aiplatform_sdk_get_experiments_data_frame_sample] diff --git a/samples/model-builder/experiment_tracking/get_experiment_data_frame_sample_test.py b/samples/model-builder/experiment_tracking/get_experiment_data_frame_sample_test.py new file mode 100644 index 0000000000..8740f5c2c7 --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_experiment_data_frame_sample_test.py @@ -0,0 +1,32 @@ +# Copyright 2022 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 +# +# https://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. + +import get_experiment_data_frame_sample + +import pytest + +import test_constants as constants + + +@pytest.mark.usefixtures("mock_sdk_init") +def test_get_experiments_data_frame_sample(mock_get_experiment_df, mock_df): + df = get_experiment_data_frame_sample.get_experiments_data_frame_sample( + experiment_name=constants.EXPERIMENT_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_get_experiment_df.assert_called_with() + + assert df is mock_df diff --git a/samples/model-builder/experiment_tracking/get_experiment_run_metrics_sample.py b/samples/model-builder/experiment_tracking/get_experiment_run_metrics_sample.py new file mode 100644 index 0000000000..ce4b36b60d --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_experiment_run_metrics_sample.py @@ -0,0 +1,34 @@ +# Copyright 2022 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 +# +# https://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 typing import Dict, Union + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_get_experiment_run_metrics_sample] +def get_experiment_run_metrics_sample( + run_name: str, + experiment: Union[str, aiplatform.Experiment], + project: str, + location: str, +) -> Dict[str, Union[float, int]]: + experiment_run = aiplatform.ExperimentRun( + run_name=run_name, experiment=experiment, project=project, location=location + ) + + return experiment_run.get_metrics() + + +# [END aiplatform_sdk_get_experiment_run_metrics_sample] diff --git a/samples/model-builder/experiment_tracking/get_experiment_run_metrics_sample_test.py b/samples/model-builder/experiment_tracking/get_experiment_run_metrics_sample_test.py new file mode 100644 index 0000000000..4ff9502392 --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_experiment_run_metrics_sample_test.py @@ -0,0 +1,34 @@ +# Copyright 2022 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 +# +# https://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. + +import get_experiment_run_metrics_sample + +import pytest + +import test_constants as constants + + +@pytest.mark.usefixtures("mock_get_run") +def test_get_experiment_run_metrics_sample(mock_get_metrics, mock_metrics): + + metrics = get_experiment_run_metrics_sample.get_experiment_run_metrics_sample( + run_name=constants.EXPERIMENT_RUN_NAME, + experiment=constants.EXPERIMENT_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_get_metrics.assert_called_with() + + assert metrics is mock_metrics diff --git a/samples/model-builder/experiment_tracking/get_experiment_run_params_sample.py b/samples/model-builder/experiment_tracking/get_experiment_run_params_sample.py new file mode 100644 index 0000000000..1bad58ebad --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_experiment_run_params_sample.py @@ -0,0 +1,34 @@ +# Copyright 2022 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 +# +# https://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 typing import Dict, Union + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_get_experiment_run_params_sample] +def get_experiment_run_params_sample( + run_name: str, + experiment: Union[str, aiplatform.Experiment], + project: str, + location: str, +) -> Dict[str, Union[float, int, str]]: + experiment_run = aiplatform.ExperimentRun( + run_name=run_name, experiment=experiment, project=project, location=location + ) + + return experiment_run.get_params() + + +# [END aiplatform_sdk_get_experiment_run_params_sample] diff --git a/samples/model-builder/experiment_tracking/get_experiment_run_params_sample_test.py b/samples/model-builder/experiment_tracking/get_experiment_run_params_sample_test.py new file mode 100644 index 0000000000..96610526a9 --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_experiment_run_params_sample_test.py @@ -0,0 +1,34 @@ +# Copyright 2022 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 +# +# https://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. + +import get_experiment_run_params_sample + +import pytest + +import test_constants as constants + + +@pytest.mark.usefixtures("mock_get_run") +def test_get_experiment_run_params_sample(mock_get_params, mock_params): + + params = get_experiment_run_params_sample.get_experiment_run_params_sample( + run_name=constants.EXPERIMENT_RUN_NAME, + experiment=constants.EXPERIMENT_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_get_params.assert_called_with() + + assert params is mock_params diff --git a/samples/model-builder/experiment_tracking/get_experiment_run_time_series_metric_data_frame_sample.py b/samples/model-builder/experiment_tracking/get_experiment_run_time_series_metric_data_frame_sample.py new file mode 100644 index 0000000000..febb92fcb6 --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_experiment_run_time_series_metric_data_frame_sample.py @@ -0,0 +1,34 @@ +# Copyright 2022 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 +# +# https://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 typing import Union + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_get_experiment_run_time_series_metric_data_frame_sample] +def get_experiment_run_time_series_metric_data_frame_sample( + run_name: str, + experiment: Union[str, aiplatform.Experiment], + project: str, + location: str, +) -> "pd.DataFrame": # noqa: F821 + experiment_run = aiplatform.ExperimentRun( + run_name=run_name, experiment=experiment, project=project, location=location + ) + + return experiment_run.get_time_series_data_frame() + + +# [END aiplatform_sdk_get_experiment_run_time_series_metric_data_frame_sample] diff --git a/samples/model-builder/experiment_tracking/get_experiment_run_time_series_metric_data_frame_sample_test.py b/samples/model-builder/experiment_tracking/get_experiment_run_time_series_metric_data_frame_sample_test.py new file mode 100644 index 0000000000..2bd40cd814 --- /dev/null +++ b/samples/model-builder/experiment_tracking/get_experiment_run_time_series_metric_data_frame_sample_test.py @@ -0,0 +1,36 @@ +# Copyright 2022 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 +# +# https://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. + +import get_experiment_run_time_series_metric_data_frame_sample + +import pytest + +import test_constants as constants + + +@pytest.mark.usefixtures("mock_get_run") +def test_get_experiment_run_time_series_metric_data_frame_sample( + mock_get_time_series_metrics, mock_time_series_metrics +): + + metrics = get_experiment_run_time_series_metric_data_frame_sample.get_experiment_run_time_series_metric_data_frame_sample( + run_name=constants.EXPERIMENT_RUN_NAME, + experiment=constants.EXPERIMENT_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_get_time_series_metrics.assert_called_with() + + assert metrics is mock_time_series_metrics diff --git a/samples/model-builder/experiment_tracking/log_metrics_sample.py b/samples/model-builder/experiment_tracking/log_metrics_sample.py new file mode 100644 index 0000000000..be139dbae1 --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_metrics_sample.py @@ -0,0 +1,35 @@ +# Copyright 2022 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 +# +# https://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 typing import Dict + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_log_metrics_sample] +def log_metrics_sample( + experiment_name: str, + run_name: str, + metrics: Dict[str, float], + project: str, + location: str, +): + aiplatform.init(experiment_name=experiment_name, project=project, location=location) + + aiplatform.start_run(run_name=run_name, resume=True) + + aiplatform.log_metrics(metrics) + + +# [END aiplatform_sdk_log_metrics_sample] diff --git a/samples/model-builder/experiment_tracking/log_metrics_sample_test.py b/samples/model-builder/experiment_tracking/log_metrics_sample_test.py new file mode 100644 index 0000000000..8cc003fa3f --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_metrics_sample_test.py @@ -0,0 +1,33 @@ +# Copyright 2022 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 +# +# https://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. + +import log_metrics_sample + +import pytest + +import test_constants as constants + + +@pytest.mark.usefixtures("mock_sdk_init", "mock_start_run") +def test_log_metrics_sample(mock_log_metrics): + + log_metrics_sample.log_metrics_sample( + experiment_name=constants.EXPERIMENT_NAME, + run_name=constants.EXPERIMENT_RUN_NAME, + metrics=constants.METRICS, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_log_metrics.assert_called_with(constants.METRICS) diff --git a/samples/model-builder/experiment_tracking/log_params_sample.py b/samples/model-builder/experiment_tracking/log_params_sample.py new file mode 100644 index 0000000000..9f07b3bdf3 --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_params_sample.py @@ -0,0 +1,35 @@ +# Copyright 2022 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 +# +# https://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 typing import Dict, Union + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_log_params_sample] +def log_params_sample( + experiment_name: str, + run_name: str, + params: Dict[str, Union[float, int, str]], + project: str, + location: str, +): + aiplatform.init(experiment_name=experiment_name, project=project, location=location) + + aiplatform.start_run(run_name=run_name, resume=True) + + aiplatform.log_params(params) + + +# [END aiplatform_sdk_log_params_sample] diff --git a/samples/model-builder/experiment_tracking/log_params_sample_test.py b/samples/model-builder/experiment_tracking/log_params_sample_test.py new file mode 100644 index 0000000000..f09673eda1 --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_params_sample_test.py @@ -0,0 +1,33 @@ +# Copyright 2022 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 +# +# https://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. + +import log_params_sample + +import pytest + +import test_constants as constants + + +@pytest.mark.usefixtures("mock_sdk_init", "mock_start_run") +def test_log_params_sample(mock_log_params): + + log_params_sample.log_params_sample( + experiment_name=constants.EXPERIMENT_NAME, + run_name=constants.EXPERIMENT_RUN_NAME, + params=constants.PARAMS, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_log_params.assert_called_with(constants.PARAMS) diff --git a/samples/model-builder/experiment_tracking/log_pipeline_job_sample.py b/samples/model-builder/experiment_tracking/log_pipeline_job_sample.py new file mode 100644 index 0000000000..467e969edd --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_pipeline_job_sample.py @@ -0,0 +1,34 @@ +# Copyright 2022 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 +# +# https://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 import aiplatform + + +# [START aiplatform_sdk_log_pipeline_job_sample] +def log_pipeline_job_sample( + experiment_name: str, + run_name: str, + pipeline_job: aiplatform.PipelineJob, + project: str, + location: str, +): + aiplatform.init(experiment_name=experiment_name, project=project, location=location) + + aiplatform.start_run(run_name=run_name, resume=True) + + aiplatform.log(pipeline_job=pipeline_job) + + +# [END aiplatform_sdk_log_pipeline_job_sample] diff --git a/samples/model-builder/experiment_tracking/log_pipeline_job_sample_test.py b/samples/model-builder/experiment_tracking/log_pipeline_job_sample_test.py new file mode 100644 index 0000000000..68f3293f1c --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_pipeline_job_sample_test.py @@ -0,0 +1,33 @@ +# Copyright 2022 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 +# +# https://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. + +import log_pipeline_job_sample + +import pytest + +import test_constants as constants + + +@pytest.mark.usefixtures("mock_sdk_init", "mock_start_run") +def test_log_pipeline_job_sample(mock_log_pipeline_job, mock_pipeline_job): + + log_pipeline_job_sample.log_pipeline_job_sample( + experiment_name=constants.EXPERIMENT_NAME, + run_name=constants.EXPERIMENT_RUN_NAME, + pipeline_job=mock_pipeline_job, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_log_pipeline_job.assert_called_with(pipeline_job=mock_pipeline_job) diff --git a/samples/model-builder/experiment_tracking/log_pipeline_job_to_experiment_sample.py b/samples/model-builder/experiment_tracking/log_pipeline_job_to_experiment_sample.py new file mode 100644 index 0000000000..495f9c2581 --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_pipeline_job_to_experiment_sample.py @@ -0,0 +1,42 @@ +# Copyright 2022 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 +# +# https://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 typing import Any, Dict, Optional + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_log_pipeline_job_to_experiment_sample] +def log_pipeline_job_to_experiment_sample( + experiment_name: str, + pipeline_job_display_name: str, + template_path: str, + pipeline_root: str, + parameter_values: Optional[Dict[str, Any]], + project: str, + location: str, +): + aiplatform.init(project=project, location=location) + + pipeline_job = aiplatform.PipelineJob( + display_name=pipeline_job_display_name, + template_path=template_path, + pipeline_root=pipeline_root, + parameter_values=parameter_values, + ) + + pipeline_job.submit(experiment=experiment_name) + + +# [END aiplatform_sdk_log_pipeline_job_to_experiment_sample] diff --git a/samples/model-builder/experiment_tracking/log_pipeline_job_to_experiment_sample_test.py b/samples/model-builder/experiment_tracking/log_pipeline_job_to_experiment_sample_test.py new file mode 100644 index 0000000000..3cb5e8cfff --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_pipeline_job_to_experiment_sample_test.py @@ -0,0 +1,45 @@ +# Copyright 2022 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 +# +# https://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. + +import log_pipeline_job_to_experiment_sample + +import test_constants as constants + + +def test_log_pipeline_job_sample( + mock_sdk_init, mock_pipeline_job_create, mock_pipeline_job_submit +): + + log_pipeline_job_to_experiment_sample.log_pipeline_job_to_experiment_sample( + experiment_name=constants.EXPERIMENT_NAME, + pipeline_job_display_name=constants.DISPLAY_NAME, + template_path=constants.TEMPLATE_PATH, + pipeline_root=constants.STAGING_BUCKET, + parameter_values=constants.PARAMS, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_sdk_init.assert_called_with( + project=constants.PROJECT, location=constants.LOCATION + ) + + mock_pipeline_job_create.assert_called_with( + display_name=constants.DISPLAY_NAME, + template_path=constants.TEMPLATE_PATH, + pipeline_root=constants.STAGING_BUCKET, + parameter_values=constants.PARAMS, + ) + + mock_pipeline_job_submit.assert_called_with(experiment=constants.EXPERIMENT_NAME) diff --git a/samples/model-builder/experiment_tracking/log_time_series_metrics_sample.py b/samples/model-builder/experiment_tracking/log_time_series_metrics_sample.py new file mode 100644 index 0000000000..2ab8d46d2f --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_time_series_metrics_sample.py @@ -0,0 +1,38 @@ +# Copyright 2022 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 +# +# https://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 typing import Dict, Optional + +from google.cloud import aiplatform +from google.protobuf import timestamp_pb2 + + +# [START aiplatform_sdk_log_time_series_metrics_sample] +def log_time_series_metrics_sample( + experiment_name: str, + run_name: str, + metrics: Dict[str, float], + step: Optional[int], + wall_time: Optional[timestamp_pb2.Timestamp], + project: str, + location: str, +): + aiplatform.init(experiment_name=experiment_name, project=project, location=location) + + aiplatform.start_run(run_name=run_name, resume=True) + + aiplatform.log_time_series_metrics(metrics=metrics, step=step, wall_time=wall_time) + + +# [END aiplatform_sdk_log_time_series_metrics_sample] diff --git a/samples/model-builder/experiment_tracking/log_time_series_metrics_sample_test.py b/samples/model-builder/experiment_tracking/log_time_series_metrics_sample_test.py new file mode 100644 index 0000000000..e85c204fb1 --- /dev/null +++ b/samples/model-builder/experiment_tracking/log_time_series_metrics_sample_test.py @@ -0,0 +1,39 @@ +# Copyright 2022 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 +# +# https://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. + +import log_time_series_metrics_sample + +import pytest + +import test_constants as constants + + +@pytest.mark.usefixtures("mock_sdk_init", "mock_start_run") +def test_log_time_series_metrics_sample(mock_log_time_series_metrics): + + log_time_series_metrics_sample.log_time_series_metrics_sample( + experiment_name=constants.EXPERIMENT_NAME, + run_name=constants.EXPERIMENT_RUN_NAME, + metrics=constants.METRICS, + step=constants.STEP, + wall_time=constants.TIMESTAMP, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_log_time_series_metrics.assert_called_with( + metrics=constants.METRICS, + step=constants.STEP, + wall_time=constants.TIMESTAMP, + ) diff --git a/samples/model-builder/experiment_tracking/resume_experiment_run_sample.py b/samples/model-builder/experiment_tracking/resume_experiment_run_sample.py new file mode 100644 index 0000000000..6ba254f07f --- /dev/null +++ b/samples/model-builder/experiment_tracking/resume_experiment_run_sample.py @@ -0,0 +1,31 @@ +# Copyright 2022 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 +# +# https://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 import aiplatform + + +# [START aiplatform_sdk_resume_experiment_run_sample] +def resume_experiment_run_sample( + experiment_name: str, + run_name: str, + project: str, + location: str, +): + aiplatform.init(experiment_name=experiment_name, project=project, location=location) + + aiplatform.start_run(run=run_name, resume=True) + + +# [END aiplatform_sdk_resume_experiment_run_sample] diff --git a/samples/model-builder/experiment_tracking/resume_experiment_run_sample_test.py b/samples/model-builder/experiment_tracking/resume_experiment_run_sample_test.py new file mode 100644 index 0000000000..964178f857 --- /dev/null +++ b/samples/model-builder/experiment_tracking/resume_experiment_run_sample_test.py @@ -0,0 +1,35 @@ +# Copyright 2022 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 +# +# https://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. + +import resume_experiment_run_sample + +import test_constants as constants + + +def test_resume_experiment_run_sample(mock_sdk_init, mock_start_run): + + resume_experiment_run_sample.resume_experiment_run_sample( + experiment_name=constants.EXPERIMENT_NAME, + run_name=constants.EXPERIMENT_RUN_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_sdk_init.assert_called_with( + experiment_name=constants.EXPERIMENT_NAME, + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_start_run.assert_called_with(run=constants.EXPERIMENT_RUN_NAME, resume=True) diff --git a/samples/model-builder/experiment_tracking/start_execution_sample.py b/samples/model-builder/experiment_tracking/start_execution_sample.py new file mode 100644 index 0000000000..8d8632bfce --- /dev/null +++ b/samples/model-builder/experiment_tracking/start_execution_sample.py @@ -0,0 +1,48 @@ +# Copyright 2022 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 +# +# https://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 typing import Any, Dict, List, Optional + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_start_execution_sample] +def start_execution_sample( + schema_title: str, + display_name: str, + input_artifacts: List[aiplatform.Artifact], + output_artifacts: List[aiplatform.Artifact], + project: str, + location: str, + resource_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + schema_version: Optional[str] = None, + resume: bool = False, +): + aiplatform.init(project=project, location=location) + + with aiplatform.start_execution( + schema_title=schema_title, + display_name=display_name, + resource_id=resource_id, + metadata=metadata, + schema_version=schema_version, + resume=resume, + ) as execution: + execution.assign_input_artifacts(input_artifacts) + execution.assign_output_artifacts(output_artifacts) + return execution + + +# [END aiplatform_sdk_start_execution_sample] diff --git a/samples/model-builder/experiment_tracking/start_execution_sample_test.py b/samples/model-builder/experiment_tracking/start_execution_sample_test.py new file mode 100644 index 0000000000..c1be84fe0f --- /dev/null +++ b/samples/model-builder/experiment_tracking/start_execution_sample_test.py @@ -0,0 +1,63 @@ +# Copyright 2022 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 +# +# https://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 import aiplatform + +import start_execution_sample + +import test_constants as constants + + +def test_start_execution_sample( + mock_sdk_init, + mock_get_execution, + mock_get_artifact, + mock_start_execution, + mock_execution, +): + + input_art = aiplatform.Artifact() + output_art = aiplatform.Artifact() + + exc = start_execution_sample.start_execution_sample( + schema_title=constants.SCHEMA_TITLE, + display_name=constants.DISPLAY_NAME, + input_artifacts=[input_art], + output_artifacts=[output_art], + project=constants.PROJECT, + location=constants.LOCATION, + resource_id=constants.RESOURCE_ID, + metadata=constants.METADATA, + schema_version=constants.SCHEMA_VERSION, + resume=True, + ) + + mock_sdk_init.assert_called_with( + project=constants.PROJECT, + location=constants.LOCATION, + ) + + mock_start_execution.assert_called_with( + schema_title=constants.SCHEMA_TITLE, + display_name=constants.DISPLAY_NAME, + resource_id=constants.RESOURCE_ID, + metadata=constants.METADATA, + schema_version=constants.SCHEMA_VERSION, + resume=True, + ) + + mock_execution.assign_input_artifacts.assert_called_with([input_art]) + mock_execution.assign_output_artifacts.assert_called_with([output_art]) + + assert exc is mock_execution diff --git a/samples/model-builder/explain_tabular_sample_test.py b/samples/model-builder/explain_tabular_sample_test.py index d088da9658..11af809632 100644 --- a/samples/model-builder/explain_tabular_sample_test.py +++ b/samples/model-builder/explain_tabular_sample_test.py @@ -32,7 +32,9 @@ def test_explain_tabular_sample( project=constants.PROJECT, location=constants.LOCATION ) - mock_get_endpoint.assert_called_once_with(constants.ENDPOINT_NAME,) + mock_get_endpoint.assert_called_once_with( + constants.ENDPOINT_NAME, + ) mock_endpoint_explain.assert_called_once_with( instances=[constants.PREDICTION_TABULAR_INSTANCE], parameters={} diff --git a/samples/model-builder/import_data_text_entity_extraction_sample_test.py b/samples/model-builder/import_data_text_entity_extraction_sample_test.py index 44ce9cc328..a3b93e9200 100644 --- a/samples/model-builder/import_data_text_entity_extraction_sample_test.py +++ b/samples/model-builder/import_data_text_entity_extraction_sample_test.py @@ -34,7 +34,9 @@ def test_import_data_text_entity_extraction_sample( project=constants.PROJECT, location=constants.LOCATION ) - mock_get_text_dataset.assert_called_once_with(constants.DATASET_NAME,) + mock_get_text_dataset.assert_called_once_with( + constants.DATASET_NAME, + ) mock_import_text_dataset.assert_called_once_with( gcs_source=constants.GCS_SOURCES, diff --git a/samples/model-builder/import_data_text_sentiment_analysis_sample_test.py b/samples/model-builder/import_data_text_sentiment_analysis_sample_test.py index 8bfd6ac0c3..2134d66b35 100644 --- a/samples/model-builder/import_data_text_sentiment_analysis_sample_test.py +++ b/samples/model-builder/import_data_text_sentiment_analysis_sample_test.py @@ -34,7 +34,9 @@ def test_import_data_text_sentiment_analysis_sample( project=constants.PROJECT, location=constants.LOCATION ) - mock_get_text_dataset.assert_called_once_with(constants.DATASET_NAME,) + mock_get_text_dataset.assert_called_once_with( + constants.DATASET_NAME, + ) mock_import_text_dataset.assert_called_once_with( gcs_source=constants.GCS_SOURCES, diff --git a/samples/model-builder/noxfile.py b/samples/model-builder/noxfile.py index 38bb0a572b..d13780d33a 100644 --- a/samples/model-builder/noxfile.py +++ b/samples/model-builder/noxfile.py @@ -171,7 +171,11 @@ def lint(session: nox.sessions.Session) -> None: def blacken(session: nox.sessions.Session) -> None: """Run black. Format code to uniform standard.""" session.install(BLACK_VERSION) - python_files = [path for path in os.listdir(".") if path.endswith(".py")] + python_files = [ + path + for path in os.listdir(".") + if path.endswith(".py") or path.endswith("experiment_tracking") + ] session.run("black", *python_files) @@ -180,6 +184,7 @@ def blacken(session: nox.sessions.Session) -> None: # format = isort + black # + @nox.session def format(session: nox.sessions.Session) -> None: """ @@ -229,9 +234,7 @@ def _session_tests( if os.path.exists("requirements-test.txt"): if os.path.exists("constraints-test.txt"): - session.install( - "-r", "requirements-test.txt", "-c", "constraints-test.txt" - ) + session.install("-r", "requirements-test.txt", "-c", "constraints-test.txt") else: session.install("-r", "requirements-test.txt") with open("requirements-test.txt") as rtfile: @@ -244,9 +247,9 @@ def _session_tests( post_install(session) if "pytest-parallel" in packages: - concurrent_args.extend(['--workers', 'auto', '--tests-per-worker', 'auto']) + concurrent_args.extend(["--workers", "auto", "--tests-per-worker", "auto"]) elif "pytest-xdist" in packages: - concurrent_args.extend(['-n', 'auto']) + concurrent_args.extend(["-n", "auto"]) session.run( "pytest", @@ -276,7 +279,7 @@ def py(session: nox.sessions.Session) -> None: def _get_repo_root() -> Optional[str]: - """ Returns the root folder of the project. """ + """Returns the root folder of the project.""" # Get root of this repository. Assume we don't have directories nested deeper than 10 items. p = Path(os.getcwd()) for i in range(10): diff --git a/samples/model-builder/noxfile_config.py b/samples/model-builder/noxfile_config.py index 024eece69f..fa5af7e8f9 100644 --- a/samples/model-builder/noxfile_config.py +++ b/samples/model-builder/noxfile_config.py @@ -31,7 +31,7 @@ # build specific Cloud project. You can also use your own string # to use your own Cloud project. # "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", - 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT', + "gcloud_project_env": "BUILD_SPECIFIC_GCLOUD_PROJECT", # If you need to use a specific version of pip, # change pip_version_override to the string representation # of the version number, for example, "20.2.4" diff --git a/samples/model-builder/predict_tabular_classification_sample.py b/samples/model-builder/predict_tabular_classification_sample.py index 0a0a1da8bd..0490a1beb1 100644 --- a/samples/model-builder/predict_tabular_classification_sample.py +++ b/samples/model-builder/predict_tabular_classification_sample.py @@ -25,14 +25,14 @@ def predict_tabular_classification_sample( endpoint_name: str, instances: List[Dict], ): - ''' + """ Args project: Your project ID or project number. location: Region where Endpoint is located. For example, 'us-central1'. endpoint_name: A fully qualified endpoint name or endpoint ID. Example: "projects/123/locations/us-central1/endpoints/456" or "456" when project and location are initialized or passed. instances: A list of one or more instances (examples) to return a prediction for. - ''' + """ aiplatform.init(project=project, location=location) endpoint = aiplatform.Endpoint(endpoint_name) diff --git a/samples/model-builder/predict_tabular_classification_sample_test.py b/samples/model-builder/predict_tabular_classification_sample_test.py index 66f2976803..c44388131e 100644 --- a/samples/model-builder/predict_tabular_classification_sample_test.py +++ b/samples/model-builder/predict_tabular_classification_sample_test.py @@ -30,4 +30,6 @@ def test_predict_tabular_classification_sample(mock_sdk_init, mock_get_endpoint) project=constants.PROJECT, location=constants.LOCATION ) - mock_get_endpoint.assert_called_once_with(constants.ENDPOINT_NAME,) + mock_get_endpoint.assert_called_once_with( + constants.ENDPOINT_NAME, + ) diff --git a/samples/model-builder/predict_tabular_regression_sample_test.py b/samples/model-builder/predict_tabular_regression_sample_test.py index abda65b3c4..572539e987 100644 --- a/samples/model-builder/predict_tabular_regression_sample_test.py +++ b/samples/model-builder/predict_tabular_regression_sample_test.py @@ -30,4 +30,6 @@ def test_predict_tabular_regression_sample(mock_sdk_init, mock_get_endpoint): project=constants.PROJECT, location=constants.LOCATION ) - mock_get_endpoint.assert_called_once_with(constants.ENDPOINT_NAME,) + mock_get_endpoint.assert_called_once_with( + constants.ENDPOINT_NAME, + ) diff --git a/samples/model-builder/predict_text_classification_single_label_sample_test.py b/samples/model-builder/predict_text_classification_single_label_sample_test.py index 789f2962c3..c446235a79 100644 --- a/samples/model-builder/predict_text_classification_single_label_sample_test.py +++ b/samples/model-builder/predict_text_classification_single_label_sample_test.py @@ -32,4 +32,6 @@ def test_predict_text_classification_single_label_sample( project=constants.PROJECT, location=constants.LOCATION ) - mock_get_endpoint.assert_called_once_with(constants.ENDPOINT_NAME,) + mock_get_endpoint.assert_called_once_with( + constants.ENDPOINT_NAME, + ) diff --git a/samples/model-builder/predict_text_entity_extraction_sample_test.py b/samples/model-builder/predict_text_entity_extraction_sample_test.py index 3b123ff148..3ca2b49b43 100644 --- a/samples/model-builder/predict_text_entity_extraction_sample_test.py +++ b/samples/model-builder/predict_text_entity_extraction_sample_test.py @@ -30,4 +30,6 @@ def test_predict_text_entity_extraction_sample(mock_sdk_init, mock_get_endpoint) project=constants.PROJECT, location=constants.LOCATION ) - mock_get_endpoint.assert_called_once_with(constants.ENDPOINT_NAME,) + mock_get_endpoint.assert_called_once_with( + constants.ENDPOINT_NAME, + ) diff --git a/samples/model-builder/predict_text_sentiment_analysis_sample_test.py b/samples/model-builder/predict_text_sentiment_analysis_sample_test.py index e3a3fad58c..c2ed180c9f 100644 --- a/samples/model-builder/predict_text_sentiment_analysis_sample_test.py +++ b/samples/model-builder/predict_text_sentiment_analysis_sample_test.py @@ -30,4 +30,6 @@ def test_predict_text_sentiment_analysis_sample(mock_sdk_init, mock_get_endpoint project=constants.PROJECT, location=constants.LOCATION ) - mock_get_endpoint.assert_called_once_with(constants.ENDPOINT_NAME,) + mock_get_endpoint.assert_called_once_with( + constants.ENDPOINT_NAME, + ) diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py index 2e50ef62e6..6f13f137bd 100644 --- a/samples/model-builder/test_constants.py +++ b/samples/model-builder/test_constants.py @@ -17,6 +17,7 @@ from google.auth import credentials from google.cloud import aiplatform +from google.protobuf import timestamp_pb2 PROJECT = "abc" LOCATION = "us-central1" @@ -255,3 +256,21 @@ QUANTILES = [0, 0.5, 1] VALIDATION_OPTIONS = "fail-pipeline" PREDEFINED_SPLIT_COLUMN_NAME = "predefined" + +TENSORBOARD_NAME = ( + f"projects/{PROJECT}/locations/{LOCATION}/tensorboards/my-tensorboard" +) + +SCHEMA_TITLE = "system.Schema" +SCHEMA_VERSION = "0.0.1" +METADATA = {} + +EXPERIMENT_RUN_NAME = "my-run" + +METRICS = {"accuracy": 0.1} +PARAMS = {"learning_rate": 0.1} + +TEMPLATE_PATH = "pipeline.json" + +STEP = 1 +TIMESTAMP = timestamp_pb2.Timestamp() diff --git a/tests/system/aiplatform/e2e_base.py b/tests/system/aiplatform/e2e_base.py index 93cdd4fe9e..26d67e9b66 100644 --- a/tests/system/aiplatform/e2e_base.py +++ b/tests/system/aiplatform/e2e_base.py @@ -17,6 +17,7 @@ import abc import importlib +import logging import os import pytest import uuid @@ -26,6 +27,7 @@ from google.api_core import exceptions from google.cloud import aiplatform from google.cloud import bigquery +from google.cloud import resourcemanager from google.cloud import storage from google.cloud.aiplatform import initializer @@ -78,9 +80,29 @@ def prepare_staging_bucket( storage_client = storage.Client(project=_PROJECT) shared_state["storage_client"] = storage_client - shared_state["bucket"] = storage_client.create_bucket( + bucket = storage_client.create_bucket( staging_bucket_name, project=_PROJECT, location=_LOCATION ) + + # TODO(#1415) Once PR Is merged, use the added utilities to + # provide create/view access to Pipeline's default service account (compute) + project_number = ( + resourcemanager.ProjectsClient() + .get_project(name=f"projects/{_PROJECT}") + .name.split("/", 1)[1] + ) + + service_account = f"{project_number}-compute@developer.gserviceaccount.com" + bucket_iam_policy = bucket.get_iam_policy() + bucket_iam_policy.setdefault("roles/storage.objectCreator", set()).add( + f"serviceAccount:{service_account}" + ) + bucket_iam_policy.setdefault("roles/storage.objectViewer", set()).add( + f"serviceAccount:{service_account}" + ) + bucket.set_iam_policy(bucket_iam_policy) + + shared_state["bucket"] = bucket yield @pytest.fixture(scope="class") @@ -177,4 +199,4 @@ def tear_down_resources(self, shared_state: Dict[str, Any]): else: resource.delete() except exceptions.GoogleAPIError as e: - print(f"Could not delete resource: {resource} due to: {e}") + logging.error(f"Could not delete resource: {resource} due to: {e}") diff --git a/tests/system/aiplatform/test_experiments.py b/tests/system/aiplatform/test_experiments.py new file mode 100644 index 0000000000..7bb66eb146 --- /dev/null +++ b/tests/system/aiplatform/test_experiments.py @@ -0,0 +1,306 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 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. +# +import tempfile + +import pytest + +from google.api_core import exceptions +from google.cloud import storage + +from google.cloud import aiplatform +from google.cloud.aiplatform.utils import rest_utils +from tests.system.aiplatform import e2e_base +from tests.system.aiplatform import test_model_upload + + +_RUN = "run-1" +_PARAMS = {"sdk-param-test-1": 0.1, "sdk-param-test-2": 0.2} +_METRICS = {"sdk-metric-test-1": 0.8, "sdk-metric-test-2": 100.0} + +_RUN_2 = "run-2" +_PARAMS_2 = {"sdk-param-test-1": 0.2, "sdk-param-test-2": 0.4} +_METRICS_2 = {"sdk-metric-test-1": 1.6, "sdk-metric-test-2": 200.0} + +_TIME_SERIES_METRIC_KEY = "accuracy" + + +@pytest.mark.usefixtures( + "prepare_staging_bucket", "delete_staging_bucket", "tear_down_resources" +) +class TestExperiments(e2e_base.TestEndToEnd): + + _temp_prefix = "tmpvrtxsdk-e2e" + + def setup_class(cls): + cls._experiment_name = cls._make_display_name("experiment")[:30] + cls._dataset_artifact_name = cls._make_display_name("ds-artifact")[:30] + cls._dataset_artifact_uri = cls._make_display_name("ds-uri") + cls._pipeline_job_id = cls._make_display_name("job-id") + + def test_create_experiment(self, shared_state): + + # Truncating the name because of resource id constraints from the service + tensorboard = aiplatform.Tensorboard.create( + project=e2e_base._PROJECT, + location=e2e_base._LOCATION, + display_name=self._experiment_name, + ) + + shared_state["resources"] = [tensorboard] + + aiplatform.init( + project=e2e_base._PROJECT, + location=e2e_base._LOCATION, + experiment=self._experiment_name, + experiment_tensorboard=tensorboard, + ) + + shared_state["resources"].append( + aiplatform.metadata.metadata._experiment_tracker.experiment + ) + + def test_get_experiment(self): + experiment = aiplatform.Experiment(experiment_name=self._experiment_name) + assert experiment.name == self._experiment_name + + def test_start_run(self): + run = aiplatform.start_run(_RUN) + assert run.name == _RUN + + def test_get_run(self): + run = aiplatform.ExperimentRun(run_name=_RUN, experiment=self._experiment_name) + assert run.name == _RUN + assert run.state == aiplatform.gapic.Execution.State.RUNNING + + def test_log_params(self): + aiplatform.log_params(_PARAMS) + run = aiplatform.ExperimentRun(run_name=_RUN, experiment=self._experiment_name) + assert run.get_params() == _PARAMS + + def test_log_metrics(self): + aiplatform.log_metrics(_METRICS) + run = aiplatform.ExperimentRun(run_name=_RUN, experiment=self._experiment_name) + assert run.get_metrics() == _METRICS + + def test_log_time_series_metrics(self): + for i in range(5): + aiplatform.log_time_series_metrics({_TIME_SERIES_METRIC_KEY: i}) + + run = aiplatform.ExperimentRun(run_name=_RUN, experiment=self._experiment_name) + + time_series_result = run.get_time_series_data_frame()[ + [_TIME_SERIES_METRIC_KEY, "step"] + ].to_dict("list") + + assert time_series_result == { + "step": list(range(1, 6)), + _TIME_SERIES_METRIC_KEY: [float(value) for value in range(5)], + } + + def test_create_artifact(self, shared_state): + ds = aiplatform.Artifact.create( + schema_title="system.Dataset", + resource_id=self._dataset_artifact_name, + uri=self._dataset_artifact_uri, + ) + + shared_state["resources"].append(ds) + assert ds.uri == self._dataset_artifact_uri + + def test_get_artifact_by_uri(self): + ds = aiplatform.Artifact.get_with_uri(uri=self._dataset_artifact_uri) + + assert ds.uri == self._dataset_artifact_uri + assert ds.name == self._dataset_artifact_name + + def test_log_execution_and_artifact(self, shared_state): + with aiplatform.start_execution( + schema_title="system.ContainerExecution", + resource_id=self._make_display_name("execution"), + ) as execution: + + shared_state["resources"].append(execution) + + ds = aiplatform.Artifact(artifact_name=self._dataset_artifact_name) + execution.assign_input_artifacts([ds]) + + model = aiplatform.Artifact.create(schema_title="system.Model") + shared_state["resources"].append(model) + + storage_client = storage.Client(project=e2e_base._PROJECT) + model_blob = storage.Blob.from_string( + uri=test_model_upload._XGBOOST_MODEL_URI, client=storage_client + ) + model_path = tempfile.mktemp() + ".my_model.xgb" + model_blob.download_to_filename(filename=model_path) + + vertex_model = aiplatform.Model.upload_xgboost_model_file( + display_name=self._make_display_name("model"), + model_file_path=model_path, + ) + shared_state["resources"].append(vertex_model) + + execution.assign_output_artifacts([model, vertex_model]) + + input_artifacts = execution.get_input_artifacts() + assert input_artifacts[0].name == ds.name + + output_artifacts = execution.get_output_artifacts() + # system.Model, google.VertexModel + output_artifacts.sort(key=lambda artifact: artifact.schema_title, reverse=True) + + shared_state["resources"].append(output_artifacts[-1]) + + assert output_artifacts[0].name == model.name + assert output_artifacts[1].uri == rest_utils.make_gcp_resource_rest_url( + resource=vertex_model + ) + + run = aiplatform.ExperimentRun(run_name=_RUN, experiment=self._experiment_name) + executions = run.get_executions() + assert executions[0].name == execution.name + + artifacts = run.get_artifacts() + + # system.Model, system.Dataset, google.VertexTensorboardRun, google.VertexModel + artifacts.sort(key=lambda artifact: artifact.schema_title, reverse=True) + assert artifacts.pop().uri == rest_utils.make_gcp_resource_rest_url( + resource=vertex_model + ) + + # tensorboard run artifact is also included + assert sorted([artifact.name for artifact in artifacts]) == sorted( + [ds.name, model.name, run._tensorboard_run_id(run.resource_id)] + ) + + def test_end_run(self): + aiplatform.end_run() + run = aiplatform.ExperimentRun(run_name=_RUN, experiment=self._experiment_name) + assert run.state == aiplatform.gapic.Execution.State.COMPLETE + + def test_run_context_manager(self): + with aiplatform.start_run(_RUN_2) as run: + run.log_params(_PARAMS_2) + run.log_metrics(_METRICS_2) + assert run.state == aiplatform.gapic.Execution.State.RUNNING + + assert run.state == aiplatform.gapic.Execution.State.COMPLETE + + def test_add_pipeline_job_to_experiment(self, shared_state): + import kfp.v2.dsl as dsl + import kfp.v2.compiler as compiler + from kfp.v2.dsl import component, Metrics, Output + + @component + def trainer( + learning_rate: float, dropout_rate: float, metrics: Output[Metrics] + ): + metrics.log_metric("accuracy", 0.8) + metrics.log_metric("mse", 1.2) + + @dsl.pipeline(name=self._make_display_name("pipeline")) + def pipeline(learning_rate: float, dropout_rate: float): + trainer(learning_rate=learning_rate, dropout_rate=dropout_rate) + + compiler.Compiler().compile( + pipeline_func=pipeline, package_path="pipeline.json" + ) + + job = aiplatform.PipelineJob( + display_name=self._make_display_name("experiment pipeline job"), + template_path="pipeline.json", + job_id=self._pipeline_job_id, + pipeline_root=f'gs://{shared_state["staging_bucket_name"]}', + parameter_values={"learning_rate": 0.1, "dropout_rate": 0.2}, + ) + + job.submit(experiment=self._experiment_name) + + shared_state["resources"].append(job) + + job.wait() + + def test_get_experiments_df(self): + df = aiplatform.get_experiment_df() + + pipelines_param_and_metrics = { + "param.dropout_rate": 0.2, + "param.learning_rate": 0.1, + "metric.accuracy": 0.8, + "metric.mse": 1.2, + } + + true_df_dict_1 = {f"metric.{key}": value for key, value in _METRICS.items()} + for key, value in _PARAMS.items(): + true_df_dict_1[f"param.{key}"] = value + + true_df_dict_1["experiment_name"] = self._experiment_name + true_df_dict_1["run_name"] = _RUN + true_df_dict_1["state"] = aiplatform.gapic.Execution.State.COMPLETE.name + true_df_dict_1["run_type"] = aiplatform.metadata.constants.SYSTEM_EXPERIMENT_RUN + true_df_dict_1[f"time_series_metric.{_TIME_SERIES_METRIC_KEY}"] = 4.0 + + true_df_dict_2 = {f"metric.{key}": value for key, value in _METRICS_2.items()} + for key, value in _PARAMS_2.items(): + true_df_dict_2[f"param.{key}"] = value + + true_df_dict_2["experiment_name"] = self._experiment_name + true_df_dict_2["run_name"] = _RUN_2 + true_df_dict_2["state"] = aiplatform.gapic.Execution.State.COMPLETE.name + true_df_dict_2["run_type"] = aiplatform.metadata.constants.SYSTEM_EXPERIMENT_RUN + true_df_dict_2[f"time_series_metric.{_TIME_SERIES_METRIC_KEY}"] = 0.0 + + # TODO(remove when running CI) + + true_df_dict_3 = { + "experiment_name": self._experiment_name, + "run_name": self._pipeline_job_id, + "run_type": aiplatform.metadata.constants.SYSTEM_PIPELINE_RUN, + "state": aiplatform.gapic.Execution.State.COMPLETE.name, + "time_series_metric.accuracy": 0.0, + } + + true_df_dict_3.update(pipelines_param_and_metrics) + + for key in pipelines_param_and_metrics.keys(): + true_df_dict_1[key] = 0.0 + true_df_dict_2[key] = 0.0 + + for key in _PARAMS.keys(): + true_df_dict_3[f"param.{key}"] = 0.0 + + for key in _METRICS.keys(): + true_df_dict_3[f"metric.{key}"] = 0.0 + + assert sorted( + [true_df_dict_1, true_df_dict_2, true_df_dict_3], + key=lambda d: d["run_name"], + ) == sorted(df.fillna(0.0).to_dict("records"), key=lambda d: d["run_name"]) + + def test_delete_run(self): + run = aiplatform.ExperimentRun(run_name=_RUN, experiment=self._experiment_name) + run.delete(delete_backing_tensorboard_run=True) + + with pytest.raises(exceptions.NotFound): + aiplatform.ExperimentRun(run_name=_RUN, experiment=self._experiment_name) + + def test_delete_experiment(self): + experiment = aiplatform.Experiment(experiment_name=self._experiment_name) + experiment.delete(delete_backing_tensorboard_runs=True) + + with pytest.raises(exceptions.NotFound): + aiplatform.Experiment(experiment_name=self._experiment_name) diff --git a/tests/system/aiplatform/test_tensorboard.py b/tests/system/aiplatform/test_tensorboard.py index a1f4634bd9..e1262640dc 100644 --- a/tests/system/aiplatform/test_tensorboard.py +++ b/tests/system/aiplatform/test_tensorboard.py @@ -41,6 +41,7 @@ def test_create_and_get_tensorboard(self, shared_state): ) shared_state["resources"] = [tb] + shared_state["tensorboard"] = tb get_tb = aiplatform.Tensorboard(tb.resource_name) @@ -50,6 +51,10 @@ def test_create_and_get_tensorboard(self, shared_state): assert len(list_tb) > 0 + def test_create_and_get_tensorboard_experiment(self, shared_state): + assert shared_state["tensorboard"] + tb = shared_state["tensorboard"] + tb_experiment = aiplatform.TensorboardExperiment.create( tensorboard_experiment_id="vertex-sdk-e2e-test-experiment", tensorboard_name=tb.resource_name, @@ -60,6 +65,7 @@ def test_create_and_get_tensorboard(self, shared_state): ) shared_state["resources"].append(tb_experiment) + shared_state["tensorboard_experiment"] = tb_experiment get_tb_experiment = aiplatform.TensorboardExperiment( tb_experiment.resource_name @@ -73,6 +79,10 @@ def test_create_and_get_tensorboard(self, shared_state): assert len(list_tb_experiment) > 0 + def test_create_and_get_tensorboard_run(self, shared_state): + assert shared_state["tensorboard_experiment"] + tb_experiment = shared_state["tensorboard_experiment"] + tb_run = aiplatform.TensorboardRun.create( tensorboard_run_id="test-run", tensorboard_experiment_name=tb_experiment.resource_name, @@ -82,6 +92,7 @@ def test_create_and_get_tensorboard(self, shared_state): ) shared_state["resources"].append(tb_run) + shared_state["tensorboard_run"] = tb_run get_tb_run = aiplatform.TensorboardRun(tb_run.resource_name) @@ -92,3 +103,39 @@ def test_create_and_get_tensorboard(self, shared_state): ) assert len(list_tb_run) > 0 + + def test_create_and_get_tensorboard_time_series(self, shared_state): + assert shared_state["tensorboard_run"] + tb_run = shared_state["tensorboard_run"] + + tb_time_series = aiplatform.TensorboardTimeSeries.create( + display_name="test-time-series", + tensorboard_run_name=tb_run.resource_name, + description="Vertex SDK Integration test run", + ) + + shared_state["resources"].append(tb_time_series) + shared_state["tensorboard_time_series"] = tb_time_series + + get_tb_time_series = aiplatform.TensorboardTimeSeries( + tb_time_series.resource_name + ) + + assert tb_time_series.resource_name == get_tb_time_series.resource_name + + list_tb_time_series = aiplatform.TensorboardTimeSeries.list( + tensorboard_run_name=tb_run.resource_name + ) + + assert len(list_tb_time_series) > 0 + + def test_write_tensorboard_scalar_data(self, shared_state): + assert shared_state["tensorboard_time_series"] + assert shared_state["tensorboard_run"] + tb_run = shared_state["tensorboard_run"] + tb_time_series = shared_state["tensorboard_time_series"] + + tb_run.write_tensorboard_scalar_data( + time_series_data={tb_time_series.display_name: 1.0}, + step=1, + ) diff --git a/tests/unit/aiplatform/test_initializer.py b/tests/unit/aiplatform/test_initializer.py index 6a31a316e1..edc78c6967 100644 --- a/tests/unit/aiplatform/test_initializer.py +++ b/tests/unit/aiplatform/test_initializer.py @@ -25,7 +25,7 @@ from google.auth import credentials from google.cloud.aiplatform import initializer -from google.cloud.aiplatform.metadata.metadata import metadata_service +from google.cloud.aiplatform.metadata.metadata import _experiment_tracker from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform import utils from google.cloud.aiplatform.utils import resource_manager_utils @@ -90,14 +90,14 @@ def test_init_location_with_invalid_location_raises(self): with pytest.raises(ValueError): initializer.global_config.init(location=_TEST_INVALID_LOCATION) - @patch.object(metadata_service, "set_experiment") + @patch.object(_experiment_tracker, "set_experiment") def test_init_experiment_sets_experiment(self, set_experiment_mock): initializer.global_config.init(experiment=_TEST_EXPERIMENT) set_experiment_mock.assert_called_once_with( - experiment=_TEST_EXPERIMENT, description=None + experiment=_TEST_EXPERIMENT, description=None, backing_tensorboard=None ) - @patch.object(metadata_service, "set_experiment") + @patch.object(_experiment_tracker, "set_experiment") def test_init_experiment_sets_experiment_with_description( self, set_experiment_mock ): @@ -105,7 +105,9 @@ def test_init_experiment_sets_experiment_with_description( experiment=_TEST_EXPERIMENT, experiment_description=_TEST_DESCRIPTION ) set_experiment_mock.assert_called_once_with( - experiment=_TEST_EXPERIMENT, description=_TEST_DESCRIPTION + experiment=_TEST_EXPERIMENT, + description=_TEST_DESCRIPTION, + backing_tensorboard=None, ) def test_init_experiment_description_fail_without_experiment(self): diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py index acbfa9098b..6fb37e20f3 100644 --- a/tests/unit/aiplatform/test_metadata.py +++ b/tests/unit/aiplatform/test_metadata.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # - +import copy from importlib import reload from unittest import mock from unittest.mock import patch, call @@ -24,27 +24,44 @@ from google.api_core import operation from google.auth import credentials +import google.cloud.aiplatform.metadata.constants from google.cloud import aiplatform from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer +from google.cloud.aiplatform.compat.types import event as gca_event +from google.cloud.aiplatform.compat.types import execution as gca_execution +from google.cloud.aiplatform.compat.types import ( + tensorboard_data as gca_tensorboard_data, +) +from google.cloud.aiplatform.compat.types import ( + tensorboard_experiment as gca_tensorboard_experiment, +) +from google.cloud.aiplatform.compat.types import tensorboard_run as gca_tensorboard_run +from google.cloud.aiplatform.compat.types import ( + tensorboard_time_series as gca_tensorboard_time_series, +) from google.cloud.aiplatform.metadata import constants +from google.cloud.aiplatform.metadata import experiment_run_resource from google.cloud.aiplatform.metadata import metadata -from google.cloud.aiplatform_v1 import ( - AddContextArtifactsAndExecutionsResponse, - Event, - LineageSubgraph, - ListExecutionsRequest, -) +from google.cloud.aiplatform.metadata import metadata_store +from google.cloud.aiplatform.metadata import utils as metadata_utils +from google.cloud.aiplatform import utils + +from google.cloud.aiplatform_v1 import AddContextArtifactsAndExecutionsResponse +from google.cloud.aiplatform_v1 import AddExecutionEventsResponse from google.cloud.aiplatform_v1 import Artifact as GapicArtifact from google.cloud.aiplatform_v1 import Context as GapicContext from google.cloud.aiplatform_v1 import Execution as GapicExecution -from google.cloud.aiplatform_v1 import ( - MetadataServiceClient, - AddExecutionEventsResponse, -) +from google.cloud.aiplatform_v1 import LineageSubgraph +from google.cloud.aiplatform_v1 import MetadataServiceClient from google.cloud.aiplatform_v1 import MetadataStore as GapicMetadataStore +from google.cloud.aiplatform_v1 import TensorboardServiceClient + +from test_pipeline_jobs import mock_pipeline_service_get # noqa: F401 +from test_pipeline_jobs import _TEST_PIPELINE_JOB_NAME # noqa: F401 -# project +import test_pipeline_jobs +import test_tensorboard _TEST_PROJECT = "test-project" _TEST_OTHER_PROJECT = "test-project-1" @@ -191,17 +208,23 @@ def get_context_not_found_mock(): yield get_context_not_found_mock +_TEST_EXPERIMENT_CONTEXT = GapicContext( + name=_TEST_CONTEXT_NAME, + display_name=_TEST_EXPERIMENT, + description=_TEST_EXPERIMENT_DESCRIPTION, + schema_title=constants.SYSTEM_EXPERIMENT, + schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_EXPERIMENT], + metadata={ + **constants.EXPERIMENT_METADATA, + constants._BACKING_TENSORBOARD_RESOURCE_KEY: test_tensorboard._TEST_NAME, + }, +) + + @pytest.fixture def update_context_mock(): with patch.object(MetadataServiceClient, "update_context") as update_context_mock: - update_context_mock.return_value = GapicContext( - name=_TEST_CONTEXT_NAME, - display_name=_TEST_EXPERIMENT, - description=_TEST_OTHER_EXPERIMENT_DESCRIPTION, - schema_title=constants.SYSTEM_EXPERIMENT, - schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_EXPERIMENT], - metadata=constants.EXPERIMENT_METADATA, - ) + update_context_mock.return_value = _TEST_EXPERIMENT_CONTEXT yield update_context_mock @@ -288,6 +311,42 @@ def list_executions_mock(): yield list_executions_mock +@pytest.fixture +def get_tensorboard_run_not_found_mock(): + with patch.object( + TensorboardServiceClient, "get_tensorboard_run" + ) as get_tensorboard_run_mock: + get_tensorboard_run_mock.side_effect = [ + exceptions.NotFound(""), + test_tensorboard._TEST_TENSORBOARD_RUN, + ] + yield get_tensorboard_run_mock + + +@pytest.fixture +def get_tensorboard_experiment_not_found_mock(): + with patch.object( + TensorboardServiceClient, "get_tensorboard_experiment" + ) as get_tensorboard_experiment_mock: + get_tensorboard_experiment_mock.side_effect = [ + exceptions.NotFound(""), + test_tensorboard._TEST_TENSORBOARD_EXPERIMENT, + ] + yield get_tensorboard_experiment_mock + + +@pytest.fixture +def get_tensorboard_time_series_not_found_mock(): + with patch.object( + TensorboardServiceClient, "get_tensorboard_time_series" + ) as get_tensorboard_time_series_mock: + get_tensorboard_time_series_mock.side_effect = [ + exceptions.NotFound(""), + # test_tensorboard._TEST_TENSORBOARD_TIME_SERIES # change to time series + ] + yield get_tensorboard_time_series_mock + + @pytest.fixture def query_execution_inputs_and_outputs_mock(): with patch.object( @@ -304,7 +363,14 @@ def query_execution_inputs_and_outputs_mock(): constants.SYSTEM_METRICS ], metadata=_TEST_METRICS, - ), + ) + ], + events=[ + gca_event.Event( + artifact=_TEST_ARTIFACT_NAME, + execution=_TEST_EXECUTION_NAME, + type_=gca_event.Event.Type.OUTPUT, + ) ], ), LineageSubgraph( @@ -319,6 +385,13 @@ def query_execution_inputs_and_outputs_mock(): metadata=_TEST_OTHER_METRICS, ), ], + events=[ + gca_event.Event( + artifact=_TEST_OTHER_ARTIFACT_NAME, + execution=_TEST_OTHER_EXECUTION_NAME, + type_=gca_event.Event.Type.OUTPUT, + ) + ], ), ] yield query_execution_inputs_and_outputs_mock @@ -336,6 +409,13 @@ def get_artifact_mock(): yield get_artifact_mock +@pytest.fixture +def get_artifact_not_found_mock(): + with patch.object(MetadataServiceClient, "get_artifact") as get_artifact_mock: + get_artifact_mock.side_effect = exceptions.NotFound("") + yield get_artifact_mock + + @pytest.fixture def get_artifact_wrong_schema_mock(): with patch.object( @@ -387,22 +467,405 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) + @pytest.mark.usefixtures("get_pipeline_context_mock") + def test_get_pipeline_df( + self, list_executions_mock, query_execution_inputs_and_outputs_mock + ): + try: + import pandas as pd + except ImportError: + raise ImportError( + "Pandas is not installed and is required to test the get_pipeline_df method. " + 'Please install the SDK using "pip install python-aiplatform[full]"' + ) + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + + pipeline_df = aiplatform.get_pipeline_df(_TEST_PIPELINE) + + expected_filter = metadata_utils._make_filter_string( + schema_title=constants.SYSTEM_RUN, in_context=[_TEST_CONTEXT_NAME] + ) + + list_executions_mock.assert_called_once_with( + request={"parent": _TEST_PARENT, "filter": expected_filter} + ) + query_execution_inputs_and_outputs_mock.assert_has_calls( + [ + call(execution=_TEST_EXECUTION_NAME), + call(execution=_TEST_OTHER_EXECUTION_NAME), + ] + ) + pipeline_df_truth = pd.DataFrame( + [ + { + "pipeline_name": _TEST_PIPELINE, + "run_name": _TEST_RUN, + "param.%s" % _TEST_PARAM_KEY_1: 0.01, + "param.%s" % _TEST_PARAM_KEY_2: 0.2, + "metric.%s" % _TEST_METRIC_KEY_1: 222, + "metric.%s" % _TEST_METRIC_KEY_2: 1, + }, + { + "pipeline_name": _TEST_PIPELINE, + "run_name": _TEST_OTHER_RUN, + "param.%s" % _TEST_PARAM_KEY_1: 0.02, + "param.%s" % _TEST_PARAM_KEY_2: 0.3, + "metric.%s" % _TEST_METRIC_KEY_2: 0.9, + }, + ] + ) + + _assert_frame_equal_with_sorted_columns(pipeline_df, pipeline_df_truth) + + @pytest.mark.usefixtures("get_context_not_found_mock") + def test_get_pipeline_df_not_exist(self): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + with pytest.raises(exceptions.NotFound): + aiplatform.get_pipeline_df(_TEST_PIPELINE) + + @pytest.mark.usefixtures("get_context_mock") + def test_get_pipeline_df_wrong_schema(self): + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + with pytest.raises(ValueError): + aiplatform.get_pipeline_df(_TEST_PIPELINE) + + +_TEST_EXPERIMENT_RUN_CONTEXT_NAME = f"{_TEST_PARENT}/contexts/{_TEST_EXECUTION_ID}" +_TEST_OTHER_EXPERIMENT_RUN_CONTEXT_NAME = ( + f"{_TEST_PARENT}/contexts/{_TEST_OTHER_EXECUTION_ID}" +) + +_EXPERIMENT_MOCK = GapicContext( + name=_TEST_CONTEXT_NAME, + display_name=_TEST_EXPERIMENT, + description=_TEST_EXPERIMENT_DESCRIPTION, + schema_title=constants.SYSTEM_EXPERIMENT, + schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_EXPERIMENT], + metadata={**constants.EXPERIMENT_METADATA}, +) + +_EXPERIMENT_RUN_MOCK = GapicContext( + name=_TEST_EXPERIMENT_RUN_CONTEXT_NAME, + display_name=_TEST_RUN, + schema_title=constants.SYSTEM_EXPERIMENT_RUN, + schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_EXPERIMENT_RUN], + metadata={ + constants._PARAM_KEY: {}, + constants._METRIC_KEY: {}, + constants._STATE_KEY: gca_execution.Execution.State.RUNNING.name, + }, +) + +_EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT = copy.deepcopy(_EXPERIMENT_RUN_MOCK) +_EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT.parent_contexts = [_TEST_CONTEXT_NAME] + + +@pytest.fixture +def get_experiment_mock(): + with patch.object(MetadataServiceClient, "get_context") as get_context_mock: + get_context_mock.return_value = _EXPERIMENT_MOCK + yield get_context_mock + + +@pytest.fixture +def get_experiment_run_run_mock(): + with patch.object(MetadataServiceClient, "get_context") as get_context_mock: + get_context_mock.side_effect = [ + _EXPERIMENT_MOCK, + _EXPERIMENT_RUN_MOCK, + _EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT, + ] + + yield get_context_mock + + +@pytest.fixture +def get_experiment_run_mock(): + with patch.object(MetadataServiceClient, "get_context") as get_context_mock: + get_context_mock.side_effect = [ + _EXPERIMENT_MOCK, + _EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT, + ] + + yield get_context_mock + + +@pytest.fixture +def create_experiment_run_context_mock(): + with patch.object(MetadataServiceClient, "create_context") as create_context_mock: + create_context_mock.side_effect = [_EXPERIMENT_RUN_MOCK] + yield create_context_mock + + +@pytest.fixture +def update_experiment_run_context_to_running(): + with patch.object(MetadataServiceClient, "update_context") as update_context_mock: + update_context_mock.side_effect = [_EXPERIMENT_RUN_MOCK] + yield update_context_mock + + +@pytest.fixture +def update_context_mock_v2(): + with patch.object(MetadataServiceClient, "update_context") as update_context_mock: + update_context_mock.side_effect = [ + # experiment run + GapicContext( + name=_TEST_EXPERIMENT_RUN_CONTEXT_NAME, + display_name=_TEST_RUN, + schema_title=constants.SYSTEM_EXPERIMENT_RUN, + schema_version=constants.SCHEMA_VERSIONS[ + constants.SYSTEM_EXPERIMENT_RUN + ], + metadata={**constants.EXPERIMENT_METADATA}, + ), + # experiment run + GapicContext( + name=_TEST_EXPERIMENT_RUN_CONTEXT_NAME, + display_name=_TEST_RUN, + schema_title=constants.SYSTEM_EXPERIMENT_RUN, + schema_version=constants.SCHEMA_VERSIONS[ + constants.SYSTEM_EXPERIMENT_RUN + ], + metadata=constants.EXPERIMENT_METADATA, + parent_contexts=[_TEST_CONTEXT_NAME], + ), + ] + + yield update_context_mock + + +@pytest.fixture +def list_contexts_mock(): + with patch.object(MetadataServiceClient, "list_contexts") as list_contexts_mock: + list_contexts_mock.return_value = [ + GapicContext( + name=_TEST_EXPERIMENT_RUN_CONTEXT_NAME, + display_name=_TEST_RUN, + schema_title=constants.SYSTEM_EXPERIMENT_RUN, + schema_version=constants.SCHEMA_VERSIONS[ + constants.SYSTEM_EXPERIMENT_RUN + ], + metadata=constants.EXPERIMENT_METADATA, + parent_contexts=[_TEST_CONTEXT_NAME], + ), + GapicContext( + name=_TEST_OTHER_EXPERIMENT_RUN_CONTEXT_NAME, + display_name=_TEST_OTHER_RUN, + schema_title=constants.SYSTEM_EXPERIMENT_RUN, + schema_version=constants.SCHEMA_VERSIONS[ + constants.SYSTEM_EXPERIMENT_RUN + ], + metadata=constants.EXPERIMENT_METADATA, + parent_contexts=[_TEST_CONTEXT_NAME], + ), + ] + yield list_contexts_mock + + +@pytest.fixture +def add_context_children_mock(): + with patch.object( + MetadataServiceClient, "add_context_children" + ) as add_context_children_mock: + yield add_context_children_mock + + +_EXPERIMENT_RUN_MOCK_POPULATED_1 = copy.deepcopy( + _EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT +) +_EXPERIMENT_RUN_MOCK_POPULATED_1.metadata[constants._PARAM_KEY].update(_TEST_PARAMS) +_EXPERIMENT_RUN_MOCK_POPULATED_1.metadata[constants._METRIC_KEY].update(_TEST_METRICS) +_EXPERIMENT_RUN_MOCK_POPULATED_2 = copy.deepcopy( + _EXPERIMENT_RUN_MOCK_WITH_PARENT_EXPERIMENT +) +_EXPERIMENT_RUN_MOCK_POPULATED_2.display_name = _TEST_OTHER_RUN +_EXPERIMENT_RUN_MOCK_POPULATED_2.metadata[constants._PARAM_KEY].update( + _TEST_OTHER_PARAMS +) +_EXPERIMENT_RUN_MOCK_POPULATED_2.metadata[constants._METRIC_KEY].update( + _TEST_OTHER_METRICS +) + +_TEST_PIPELINE_RUN_ID = "test-pipeline-run" +_TEST_PIPELINE_RUN_CONTEXT_NAME = f"{_TEST_PARENT}/contexts/{_TEST_PIPELINE_RUN_ID}" + +_TEST_PIPELINE_CONTEXT = GapicContext( + name=_TEST_PIPELINE_RUN_CONTEXT_NAME, + display_name=_TEST_PIPELINE_RUN_ID, + schema_title=constants.SYSTEM_PIPELINE_RUN, + parent_contexts=[_TEST_CONTEXT_NAME], +) + + +@pytest.fixture() +def list_context_mock_for_experiment_dataframe_mock(): + with patch.object(MetadataServiceClient, "list_contexts") as list_context_mock: + list_context_mock.side_effect = [ + # experiment runs + [ + _EXPERIMENT_RUN_MOCK_POPULATED_1, + _EXPERIMENT_RUN_MOCK_POPULATED_2, + _TEST_PIPELINE_CONTEXT, + ], + # pipeline runs + [], + ] + yield list_context_mock + + +_TEST_LEGACY_METRIC_ARTIFACT = GapicArtifact( + name=_TEST_ARTIFACT_NAME, + schema_title=constants.SYSTEM_METRICS, + metadata=_TEST_METRICS, +) + +_TEST_PIPELINE_METRIC_ARTIFACT = GapicArtifact( + name=_TEST_ARTIFACT_NAME, + schema_title=constants.SYSTEM_METRICS, + metadata={key: value + 1 for key, value in _TEST_METRICS.items()}, +) + + +@pytest.fixture() +def list_artifact_mock_for_experiment_dataframe(): + with patch.object(MetadataServiceClient, "list_artifacts") as list_artifacts_mock: + list_artifacts_mock.side_effect = [ + # pipeline run metric artifact + [_TEST_PIPELINE_METRIC_ARTIFACT], + ] + yield list_artifacts_mock + + +_TEST_PIPELINE_SYSTEM_RUN_EXECUTION = GapicExecution( + name=_TEST_EXECUTION_NAME, + schema_title=constants.SYSTEM_RUN, + state=gca_execution.Execution.State.RUNNING, + metadata={f"input:{key}": value + 1 for key, value in _TEST_PARAMS.items()}, +) + +_TEST_LEGACY_SYSTEM_RUN_EXECUTION = GapicExecution( + name=_TEST_EXECUTION_NAME, + display_name=_TEST_RUN, + schema_title=constants.SYSTEM_RUN, + schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_RUN], + metadata=_TEST_PARAMS, +) + + +# backward compatibility +@pytest.fixture() +def list_executions_mock_for_experiment_dataframe(): + with patch.object(MetadataServiceClient, "list_executions") as list_executions_mock: + list_executions_mock.side_effect = [ + # legacy system.run execution + [_TEST_LEGACY_SYSTEM_RUN_EXECUTION], + # pipeline system.run execution + [_TEST_PIPELINE_SYSTEM_RUN_EXECUTION], + ] + yield list_executions_mock + + +@pytest.fixture +def get_tensorboard_run_artifact_not_found_mock(): + with patch.object(MetadataServiceClient, "get_artifact") as get_artifact_mock: + get_artifact_mock.side_effect = exceptions.NotFound("") + yield get_artifact_mock + + +_TEST_LEGACY_METRIC_ARTIFACT + +_TEST_TENSORBOARD_RUN_ARTIFACT = GapicArtifact( + name=experiment_run_resource.ExperimentRun._tensorboard_run_id( + _TEST_EXPERIMENT_RUN_CONTEXT_NAME + ), + uri="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/tensorboards/1028944691210842416/experiments/test-experiment/runs/test-run", + schema_title=google.cloud.aiplatform.metadata.constants._TENSORBOARD_RUN_REFERENCE_ARTIFACT.schema_title, + schema_version=google.cloud.aiplatform.metadata.constants._TENSORBOARD_RUN_REFERENCE_ARTIFACT.schema_version, + state=GapicArtifact.State.LIVE, + metadata={ + google.cloud.aiplatform.metadata.constants._VERTEX_EXPERIMENT_TRACKING_LABEL: True, + constants.GCP_ARTIFACT_RESOURCE_NAME_KEY: test_tensorboard._TEST_TENSORBOARD_RUN_NAME, + }, +) + +get_tensorboard_mock = test_tensorboard.get_tensorboard_mock +create_tensorboard_experiment_mock = test_tensorboard.create_tensorboard_experiment_mock +create_tensorboard_run_mock = test_tensorboard.create_tensorboard_run_mock +write_tensorboard_run_data_mock = test_tensorboard.write_tensorboard_run_data_mock +create_tensorboard_time_series_mock = ( + test_tensorboard.create_tensorboard_time_series_mock +) + +get_tensorboard_run_mock = test_tensorboard.get_tensorboard_run_mock +list_tensorboard_time_series_mock = test_tensorboard.list_tensorboard_time_series_mock +batch_read_tensorboard_time_series_mock = ( + test_tensorboard.batch_read_tensorboard_time_series_mock +) +get_pipeline_job_mock = test_pipeline_jobs.mock_pipeline_service_get + + +@pytest.fixture +def list_tensorboard_time_series_mock_empty(): + with patch.object( + TensorboardServiceClient, + "list_tensorboard_time_series", + ) as list_tensorboard_time_series_mock: + list_tensorboard_time_series_mock.side_effect = [ + [], # initially empty + [], + [test_tensorboard._TEST_TENSORBOARD_TIME_SERIES], + ] + yield list_tensorboard_time_series_mock + + +@pytest.fixture +def create_tensorboard_run_artifact_mock(): + with patch.object(MetadataServiceClient, "create_artifact") as create_artifact_mock: + create_artifact_mock.side_effect = [_TEST_TENSORBOARD_RUN_ARTIFACT] + yield create_artifact_mock + + +@pytest.fixture +def get_tensorboard_run_artifact_mock(): + with patch.object(MetadataServiceClient, "get_artifact") as get_artifact_mock: + get_artifact_mock.side_effect = [ + _TEST_TENSORBOARD_RUN_ARTIFACT, + exceptions.NotFound(""), + _TEST_LEGACY_METRIC_ARTIFACT, + ] + yield get_artifact_mock + + +@pytest.mark.usefixtures("google_auth_mock") +class TestExperiments: + def setup_method(self): + reload(initializer) + reload(metadata) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + def test_init_experiment_with_existing_metadataStore_and_context( - self, get_metadata_store_mock, get_context_mock + self, get_metadata_store_mock, get_experiment_run_run_mock ): aiplatform.init( - project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT + project=_TEST_PROJECT, + location=_TEST_LOCATION, + experiment=_TEST_EXPERIMENT, ) get_metadata_store_mock.assert_called_once_with( name=_TEST_METADATASTORE, retry=base._DEFAULT_RETRY ) - get_context_mock.assert_called_once_with( + get_experiment_run_run_mock.assert_called_once_with( name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY ) def test_init_experiment_with_credentials( - self, get_metadata_store_mock, get_context_mock + self, + get_metadata_store_mock, + get_experiment_run_run_mock, ): creds = credentials.AnonymousCredentials() @@ -414,14 +877,14 @@ def test_init_experiment_with_credentials( ) assert ( - metadata.metadata_service._experiment.api_client._transport._credentials + metadata._experiment_tracker._experiment._metadata_context.api_client._transport._credentials == creds ) get_metadata_store_mock.assert_called_once_with( name=_TEST_METADATASTORE, retry=base._DEFAULT_RETRY ) - get_context_mock.assert_called_once_with( + get_experiment_run_run_mock.assert_called_once_with( name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY ) @@ -434,7 +897,7 @@ def test_init_and_get_metadata_store_with_credentials( project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=creds ) - store = metadata._MetadataStore.get_or_create() + store = metadata_store._MetadataStore.get_or_create() assert store.api_client._transport._credentials == creds @@ -442,19 +905,21 @@ def test_init_and_get_metadata_store_with_credentials( "get_metadata_store_mock_raise_not_found_exception", "create_metadata_store_mock", ) - def test_init_and_get_then_create_metadata_store_with_credentials(self): + def test_init_and_get_then_create_metadata_store_with_credentials( + self, + ): creds = credentials.AnonymousCredentials() aiplatform.init( project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=creds ) - store = metadata._MetadataStore.get_or_create() + store = metadata_store._MetadataStore.get_or_create() assert store.api_client._transport._credentials == creds def test_init_experiment_with_existing_description( - self, get_metadata_store_mock, get_context_mock + self, get_metadata_store_mock, get_experiment_run_run_mock ): aiplatform.init( project=_TEST_PROJECT, @@ -466,13 +931,15 @@ def test_init_experiment_with_existing_description( get_metadata_store_mock.assert_called_once_with( name=_TEST_METADATASTORE, retry=base._DEFAULT_RETRY ) - get_context_mock.assert_called_once_with( + get_experiment_run_run_mock.assert_called_once_with( name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY ) - @pytest.mark.usefixtures("get_metadata_store_mock") - @pytest.mark.usefixtures("get_context_mock") - def test_init_experiment_without_existing_description(self, update_context_mock): + @pytest.mark.usefixtures("get_metadata_store_mock", "get_experiment_run_run_mock") + def test_init_experiment_without_existing_description( + self, + update_context_mock, + ): aiplatform.init( project=_TEST_PROJECT, location=_TEST_LOCATION, @@ -491,8 +958,31 @@ def test_init_experiment_without_existing_description(self, update_context_mock) update_context_mock.assert_called_once_with(context=experiment_context) - @pytest.mark.usefixtures("get_metadata_store_mock") - @pytest.mark.usefixtures("get_context_wrong_schema_mock") + @pytest.mark.usefixtures( + "get_metadata_store_mock", + "get_experiment_run_mock", + "update_experiment_run_context_to_running", + "get_tensorboard_run_artifact_not_found_mock", + ) + def test_init_experiment_reset(self): + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + experiment=_TEST_EXPERIMENT, + ) + aiplatform.start_run(_TEST_RUN, resume=True) + + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + + assert metadata._experiment_tracker.experiment_name == _TEST_EXPERIMENT + assert metadata._experiment_tracker.experiment_run.name == _TEST_RUN + + aiplatform.init(project=_TEST_OTHER_PROJECT, location=_TEST_LOCATION) + + assert metadata._experiment_tracker.experiment_name is None + assert metadata._experiment_tracker.experiment_run is None + + @pytest.mark.usefixtures("get_metadata_store_mock", "get_context_wrong_schema_mock") def test_init_experiment_wrong_schema(self): with pytest.raises(ValueError): aiplatform.init( @@ -502,193 +992,377 @@ def test_init_experiment_wrong_schema(self): ) @pytest.mark.usefixtures("get_metadata_store_mock") - @pytest.mark.usefixtures("get_context_mock") - @pytest.mark.usefixtures("get_execution_mock") - @pytest.mark.usefixtures("add_context_artifacts_and_executions_mock") - @pytest.mark.usefixtures("get_artifact_mock") - @pytest.mark.usefixtures("add_execution_events_mock") - def test_init_experiment_reset(self): + @pytest.mark.usefixtures() + def test_start_run( + self, + get_experiment_mock, + create_experiment_run_context_mock, + add_context_children_mock, + ): + aiplatform.init( - project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT + project=_TEST_PROJECT, + location=_TEST_LOCATION, + experiment=_TEST_EXPERIMENT, ) aiplatform.start_run(_TEST_RUN) - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + get_experiment_mock.assert_called_with( + name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY + ) - assert metadata.metadata_service.experiment_name == _TEST_EXPERIMENT - assert metadata.metadata_service.run_name == _TEST_RUN + _TRUE_CONTEXT = copy.deepcopy(_EXPERIMENT_RUN_MOCK) + _TRUE_CONTEXT.name = None - aiplatform.init(project=_TEST_OTHER_PROJECT, location=_TEST_LOCATION) + create_experiment_run_context_mock.assert_called_with( + parent=_TEST_METADATASTORE, + context=_TRUE_CONTEXT, + context_id=_EXPERIMENT_RUN_MOCK.name.split("/")[-1], + ) - assert metadata.metadata_service.experiment_name is None - assert metadata.metadata_service.run_name is None + add_context_children_mock.assert_called_with( + context=_EXPERIMENT_MOCK.name, child_contexts=[_EXPERIMENT_RUN_MOCK.name] + ) - @pytest.mark.usefixtures("get_metadata_store_mock") - @pytest.mark.usefixtures("get_context_mock") - def test_start_run_with_existing_execution_and_artifact( + @pytest.mark.usefixtures( + "get_metadata_store_mock", + "get_experiment_mock", + "create_experiment_run_context_mock", + "add_context_children_mock", + ) + def test_log_params( + self, + update_context_mock, + ): + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + experiment=_TEST_EXPERIMENT, + ) + aiplatform.start_run(_TEST_RUN) + aiplatform.log_params(_TEST_PARAMS) + + _TRUE_CONTEXT = copy.deepcopy(_EXPERIMENT_RUN_MOCK) + _TRUE_CONTEXT.metadata[constants._PARAM_KEY].update(_TEST_PARAMS) + + update_context_mock.assert_called_once_with(context=_TRUE_CONTEXT) + + @pytest.mark.usefixtures( + "get_metadata_store_mock", + "get_experiment_mock", + "create_experiment_run_context_mock", + "add_context_children_mock", + ) + def test_log_metrics(self, update_context_mock): + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + experiment=_TEST_EXPERIMENT, + ) + aiplatform.start_run(_TEST_RUN) + aiplatform.log_metrics(_TEST_METRICS) + + _TRUE_CONTEXT = copy.deepcopy(_EXPERIMENT_RUN_MOCK) + _TRUE_CONTEXT.metadata[constants._METRIC_KEY].update(_TEST_METRICS) + + update_context_mock.assert_called_once_with(context=_TRUE_CONTEXT) + + @pytest.mark.usefixtures( + "get_metadata_store_mock", + "get_experiment_mock", + "create_experiment_run_context_mock", + "add_context_children_mock", + "get_tensorboard_mock", + "get_tensorboard_run_not_found_mock", + "get_tensorboard_experiment_not_found_mock", + "get_artifact_not_found_mock", + "get_tensorboard_time_series_not_found_mock", + "list_tensorboard_time_series_mock_empty", + ) + def test_log_time_series_metrics( self, - get_execution_mock, + update_context_mock, + create_tensorboard_experiment_mock, + create_tensorboard_run_mock, + create_tensorboard_run_artifact_mock, add_context_artifacts_and_executions_mock, - get_artifact_mock, - add_execution_events_mock, + create_tensorboard_time_series_mock, + batch_read_tensorboard_time_series_mock, + write_tensorboard_run_data_mock, ): + tb = aiplatform.Tensorboard(test_tensorboard._TEST_NAME) + aiplatform.init( - project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT + project=_TEST_PROJECT, + location=_TEST_LOCATION, + experiment=_TEST_EXPERIMENT, + experiment_tensorboard=tb, ) + + update_context_mock.assert_called_once_with(context=_TEST_EXPERIMENT_CONTEXT) + aiplatform.start_run(_TEST_RUN) + timestamp = utils.get_timestamp_proto() + aiplatform.log_time_series_metrics(_TEST_OTHER_METRICS, wall_time=timestamp) + + create_tensorboard_experiment_mock.assert_called_once_with( + parent=test_tensorboard._TEST_NAME, + tensorboard_experiment_id=_TEST_CONTEXT_ID, + tensorboard_experiment=gca_tensorboard_experiment.TensorboardExperiment( + display_name=experiment_run_resource.ExperimentRun._format_tensorboard_experiment_display_name( + _TEST_CONTEXT_ID + ), + ), + metadata=(), + timeout=None, + ) - get_execution_mock.assert_called_once_with( - name=_TEST_EXECUTION_NAME, retry=base._DEFAULT_RETRY + create_tensorboard_run_mock.assert_called_once_with( + parent=test_tensorboard._TEST_TENSORBOARD_EXPERIMENT_NAME, + tensorboard_run_id=_TEST_RUN, + tensorboard_run=gca_tensorboard_run.TensorboardRun( + display_name=_TEST_RUN, + ), + metadata=(), + timeout=None, ) + + true_tb_run_artifact = copy.deepcopy(_TEST_TENSORBOARD_RUN_ARTIFACT) + true_tb_run_artifact.name = None + + create_tensorboard_run_artifact_mock.assert_called_once_with( + parent=_TEST_PARENT, + artifact=true_tb_run_artifact, + artifact_id=experiment_run_resource.ExperimentRun._tensorboard_run_id( + _TEST_EXECUTION_ID + ), + ) + add_context_artifacts_and_executions_mock.assert_called_once_with( - context=_TEST_CONTEXT_NAME, - artifacts=None, - executions=[_TEST_EXECUTION_NAME], + context=_TEST_EXPERIMENT_RUN_CONTEXT_NAME, + artifacts=[_TEST_TENSORBOARD_RUN_ARTIFACT.name], + executions=None, ) - get_artifact_mock.assert_called_once_with( - name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY + + create_tensorboard_time_series_mock.assert_called_with( + parent=test_tensorboard._TEST_TENSORBOARD_RUN_NAME, + tensorboard_time_series=gca_tensorboard_time_series.TensorboardTimeSeries( + display_name=list(_TEST_OTHER_METRICS.keys())[0], + value_type="SCALAR", + plugin_name="scalars", + ), ) - add_execution_events_mock.assert_called_once_with( - execution=_TEST_EXECUTION_NAME, - events=[Event(artifact=_TEST_ARTIFACT_NAME, type_=Event.Type.OUTPUT)], + + ts_data = [ + gca_tensorboard_data.TimeSeriesData( + tensorboard_time_series_id=test_tensorboard._TEST_TENSORBOARD_TIME_SERIES_ID, + value_type=gca_tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR, + values=[ + gca_tensorboard_data.TimeSeriesDataPoint( + scalar=gca_tensorboard_data.Scalar(value=value), + wall_time=timestamp, + step=2, + ) + ], + ) + for value in _TEST_OTHER_METRICS.values() + ] + + write_tensorboard_run_data_mock.assert_called_once_with( + tensorboard_run=test_tensorboard._TEST_TENSORBOARD_RUN_NAME, + time_series_data=ts_data, ) - @pytest.mark.usefixtures("get_metadata_store_mock") - @pytest.mark.usefixtures("get_context_mock") - @pytest.mark.usefixtures("get_execution_wrong_schema_mock") - def test_start_run_with_wrong_run_execution_schema( - self, - ): + @pytest.mark.usefixtures( + "get_metadata_store_mock", + "get_experiment_mock", + "create_experiment_run_context_mock", + "add_context_children_mock", + ) + def test_log_metrics_nest_value_raises_error(self): aiplatform.init( project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT ) - with pytest.raises(ValueError): - aiplatform.start_run(_TEST_RUN) + aiplatform.start_run(_TEST_RUN) + with pytest.raises(TypeError): + aiplatform.log_metrics({"test": {"nested": "string"}}) - @pytest.mark.usefixtures("get_metadata_store_mock") - @pytest.mark.usefixtures("get_context_mock") - @pytest.mark.usefixtures("get_execution_mock") - @pytest.mark.usefixtures("add_context_artifacts_and_executions_mock") - @pytest.mark.usefixtures("get_artifact_wrong_schema_mock") - def test_start_run_with_wrong_metrics_artifact_schema( - self, - ): + @pytest.mark.usefixtures( + "get_metadata_store_mock", + "get_experiment_mock", + "create_experiment_run_context_mock", + "add_context_children_mock", + ) + def test_log_params_nest_value_raises_error(self): aiplatform.init( project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT ) - with pytest.raises(ValueError): - aiplatform.start_run(_TEST_RUN) + aiplatform.start_run(_TEST_RUN) + with pytest.raises(TypeError): + aiplatform.log_params({"test": {"nested": "string"}}) - @pytest.mark.usefixtures("get_metadata_store_mock") - @pytest.mark.usefixtures("get_context_mock") - @pytest.mark.usefixtures("get_execution_mock") - @pytest.mark.usefixtures("add_context_artifacts_and_executions_mock") - @pytest.mark.usefixtures("get_artifact_mock") - @pytest.mark.usefixtures("add_execution_events_mock") - def test_log_params( + @pytest.mark.usefixtures( + "get_metadata_store_mock", + "get_experiment_mock", + "create_experiment_run_context_mock", + "add_context_children_mock", + ) + def test_end_run( self, - update_execution_mock, + update_context_mock, ): aiplatform.init( - project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT + project=_TEST_PROJECT, + location=_TEST_LOCATION, + experiment=_TEST_EXPERIMENT, ) aiplatform.start_run(_TEST_RUN) - aiplatform.log_params(_TEST_PARAMS) + aiplatform.end_run() - updated_execution = GapicExecution( - name=_TEST_EXECUTION_NAME, - display_name=_TEST_RUN, - schema_title=constants.SYSTEM_RUN, - schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_RUN], - metadata=_TEST_PARAMS, - ) + _TRUE_CONTEXT = copy.deepcopy(_EXPERIMENT_RUN_MOCK) + _TRUE_CONTEXT.metadata[ + constants._STATE_KEY + ] = gca_execution.Execution.State.COMPLETE.name - update_execution_mock.assert_called_once_with(execution=updated_execution) + update_context_mock.assert_called_once_with(context=_TRUE_CONTEXT) - @pytest.mark.usefixtures("get_metadata_store_mock") - @pytest.mark.usefixtures("get_context_mock") - @pytest.mark.usefixtures("get_execution_mock") - @pytest.mark.usefixtures("add_context_artifacts_and_executions_mock") - @pytest.mark.usefixtures("get_artifact_mock") - @pytest.mark.usefixtures("add_execution_events_mock") - def test_log_metrics( + @pytest.mark.usefixtures( + "get_metadata_store_mock", + "get_experiment_mock", + "create_experiment_run_context_mock", + "get_pipeline_job_mock", + ) + def test_log_pipeline_job( self, - update_artifact_mock, + add_context_children_mock, ): aiplatform.init( - project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT + project=_TEST_PROJECT, + location=_TEST_LOCATION, + experiment=_TEST_EXPERIMENT, ) aiplatform.start_run(_TEST_RUN) - aiplatform.log_metrics(_TEST_METRICS) - updated_artifact = GapicArtifact( - name=_TEST_ARTIFACT_NAME, - display_name=_TEST_ARTIFACT_ID, - schema_title=constants.SYSTEM_METRICS, - schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_METRICS], - metadata=_TEST_METRICS, + pipeline_job = aiplatform.PipelineJob.get( + test_pipeline_jobs._TEST_PIPELINE_JOB_ID ) + pipeline_job.wait() - update_artifact_mock.assert_called_once_with(artifact=updated_artifact) + aiplatform.log(pipeline_job=pipeline_job) - @pytest.mark.usefixtures("get_metadata_store_mock") - @pytest.mark.usefixtures("get_context_mock") - @pytest.mark.usefixtures("get_execution_mock") - @pytest.mark.usefixtures("add_context_artifacts_and_executions_mock") - @pytest.mark.usefixtures("get_artifact_mock") - @pytest.mark.usefixtures("add_execution_events_mock") - def test_log_metrics_string_value_raise_error(self): - aiplatform.init( - project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT + add_context_children_mock.assert_called_with( + context=_EXPERIMENT_RUN_MOCK.name, + child_contexts=[ + pipeline_job.gca_resource.job_detail.pipeline_run_context.name + ], ) - aiplatform.start_run(_TEST_RUN) - with pytest.raises(TypeError): - aiplatform.log_metrics({"test": "string"}) - @pytest.mark.usefixtures("get_context_mock") + @pytest.mark.usefixtures( + "get_experiment_mock", + "list_tensorboard_time_series_mock", + "batch_read_tensorboard_time_series_mock", + ) def test_get_experiment_df( - self, list_executions_mock, query_execution_inputs_and_outputs_mock + self, + list_context_mock_for_experiment_dataframe_mock, + list_artifact_mock_for_experiment_dataframe, + list_executions_mock_for_experiment_dataframe, + get_tensorboard_run_artifact_mock, + get_tensorboard_run_mock, ): - try: - import pandas as pd - except ImportError: - raise ImportError( - "Pandas is not installed and is required to test the get_experiment_df method. " - 'Please install the SDK using "pip install python-aiplatform[full]"' - ) + import pandas as pd + aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) experiment_df = aiplatform.get_experiment_df(_TEST_EXPERIMENT) - expected_filter = f'schema_title="{constants.SYSTEM_RUN}" AND in_context("{_TEST_CONTEXT_NAME}")' - list_executions_mock.assert_called_once_with( - request=ListExecutionsRequest( - parent=_TEST_PARENT, - filter=expected_filter, - ) + expected_filter = metadata_utils._make_filter_string( + parent_contexts=[_TEST_CONTEXT_NAME], + schema_title=[ + constants.SYSTEM_EXPERIMENT_RUN, + constants.SYSTEM_PIPELINE_RUN, + ], ) - query_execution_inputs_and_outputs_mock.assert_has_calls( - [ - call(execution=_TEST_EXECUTION_NAME), - call(execution=_TEST_OTHER_EXECUTION_NAME), - ] + + list_context_mock_for_experiment_dataframe_mock.assert_called_once_with( + request=dict(parent=_TEST_PARENT, filter=expected_filter) + ) + + expected_legacy_filter = metadata_utils._make_filter_string( + in_context=[_TEST_CONTEXT_NAME], schema_title=[constants.SYSTEM_RUN] + ) + expected_pipeline_filter = metadata_utils._make_filter_string( + in_context=[_TEST_PIPELINE_CONTEXT.name], schema_title=constants.SYSTEM_RUN + ) + + list_executions_mock_for_experiment_dataframe.assert_has_calls( + calls=[ + call(request=dict(parent=_TEST_PARENT, filter=expected_legacy_filter)), + call( + request=dict(parent=_TEST_PARENT, filter=expected_pipeline_filter) + ), + ], + any_order=False, + ) + + expected_filter = metadata_utils._make_filter_string( + in_context=[_TEST_PIPELINE_CONTEXT.name], + schema_title=constants.SYSTEM_METRICS, ) + + list_artifact_mock_for_experiment_dataframe.assert_has_calls( + calls=[call(request=dict(parent=_TEST_PARENT, filter=expected_filter))], + any_order=False, + ) + experiment_df_truth = pd.DataFrame( [ { "experiment_name": _TEST_EXPERIMENT, + "run_type": constants.SYSTEM_EXPERIMENT_RUN, + "state": gca_execution.Execution.State.RUNNING.name, "run_name": _TEST_RUN, - "param.%s" % _TEST_PARAM_KEY_1: 0.01, - "param.%s" % _TEST_PARAM_KEY_2: 0.2, - "metric.%s" % _TEST_METRIC_KEY_1: 222, - "metric.%s" % _TEST_METRIC_KEY_2: 1, + "param.%s" % _TEST_PARAM_KEY_1: _TEST_PARAMS[_TEST_PARAM_KEY_1], + "param.%s" % _TEST_PARAM_KEY_2: _TEST_PARAMS[_TEST_PARAM_KEY_2], + "metric.%s" % _TEST_METRIC_KEY_1: _TEST_METRICS[_TEST_METRIC_KEY_1], + "metric.%s" % _TEST_METRIC_KEY_2: _TEST_METRICS[_TEST_METRIC_KEY_2], + "time_series_metric.accuracy": test_tensorboard._TEST_TENSORBOARD_TIME_SERIES_DATA.values[ + 0 + ].scalar.value, }, { "experiment_name": _TEST_EXPERIMENT, + "run_type": constants.SYSTEM_EXPERIMENT_RUN, + "state": gca_execution.Execution.State.RUNNING.name, "run_name": _TEST_OTHER_RUN, - "param.%s" % _TEST_PARAM_KEY_1: 0.02, - "param.%s" % _TEST_PARAM_KEY_2: 0.3, - "metric.%s" % _TEST_METRIC_KEY_2: 0.9, + "param.%s" + % _TEST_PARAM_KEY_1: _TEST_OTHER_PARAMS[_TEST_PARAM_KEY_1], + "param.%s" + % _TEST_PARAM_KEY_2: _TEST_OTHER_PARAMS[_TEST_PARAM_KEY_2], + "metric.%s" + % _TEST_METRIC_KEY_2: _TEST_OTHER_METRICS[_TEST_METRIC_KEY_2], + }, + { + "experiment_name": _TEST_EXPERIMENT, + "run_type": constants.SYSTEM_PIPELINE_RUN, + "state": gca_execution.Execution.State.RUNNING.name, + "run_name": _TEST_PIPELINE_RUN_ID, + "param.%s" % _TEST_PARAM_KEY_1: _TEST_PARAMS[_TEST_PARAM_KEY_1] + 1, + "param.%s" % _TEST_PARAM_KEY_2: _TEST_PARAMS[_TEST_PARAM_KEY_2] + 1, + "metric.%s" % _TEST_METRIC_KEY_1: _TEST_METRICS[_TEST_METRIC_KEY_1] + + 1, + "metric.%s" % _TEST_METRIC_KEY_2: _TEST_METRICS[_TEST_METRIC_KEY_2] + + 1, + }, + { + "experiment_name": _TEST_EXPERIMENT, + "run_type": constants.SYSTEM_RUN, + "state": gca_execution.Execution.State.STATE_UNSPECIFIED.name, + "run_name": _TEST_RUN, + "param.%s" % _TEST_PARAM_KEY_1: _TEST_PARAMS[_TEST_PARAM_KEY_1], + "param.%s" % _TEST_PARAM_KEY_2: _TEST_PARAMS[_TEST_PARAM_KEY_2], + "metric.%s" % _TEST_METRIC_KEY_1: _TEST_METRICS[_TEST_METRIC_KEY_1], + "metric.%s" % _TEST_METRIC_KEY_2: _TEST_METRICS[_TEST_METRIC_KEY_2], }, ] ) @@ -706,65 +1380,3 @@ def test_get_experiment_df_wrong_schema(self): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) with pytest.raises(ValueError): aiplatform.get_experiment_df(_TEST_EXPERIMENT) - - @pytest.mark.usefixtures("get_pipeline_context_mock") - def test_get_pipeline_df( - self, list_executions_mock, query_execution_inputs_and_outputs_mock - ): - try: - import pandas as pd - except ImportError: - raise ImportError( - "Pandas is not installed and is required to test the get_pipeline_df method. " - 'Please install the SDK using "pip install python-aiplatform[full]"' - ) - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - - pipeline_df = aiplatform.get_pipeline_df(_TEST_PIPELINE) - - expected_filter = f'schema_title="{constants.SYSTEM_RUN}" AND in_context("{_TEST_CONTEXT_NAME}")' - list_executions_mock.assert_called_once_with( - request=ListExecutionsRequest( - parent=_TEST_PARENT, - filter=expected_filter, - ) - ) - query_execution_inputs_and_outputs_mock.assert_has_calls( - [ - call(execution=_TEST_EXECUTION_NAME), - call(execution=_TEST_OTHER_EXECUTION_NAME), - ] - ) - pipeline_df_truth = pd.DataFrame( - [ - { - "pipeline_name": _TEST_PIPELINE, - "run_name": _TEST_RUN, - "param.%s" % _TEST_PARAM_KEY_1: 0.01, - "param.%s" % _TEST_PARAM_KEY_2: 0.2, - "metric.%s" % _TEST_METRIC_KEY_1: 222, - "metric.%s" % _TEST_METRIC_KEY_2: 1, - }, - { - "pipeline_name": _TEST_PIPELINE, - "run_name": _TEST_OTHER_RUN, - "param.%s" % _TEST_PARAM_KEY_1: 0.02, - "param.%s" % _TEST_PARAM_KEY_2: 0.3, - "metric.%s" % _TEST_METRIC_KEY_2: 0.9, - }, - ] - ) - - _assert_frame_equal_with_sorted_columns(pipeline_df, pipeline_df_truth) - - @pytest.mark.usefixtures("get_context_not_found_mock") - def test_get_pipeline_df_not_exist(self): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - with pytest.raises(exceptions.NotFound): - aiplatform.get_pipeline_df(_TEST_PIPELINE) - - @pytest.mark.usefixtures("get_context_mock") - def test_get_pipeline_df_wrong_schema(self): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - with pytest.raises(ValueError): - aiplatform.get_pipeline_df(_TEST_PIPELINE) diff --git a/tests/unit/aiplatform/test_metadata_resources.py b/tests/unit/aiplatform/test_metadata_resources.py index 2a7180adfd..2544767660 100644 --- a/tests/unit/aiplatform/test_metadata_resources.py +++ b/tests/unit/aiplatform/test_metadata_resources.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,6 +24,7 @@ from google.cloud import aiplatform from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer +from google.cloud.aiplatform.compat.types import event as gca_event from google.cloud.aiplatform.metadata import artifact from google.cloud.aiplatform.metadata import context from google.cloud.aiplatform.metadata import execution @@ -36,9 +37,6 @@ MetadataServiceClient, AddExecutionEventsResponse, Event, - ListExecutionsRequest, - ListArtifactsRequest, - ListContextsRequest, ) # project @@ -198,6 +196,7 @@ def create_execution_mock(): schema_version=_TEST_SCHEMA_VERSION, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, + state=GapicExecution.State.RUNNING, ) yield create_execution_mock @@ -241,14 +240,13 @@ def query_execution_inputs_and_outputs_mock(): description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, ), - GapicArtifact( - name=_TEST_ARTIFACT_NAME, - display_name=_TEST_DISPLAY_NAME, - schema_title=_TEST_SCHEMA_TITLE, - schema_version=_TEST_SCHEMA_VERSION, - description=_TEST_DESCRIPTION, - metadata=_TEST_METADATA, - ), + ], + events=[ + gca_event.Event( + artifact=_TEST_ARTIFACT_NAME, + execution=_TEST_EXECUTION_NAME, + type_=gca_event.Event.Type.OUTPUT, + ) ], ) yield query_execution_inputs_and_outputs_mock @@ -266,6 +264,7 @@ def update_execution_mock(): schema_version=_TEST_SCHEMA_VERSION, description=_TEST_DESCRIPTION, metadata=_TEST_UPDATED_METADATA, + state=GapicExecution.State.RUNNING, ) yield update_execution_mock @@ -314,6 +313,7 @@ def create_artifact_mock(): schema_version=_TEST_SCHEMA_VERSION, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, + state=GapicArtifact.State.STATE_UNSPECIFIED, ) yield create_artifact_mock @@ -352,6 +352,7 @@ def update_artifact_mock(): schema_version=_TEST_SCHEMA_VERSION, description=_TEST_DESCRIPTION, metadata=_TEST_UPDATED_METADATA, + state=GapicArtifact.State.STATE_UNSPECIFIED, ) yield update_artifact_mock @@ -462,10 +463,10 @@ def test_list_contexts(self, list_contexts_mock): ) list_contexts_mock.assert_called_once_with( - request=ListContextsRequest( - parent=_TEST_PARENT, - filter=filter, - ) + request={ + "parent": _TEST_PARENT, + "filter": filter, + } ) assert len(context_list) == 2 assert context_list[0]._gca_resource == expected_context @@ -552,15 +553,15 @@ def teardown_method(self): def test_init_execution(self, get_execution_mock): aiplatform.init(project=_TEST_PROJECT) - execution._Execution(resource_name=_TEST_EXECUTION_NAME) + execution.Execution(execution_name=_TEST_EXECUTION_NAME) get_execution_mock.assert_called_once_with( name=_TEST_EXECUTION_NAME, retry=base._DEFAULT_RETRY ) def test_init_execution_with_id(self, get_execution_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - execution._Execution( - resource_name=_TEST_EXECUTION_ID, metadata_store_id=_TEST_METADATA_STORE + execution.Execution( + execution_name=_TEST_EXECUTION_ID, metadata_store_id=_TEST_METADATA_STORE ) get_execution_mock.assert_called_once_with( name=_TEST_EXECUTION_NAME, retry=base._DEFAULT_RETRY @@ -571,7 +572,7 @@ def test_get_or_create_execution( ): aiplatform.init(project=_TEST_PROJECT) - my_execution = execution._Execution.get_or_create( + my_execution = execution.Execution.get_or_create( resource_id=_TEST_EXECUTION_ID, schema_title=_TEST_SCHEMA_TITLE, display_name=_TEST_DISPLAY_NAME, @@ -587,6 +588,7 @@ def test_get_or_create_execution( display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, + state=GapicExecution.State.RUNNING, ) get_execution_for_get_or_create_mock.assert_called_once_with( name=_TEST_EXECUTION_NAME, retry=base._DEFAULT_RETRY @@ -605,7 +607,7 @@ def test_get_or_create_execution( def test_update_execution(self, update_execution_mock): aiplatform.init(project=_TEST_PROJECT) - my_execution = execution._Execution._create( + my_execution = execution.Execution._create( resource_id=_TEST_EXECUTION_ID, schema_title=_TEST_SCHEMA_TITLE, display_name=_TEST_DISPLAY_NAME, @@ -614,7 +616,7 @@ def test_update_execution(self, update_execution_mock): metadata=_TEST_METADATA, metadata_store_id=_TEST_METADATA_STORE, ) - my_execution.update(_TEST_UPDATED_METADATA) + my_execution.update(metadata=_TEST_UPDATED_METADATA) updated_execution = GapicExecution( name=_TEST_EXECUTION_NAME, @@ -623,6 +625,7 @@ def test_update_execution(self, update_execution_mock): display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, metadata=_TEST_UPDATED_METADATA, + state=GapicExecution.State.RUNNING, ) update_execution_mock.assert_called_once_with(execution=updated_execution) @@ -633,7 +636,7 @@ def test_list_executions(self, list_executions_mock): aiplatform.init(project=_TEST_PROJECT) filter = "test-filter" - execution_list = execution._Execution.list( + execution_list = execution.Execution.list( filter=filter, metadata_store_id=_TEST_METADATA_STORE ) @@ -647,7 +650,7 @@ def test_list_executions(self, list_executions_mock): ) list_executions_mock.assert_called_once_with( - request=ListExecutionsRequest( + request=dict( parent=_TEST_PARENT, filter=filter, ) @@ -656,11 +659,11 @@ def test_list_executions(self, list_executions_mock): assert execution_list[0]._gca_resource == expected_execution assert execution_list[1]._gca_resource == expected_execution - @pytest.mark.usefixtures("get_execution_mock") + @pytest.mark.usefixtures("get_execution_mock", "get_artifact_mock") def test_add_artifact(self, add_execution_events_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - my_execution = execution._Execution.get_or_create( + my_execution = execution.Execution.get_or_create( resource_id=_TEST_EXECUTION_ID, schema_title=_TEST_SCHEMA_TITLE, display_name=_TEST_DISPLAY_NAME, @@ -669,10 +672,9 @@ def test_add_artifact(self, add_execution_events_mock): metadata=_TEST_METADATA, metadata_store_id=_TEST_METADATA_STORE, ) - my_execution.add_artifact( - artifact_resource_name=_TEST_ARTIFACT_NAME, - input=False, - ) + + my_artifact = aiplatform.Artifact(_TEST_ARTIFACT_ID) + my_execution.assign_output_artifacts(artifacts=[my_artifact]) add_execution_events_mock.assert_called_once_with( execution=_TEST_EXECUTION_NAME, events=[Event(artifact=_TEST_ARTIFACT_NAME, type_=Event.Type.OUTPUT)], @@ -683,7 +685,7 @@ def test_query_input_and_output_artifacts( self, query_execution_inputs_and_outputs_mock ): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - my_execution = execution._Execution.get_or_create( + my_execution = execution.Execution.get_or_create( resource_id=_TEST_EXECUTION_ID, schema_title=_TEST_SCHEMA_TITLE, display_name=_TEST_DISPLAY_NAME, @@ -693,7 +695,7 @@ def test_query_input_and_output_artifacts( metadata_store_id=_TEST_METADATA_STORE, ) - artifact_list = my_execution.query_input_and_output_artifacts() + artifact_list = my_execution.get_output_artifacts() expected_artifact = GapicArtifact( name=_TEST_ARTIFACT_NAME, @@ -707,9 +709,8 @@ def test_query_input_and_output_artifacts( query_execution_inputs_and_outputs_mock.assert_called_once_with( execution=_TEST_EXECUTION_NAME, ) - assert len(artifact_list) == 2 + assert len(artifact_list) == 1 assert artifact_list[0]._gca_resource == expected_artifact - assert artifact_list[1]._gca_resource == expected_artifact class TestArtifact: @@ -722,15 +723,15 @@ def teardown_method(self): def test_init_artifact(self, get_artifact_mock): aiplatform.init(project=_TEST_PROJECT) - artifact._Artifact(resource_name=_TEST_ARTIFACT_NAME) + artifact.Artifact(artifact_name=_TEST_ARTIFACT_NAME) get_artifact_mock.assert_called_once_with( name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY ) def test_init_artifact_with_id(self, get_artifact_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) - artifact._Artifact( - resource_name=_TEST_ARTIFACT_ID, metadata_store_id=_TEST_METADATA_STORE + artifact.Artifact( + artifact_name=_TEST_ARTIFACT_ID, metadata_store_id=_TEST_METADATA_STORE ) get_artifact_mock.assert_called_once_with( name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY @@ -741,7 +742,7 @@ def test_get_or_create_artifact( ): aiplatform.init(project=_TEST_PROJECT) - my_artifact = artifact._Artifact.get_or_create( + my_artifact = artifact.Artifact.get_or_create( resource_id=_TEST_ARTIFACT_ID, schema_title=_TEST_SCHEMA_TITLE, display_name=_TEST_DISPLAY_NAME, @@ -757,6 +758,7 @@ def test_get_or_create_artifact( display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, + state=GapicArtifact.State.STATE_UNSPECIFIED, ) get_artifact_for_get_or_create_mock.assert_called_once_with( name=_TEST_ARTIFACT_NAME, retry=base._DEFAULT_RETRY @@ -775,7 +777,7 @@ def test_get_or_create_artifact( def test_update_artifact(self, update_artifact_mock): aiplatform.init(project=_TEST_PROJECT) - my_artifact = artifact._Artifact._create( + my_artifact = artifact.Artifact._create( resource_id=_TEST_ARTIFACT_ID, schema_title=_TEST_SCHEMA_TITLE, display_name=_TEST_DISPLAY_NAME, @@ -784,7 +786,7 @@ def test_update_artifact(self, update_artifact_mock): metadata=_TEST_METADATA, metadata_store_id=_TEST_METADATA_STORE, ) - my_artifact.update(_TEST_UPDATED_METADATA) + my_artifact.update(metadata=_TEST_UPDATED_METADATA) updated_artifact = GapicArtifact( name=_TEST_ARTIFACT_NAME, @@ -793,6 +795,7 @@ def test_update_artifact(self, update_artifact_mock): display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, metadata=_TEST_UPDATED_METADATA, + state=GapicArtifact.State.STATE_UNSPECIFIED, ) update_artifact_mock.assert_called_once_with(artifact=updated_artifact) @@ -803,7 +806,7 @@ def test_list_artifacts(self, list_artifacts_mock): aiplatform.init(project=_TEST_PROJECT) filter = "test-filter" - artifact_list = artifact._Artifact.list( + artifact_list = artifact.Artifact.list( filter=filter, metadata_store_id=_TEST_METADATA_STORE ) @@ -817,7 +820,7 @@ def test_list_artifacts(self, list_artifacts_mock): ) list_artifacts_mock.assert_called_once_with( - request=ListArtifactsRequest( + request=dict( parent=_TEST_PARENT, filter=filter, ) diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index 1f6f2bb50c..9bfda28353 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.compat.types import ( pipeline_job as gca_pipeline_job, pipeline_state as gca_pipeline_state, + context as gca_context, ) _TEST_PROJECT = "test-project" @@ -207,6 +208,11 @@ def make_pipeline_job(state): create_time=_TEST_PIPELINE_CREATE_TIME, service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK, + job_detail=gca_pipeline_job.PipelineJobDetail( + pipeline_run_context=gca_context.Context( + name=_TEST_PIPELINE_JOB_NAME, + ) + ), ) @@ -285,15 +291,6 @@ def mock_load_yaml_and_json(job_spec): @pytest.mark.usefixtures("google_auth_mock") class TestPipelineJob: - class FakePipelineJob(pipeline_jobs.PipelineJob): - - _resource_noun = "fakePipelineJobs" - _getter_method = _TEST_PIPELINE_GET_METHOD_NAME - _list_method = _TEST_PIPELINE_LIST_METHOD_NAME - _cancel_method = _TEST_PIPELINE_CANCEL_METHOD_NAME - _delete_method = _TEST_PIPELINE_DELETE_METHOD_NAME - resource_name = _TEST_PIPELINE_RESOURCE_NAME - def setup_method(self): reload(initializer) reload(aiplatform) @@ -457,7 +454,7 @@ def test_run_call_pipeline_service_create_with_timeout( # ) # assert job._gca_resource == make_pipeline_job( - # gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + # gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED # ) @pytest.mark.parametrize( diff --git a/tests/unit/aiplatform/test_tensorboard.py b/tests/unit/aiplatform/test_tensorboard.py index 37e3376875..b73b805202 100644 --- a/tests/unit/aiplatform/test_tensorboard.py +++ b/tests/unit/aiplatform/test_tensorboard.py @@ -15,7 +15,6 @@ # limitations under the License. # - import pytest from unittest import mock @@ -28,6 +27,7 @@ from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform import tensorboard +from google.cloud.aiplatform import utils from google.cloud.aiplatform.compat.services import ( tensorboard_service_client, @@ -36,9 +36,11 @@ from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec, tensorboard as gca_tensorboard, + tensorboard_data as gca_tensorboard_data, tensorboard_experiment as gca_tensorboard_experiment, tensorboard_run as gca_tensorboard_run, tensorboard_service as gca_tensorboard_service, + tensorboard_time_series as gca_tensorboard_time_series, ) from google.protobuf import field_mask_pb2 @@ -75,6 +77,11 @@ f"{_TEST_TENSORBOARD_EXPERIMENT_NAME}/runs/{_TEST_TENSORBOARD_RUN_ID}" ) +_TEST_TENSORBOARD_TIME_SERIES_ID = "test-time-series" +_TEST_TENSORBOARD_TIME_SERIES_NAME = ( + f"{_TEST_TENSORBOARD_RUN_NAME}/timeSeries/{_TEST_TENSORBOARD_TIME_SERIES_ID}" +) + # request_metadata _TEST_REQUEST_METADATA = () @@ -84,6 +91,8 @@ kms_key_name=_TEST_ENCRYPTION_KEY_NAME ) +_TEST_TIME_SERIES_DISPLAY_NAME = "accuracy" + @pytest.fixture def get_tensorboard_mock(): @@ -143,18 +152,19 @@ def delete_tensorboard_mock(): yield delete_tensorboard_mock +_TEST_TENSORBOARD_EXPERIMENT = gca_tensorboard_experiment.TensorboardExperiment( + name=_TEST_TENSORBOARD_EXPERIMENT_NAME, + display_name=_TEST_DISPLAY_NAME, +) + + @pytest.fixture def get_tensorboard_experiment_mock(): with patch.object( tensorboard_service_client.TensorboardServiceClient, "get_tensorboard_experiment", ) as get_tensorboard_experiment__mock: - get_tensorboard_experiment__mock.return_value = ( - gca_tensorboard_experiment.TensorboardExperiment( - name=_TEST_TENSORBOARD_EXPERIMENT_NAME, - display_name=_TEST_DISPLAY_NAME, - ) - ) + get_tensorboard_experiment__mock.return_value = _TEST_TENSORBOARD_EXPERIMENT yield get_tensorboard_experiment__mock @@ -164,12 +174,7 @@ def create_tensorboard_experiment_mock(): tensorboard_service_client.TensorboardServiceClient, "create_tensorboard_experiment", ) as create_tensorboard_experiment_mock: - create_tensorboard_experiment_mock.return_value = ( - gca_tensorboard_experiment.TensorboardExperiment( - name=_TEST_TENSORBOARD_EXPERIMENT_NAME, - display_name=_TEST_DISPLAY_NAME, - ) - ) + create_tensorboard_experiment_mock.return_value = _TEST_TENSORBOARD_EXPERIMENT yield create_tensorboard_experiment_mock @@ -197,25 +202,23 @@ def list_tensorboard_experiment_mock(): tensorboard_service_client.TensorboardServiceClient, "list_tensorboard_experiments", ) as list_tensorboard_experiment_mock: - list_tensorboard_experiment_mock.return_value = [ - gca_tensorboard_experiment.TensorboardExperiment( - name=_TEST_TENSORBOARD_EXPERIMENT_NAME, - display_name=_TEST_DISPLAY_NAME, - ) - ] + list_tensorboard_experiment_mock.return_value = [_TEST_TENSORBOARD_EXPERIMENT] yield list_tensorboard_experiment_mock +_TEST_TENSORBOARD_RUN = gca_tensorboard_run.TensorboardRun( + name=_TEST_TENSORBOARD_RUN_NAME, + display_name=_TEST_DISPLAY_NAME, +) + + @pytest.fixture def get_tensorboard_run_mock(): with patch.object( tensorboard_service_client.TensorboardServiceClient, "get_tensorboard_run", ) as get_tensorboard_run_mock: - get_tensorboard_run_mock.return_value = gca_tensorboard_run.TensorboardRun( - name=_TEST_TENSORBOARD_RUN_NAME, - display_name=_TEST_DISPLAY_NAME, - ) + get_tensorboard_run_mock.return_value = _TEST_TENSORBOARD_RUN yield get_tensorboard_run_mock @@ -225,10 +228,7 @@ def create_tensorboard_run_mock(): tensorboard_service_client.TensorboardServiceClient, "create_tensorboard_run", ) as create_tensorboard_run_mock: - create_tensorboard_run_mock.return_value = gca_tensorboard_run.TensorboardRun( - name=_TEST_TENSORBOARD_RUN_NAME, - display_name=_TEST_DISPLAY_NAME, - ) + create_tensorboard_run_mock.return_value = _TEST_TENSORBOARD_RUN yield create_tensorboard_run_mock @@ -263,6 +263,97 @@ def list_tensorboard_run_mock(): yield list_tensorboard_run_mock +@pytest.fixture +def write_tensorboard_run_data_mock(): + with patch.object( + tensorboard_service_client.TensorboardServiceClient, + "write_tensorboard_run_data", + ) as write_tensorboard_run_data_mock: + yield write_tensorboard_run_data_mock + + +_TEST_TENSORBOARD_TIME_SERIES = gca_tensorboard_time_series.TensorboardTimeSeries( + name=_TEST_TENSORBOARD_TIME_SERIES_NAME, + display_name=_TEST_TIME_SERIES_DISPLAY_NAME, + value_type=gca_tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR, +) + + +@pytest.fixture +def get_tensorboard_time_series_mock(): + with patch.object( + tensorboard_service_client.TensorboardServiceClient, + "get_tensorboard_time_series", + ) as get_tensorboard_time_series_mock: + get_tensorboard_time_series_mock.return_value = _TEST_TENSORBOARD_TIME_SERIES + yield get_tensorboard_time_series_mock + + +@pytest.fixture +def create_tensorboard_time_series_mock(): + with patch.object( + tensorboard_service_client.TensorboardServiceClient, + "create_tensorboard_time_series", + ) as create_tensorboard_time_series_mock: + create_tensorboard_time_series_mock.return_value = _TEST_TENSORBOARD_TIME_SERIES + yield create_tensorboard_time_series_mock + + +@pytest.fixture +def delete_tensorboard_time_series_mock(): + with mock.patch.object( + tensorboard_service_client.TensorboardServiceClient, + "delete_tensorboard_time_series", + ) as delete_tensorboard_time_series_mock: + delete_tensorboard_lro_time_series_mock = mock.Mock(operation.Operation) + delete_tensorboard_lro_time_series_mock.result.return_value = ( + gca_tensorboard_service.DeleteTensorboardTimeSeriesRequest( + name=_TEST_TENSORBOARD_TIME_SERIES_NAME, + ) + ) + delete_tensorboard_time_series_mock.return_value = ( + delete_tensorboard_lro_time_series_mock + ) + yield delete_tensorboard_time_series_mock + + +@pytest.fixture +def list_tensorboard_time_series_mock(): + with patch.object( + tensorboard_service_client.TensorboardServiceClient, + "list_tensorboard_time_series", + ) as list_tensorboard_time_series_mock: + list_tensorboard_time_series_mock.return_value = [_TEST_TENSORBOARD_TIME_SERIES] + yield list_tensorboard_time_series_mock + + +_TEST_TENSORBOARD_TIME_SERIES_DATA = gca_tensorboard_data.TimeSeriesData( + tensorboard_time_series_id=_TEST_TENSORBOARD_TIME_SERIES_ID, + value_type=gca_tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR, + values=[ + gca_tensorboard_data.TimeSeriesDataPoint( + scalar=gca_tensorboard_data.Scalar(value=1.0), + step=1, + wall_time=utils.get_timestamp_proto(), + ) + ], +) + + +@pytest.fixture +def batch_read_tensorboard_time_series_mock(): + with patch.object( + tensorboard_service_client.TensorboardServiceClient, + "batch_read_tensorboard_time_series_data", + ) as batch_read_tensorboard_time_series_data_mock: + batch_read_tensorboard_time_series_data_mock.return_value = ( + gca_tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse( + time_series_data=[_TEST_TENSORBOARD_TIME_SERIES_DATA] + ) + ) + yield batch_read_tensorboard_time_series_data_mock + + @pytest.mark.usefixtures("google_auth_mock") class TestTensorboard: def setup_method(self): @@ -490,6 +581,7 @@ def test_update_tensorboard_encryption_spec(self, update_tensorboard_mock): ) +@pytest.mark.usefixtures("google_auth_mock") class TestTensorboardExperiment: def setup_method(self): reload(initializer) @@ -647,6 +739,7 @@ def test_list_tensorboard_experiments(self, list_tensorboard_experiment_mock): ) +@pytest.mark.usefixtures("google_auth_mock") class TestTensorboardRun: def setup_method(self): reload(initializer) @@ -655,6 +748,7 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) + @pytest.mark.usefixtures("list_tensorboard_time_series_mock") def test_init_tensorboard_run(self, get_tensorboard_run_mock): aiplatform.init(project=_TEST_PROJECT) tensorboard.TensorboardRun(tensorboard_run_name=_TEST_TENSORBOARD_RUN_NAME) @@ -662,6 +756,7 @@ def test_init_tensorboard_run(self, get_tensorboard_run_mock): name=_TEST_TENSORBOARD_RUN_NAME, retry=base._DEFAULT_RETRY ) + @pytest.mark.usefixtures("list_tensorboard_time_series_mock") def test_init_tensorboard_run_with_tensorboard_and_experiment( self, get_tensorboard_run_mock ): @@ -676,7 +771,7 @@ def test_init_tensorboard_run_with_tensorboard_and_experiment( ) def test_init_tensorboard_run_with_id_only_with_project_and_location( - self, get_tensorboard_run_mock + self, get_tensorboard_run_mock, list_tensorboard_time_series_mock ): aiplatform.init(project=_TEST_PROJECT) tensorboard.TensorboardRun( @@ -689,9 +784,15 @@ def test_init_tensorboard_run_with_id_only_with_project_and_location( get_tensorboard_run_mock.assert_called_once_with( name=_TEST_TENSORBOARD_RUN_NAME, retry=base._DEFAULT_RETRY ) + list_tensorboard_time_series_mock.assert_called_once_with( + request={"parent": _TEST_TENSORBOARD_RUN_NAME, "filter": None} + ) + @pytest.mark.usefixtures("list_tensorboard_time_series_mock") def test_create_tensorboard_run( - self, create_tensorboard_run_mock, get_tensorboard_run_mock + self, + create_tensorboard_run_mock, + get_tensorboard_run_mock, ): aiplatform.init( @@ -720,6 +821,7 @@ def test_create_tensorboard_run( name=_TEST_TENSORBOARD_RUN_NAME, retry=base._DEFAULT_RETRY ) + @pytest.mark.usefixtures("list_tensorboard_time_series_mock") def test_create_tensorboard_run_with_timeout( self, create_tensorboard_run_mock, get_tensorboard_run_mock ): @@ -746,6 +848,7 @@ def test_create_tensorboard_run_with_timeout( timeout=180.0, ) + @pytest.mark.usefixtures("list_tensorboard_time_series_mock") def test_create_tensorboard_run_with_timeout_not_explicitly_set( self, create_tensorboard_run_mock, get_tensorboard_run_mock ): @@ -771,7 +874,9 @@ def test_create_tensorboard_run_with_timeout_not_explicitly_set( timeout=None, ) - @pytest.mark.usefixtures("get_tensorboard_run_mock") + @pytest.mark.usefixtures( + "get_tensorboard_run_mock", "list_tensorboard_time_series_mock" + ) def test_delete_tensorboard_run(self, delete_tensorboard_run_mock): aiplatform.init(project=_TEST_PROJECT) @@ -785,7 +890,9 @@ def test_delete_tensorboard_run(self, delete_tensorboard_run_mock): name=my_tensorboard_run.resource_name ) - def test_list_tensorboard_runs(self, list_tensorboard_run_mock): + def test_list_tensorboard_runs( + self, list_tensorboard_run_mock, list_tensorboard_time_series_mock + ): aiplatform.init(project=_TEST_PROJECT) tensorboard.TensorboardRun.list( @@ -795,3 +902,167 @@ def test_list_tensorboard_runs(self, list_tensorboard_run_mock): list_tensorboard_run_mock.assert_called_once_with( request={"parent": _TEST_TENSORBOARD_EXPERIMENT_NAME, "filter": None} ) + + list_tensorboard_time_series_mock.assert_called_once_with( + request={"parent": _TEST_TENSORBOARD_RUN_NAME, "filter": None} + ) + + @pytest.mark.usefixtures( + "get_tensorboard_run_mock", "list_tensorboard_time_series_mock" + ) + def test_write_tensorboard_run_data(self, write_tensorboard_run_data_mock): + aiplatform.init(project=_TEST_PROJECT) + + tb_run = tensorboard.TensorboardRun( + tensorboard_run_name=_TEST_TENSORBOARD_RUN_NAME + ) + + timestamp = utils.get_timestamp_proto() + tb_run.write_tensorboard_scalar_data( + time_series_data={"accuracy": 0.9}, step=1, wall_time=timestamp + ) + + expected_time_series_data = [ + gca_tensorboard_data.TimeSeriesData( + tensorboard_time_series_id=_TEST_TENSORBOARD_TIME_SERIES_ID, + value_type=gca_tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR, + values=[ + gca_tensorboard_data.TimeSeriesDataPoint( + scalar=gca_tensorboard_data.Scalar(value=0.9), + wall_time=timestamp, + step=1, + ) + ], + ), + ] + + write_tensorboard_run_data_mock.assert_called_once_with( + tensorboard_run=_TEST_TENSORBOARD_RUN_NAME, + time_series_data=expected_time_series_data, + ) + + @pytest.mark.usefixtures( + "get_tensorboard_run_mock", "list_tensorboard_time_series_mock" + ) + def test_read_tensorboard_time_series( + self, batch_read_tensorboard_time_series_mock + ): + aiplatform.init(project=_TEST_PROJECT) + + tb_run = tensorboard.TensorboardRun( + tensorboard_run_name=_TEST_TENSORBOARD_RUN_NAME + ) + + ts_data = tb_run.read_time_series_data() + + true_ts_data = { + _TEST_TIME_SERIES_DISPLAY_NAME: _TEST_TENSORBOARD_TIME_SERIES_DATA + } + + batch_read_tensorboard_time_series_mock.assert_called_once_with( + request=gca_tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest( + tensorboard=_TEST_NAME, + time_series=[_TEST_TENSORBOARD_TIME_SERIES_NAME], + ) + ) + + assert ts_data == true_ts_data + + +@pytest.mark.usefixtures("google_auth_mock") +class TestTensorboardTimeSeries: + def setup_method(self): + reload(initializer) + reload(aiplatform) + + def teardown_method(self): + initializer.global_pool.shutdown(wait=True) + + def test_init_tensorboard_time_series(self, get_tensorboard_time_series_mock): + aiplatform.init(project=_TEST_PROJECT) + tensorboard.TensorboardTimeSeries( + tensorboard_time_series_name=_TEST_TENSORBOARD_TIME_SERIES_NAME + ) + get_tensorboard_time_series_mock.assert_called_once_with( + name=_TEST_TENSORBOARD_TIME_SERIES_NAME, retry=base._DEFAULT_RETRY + ) + + def test_init_tensorboard_time_series_with_tensorboard_and_experiment_and_run( + self, get_tensorboard_time_series_mock + ): + aiplatform.init(project=_TEST_PROJECT) + tensorboard.TensorboardTimeSeries( + tensorboard_time_series_name=_TEST_TENSORBOARD_TIME_SERIES_ID, + tensorboard_run_id=_TEST_TENSORBOARD_RUN_ID, + tensorboard_experiment_id=_TEST_TENSORBOARD_EXPERIMENT_ID, + tensorboard_id=_TEST_ID, + ) + get_tensorboard_time_series_mock.assert_called_once_with( + name=_TEST_TENSORBOARD_TIME_SERIES_NAME, retry=base._DEFAULT_RETRY + ) + + def test_init_tensorboard_time_series_with_id_only_with_project_and_location( + self, get_tensorboard_time_series_mock + ): + aiplatform.init(project=_TEST_PROJECT) + tensorboard.TensorboardTimeSeries( + tensorboard_time_series_name=_TEST_TENSORBOARD_TIME_SERIES_ID, + tensorboard_run_id=_TEST_TENSORBOARD_RUN_ID, + tensorboard_experiment_id=_TEST_TENSORBOARD_EXPERIMENT_ID, + tensorboard_id=_TEST_ID, + project=_TEST_PROJECT, + location=_TEST_LOCATION, + ) + get_tensorboard_time_series_mock.assert_called_once_with( + name=_TEST_TENSORBOARD_TIME_SERIES_NAME, retry=base._DEFAULT_RETRY + ) + + def test_create_tensorboard_time_series( + self, + create_tensorboard_time_series_mock, + ): + + aiplatform.init( + project=_TEST_PROJECT, + ) + + tensorboard.TensorboardTimeSeries.create( + display_name=_TEST_TIME_SERIES_DISPLAY_NAME, + tensorboard_run_name=_TEST_TENSORBOARD_RUN_NAME, + ) + + expected_tensorboard_time_series = gca_tensorboard_time_series.TensorboardTimeSeries( + display_name=_TEST_TIME_SERIES_DISPLAY_NAME, + value_type=gca_tensorboard_time_series.TensorboardTimeSeries.ValueType.SCALAR, + plugin_name="scalars", + ) + + create_tensorboard_time_series_mock.assert_called_once_with( + parent=_TEST_TENSORBOARD_RUN_NAME, + tensorboard_time_series=expected_tensorboard_time_series, + ) + + @pytest.mark.usefixtures("get_tensorboard_time_series_mock") + def test_delete_tensorboard_time_series(self, delete_tensorboard_time_series_mock): + aiplatform.init(project=_TEST_PROJECT) + + my_tensorboard_time_series = tensorboard.TensorboardTimeSeries( + tensorboard_time_series_name=_TEST_TENSORBOARD_TIME_SERIES_NAME + ) + + my_tensorboard_time_series.delete() + + delete_tensorboard_time_series_mock.assert_called_once_with( + name=my_tensorboard_time_series.resource_name + ) + + def test_list_tensorboard_time_series(self, list_tensorboard_time_series_mock): + aiplatform.init(project=_TEST_PROJECT) + + tensorboard.TensorboardTimeSeries.list( + tensorboard_run_name=_TEST_TENSORBOARD_RUN_NAME + ) + + list_tensorboard_time_series_mock.assert_called_once_with( + request={"parent": _TEST_TENSORBOARD_RUN_NAME, "filter": None} + ) diff --git a/tests/unit/aiplatform/test_utils.py b/tests/unit/aiplatform/test_utils.py index c700271590..c467a28b05 100644 --- a/tests/unit/aiplatform/test_utils.py +++ b/tests/unit/aiplatform/test_utils.py @@ -71,9 +71,11 @@ def test_invalid_region_does_not_raise_with_valid_region(): ( "contexts", "123456", - aiplatform.metadata._Context._parse_resource_name, - aiplatform.metadata._Context._format_resource_name, - {aiplatform.metadata._MetadataStore._resource_noun: "default"}, + aiplatform.metadata.context._Context._parse_resource_name, + aiplatform.metadata.context._Context._format_resource_name, + { + aiplatform.metadata.metadata_store._MetadataStore._resource_noun: "default" + }, "europe-west4", "projects/857392/locations/us-central1/metadataStores/default/contexts/123", ), @@ -142,9 +144,11 @@ def test_full_resource_name_with_full_name( ( "123", "contexts", - aiplatform.metadata._Context._parse_resource_name, - aiplatform.metadata._Context._format_resource_name, - {aiplatform.metadata._MetadataStore._resource_noun: "default"}, + aiplatform.metadata.context._Context._parse_resource_name, + aiplatform.metadata.context._Context._format_resource_name, + { + aiplatform.metadata.metadata_store._MetadataStore._resource_noun: "default" + }, "857392", "us-central1", "projects/857392/locations/us-central1/metadataStores/default/contexts/123",