Add notifications, schedules to new flytekit

flytekit/annotated/launch_plan.py

@@ -45,6 +45,8 @@ def create(
  workflow: _annotated_workflow.Workflow,
  default_inputs: Dict[str, Any] = None,
  fixed_inputs: Dict[str, Any] = None,
+ schedule: _schedule_model.Schedule = None,
+ notifications: List[_common_models.Notification] = None,
  ) -> LaunchPlan:
  ctx = FlyteContext.current_context()
  default_inputs = default_inputs or {}
@@ -72,7 +74,14 @@ def create(
  )
  fixed_lm = _literal_models.LiteralMap(literals=fixed_literals)
 
- lp = cls(name=name, workflow=workflow, parameters=wf_signature_parameters, fixed_inputs=fixed_lm)
+ lp = cls(
+ name=name,
+ workflow=workflow,
+ parameters=wf_signature_parameters,
+ fixed_inputs=fixed_lm,
+ schedule=schedule,
+ notifications=notifications,
+ )
 
  # This is just a convenience - we'll need the fixed inputs LiteralMap for when serializing the Launch Plan out
  # to protobuf, but for local execution and such, why not save the original Python native values as well so

flytekit/annotated/notification.py

@@ -0,0 +1,69 @@
+from typing import List
+
+from flytekit.models import common as _common_model
+from flytekit.models.core import execution as _execution_model
+
+
+# Duplicates flytekit.common.notifications.Notification to avoid using the ExtendedSdkType metaclass.
+class Notification(_common_model.Notification):
+
+ VALID_PHASES = {
+ _execution_model.WorkflowExecutionPhase.ABORTED,
+ _execution_model.WorkflowExecutionPhase.FAILED,
+ _execution_model.WorkflowExecutionPhase.SUCCEEDED,
+ _execution_model.WorkflowExecutionPhase.TIMED_OUT,
+ }
+
+ def __init__(
+ self,
+ phases: List[int],
+ email: _common_model.EmailNotification = None,
+ pager_duty: _common_model.PagerDutyNotification = None,
+ slack: _common_model.SlackNotification = None,
+ ):
+ """
+ :param list[int] phases: A required list of phases for which to fire the event. Events can only be fired for
+ terminal phases. Phases should be as defined in: flytekit.models.core.execution.WorkflowExecutionPhase
+ """
+ self._validate_phases(phases)
+ super(Notification, self).__init__(phases, email=email, pager_duty=pager_duty, slack=slack)
+
+ def _validate_phases(self, phases: List[int]):
+ """
+ :param list[int] phases:
+ """
+ if len(phases) == 0:
+ raise AssertionError("You must specify at least one phase for a notification.")
+ for phase in phases:
+ if phase not in self.VALID_PHASES:
+ raise AssertionError(f"Invalid phase: {phase}. only terminal states are permitted for notifications")
+
+
+class PagerDuty(Notification):
+ def __init__(self, phases: List[int], recipients_email: List[str]):
+ """
+ :param list[int] phases: A required list of phases for which to fire the event. Events can only be fired for
+ terminal phases. Phases should be as defined in: flytekit.models.core.execution.WorkflowExecutionPhase
+ :param list[str] recipients_email: A required non-empty list of recipients for the notification.
+ """
+ super(PagerDuty, self).__init__(phases, pager_duty=_common_model.PagerDutyNotification(recipients_email))
+
+
+class Email(Notification):
+ def __init__(self, phases: List[int], recipients_email: List[str]):
+ """
+ :param list[int] phases: A required list of phases for which to fire the event. Events can only be fired for
+ terminal phases. Phases should be as defined in: flytekit.models.core.execution.WorkflowExecutionPhase
+ :param list[str] recipients_email: A required non-empty list of recipients for the notification.
+ """
+ super(Email, self).__init__(phases, email=_common_model.EmailNotification(recipients_email))
+
+
+class Slack(Notification):
+ def __init__(self, phases: List[int], recipients_email: List[str]):
+ """
+ :param list[int] phases: A required list of phases for which to fire the event. Events can only be fired for
+ terminal phases. Phases should be as defined in: flytekit.models.core.execution.WorkflowExecutionPhase
+ :param list[str] recipients_email: A required non-empty list of recipients for the notification.
+ """
+ super(Slack, self).__init__(phases, slack=_common_model.SlackNotification(recipients_email))

flytekit/annotated/schedule.py

