from metaflow import step, FlowSpec, catch, retry, timeout
from time import sleep

class SFNTimeoutFlow(FlowSpec):
    @step
    def start(self):
        print("Starting 👋")
        self.next(self.schedule_workers)

    @step
    def schedule_workers(self):
        self.worker_inputs = list(range(1,40))
        self.next(self.run_worker, foreach='worker_inputs')

    @catch(var='timeout')
    @retry(times=0)
    @timeout(seconds=60)
    @step
    def run_worker(self):
        if int(self.input) % 2 == 0:
            print("failing due to timeout")
            sleep(360)
        self.next(self.join)
    
    @step
    def join(self, inputs):
        self.next(self.end)


    @step
    def end(self):
        print("Done! 🏁")


if __name__ == "__main__":
    SFNTimeoutFlow()