Add pipeline_start_post_processing to Orchestrator env

PiperOrigin-RevId: 620324534

tfx/orchestration/experimental/core/env.py

@@ -71,6 +71,16 @@ def set_health_status(self, status: status_lib.Status) -> None:
  def check_if_can_orchestrate(self, pipeline: pipeline_pb2.Pipeline) -> None:
  """Check if this orchestrator is capable of orchestrating the pipeline."""
 
+ @abc.abstractmethod
+ def pipeline_start_post_process(self, pipeline: pipeline_pb2.Pipeline):
+ """Method for processing a pipeline at the end of its initialization, before it starts running.
+
+ This *will* mutate the provided IR in-place.
+
+ Args:
+ pipeline: The pipeline IR to process.
+ """
+
 
 class _DefaultEnv(Env):
  """Default environment."""
@@ -104,6 +114,9 @@ def set_health_status(self, status: status_lib.Status) -> None:
  def check_if_can_orchestrate(self, pipeline: pipeline_pb2.Pipeline) -> None:
  pass
 
+ def pipeline_start_post_process(self, pipeline: pipeline_pb2.Pipeline):
+ pass
+
 
 _ENV = _DefaultEnv()
 

tfx/orchestration/experimental/core/env_test.py

@@ -16,6 +16,7 @@
 import tensorflow as tf
 from tfx.orchestration.experimental.core import env
 from tfx.orchestration.experimental.core import test_utils
+from tfx.proto.orchestration import pipeline_pb2
 from tfx.utils import status as status_lib
 
 
@@ -45,6 +46,9 @@ def set_health_status(self, status: status_lib.Status) -> None:
  def check_if_can_orchestrate(self, pipeline) -> None:
  raise NotImplementedError()
 
+ def pipeline_start_post_process(self, pipeline: pipeline_pb2.Pipeline):
+ raise NotImplementedError()
+
 
 class EnvTest(test_utils.TfxTest):