| 409 | self._pool.shutdown() |
| 410 | |
| 411 | class _RayPreprocessMultiProcessingDataLoaderIter(_BaseDataLoaderIter): |
| 412 | |
| 413 | def __init__(self, loader, preprocess: Preprocess=None, queue_max_size=10): |
| 414 | self._preprocess = preprocess |
| 415 | self._stopped = False |
| 416 | self._send_idx = 0 # idx of the next task to be sent to workers |
| 417 | self._rcvd_idx = 0 # idx of the next task to be returned in __next__ |
| 418 | |
| 419 | super(_RayPreprocessMultiProcessingDataLoaderIter, self).__init__(loader) |
| 420 | |
| 421 | assert self._num_workers > 0 |
| 422 | assert self._prefetch_factor > 0 |
| 423 | |
| 424 | |
| 425 | |
| 426 | # Adds forward compatibilities so classic DataLoader can work with DataPipes: |
| 427 | # Taking care of distributed sharding |
| 428 | # if isinstance(self._dataset, (IterDataPipe, MapDataPipe)): |
| 429 | # torch.utils.data.graph_settings.apply_sharding(self._dataset, self._world_size, self._rank) |
| 430 | |
| 431 | if self._preprocess is not None and self._preprocess.has_cpu_preprocess(): |
| 432 | org_create_fetcher = _DatasetKind.create_fetcher |
| 433 | _DatasetKind.create_fetcher = _create_fetcher_proxy(org_create_fetcher, self._preprocess) |
| 434 | self._dataset_fetcher = _DatasetKind.create_fetcher( |
| 435 | self._dataset_kind, self._dataset, self._auto_collation, self._collate_fn, self._drop_last) |
| 436 | if self._preprocess is not None and self._preprocess.has_cpu_preprocess(): |
| 437 | _DatasetKind.create_fetcher = org_create_fetcher |
| 438 | # if self._preprocess is not None and self._preprocess.has_gpu_preprocess(): |
| 439 | self._queue_max_size = queue_max_size |
| 440 | self._stream = torch.cuda.Stream() |
| 441 | self._data_queue = queue.Queue(maxsize=self._queue_max_size) |
| 442 | self._device_id = torch.cuda.current_device() |
| 443 | self._timeout = self._timeout if self._timeout > 0 else (12 * _utils.MP_STATUS_CHECK_INTERVAL) |
| 444 | print("init _RayPreprocessMultiProcessingDataLoaderIter") |
| 445 | |
| 446 | |
| 447 | |
| 448 | def _start(self): |
| 449 | print("begin to start ray actor") |
| 450 | |
| 451 | if not ray.is_initialized(): |
| 452 | ray.init(address="auto") |
| 453 | |
| 454 | node_id = ray.get_runtime_context().get_node_id() |
| 455 | scheduling_strategy = NodeAffinitySchedulingStrategy(node_id=node_id, soft=False) |
| 456 | actor_options = {"num_cpus": 1, "scheduling_strategy": scheduling_strategy} |
| 457 | # worker_actor_options = {"num_cpus": 1, "scheduling_strategy": scheduling_strategy} |
| 458 | nodes = ray.nodes() |
| 459 | node_cpu = 0 |
| 460 | node_gpu = 0 |
| 461 | for node in nodes: |
| 462 | if node['NodeID'] == node_id: |
| 463 | node_cpu = int(node['Resources']['CPU']) |
| 464 | node_gpu = int(node['Resources']['GPU']) |
| 465 | break |
| 466 | |
| 467 | num_queue_actors = node_cpu // int(max(1, node_gpu)) |
| 468 | print(f"node_id = {node_id}, node_cpu = {node_cpu}, node_gpu = {node_gpu}, num_queue_actors = {num_queue_actors}") |