wehub-resource-sync
71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, List
|
|
|
|
import ray
|
|
from ray._common.network_utils import build_address
|
|
from ray.train._internal.base_worker_group import BaseWorkerGroup
|
|
from ray.train._internal.utils import get_address_and_port
|
|
from ray.train.backend import Backend, BackendConfig
|
|
from ray.train.v2._internal.util import TrainingFramework
|
|
from ray.util import PublicAPI
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@PublicAPI(stability="beta")
|
|
@dataclass
|
|
class TensorflowConfig(BackendConfig):
|
|
@property
|
|
def backend_cls(self):
|
|
return _TensorflowBackend
|
|
|
|
@property
|
|
def framework(self):
|
|
return TrainingFramework.TENSORFLOW
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {}
|
|
|
|
|
|
def _setup_tensorflow_environment(worker_addresses: List[str], index: int):
|
|
"""Set up distributed Tensorflow training information.
|
|
|
|
This function should be called on each worker.
|
|
|
|
Args:
|
|
worker_addresses: Addresses of all the workers.
|
|
index: Index (i.e. world rank) of the current worker.
|
|
"""
|
|
tf_config = {
|
|
"cluster": {"worker": worker_addresses},
|
|
"task": {"type": "worker", "index": index},
|
|
}
|
|
os.environ["TF_CONFIG"] = json.dumps(tf_config)
|
|
os.environ["TF_USE_LEGACY_KERAS"] = "1"
|
|
|
|
|
|
class _TensorflowBackend(Backend):
|
|
def on_start(self, worker_group: BaseWorkerGroup, backend_config: TensorflowConfig):
|
|
# Compute URL for initializing distributed setup.
|
|
def get_url():
|
|
address, port = get_address_and_port()
|
|
return build_address(address, port)
|
|
|
|
urls = worker_group.execute(get_url)
|
|
|
|
# Get setup tasks in order to throw errors on failure.
|
|
setup_futures = []
|
|
for i in range(len(worker_group)):
|
|
setup_futures.append(
|
|
worker_group.execute_single_async(
|
|
i,
|
|
_setup_tensorflow_environment,
|
|
worker_addresses=urls,
|
|
index=i,
|
|
)
|
|
)
|
|
ray.get(setup_futures)
|