| 292 | """ |
| 293 | |
| 294 | def _start(self, |
| 295 | dist_start_cmd, |
| 296 | func, |
| 297 | num_gpus, |
| 298 | assert_callback=None, |
| 299 | save_all_ranks=False, |
| 300 | *args, |
| 301 | **kwargs): |
| 302 | script_path = func.__code__.co_filename |
| 303 | script_dir, script_name = os.path.split(script_path) |
| 304 | script_name = os.path.splitext(script_name)[0] |
| 305 | func_name = func.__qualname__ |
| 306 | |
| 307 | func_params = [] |
| 308 | for arg in args: |
| 309 | if isinstance(arg, str): |
| 310 | arg = ('\'{}\''.format(arg)) |
| 311 | func_params.append(str(arg)) |
| 312 | |
| 313 | for k, v in kwargs.items(): |
| 314 | if isinstance(v, str): |
| 315 | v = ('\'{}\''.format(v)) |
| 316 | func_params.append('{}={}'.format(k, v)) |
| 317 | |
| 318 | func_params = ','.join(func_params).strip(',') |
| 319 | |
| 320 | tmp_run_file = tempfile.NamedTemporaryFile(suffix='.py').name |
| 321 | tmp_res_file = tempfile.NamedTemporaryFile(suffix='.pkl').name |
| 322 | |
| 323 | with open(tmp_run_file, 'w') as f: |
| 324 | print('save temporary run file to : {}'.format(tmp_run_file)) |
| 325 | print('save results to : {}'.format(tmp_res_file)) |
| 326 | run_file_content = _DIST_SCRIPT_TEMPLATE.format( |
| 327 | script_name, script_name, func_name, func_params) |
| 328 | f.write(run_file_content) |
| 329 | |
| 330 | tmp_res_files = [] |
| 331 | if save_all_ranks: |
| 332 | for i in range(num_gpus): |
| 333 | tmp_res_files.append(tmp_res_file + str(i)) |
| 334 | else: |
| 335 | tmp_res_files = [tmp_res_file] |
| 336 | self.addCleanup(self.clean_tmp, [tmp_run_file] + tmp_res_files) |
| 337 | |
| 338 | tmp_env = copy.deepcopy(os.environ) |
| 339 | tmp_env['PYTHONPATH'] = ':'.join( |
| 340 | (tmp_env.get('PYTHONPATH', ''), script_dir)).lstrip(':') |
| 341 | # avoid distributed test hang |
| 342 | tmp_env['NCCL_P2P_DISABLE'] = '1' |
| 343 | script_params = '--save_all_ranks=%s --save_file=%s' % (save_all_ranks, |
| 344 | tmp_res_file) |
| 345 | script_cmd = '%s %s %s' % (dist_start_cmd, tmp_run_file, script_params) |
| 346 | print('script command: %s' % script_cmd) |
| 347 | res = subprocess.call(script_cmd, shell=True, env=tmp_env) |
| 348 | |
| 349 | script_res = [] |
| 350 | for res_file in tmp_res_files: |
| 351 | with open(res_file, 'rb') as f: |