A local work manager for inference job.
| 452 | |
| 453 | |
| 454 | class LocalWorkMgr(object): |
| 455 | """A local work manager for inference job.""" |
| 456 | def __init__(self, job_name, task_index, restore_works_dir, name=None): |
| 457 | """Constructs a LocalWorkMgr. |
| 458 | |
| 459 | Args: |
| 460 | job_name: name of current tf-worker. |
| 461 | task_index: index of current tf-worker. |
| 462 | restore_works_dir: a directory that restore works for |
| 463 | WorkQueue when failover. |
| 464 | name: (Optional.) Name of the LocalWorkMgr. |
| 465 | |
| 466 | Raises: |
| 467 | ValueError: If one of the arguments is invalid. |
| 468 | """ |
| 469 | |
| 470 | self._job_name = job_name |
| 471 | self._task_index = task_index |
| 472 | self._restore_works_dir = restore_works_dir |
| 473 | self._name = name or 'local_workqueue_mgr' |
| 474 | assert job_name in ['worker', 'chief'], \ |
| 475 | 'support chief or worker role, not {}'.format(job_name) |
| 476 | assert task_index >= 0, 'task_index must be >= 0, not {}'.format(task_index) |
| 477 | assert gfile.IsDirectory(restore_works_dir), \ |
| 478 | '{} is not a directory.'.format(restore_works_dir) |
| 479 | self._get_local_works() |
| 480 | self._local_device = control_flow_ops.no_op().device |
| 481 | |
| 482 | @property |
| 483 | def job_name(self): |
| 484 | return self._job_name |
| 485 | |
| 486 | @property |
| 487 | def task_index(self): |
| 488 | return self._task_index |
| 489 | |
| 490 | @property |
| 491 | def restore_works_dir(self): |
| 492 | return self._restore_works_dir |
| 493 | |
| 494 | def _get_local_works(self): |
| 495 | """Get the local works that needs to be restored.""" |
| 496 | self._restore_works = [] |
| 497 | restore_work_file_dir = os.path.join( |
| 498 | self._restore_works_dir, |
| 499 | '{}_{}'.format(self._job_name, self._task_index)) |
| 500 | if gfile.IsDirectory(restore_work_file_dir): |
| 501 | restore_work_files = gfile.ListDirectory(restore_work_file_dir) |
| 502 | for work_file in restore_work_files: |
| 503 | work_file_path = os.path.join(restore_work_file_dir, work_file) |
| 504 | with gfile.GFile(work_file_path, 'r') as rfile: |
| 505 | for line in rfile.readlines(): |
| 506 | line = line.strip() |
| 507 | if not re.match(r'(.)*(?:\?start=\d+&end=\d+)$', line): |
| 508 | logging.error('Invalid format: {}'.format(line)) |
| 509 | continue |
| 510 | self._restore_works.append(line) |
| 511 | self._restore_works.sort( |