Add pipeline_start_post_processing to Orchestrator env

PiperOrigin-RevId: 620324534

tfx/orchestration/experimental/core/env.py

@@ -71,6 +71,21 @@ 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_processing(
+ self, pipeline: pipeline_pb2.Pipeline
+ ) -> pipeline_pb2.Pipeline:
+ """Method for processing a pipeline at the end of its initialization, before it starts running.
+
+ Does not mutate the input pipeline.
+
+ Args:
+ pipeline: The pipeline IR to process.
+
+ Returns:
+ The updated pipeline IR.
+ """
+
 
 class _DefaultEnv(Env):
  """Default environment."""
@@ -104,6 +119,11 @@ 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_processing(
+ self, pipeline: pipeline_pb2.Pipeline
+ ) -> pipeline_pb2.Pipeline:
+ return pipeline
+
 
 _ENV = _DefaultEnv()