import pathlib
from time import sleep

import torch

import wandb

num_bytes = 1 * 1024 * 1024
num_floats = num_bytes // 4
num_steps = 50
save_interval = 5
num_files = 0

with wandb.init(project="file-sync-test") as run:
    for i in range(num_steps):
        if (i + 1) % save_interval == 0:
            num_files += 1
            print(pathlib.Path(__file__).parent / f"data_{i + 1}.save")
            torch.save(
                torch.rand(num_floats, dtype=torch.float32),
                f"{wandb.run.dir}/data_{i + 1}.save",
                # pathlib.Path(__file__).parent / f"data_{i + 1}.save",
            )
            print(f"{i + 1}/{num_steps} -> Saved file")
        else:
            print(f"{i + 1}/{num_steps}")
        run.save("*.save")  # sync any new files of interest
        sleep(2.0)


print(f"Total traffic should be: ~{num_files * num_bytes / 1024 / 1024} MB")
print(f"Total saved files: {num_files}")