diff --git a/tfx/orchestration/experimental/core/env.py b/tfx/orchestration/experimental/core/env.py index 464ade36a5..2a5a2c2287 100644 --- a/tfx/orchestration/experimental/core/env.py +++ b/tfx/orchestration/experimental/core/env.py @@ -47,6 +47,12 @@ def get_orchestration_options( def get_base_dir(self) -> Optional[str]: """Returns the base directory for the pipeline.""" + @abc.abstractmethod + def label_and_tag_pipeline_run( + self, mlmd_handle, pipeline_id, pipeline_run_id, labels, tags + ) -> None: + """Labels and tags the pipeline run after it starts.""" + @abc.abstractmethod def max_mlmd_str_value_length(self) -> Optional[int]: """Returns max size of a string value in MLMD db, `None` if unlimited.""" @@ -111,6 +117,11 @@ def get_orchestration_options( def get_base_dir(self) -> Optional[str]: return None + def label_and_tag_pipeline_run( + self, mlmd_handle, pipeline_id, pipeline_run_id, labels, tags + ) -> None: + return None + def max_mlmd_str_value_length(self) -> Optional[int]: return None diff --git a/tfx/orchestration/experimental/core/env_test.py b/tfx/orchestration/experimental/core/env_test.py index de7e33ed36..c34f9621f3 100644 --- a/tfx/orchestration/experimental/core/env_test.py +++ b/tfx/orchestration/experimental/core/env_test.py @@ -32,6 +32,11 @@ def get_orchestration_options(self, pipeline): def get_base_dir(self): raise NotImplementedError() + def label_and_tag_pipeline_run( + self, mlmd_handle, pipeline_id, pipeline_run_id, labels, tags + ): + raise NotImplementedError() + def max_mlmd_str_value_length(self): raise NotImplementedError()