Package: wandb-osh
Wandb Offline Sync HookA convenient way to trigger synchronizations to wandb if your compute nodes don't have internet!🤔 What is this?✅ You usewandb/Weights & Biasesto record your machine learning trials?✅ Your ML experiments run on compute nodes without internet access (for example, using a batch system)?✅ Your compute nodes and your head/login node (with internet) have access to a shared file system?Then this package might be useful. For alternatives, seebelow.What you might have been doing so farYou probably have been usingexport WANDB_MODE="offline"on the compute nodes and then ran something likecd/.../result_dir/fordin$(ls-t-d*/);docd$d;wandbsync--sync-all;cd..;donefrom your head node (with internet access) every now and then.
However, obviously this is not very satisfying as it doesn't update live.
Sure, you could throw this in awhile Trueloop, but if you have a lot of trials in your directory, this will take forever,cause unnecessary network trafficand it's just not very elegant.How doeswandb-oshsolve the problem?You add a hook that is called every time an epoch concludes (that is, when we want to trigger a sync).You start thewandb-oshscript in your head node with internet access. This script will now triggerwandb syncupon request from one of the compute nodes.How is this implemented?Very simple: Every time an epoch concludes, the hook gets called and creates a file in thecommunication directory(~/.wandb_osh_communicationby default).
Thewandb-oshscript that is running on the head node (with internet) reads these files and performs the synchronization.What alternatives are there?Withray tune, you can use your ray head node as the place to synchronize from (rather than deploying it via the batch system as well, as thecurrent docssuggest). See the note below or mydemo repository.
Similar strategies might be possible forwandbas well (let me know!).📦 Installationpip3installwandb-oshFor completeness, the extra dependencieslightningandrayare given, but they only ensure that the corresponding package is installed.
For examplepip3install'wandb-osh[lightning]'also installs pytorch lightning if it is not already present, but has no other effect.For development, make sure also to include thetestingextra requirement.pip3install--editable'.[testing]'🔥 Running it!Two steps: Set up the hook, then run the script from your head node.Step 1: Setting up the hookWith pure wandbLet's adapt thesimple pytorch examplefrom the wandb docs (it only takes 3 lines!):importwandbfromwandb_osh.hooksimportTriggerWandbSyncHook# <-- New!trigger_sync=TriggerWandbSyncHook()# <--- New!wandb.init(config=args,mode="offline")model=...# set up your model# Magicwandb.watch(model,log_freq=100)model.train()forbatch_idx,(data,target)inenumerate(train_loader):output=model(data)loss=F.nll_loss(output,target)loss.backward()optimizer.step()ifbatch_idx%args.log_interval==0:wandb.log({"loss":loss})trigger_sync()# <-- New!With pytorch lightningSimply add theTriggerWandbSyncLightningCallbackto your list of callbacks and you're good to go!fromwandb_osh.lightning_hooksimportTriggerWandbSyncLightningCallback# <-- New!frompytorch_lightning.loggersimportWandbLoggerfrompytorch_lightningimportTrainerlogger=WandbLogger(project="project",group="group",offline=True,)model=MyLightningModule()trainer=Trainer(logger=logger,callbacks=[TriggerWandbSyncLightningCallback()]# <-- New!)trainer.fit(model,train_dataloader,val_dataloader)With ray tuneNoteWith ray tune, you might not need this package! While the approach suggested in theray tune SLURM docsdeploys the ray head on a worker node as well (so it doesn't
have internet), this actually isn't needed. Instead, you can run the ray head and the
tuning script on the head node and only submit batch jobs for your workers.
In this way,wandbwill be called from the head node and internet access is no
problem there.
For more information on this approach, take a look at mydemo repository.You probably already use theWandbLoggerCallbackcallback. We simply add a second callback forwandb-osh(it only takes two new lines!):importosfromwandb_osh.ray_hooksimportTriggerWandbSyncRayHook# <-- New!os.environ["WANDB_MODE"]="offline"callbacks=[WandbLoggerCallback(...),# <-- ray tune documentation tells you about thisTriggerWandbSyncRayHook(),# <-- New!]tuner=tune.Tuner(trainable,tune_config=...,run_config=RunConfig(...,callbacks=callbacks,),)With anything elseSimply take theTriggerWandbSyncHookclass and use it as a callback in your training
loop (as in thewandbexample above), passing the directory thatwandbis syncing
to as an argument.Step 2: Running the script on the head nodeAfter installation, you should have awandb-oshscript in your$PATH. Simply call it like this:wandb-oshThe output will look something like thisINFO: Starting to watch /home/kl5675/.wandb_osh_command_dir
INFO: Syncing /home/kl5675/ray_results/tcn-perfect-test-sync/DynamicTCNTrainable_b1f60706_4_attr_pt_thld=0.0273,batch_size=1,focal_alpha=0.2500,focal_gamma=2.0000,gnn_tracking_experiments_has_2022-11-03_17-08-42
Find logs at: /home/kl5675/ray_results/tcn-perfect-test-sync/DynamicTCNTrainable_b1f60706_4_attr_pt_thld=0.0273,batch_size=1,focal_alpha=0.2500,focal_gamma=2.0000,gnn_tracking_experiments_has_2022-11-03_17-08-42/wandb/debug-cli.kl5675.log
Syncing: https://wandb.ai/gnn_tracking/gnn_tracking/runs/b1f60706 ... done.
INFO: Finished syncing /home/kl5675/ray_results/tcn-perfect-test-sync/DynamicTCNTrainable_b1f60706_4_attr_pt_thld=0.0273,batch_size=1,focal_alpha=0.2500,focal_gamma=2.0000,gnn_tracking_experiments_has_2022-11-03_17-08-42
INFO: Syncing /home/kl5675/ray_results/tcn-perfect-test-sync/DynamicTCNTrainable_92a3ef1b_1_attr_pt_thld=0.0225,batch_size=1,focal_alpha=0.2500,focal_gamma=2.0000,gnn_tracking_experiments_has_2022-11-03_17-07-49
Find logs at: /home/kl5675/ray_results/tcn-perfect-test-sync/DynamicTCNTrainable_92a3ef1b_1_attr_pt_thld=0.0225,batch_size=1,focal_alpha=0.2500,focal_gamma=2.0000,gnn_tracking_experiments_has_2022-11-03_17-07-49/wandb/debug-cli.kl5675.log
Syncing: https://wandb.ai/gnn_tracking/gnn_tracking/runs/92a3ef1b ... done.
INFO: Finished syncing /home/kl5675/ray_results/tcn-perfect-test-sync/DynamicTCNTrainable_92a3ef1b_1_attr_pt_thld=0.0225,batch_size=1,focal_alpha=0.2500,focal_gamma=2.0000,gnn_tracking_experiments_has_2022-11-03_17-07-49
INFO: Syncing /home/kl5675/ray_results/tcn-perfect-test-sync/DynamicTCNTrainable_a2caa9c0_2_attr_pt_thld=0.0092,batch_size=1,focal_alpha=0.2500,focal_gamma=2.0000,gnn_tracking_experiments_has_2022-11-03_17-08-17
Find logs at: /home/kl5675/ray_results/tcn-perfect-test-sync/DynamicTCNTrainable_a2caa9c0_2_attr_pt_thld=0.0092,batch_size=1,focal_alpha=0.2500,focal_gamma=2.0000,gnn_tracking_experiments_has_2022-11-03_17-08-17/wandb/debug-cli.kl5675.log
Syncing: https://wandb.ai/gnn_tracking/gnn_tracking/runs/a2caa9c0 ... done.Take a look atwandb-osh --helpor checkthe documentationfor all command line options.
You can add options to thewandb synccall by placing them after--. For examplewandb-osh----sync-all❓ Q & AI get the warning "wandb: NOTE: use wandb sync --sync-all to sync 1 unsynced runs from local directory."You can startwandb-oshwithwandb-osh -- --sync-allto always synchronize
all available runs.How can I suppress logging messages (e.g., warnings about the syncing not being fast enough)importwandb_osh# for wandb_osh.__version__ >= 1.2.0wandb_osh.set_log_level("ERROR")🧰 Development setuppip3installpre-commit
pre-commitinstall💖 ContributingYour help is greatly appreciated! Suggestions, bug reports and feature requests are best opened asgithub issues. You are also very welcome to submit apull request!Bug reports and pull requests are credited with the help of theallcontributors bot.Barthelemy Meynard-Piganeau🐛MoH-assan🐛