support timezone when scheduling with Argo workflows

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}},
             },
         }

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))

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"]

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