Skip to content

Commit bee8e93

Browse files
SinaChavoshiSinaChavoshi
committed
remove start_execution from this pr and move to a separate PR
1 parent 1778a34 commit bee8e93

File tree

4 files changed

+14
-111
lines changed

4 files changed

+14
-111
lines changed

‎google/cloud/aiplatform/metadata/metadata.py

+14-30
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
from google.cloud.aiplatform.metadata import execution
3232
from google.cloud.aiplatform.metadata import experiment_resources
3333
from google.cloud.aiplatform.metadata import experiment_run_resource
34-
from google.cloud.aiplatform.metadata.schema import base_execution
3534
from google.cloud.aiplatform.tensorboard import tensorboard_resource
3635

3736
_LOGGER = base.Logger(__name__)
@@ -513,12 +512,10 @@ def start_execution(
513512
metadata: Optional[Dict[str, Any]] = None,
514513
schema_version: Optional[str] = None,
515514
description: Optional[str] = None,
516-
metadata_store_id: Optional[str] = "default",
517515
resume: bool = False,
518516
project: Optional[str] = None,
519517
location: Optional[str] = None,
520518
credentials: Optional[auth_credentials.Credentials] = None,
521-
base_execution_schema: Optional["base_execution.BaseExecutionSchema"] = None,
522519
) -> execution.Execution:
523520
"""
524521
Create and starts a new Metadata Execution or resumes a previously created Execution.
@@ -568,9 +565,6 @@ def start_execution(
568565
credentials (auth_credentials.Credentials):
569566
Optional. Custom credentials used to create this Execution. Overrides
570567
credentials set in aiplatform.init.
571-
base_execution_schema (BaseExecutionSchema):
572-
Optional. An instance of the BaseExecutionSchema class that can be provided instead of providing schema specific parameters. It overrides
573-
the values provided for schema_title, resource_id, state, display_name, schema_version, description, and metadata.
574568
575569
Returns:
576570
Execution: Instantiated representation of the managed Metadata Execution.
@@ -611,32 +605,22 @@ def start_execution(
611605

612606
run_execution.update(state=gca_execution.Execution.State.RUNNING)
613607
else:
614-
if base_execution_schema:
615-
run_execution = execution.Execution.create_from_base_execution_schema(
616-
base_execution_schema=base_execution_schema,
617-
metadata_store_id=metadata_store_id,
618-
project=project,
619-
location=location,
620-
credentials=credentials,
608+
if not schema_title:
609+
raise ValueError(
610+
"schema_title must be provided when starting a new Execution"
621611
)
622-
else:
623-
if not schema_title:
624-
raise ValueError(
625-
"schema_title or base_execution_schema must be provided when starting a new Execution"
626-
)
627612

628-
run_execution = execution.Execution.create(
629-
display_name=display_name,
630-
schema_title=schema_title,
631-
schema_version=schema_version,
632-
metadata=metadata,
633-
description=description,
634-
metadata_store_id=metadata_store_id,
635-
resource_id=resource_id,
636-
project=project,
637-
location=location,
638-
credentials=credentials,
639-
)
613+
run_execution = execution.Execution.create(
614+
display_name=display_name,
615+
schema_title=schema_title,
616+
schema_version=schema_version,
617+
metadata=metadata,
618+
description=description,
619+
resource_id=resource_id,
620+
project=project,
621+
location=location,
622+
credentials=credentials,
623+
)
640624

641625
if self.experiment_run:
642626
if self.experiment_run._is_legacy_experiment_run():

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

-38
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
from google.cloud.aiplatform.compat.types import execution as gca_execution
2525
from google.cloud.aiplatform.metadata import constants
2626
from google.cloud.aiplatform.metadata import execution
27-
from google.cloud.aiplatform.metadata import metadata
2827

2928

3029
class BaseExecutionSchema(metaclass=abc.ABCMeta):
@@ -109,40 +108,3 @@ def create(
109108
credentials=credentials,
110109
)
111110
return self.execution
112-
113-
def start_execution(
114-
self,
115-
metadata_store_id: Optional[str] = "default",
116-
project: Optional[str] = None,
117-
location: Optional[str] = None,
118-
credentials: Optional[auth_credentials.Credentials] = None,
119-
) -> "execution.Execution":
120-
"""Create and starts a new Metadata Execution.
121-
122-
Args:
123-
metadata_store_id (str):
124-
Optional. The <metadata_store_id> portion of the resource name with
125-
the format:
126-
projects/123/locations/us-central1/metadataStores/<metadata_store_id>/executions/<resource_id>
127-
If not provided, the MetadataStore's ID will be set to "default".
128-
project (str):
129-
Optional. Project used to create this Execution. Overrides project set in
130-
aiplatform.init.
131-
location (str):
132-
Optional. Location used to create this Execution. Overrides location set in
133-
aiplatform.init.
134-
credentials (auth_credentials.Credentials):
135-
Optional. Custom credentials used to create this Execution. Overrides
136-
credentials set in aiplatform.init.
137-
Returns:
138-
Execution: Instantiated representation of the managed Metadata Execution.
139-
140-
"""
141-
return metadata._ExperimentTracker().start_execution(
142-
base_execution_schema=self,
143-
resume=False,
144-
metadata_store_id=metadata_store_id,
145-
project=project,
146-
location=location,
147-
credentials=credentials,
148-
)

‎tests/system/aiplatform/test_e2e_metadata_schema.py

-17
Original file line numberDiff line numberDiff line change
@@ -143,20 +143,3 @@ def test_execution_create_using_system_schema_class(self):
143143
assert execution.schema_title == "system.CustomJobExecution"
144144
assert execution.description == self.execution_description
145145
assert "/metadataStores/default/executions/" in execution.resource_name
146-
147-
def test_execution_start_execution_using_system_schema_class(self):
148-
149-
aiplatform.init(
150-
project=e2e_base._PROJECT,
151-
location=e2e_base._LOCATION,
152-
)
153-
154-
execution = system_execution_schema.ContainerExecution(
155-
display_name=self.execution_display_name,
156-
description=self.execution_description,
157-
).start_execution()
158-
159-
assert execution.display_name == self.execution_display_name
160-
assert execution.schema_title == "system.ContainerExecution"
161-
assert execution.description == self.execution_description
162-
assert "/metadataStores/default/executions/" in execution.resource_name

‎tests/unit/aiplatform/test_metadata_schema.py

-26
Original file line numberDiff line numberDiff line change
@@ -242,32 +242,6 @@ class TestExecution(base_execution.BaseExecutionSchema):
242242
assert kwargs["execution"].description == _TEST_DESCRIPTION
243243
assert kwargs["execution"].metadata == _TEST_UPDATED_METADATA
244244

245-
@pytest.mark.usefixtures("create_execution_mock")
246-
def test_start_execution_method_calls_gapic_library_with_correct_parameters(
247-
self, create_execution_mock
248-
):
249-
aiplatform.init(project=_TEST_PROJECT)
250-
251-
class TestExecution(base_execution.BaseExecutionSchema):
252-
schema_title = _TEST_SCHEMA_TITLE
253-
254-
execution = TestExecution(
255-
state=_TEST_EXECUTION_STATE,
256-
display_name=_TEST_DISPLAY_NAME,
257-
description=_TEST_DESCRIPTION,
258-
metadata=_TEST_UPDATED_METADATA,
259-
)
260-
execution.start_execution(metadata_store_id=_TEST_METADATA_STORE)
261-
create_execution_mock.assert_called_once_with(
262-
parent=_TEST_PARENT, execution=mock.ANY, execution_id=None
263-
)
264-
_, _, kwargs = create_execution_mock.mock_calls[0]
265-
assert kwargs["execution"].schema_title == _TEST_SCHEMA_TITLE
266-
assert kwargs["execution"].state == _TEST_EXECUTION_STATE
267-
assert kwargs["execution"].display_name == _TEST_DISPLAY_NAME
268-
assert kwargs["execution"].description == _TEST_DESCRIPTION
269-
assert kwargs["execution"].metadata == _TEST_UPDATED_METADATA
270-
271245

272246
class TestMetadataGoogleArtifactSchema:
273247
def setup_method(self):

0 commit comments

Comments
 (0)