| 167 | |
| 168 | |
| 169 | def _pick_tqdm_interval(file): |
| 170 | # Heuristics to pick a update interval for progress bar that's nice-looking for users. |
| 171 | isatty = file.isatty() |
| 172 | # Jupyter notebook should be recognized as tty. |
| 173 | # Wait for https://github.com/ipython/ipykernel/issues/268 |
| 174 | try: |
| 175 | from ipykernel import iostream |
| 176 | if isinstance(file, iostream.OutStream): |
| 177 | isatty = True |
| 178 | except ImportError: |
| 179 | pass |
| 180 | |
| 181 | if isatty: |
| 182 | return 0.5 |
| 183 | else: |
| 184 | # When run under mpirun/slurm, isatty is always False. |
| 185 | # Here we apply some hacky heuristics for slurm. |
| 186 | if 'SLURM_JOB_ID' in os.environ: |
| 187 | if int(os.environ.get('SLURM_JOB_NUM_NODES', 1)) > 1: |
| 188 | # multi-machine job, probably not interactive |
| 189 | return 60 |
| 190 | else: |
| 191 | # possibly interactive, so let's be conservative |
| 192 | return 15 |
| 193 | |
| 194 | if 'OMPI_COMM_WORLD_SIZE' in os.environ: |
| 195 | return 60 |
| 196 | |
| 197 | # If not a tty, don't refresh progress bar that often |
| 198 | return 180 |
| 199 | |
| 200 | |
| 201 | def get_tqdm_kwargs(**kwargs): |