import sys
import time
import traceback
from functools import wraps

from metaflow import FlowSpec, step, retry
from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY

def platform_retry(f):
    @wraps(f)
    def wrapper(self):
        try:
            f(self)
        except:
            traceback.print_exc()
            sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
    return retry(wrapper)

class PlatformRetryFlow(FlowSpec):

    @platform_retry
    @step
    def start(self):
        time.sleep(10)
        print('fail', 1 / 0)
        self.next(self.end)

    @platform_retry
    @step
    def end(self):
        print("done!")

if __name__ == '__main__':
    PlatformRetryFlow()