@@ -0,0 +1,151 @@
+import datetime
+import re as _re
+
+import croniter as _croniter
+
+from flytekit.models import schedule as _schedule_models
+
+
+# Duplicates flytekit.common.schedules.Schedule to avoid using the ExtendedSdkType metaclass.
+class CronSchedule(_schedule_models.Schedule):
+ _VALID_CRON_ALIASES = [
+ "hourly",
+ "hours",
+ "@hourly",
+ "daily",
+ "days",
+ "@daily",
+ "weekly",
+ "weeks",
+ "@weekly",
+ "monthly",
+ "months",
+ "@monthly",
+ "annually",
+ "@annually",
+ "yearly",
+ "years",
+ "@yearly",
+ ]
+
+ # Not a perfect regex but good enough and simple to reason about
+ _OFFSET_PATTERN = _re.compile("([-+]?)P([-+0-9YMWD]+)?(T([-+0-9HMS.,]+)?)?")
+
+ def __init__(
+ self, cron_expression: str = None, schedule: str = None, offset: str = None, kickoff_time_input_arg: str = None
+ ):
+ """
+ :param str cron_expression:
+ :param str schedule:
+ :param str offset:
+ :param str kickoff_time_input_arg:
+ """
+ if cron_expression is None and schedule is None:
+ raise AssertionError("Either `cron_expression` or `schedule` should be specified.")
+
+ if cron_expression is not None and offset is not None:
+ raise AssertionError("Only `schedule` is supported when specifying `offset`.")
+
+ if cron_expression is not None:
+ CronSchedule._validate_expression(cron_expression)
+
+ if schedule is not None:
+ CronSchedule._validate_schedule(schedule)
+
+ if offset is not None:
+ CronSchedule._validate_offset(offset)
+
+ super(CronSchedule, self).__init__(
+ kickoff_time_input_arg,
+ cron_expression=cron_expression,
+ cron_schedule=_schedule_models.Schedule.CronSchedule(schedule, offset) if schedule is not None else None,
+ )
+
+ @staticmethod
+ def _validate_expression(cron_expression: str):
+ """
+ Ensures that the set value is a valid cron string. We use the format used in Cloudwatch and the best
+ explanation can be found here:
+ https://docs.aws.amazon.com/AmazonCloudWatch/latest/events/ScheduledEvents.html#CronExpressions
+ :param str cron_expression: cron expression
+ """
+ # We use the croniter lib to validate our cron expression. Since on the admin side we use Cloudwatch,
+ # we have a couple checks in order to line up Cloudwatch with Croniter.
+ tokens = cron_expression.split()
+ if len(tokens) != 6:
+ raise ValueError(
+ "Cron expression is invalid. A cron expression must have 6 fields. Cron expressions are in the "
+ "format of: `minute hour day-of-month month day-of-week year`. "
+ "Use `schedule` for 5 fields cron expression. Received: `{}`".format(cron_expression)
+ )
+
+ if tokens[2] != "?" and tokens[4] != "?":
+ raise ValueError(
+ "Scheduled string is invalid. A cron expression must have a '?' for either day-of-month or "
+ "day-of-week. Please specify '?' for one of those fields. Cron expressions are in the format of: "
+ "minute hour day-of-month month day-of-week year.\n\n"
+ "For more information: "
+ "https://docs.aws.amazon.com/AmazonCloudWatch/latest/events/ScheduledEvents.html#CronExpressions"
+ )
+
+ try:
+ # Cut to 5 fields and just assume year field is good because croniter treats the 6th field as seconds.
+ # TODO: Parse this field ourselves and check
+ _croniter.croniter(" ".join(cron_expression.replace("?", "*").split()[:5]))
+ except Exception:
+ raise ValueError(
+ "Scheduled string is invalid. The cron expression was found to be invalid."
+ f" Provided cron expr: {cron_expression}"
+ )
+
+ @staticmethod
+ def _validate_schedule(schedule: str):
+ if schedule.lower() not in CronSchedule._VALID_CRON_ALIASES:
+ try:
+ _croniter.croniter(schedule)
+ except Exception:
+ raise ValueError(
+ "Schedule is invalid. It must be set to either a cron alias or valid cron expression."
+ f" Provided schedule: {schedule}"
+ )
+
+ @staticmethod
+ def _validate_offset(offset: str):
+ if CronSchedule._OFFSET_PATTERN.fullmatch(offset) is None:
+ raise ValueError("Offset is invalid. It must be an ISO 8601 duration. Provided offset: {}".format(offset))
+
+
+class FixedRate(_schedule_models.Schedule):
+ def __init__(self, duration: datetime.timedelta, kickoff_time_input_arg: str = None):
+ """
+ :param datetime.timedelta duration:
+ :param str kickoff_time_input_arg:
+ """
+ super(FixedRate, self).__init__(kickoff_time_input_arg, rate=self._translate_duration(duration))
+
+ @staticmethod
+ def _translate_duration(duration: datetime.timedelta):
+ """
+ :param datetime.timedelta duration: timedelta between runs
+ :rtype: flytekit.models.schedule.Schedule.FixedRate
+ """
+ _SECONDS_TO_MINUTES = 60
+ _SECONDS_TO_HOURS = _SECONDS_TO_MINUTES * 60
+ _SECONDS_TO_DAYS = _SECONDS_TO_HOURS * 24
+
+ if duration.microseconds != 0 or duration.seconds % _SECONDS_TO_MINUTES != 0:
+ raise AssertionError(
+ f"Granularity of less than a minute is not supported for FixedRate schedules. Received: {duration}"
+ )
+ elif int(duration.total_seconds()) % _SECONDS_TO_DAYS == 0:
+ return _schedule_models.Schedule.FixedRate(
+ int(duration.total_seconds() / _SECONDS_TO_DAYS), _schedule_models.Schedule.FixedRateUnit.DAY,
+ )
+ elif int(duration.total_seconds()) % _SECONDS_TO_HOURS == 0:
+ return _schedule_models.Schedule.FixedRate(
+ int(duration.total_seconds() / _SECONDS_TO_HOURS), _schedule_models.Schedule.FixedRateUnit.HOUR,
+ )
+ else:
+ return _schedule_models.Schedule.FixedRate(
+ int(duration.total_seconds() / _SECONDS_TO_MINUTES), _schedule_models.Schedule.FixedRateUnit.MINUTE,
+ )

