Inconsistent ExitHandler behavior

[logicbomb421]

This is an expansion on #999.

---

What happened:
It appears currently the only valid way to create an ExitHandler is using dsl.ContainerOp, with is_exit_handler set to True. When the pipeline runs, duplicate nodes are spawned at the beginning of the workflow and execute the exit handler logic immediately. Attempting to set is_exit_handler manually on an ad-hoc component created with func_to_container_op, or one loaded via load_component[*], appears to do nothing.

What did you expect to happen:
Regardless of where the component came from, marking it as an exit handler should inform Kubeflow the component is an exit handler, and said handler should only be invoked once upon exiting the pipeline.

What steps did you take:

func_to_container_op exit handler

import kfp.dsl as dsl
import kfp.components as components

def print_message(message: str) -> None:
  print(message or 'NO MESSAGE SUPPLIED!')
print_op = components.func_to_container_op(print_message, base_image='python:3.7-alpine')

@dsl.pipeline(name='simple pipeline')
def pipeline_func():
  exit_handler = print_op('EXITED')
  exit_handler.is_exit_handler=True
  with dsl.ExitHandler(exit_handler):
    first = print_op(message='first')
    second = print_op(message='second').after(first)
    third = print_op(message='third').after(second)

This design produces the following result:

load_component_from_text exit handler

import kfp.components as components
from kfp import dsl

exit_op = components.load_component_from_text("""
name: 'Print Message Component'
description: Prints a message.
inputs:
  - {name: message, type: String, description: 'Required. The message to print.'}
implementation:
  container:
    image: python:3.7-alpine
    command: [python]
    args: [
      -c,
      import sys; print(sys.argv.pop() or 'FALLBACK MESSAGE'),
      {inputValue: message},
    ]
""")

def print_message(message: str) -> None:
  print(message or 'NO MESSAGE SUPPLIED!')
print_op = components.func_to_container_op(print_message, base_image='python:3.7-alpine')

@dsl.pipeline(name='simple pipeline')
def pipeline_func():
  exit_handler = exit_op(message='EXITED')
  exit_handler.is_exit_handler=True
  with dsl.ExitHandler(exit_handler):
    first = print_op(message='first')
    second = print_op(message='second').after(first)
    third = print_op(message='third').after(second)

This design produces the following result:

pure dsl.ContainerOp exit handler

import kfp.dsl as dsl
import kfp.components as components

def print_message(message: str) -> None:
  print(message or 'NO MESSAGE SUPPLIED!')
print_op = components.func_to_container_op(print_message, base_image='python:3.7-alpine')

@dsl.pipeline(name='simple pipeline')
def pipeline_func():
  exit_handler = dsl.ContainerOp(
    name='Exit Handler',
    is_exit_handler=True,
    image='python:3.7-alpine',
    command='python',
    arguments=[ 
      '-c',
      "print('EXITED')"
    ]
  )
  with dsl.ExitHandler(exit_handler):
    first = print_op(message='first')
    second = print_op(message='second').after(first)
    third = print_op(message='third').after(second)

This design correctly produces the following result:

Anything else you would like to add:
Ideally this boolean value would be supplied with all calls that return a component factory function (to avoid needing to set it explicitly on the resulting operation).

[Ark-kun]

This is indeed an issue - is_exit_handler is only evaluated at construction time.
We should probably change the dsl.ExitHandler context manager so that it fixes the exit handler task (removes it from the op groups).

[Ark-kun]

Fixed by #2411