wandb_resume.py

import time
import os
import pickle
import wandb

class Model:
    def __init__(self):
        self.var = 0

    def step(self):
        time.sleep(1)
        self.var += 1

wandb.init(project="test", resume="test13")

_ = wandb.run.history

# Attempt to restore model
if wandb.run.resumed:
    print("restoring model")
    wandb.restore("model.pickle")
    path = os.path.join(wandb.run.dir, "model.pickle")
    with open(path) as f:
        model = pickle.load(f)
    i = wandb.run.step
else:
    print("no existing run found")
    model = Model()
    i = 0


# "Train" the model
while True:
    print(f"i is {i}")
    model.step()
    print(f"model var is {model.var}")
    wandb.log({"var": model.var}, step=i)

    if i % 10 == 0:
        # Save model to file
        path = os.path.join(wandb.run.dir, "model.pickle")
        with open(path, 'wb') as f:
            pickle.dump(model, f)
            wandb.save("model.pickle")

    i += 1