feat: add get_associated_experiment method to pipeline_jobs #1476
Conversation
product-auto-label
bot
added
size: m
| @@ -776,3 +776,31 @@ def clone( | |||
| ) | |||
|
|
|||
| return cloned | |||
|
|
|||
| def get_associated_experiment(self) -> Optional["experiment_resources.Experiment"]: | |||
There was a problem hiding this comment.
| def get_associated_experiment(self) -> Optional["experiment_resources.Experiment"]: | |
| def get_associated_experiment(self) -> Optional["aiplatform.Experiment"]: |
|
|
||
| def get_associated_experiment(self) -> Optional["experiment_resources.Experiment"]: | ||
| """Returns the aiplatform.Experiment associated with this PipelineJob, | ||
| or None if this PipelineJob is not associated with an experiment.""" |
| pipeline_experiment_resource_names = [] | ||
|
|
||
| for i in pipeline_context_resources: | ||
| if i.schema_title == "system.Experiment": |
There was a problem hiding this comment.
|
|
||
| pipeline_parent_contexts = ( | ||
| self._gca_resource.job_detail.pipeline_run_context.parent_contexts | ||
| ) |
There was a problem hiding this comment.
You can locally filter out the pipeline_context to avoid an additional GET: https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform_v1/types/pipeline_job.py#L309
ie:
pipeline_parent_context = [c for c in pipeline_parent_contexts if c.name != self._gca_resource.job_detail.pipeline_context.name|
|
||
| if len(pipeline_experiment_resource_names) > 1: | ||
| _LOGGER.warning( | ||
| "There are more than experiments associated with this pipeline, returning only one..." |
There was a problem hiding this comment.
More informative logging message, logging all the found experiments and and experiment that is being returned.
|
|
||
| for i in pipeline_context_resources: | ||
| if i.schema_title == "system.Experiment": | ||
| pipeline_experiment_resource_names.append(i.name.split("contexts/")[1]) |
There was a problem hiding this comment.
I think you can use context.name
| @@ -29,6 +29,10 @@ | |||
| from google.cloud import aiplatform | |||
There was a problem hiding this comment.
Please also update the experiment system tests to test this functinality. Add a test right after this test: https://github.com/googleapis/python-aiplatform/blob/main/tests/system/aiplatform/test_experiments.py#L263
It should be able to get the pipeline_job using job = aiplatform.PipelineJob(self._pipeline_job_id and then verify the associated experiment is self._experiment_name.
There was a problem hiding this comment.
I added a call to get_associated_experiment to test_add_pipeline_job_to_experiment. Verified it passed on my test project and just kicked off our system tests from my branch.
…is#1476) * feat: add get_associated_experiment method to pipeline_jobs * updates from reviewer feedback * clean up system test * re-add check for experiment schema title
…oogleapis#1476)" This reverts commit b2c9939.
Added a
get_associated_experimentmethod toPipelineJobto return the Experiment associated with the provided pipeline job.