flytekit/models/common.py

@@ -242,7 +242,13 @@ def from_flyte_idl(cls, pb2_object):
 
 
 class Notification(FlyteIdlEntity):
- def __init__(self, phases, email=None, pager_duty=None, slack=None):
+ def __init__(
+ self,
+ phases,
+ email: EmailNotification = None,
+ pager_duty: PagerDutyNotification = None,
+ slack: SlackNotification = None,
+ ):
  """
  Represents a structure for notifications based on execution status.
 

flytekit/models/schedule.py

@@ -158,7 +158,8 @@ def from_flyte_idl(cls, pb2_object):
  :param flyteidl.admin.schedule_pb2.Schedule pb2_object:
  :rtype: Schedule
  """
- return cls(
+ # Explicitly instantiate a Schedule model rather than a potential sub-class.
+ return Schedule(
  pb2_object.kickoff_time_input_arg,
  cron_expression=pb2_object.cron_expression if pb2_object.HasField("cron_expression") else None,
  rate=Schedule.FixedRate.from_flyte_idl(pb2_object.rate) if pb2_object.HasField("rate") else None,

tests/flytekit/unit/annotated/test_notifications.py

@@ -0,0 +1,65 @@
+from flyteidl.admin import common_pb2 as _common_pb2
+
+from flytekit.annotated import notification
+from flytekit.annotated.launch_plan import LaunchPlan
+from flytekit.annotated.task import task
+from flytekit.annotated.workflow import workflow
+from flytekit.models import common as _common_model
+from flytekit.models.core import execution as _execution_model
+
+_workflow_execution_succeeded = _execution_model.WorkflowExecutionPhase.SUCCEEDED
+
+
+def test_pager_duty_notification():
+ pager_duty_notif = notification.PagerDuty(
+ phases=[_workflow_execution_succeeded], recipients_email=["my-team@pagerduty.com"]
+ )
+ assert pager_duty_notif.to_flyte_idl() == _common_pb2.Notification(
+ phases=[_workflow_execution_succeeded],
+ email=None,
+ pager_duty=_common_model.PagerDutyNotification(["my-team@pagerduty.com"]).to_flyte_idl(),
+ slack=None,
+ )
+
+
+def test_slack_notification():
+ slack_notif = notification.Slack(phases=[_workflow_execution_succeeded], recipients_email=["my-team@slack.com"])
+ assert slack_notif.to_flyte_idl() == _common_pb2.Notification(
+ phases=[_workflow_execution_succeeded],
+ email=None,
+ pager_duty=None,
+ slack=_common_model.SlackNotification(["my-team@slack.com"]).to_flyte_idl(),
+ )
+
+
+def test_email_notification():
+ email_notif = notification.Email(phases=[_workflow_execution_succeeded], recipients_email=["my-team@email.com"])
+ assert email_notif.to_flyte_idl() == _common_pb2.Notification(
+ phases=[_workflow_execution_succeeded],
+ email=_common_model.EmailNotification(["my-team@email.com"]).to_flyte_idl(),
+ pager_duty=None,
+ slack=None,
+ )
+
+
+def test_with_launch_plan():
+ @task
+ def double(a: int) -> int:
+ return a * 2
+
+ @workflow
+ def quadruple(a: int) -> int:
+ b = double(a=a)
+ c = double(a=b)
+ return c
+
+ lp = LaunchPlan.create(
+ "notif_test",
+ quadruple,
+ notifications=[
+ notification.Email(phases=[_workflow_execution_succeeded], recipients_email=["my-team@email.com"])
+ ],
+ )
+ assert lp.notifications == [
+ notification.Email(phases=[_workflow_execution_succeeded], recipients_email=["my-team@email.com"])
+ ]