I m looking for a way to implement retry functionality for t Outerbounds #ask-metaflow

I’m looking for a way to implement retry functiona...

elegant-beach-10818

04/05/2023, 7:01 PM

I’m looking for a way to implement retry functionality for training a model because we’re leveraging AWS spot instances for 6+ hr training sessions. Suppose there are 10 epochs, I’d like to save the model after each epoch. Then if there’s a failure, I’d like to load the latest trained model so I can resume training. I’ve attempted the following, but unfortunately I discovered that artifacts are not saved in the event of a failure:

Copy code

class ResumeTraining(FlowSpec):
    @step
    def start(self):
        self.finished_epoch = 0
        self.next(self.looping_training)
        
    @retry(times=1)
    @step
    def looping_training(self):
        import time
        print(f"finished_epoch: {self.finished_epoch}")
        if self.finished_epoch > 0:
            start =  self.finished_epoch + 1
        else:
            start = 1
        
        for epoch in range(start, 10):
            if epoch == 5:
                self.finished_epoch = epoch
                raise Exception('fake failure')

            time.sleep(1)
            
            self.model = f"model iteration {epoch}"
            print(self.model)
            self.finished_epoch = epoch
            print(f"in loop finished_epoch: {self.finished_epoch}")
        
        self.next(self.end)

✅ 1

square-wire-39606

04/05/2023, 7:21 PM

You can use

metaflow.s3

to explicitly save the model after every epoch and resume from the stored model at the beginning of the step. It's a good idea to support this pattern out of the box soon.

➕ 1

elegant-beach-10818

04/05/2023, 9:14 PM

That’s what I needed! Thanks!

elegant-beach-10818

04/05/2023, 9:42 PM

For posterity should this come up in a search:

Copy code

from metaflow import FlowSpec, step, retry, S3

class ResumeTraining(FlowSpec):
    @step
    def start(self):
        self.next(self.looping_training)
        
    @retry(times=1)
    @step
    def looping_training(self):
        import time
        s3 = S3(run=self)

        if <http://s3.info|s3.info>("model", return_missing=True).exists == False:
            print("model is none, loading model")
        else:
            print("loading model")

        if <http://s3.info|s3.info>("finished_epoch", return_missing=True).exists == False:
            start = 1
        else:
            start = int(s3.get('finished_epoch').text) + 1
        
        for epoch in range(start, 10):
            if epoch == 5:
                s3.put('finished_epoch', str(epoch))
                raise Exception('fake failure')

            time.sleep(1)
            s3.put("model", f"model_{epoch}")
            s3.put("finished_epoch", str(epoch))
            
            print(s3.get("finished_epoch").text)
        
        self.next(self.end)

    @step
    def end(self):
        print('finished')

🙌 2

2 Views

Open in Slack

Previous Next