remove start_execution from this pr and move to a separate PR

google/cloud/aiplatform/metadata/metadata.py

@@ -31,7 +31,6 @@
 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.metadata.schema import base_execution
 from google.cloud.aiplatform.tensorboard import tensorboard_resource
 
 _LOGGER = base.Logger(__name__)
@@ -513,12 +512,10 @@ def start_execution(
         metadata: Optional[Dict[str, Any]] = None,
         schema_version: Optional[str] = None,
         description: Optional[str] = None,
-        metadata_store_id: Optional[str] = "default",
         resume: bool = False,
         project: Optional[str] = None,
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
-        base_execution_schema: Optional["base_execution.BaseExecutionSchema"] = None,
     ) -> execution.Execution:
         """
         Create and starts a new Metadata Execution or resumes a previously created Execution.
@@ -568,9 +565,6 @@ def start_execution(
             credentials (auth_credentials.Credentials):
                 Optional. Custom credentials used to create this Execution. Overrides
                 credentials set in aiplatform.init.
-            base_execution_schema (BaseExecutionSchema):
-                Optional. An instance of the BaseExecutionSchema class that can be provided instead of providing schema specific parameters. It overrides
-                the values provided for schema_title, resource_id, state, display_name, schema_version, description, and metadata.
 
         Returns:
             Execution: Instantiated representation of the managed Metadata Execution.
@@ -611,32 +605,22 @@ def start_execution(
 
             run_execution.update(state=gca_execution.Execution.State.RUNNING)
         else:
-            if base_execution_schema:
-                run_execution = execution.Execution.create_from_base_execution_schema(
-                    base_execution_schema=base_execution_schema,
-                    metadata_store_id=metadata_store_id,
-                    project=project,
-                    location=location,
-                    credentials=credentials,
+            if not schema_title:
+                raise ValueError(
+                    "schema_title must be provided when starting a new Execution"
                 )
-            else:
-                if not schema_title:
-                    raise ValueError(
-                        "schema_title or base_execution_schema must be provided when starting a new Execution"
-                    )
 
-                run_execution = execution.Execution.create(
-                    display_name=display_name,
-                    schema_title=schema_title,
-                    schema_version=schema_version,
-                    metadata=metadata,
-                    description=description,
-                    metadata_store_id=metadata_store_id,
-                    resource_id=resource_id,
-                    project=project,
-                    location=location,
-                    credentials=credentials,
-                )
+            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,
+            )
 
         if self.experiment_run:
             if self.experiment_run._is_legacy_experiment_run():

google/cloud/aiplatform/metadata/schema/base_execution.py

@@ -24,7 +24,6 @@
 from google.cloud.aiplatform.compat.types import execution as gca_execution
 from google.cloud.aiplatform.metadata import constants
 from google.cloud.aiplatform.metadata import execution
-from google.cloud.aiplatform.metadata import metadata
 
 
 class BaseExecutionSchema(metaclass=abc.ABCMeta):
@@ -109,40 +108,3 @@ def create(
             credentials=credentials,
         )
         return self.execution
-
-    def start_execution(
-        self,
-        metadata_store_id: Optional[str] = "default",
-        project: Optional[str] = None,
-        location: Optional[str] = None,
-        credentials: Optional[auth_credentials.Credentials] = None,
-    ) -> "execution.Execution":
-        """Create and starts a new Metadata Execution.
-
-        Args:
-            metadata_store_id (str):
-                Optional. The <metadata_store_id> portion of the resource name with
-                the format:
-                projects/123/locations/us-central1/metadataStores/<metadata_store_id>/executions/<resource_id>
-                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.
-
-        """
-        return metadata._ExperimentTracker().start_execution(
-            base_execution_schema=self,
-            resume=False,
-            metadata_store_id=metadata_store_id,
-            project=project,
-            location=location,
-            credentials=credentials,
-        )

tests/system/aiplatform/test_e2e_metadata_schema.py

@@ -143,20 +143,3 @@ def test_execution_create_using_system_schema_class(self):
         assert execution.schema_title == "system.CustomJobExecution"
         assert execution.description == self.execution_description
         assert "/metadataStores/default/executions/" in execution.resource_name
-
-    def test_execution_start_execution_using_system_schema_class(self):
-
-        aiplatform.init(
-            project=e2e_base._PROJECT,
-            location=e2e_base._LOCATION,
-        )
-
-        execution = system_execution_schema.ContainerExecution(
-            display_name=self.execution_display_name,
-            description=self.execution_description,
-        ).start_execution()
-
-        assert execution.display_name == self.execution_display_name
-        assert execution.schema_title == "system.ContainerExecution"
-        assert execution.description == self.execution_description
-        assert "/metadataStores/default/executions/" in execution.resource_name

tests/unit/aiplatform/test_metadata_schema.py

@@ -242,32 +242,6 @@ class TestExecution(base_execution.BaseExecutionSchema):
         assert kwargs["execution"].description == _TEST_DESCRIPTION
         assert kwargs["execution"].metadata == _TEST_UPDATED_METADATA
 
-    @pytest.mark.usefixtures("create_execution_mock")
-    def test_start_execution_method_calls_gapic_library_with_correct_parameters(
-        self, create_execution_mock
-    ):
-        aiplatform.init(project=_TEST_PROJECT)
-
-        class TestExecution(base_execution.BaseExecutionSchema):
-            schema_title = _TEST_SCHEMA_TITLE
-
-        execution = TestExecution(
-            state=_TEST_EXECUTION_STATE,
-            display_name=_TEST_DISPLAY_NAME,
-            description=_TEST_DESCRIPTION,
-            metadata=_TEST_UPDATED_METADATA,
-        )
-        execution.start_execution(metadata_store_id=_TEST_METADATA_STORE)
-        create_execution_mock.assert_called_once_with(
-            parent=_TEST_PARENT, execution=mock.ANY, execution_id=None
-        )
-        _, _, kwargs = create_execution_mock.mock_calls[0]
-        assert kwargs["execution"].schema_title == _TEST_SCHEMA_TITLE
-        assert kwargs["execution"].state == _TEST_EXECUTION_STATE
-        assert kwargs["execution"].display_name == _TEST_DISPLAY_NAME
-        assert kwargs["execution"].description == _TEST_DESCRIPTION
-        assert kwargs["execution"].metadata == _TEST_UPDATED_METADATA
-
 
 class TestMetadataGoogleArtifactSchema:
     def setup_method(self):