From 2c7752af37c055a0a4db50cb0218f344614361ed Mon Sep 17 00:00:00 2001 From: Adam Merberg Date: Fri, 27 Jan 2023 16:42:01 -0500 Subject: [PATCH 1/9] support timezone when scheduling with Argo workflows --- metaflow/plugins/argo/argo_client.py | 3 ++- metaflow/plugins/argo/argo_workflows.py | 7 ++++++- metaflow/plugins/aws/step_functions/schedule_decorator.py | 4 +++- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/metaflow/plugins/argo/argo_client.py b/metaflow/plugins/argo/argo_client.py index bfe70cff5b9..0f64ce0b02a 100644 --- a/metaflow/plugins/argo/argo_client.py +++ b/metaflow/plugins/argo/argo_client.py @@ -116,7 +116,7 @@ def trigger_workflow_template(self, name, parameters={}): json.loads(e.body)["message"] if e.body is not None else e.reason ) - def schedule_workflow_template(self, name, schedule=None): + def schedule_workflow_template(self, name, schedule=None, timezone=None): # Unfortunately, Kubernetes client does not handle optimistic # concurrency control by itself unlike kubectl client = self._kubernetes_client.get() @@ -127,6 +127,7 @@ def schedule_workflow_template(self, name, schedule=None): "spec": { "suspend": schedule is None, "schedule": schedule, + "timezone": timezone, "workflowSpec": {"workflowTemplateRef": {"name": name}}, }, } diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index 071e097baf0..46fdf98a53d 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -120,6 +120,7 @@ def __init__( self.parameters = self._process_parameters() self._workflow_template = self._compile() self._cron = self._cron() + self._timezone = self._timezone() def __str__(self): return str(self._workflow_template) @@ -179,10 +180,14 @@ def _cron(self): return " ".join(schedule.schedule.split()[:5]) return None + def _timezone(self): + schedule = self.flow._flow_decorators.get("schedule") + return schedule.timezone if schedule else None + def schedule(self): try: ArgoClient(namespace=KUBERNETES_NAMESPACE).schedule_workflow_template( - self.name, self._cron + self.name, self._cron, self._timezone ) except Exception as e: raise ArgoWorkflowsSchedulingException(str(e)) diff --git a/metaflow/plugins/aws/step_functions/schedule_decorator.py b/metaflow/plugins/aws/step_functions/schedule_decorator.py index b3bc174bc2a..cfa5fb2865e 100644 --- a/metaflow/plugins/aws/step_functions/schedule_decorator.py +++ b/metaflow/plugins/aws/step_functions/schedule_decorator.py @@ -21,7 +21,7 @@ class ScheduleDecorator(FlowDecorator): """ name = "schedule" - defaults = {"cron": None, "weekly": False, "daily": True, "hourly": False} + defaults = {"cron": None, "weekly": False, "daily": True, "hourly": False, "timezone": None} def flow_init( self, flow, graph, environment, flow_datastore, metadata, logger, echo, options @@ -38,3 +38,5 @@ def flow_init( self.schedule = "0 0 * * ? *" else: self.schedule = None + + self.timezone = self.attributes["timezone"] From 3389d8472529af76c9130ba74bc99dc9597f4993 Mon Sep 17 00:00:00 2001 From: Adam Merberg Date: Fri, 27 Jan 2023 17:22:49 -0500 Subject: [PATCH 2/9] black formatting --- metaflow/plugins/aws/step_functions/schedule_decorator.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/metaflow/plugins/aws/step_functions/schedule_decorator.py b/metaflow/plugins/aws/step_functions/schedule_decorator.py index cfa5fb2865e..2e4af346246 100644 --- a/metaflow/plugins/aws/step_functions/schedule_decorator.py +++ b/metaflow/plugins/aws/step_functions/schedule_decorator.py @@ -21,7 +21,13 @@ class ScheduleDecorator(FlowDecorator): """ name = "schedule" - defaults = {"cron": None, "weekly": False, "daily": True, "hourly": False, "timezone": None} + defaults = { + "cron": None, + "weekly": False, + "daily": True, + "hourly": False, + "timezone": None, + } def flow_init( self, flow, graph, environment, flow_datastore, metadata, logger, echo, options From 5057a3e2bda0cd510d15a17651f957659531cc01 Mon Sep 17 00:00:00 2001 From: Adam Merberg Date: Fri, 27 Jan 2023 17:39:20 -0500 Subject: [PATCH 3/9] roll timezone into _cron --- metaflow/plugins/argo/argo_workflows.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index 46fdf98a53d..1730b493248 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -120,7 +120,6 @@ def __init__( self.parameters = self._process_parameters() self._workflow_template = self._compile() self._cron = self._cron() - self._timezone = self._timezone() def __str__(self): return str(self._workflow_template) @@ -177,17 +176,14 @@ def _cron(self): schedule = self.flow._flow_decorators.get("schedule") if schedule: # Remove the field "Year" if it exists - return " ".join(schedule.schedule.split()[:5]) + return " ".join(schedule.schedule.split()[:5]), schedule.timezone return None - def _timezone(self): - schedule = self.flow._flow_decorators.get("schedule") - return schedule.timezone if schedule else None - def schedule(self): try: + cron, timezone = self._cron ArgoClient(namespace=KUBERNETES_NAMESPACE).schedule_workflow_template( - self.name, self._cron, self._timezone + self.name, cron, timezone ) except Exception as e: raise ArgoWorkflowsSchedulingException(str(e)) From 228c1ebb1b9dae4e1f262eb3f609ae94b8e5cb55 Mon Sep 17 00:00:00 2001 From: Adam Merberg Date: Fri, 27 Jan 2023 18:03:51 -0500 Subject: [PATCH 4/9] document supported timezone format --- metaflow/plugins/aws/step_functions/schedule_decorator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/metaflow/plugins/aws/step_functions/schedule_decorator.py b/metaflow/plugins/aws/step_functions/schedule_decorator.py index 2e4af346246..8cd2b83b208 100644 --- a/metaflow/plugins/aws/step_functions/schedule_decorator.py +++ b/metaflow/plugins/aws/step_functions/schedule_decorator.py @@ -45,4 +45,5 @@ def flow_init( else: self.schedule = None + # Argo Workflows supports the IANA timezone standard, e.g. America/Los_Angeles self.timezone = self.attributes["timezone"] From 5e2171a93c49e469e37206f4a7318d5170bdcd82 Mon Sep 17 00:00:00 2001 From: Adam Merberg Date: Fri, 27 Jan 2023 18:06:28 -0500 Subject: [PATCH 5/9] fail if timezone is provided with AWS Step Functions --- metaflow/plugins/aws/step_functions/step_functions.py | 1 + 1 file changed, 1 insertion(+) diff --git a/metaflow/plugins/aws/step_functions/step_functions.py b/metaflow/plugins/aws/step_functions/step_functions.py index 342fbff6b2a..f95a5cca833 100644 --- a/metaflow/plugins/aws/step_functions/step_functions.py +++ b/metaflow/plugins/aws/step_functions/step_functions.py @@ -328,6 +328,7 @@ def _visit(node, workflow, exit_node=None): def _cron(self): schedule = self.flow._flow_decorators.get("schedule") + assert schedule.timezone is None, "Step Functions does not support scheduling with a timezone." if schedule: return schedule.schedule return None From 5cfc254390109fcdaa21b1182a6d32d178d1ffe7 Mon Sep 17 00:00:00 2001 From: Adam Merberg Date: Fri, 27 Jan 2023 18:08:00 -0500 Subject: [PATCH 6/9] move assert into conditional --- metaflow/plugins/aws/step_functions/step_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/metaflow/plugins/aws/step_functions/step_functions.py b/metaflow/plugins/aws/step_functions/step_functions.py index f95a5cca833..597a242bd12 100644 --- a/metaflow/plugins/aws/step_functions/step_functions.py +++ b/metaflow/plugins/aws/step_functions/step_functions.py @@ -328,8 +328,8 @@ def _visit(node, workflow, exit_node=None): def _cron(self): schedule = self.flow._flow_decorators.get("schedule") - assert schedule.timezone is None, "Step Functions does not support scheduling with a timezone." if schedule: + assert schedule.timezone is None, "Step Functions does not support scheduling with a timezone." return schedule.schedule return None From 615e30e74e0f2b8871ae61c4ffca3bef3bc9489d Mon Sep 17 00:00:00 2001 From: Adam Merberg Date: Fri, 27 Jan 2023 18:10:40 -0500 Subject: [PATCH 7/9] black formatting --- metaflow/plugins/aws/step_functions/step_functions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/metaflow/plugins/aws/step_functions/step_functions.py b/metaflow/plugins/aws/step_functions/step_functions.py index 597a242bd12..bc6f3dff186 100644 --- a/metaflow/plugins/aws/step_functions/step_functions.py +++ b/metaflow/plugins/aws/step_functions/step_functions.py @@ -329,7 +329,9 @@ def _visit(node, workflow, exit_node=None): def _cron(self): schedule = self.flow._flow_decorators.get("schedule") if schedule: - assert schedule.timezone is None, "Step Functions does not support scheduling with a timezone." + assert ( + schedule.timezone is None + ), "Step Functions does not support scheduling with a timezone." return schedule.schedule return None From 768d2d9e5b16350c80491a223e1089d7fe68b850 Mon Sep 17 00:00:00 2001 From: Adam Merberg Date: Fri, 27 Jan 2023 21:28:55 -0500 Subject: [PATCH 8/9] raise StepFunctionsException if timezone is provided with step functions --- metaflow/plugins/aws/step_functions/step_functions.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/metaflow/plugins/aws/step_functions/step_functions.py b/metaflow/plugins/aws/step_functions/step_functions.py index bc6f3dff186..8c9584ed706 100644 --- a/metaflow/plugins/aws/step_functions/step_functions.py +++ b/metaflow/plugins/aws/step_functions/step_functions.py @@ -329,9 +329,10 @@ def _visit(node, workflow, exit_node=None): def _cron(self): schedule = self.flow._flow_decorators.get("schedule") if schedule: - assert ( - schedule.timezone is None - ), "Step Functions does not support scheduling with a timezone." + if schedule.timezone is not None: + raise StepFunctionsException( + "Step Functions does not support scheduling with a timezone." + ) return schedule.schedule return None From f6c632022d94f11313453084264a6fb9ae16ef28 Mon Sep 17 00:00:00 2001 From: Adam Merberg Date: Mon, 30 Jan 2023 14:25:43 -0500 Subject: [PATCH 9/9] add timezone to schedule decorator docstring --- metaflow/plugins/aws/step_functions/schedule_decorator.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/metaflow/plugins/aws/step_functions/schedule_decorator.py b/metaflow/plugins/aws/step_functions/schedule_decorator.py index 8cd2b83b208..8fd1103115b 100644 --- a/metaflow/plugins/aws/step_functions/schedule_decorator.py +++ b/metaflow/plugins/aws/step_functions/schedule_decorator.py @@ -18,6 +18,9 @@ class ScheduleDecorator(FlowDecorator): cron : str Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions) specified by this expression. + timezone : str + Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows, + which accepts timezones in [IANA format](https://nodatime.org/TimeZones). """ name = "schedule"