feat: Vertex AI Experiments GA (#1410)

Author: sasha-gitg

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",

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."""

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

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,

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,