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..1730b493248 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -176,13 +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 schedule(self): try: + cron, timezone = self._cron ArgoClient(namespace=KUBERNETES_NAMESPACE).schedule_workflow_template( - self.name, self._cron + self.name, cron, 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..8fd1103115b 100644 --- a/metaflow/plugins/aws/step_functions/schedule_decorator.py +++ b/metaflow/plugins/aws/step_functions/schedule_decorator.py @@ -18,10 +18,19 @@ 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" - 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 +47,6 @@ def flow_init( self.schedule = "0 0 * * ? *" else: self.schedule = None + + # Argo Workflows supports the IANA timezone standard, e.g. America/Los_Angeles + self.timezone = self.attributes["timezone"] diff --git a/metaflow/plugins/aws/step_functions/step_functions.py b/metaflow/plugins/aws/step_functions/step_functions.py index 342fbff6b2a..8c9584ed706 100644 --- a/metaflow/plugins/aws/step_functions/step_functions.py +++ b/metaflow/plugins/aws/step_functions/step_functions.py @@ -329,6 +329,10 @@ def _visit(node, workflow, exit_node=None): def _cron(self): schedule = self.flow._flow_decorators.get("schedule") if schedule: + if schedule.timezone is not None: + raise StepFunctionsException( + "Step Functions does not support scheduling with a timezone." + ) return schedule.schedule return None