(neox_args)
| 149 | |
| 150 | |
| 151 | def init_wandb(neox_args): |
| 152 | # Wandb. (one worker per machine) |
| 153 | if neox_args.use_wandb == False: |
| 154 | return |
| 155 | |
| 156 | if not neox_args.wandb_init_all_ranks: |
| 157 | use_wandb = is_local_main() and ( |
| 158 | get_wandb_api_key(neox_args=neox_args) is not None |
| 159 | ) |
| 160 | neox_args.update_value("use_wandb", use_wandb) |
| 161 | if neox_args.use_wandb: |
| 162 | group_name = neox_args.wandb_group |
| 163 | name = f"{socket.gethostname()}-{local_rank()}" if group_name else None |
| 164 | try: |
| 165 | wandb.init( |
| 166 | project=neox_args.wandb_project, |
| 167 | group=group_name, |
| 168 | name=name, |
| 169 | save_code=False, |
| 170 | force=False, |
| 171 | entity=neox_args.wandb_team, |
| 172 | ) |
| 173 | except wandb.UsageError as e: |
| 174 | neox_args.update_value("use_wandb", False) |
| 175 | print(e) |
| 176 | print( |
| 177 | "Skipping wandb. Execute `wandb login` on local or main node machine to enable.", |
| 178 | flush=True, |
| 179 | ) |
| 180 | wandb.config.update(neox_args.all_config) |
| 181 | |
| 182 | |
| 183 | def obtain_resource_pool( |
no test coverage detected