the-algorithm-ml/reader/dds.py

181 lines
6.1 KiB
Python
Raw Normal View History

"""Dataset service orchestrated by a TFJob
"""
from typing import Optional
import uuid
from tml.ml_logging.torch_logging import logging
import tml.machines.environment as env
import packaging.version
import tensorflow as tf
try:
import tensorflow_io as tfio
except:
pass
from tensorflow.python.data.experimental.ops.data_service_ops import (
_from_dataset_id,
_register_dataset,
)
import torch.distributed as dist
def maybe_start_dataset_service():
2023-09-13 07:52:13 +02:00
"""
Start the dataset service if readers are available and required dependencies are met.
This function checks if readers are available and if the required TensorFlow version is >= 2.5.
If both conditions are met and the current environment is the dispatcher or reader, it starts
the TensorFlow dataset service.
Raises:
Exception: If the required TensorFlow version is not met (>= 2.5).
"""
if not env.has_readers():
return
if packaging.version.parse(tf.__version__) < packaging.version.parse("2.5"):
raise Exception(f"maybe_distribute_dataset requires TF >= 2.5; got {tf.__version__}")
if env.is_dispatcher():
logging.info(f"env.get_reader_port() = {env.get_reader_port()}")
logging.info(f"env.get_dds_journaling_dir() = {env.get_dds_journaling_dir()}")
work_dir = env.get_dds_journaling_dir()
server = tf.data.experimental.service.DispatchServer(
tf.data.experimental.service.DispatcherConfig(
port=env.get_reader_port(),
protocol="grpc",
work_dir=work_dir,
fault_tolerant_mode=bool(work_dir),
)
)
server.join()
elif env.is_reader():
logging.info(f"env.get_reader_port() = {env.get_reader_port()}")
logging.info(f"env.get_dds_dispatcher_address() = {env.get_dds_dispatcher_address()}")
logging.info(f"env.get_dds_worker_address() = {env.get_dds_worker_address()}")
server = tf.data.experimental.service.WorkerServer(
tf.data.experimental.service.WorkerConfig(
port=env.get_reader_port(),
dispatcher_address=env.get_dds_dispatcher_address(),
worker_address=env.get_dds_worker_address(),
protocol="grpc",
)
)
server.join()
def register_dataset(
dataset: tf.data.Dataset, dataset_service: str, compression: Optional[str] = "AUTO"
):
2023-09-13 07:52:13 +02:00
"""
Register a dataset with the distributed dataset service.
This function registers a dataset with the distributed dataset service and broadcasts the dataset ID
and job name to all processes in the distributed environment.
Args:
dataset (tf.data.Dataset): The dataset to be registered.
dataset_service (str): The name of the dataset service.
compression (Optional[str]): The compression type for the dataset (default is "AUTO").
Returns:
Tuple[int, str]: A tuple containing the dataset ID and job name.
Note:
This function should be called on the rank 0 process.
"""
if dist.get_rank() == 0:
dataset_id = _register_dataset(
service=dataset_service,
dataset=dataset,
compression=compression,
)
job_name = uuid.uuid4().hex[:8]
id_and_job = [dataset_id.numpy(), job_name]
logging.info(f"rank{dist.get_rank()}: Created dds job with {dataset_id.numpy()}, {job_name}")
else:
id_and_job = [None, None]
dist.broadcast_object_list(id_and_job, src=0)
return tuple(id_and_job)
def distribute_from_dataset_id(
dataset_service: str,
dataset_id: int,
job_name: Optional[str],
compression: Optional[str] = "AUTO",
prefetch: Optional[int] = tf.data.experimental.AUTOTUNE,
) -> tf.data.Dataset:
2023-09-13 07:52:13 +02:00
"""
Distribute a dataset from a registered dataset ID.
This function consumes a dataset from the distributed dataset service using the provided dataset ID
and job name. It also supports prefetching for improved performance.
Args:
dataset_service (str): The name of the dataset service.
dataset_id (int): The ID of the dataset to be consumed.
job_name (Optional[str]): The name of the job associated with the dataset (optional).
compression (Optional[str]): The compression type for the dataset (default is "AUTO").
prefetch (Optional[int]): The number of elements to prefetch (default is tf.data.experimental.AUTOTUNE).
Returns:
tf.data.Dataset: The distributed dataset.
"""
logging.info(f"rank{dist.get_rank()}: Consuming dds job with {dataset_id}, {job_name}")
dataset = _from_dataset_id(
processing_mode="parallel_epochs",
service=dataset_service,
dataset_id=dataset_id,
job_name=job_name,
element_spec=None,
compression=compression,
)
if prefetch is not None:
dataset = dataset.prefetch(prefetch)
return dataset
def maybe_distribute_dataset(dataset: tf.data.Dataset) -> tf.data.Dataset:
2023-09-13 07:52:13 +02:00
"""
Distribute a TensorFlow dataset for Torch-compatible and distributed training-aware consumption.
2023-09-13 07:52:13 +02:00
This function is used to distribute a dataset in a distributed training environment. It performs the
following steps:
- On the rank 0 process, it registers the given dataset with the distributed dataset service.
- It broadcasts the job name and dataset ID to all rank processes.
- All rank processes then consume the same dataset from the distributed dataset service.
2023-09-13 07:52:13 +02:00
Args:
dataset (tf.data.Dataset): The TensorFlow dataset to be distributed.
2023-09-13 07:52:13 +02:00
Returns:
tf.data.Dataset: The distributed TensorFlow dataset.
Note:
- If there are no reader processes in the distributed environment, the original dataset is returned
without any distribution.
- This function is intended for use in distributed training environments to prevent out-of-memory (OOM)
issues caused by each rank process trying to serve one job.
"""
if not env.has_readers():
return dataset
dataset_service = env.get_dds()
logging.info(f"using DDS = {dataset_service}")
dataset_id, job_name = register_dataset(dataset=dataset, dataset_service=dataset_service)
dataset = distribute_from_dataset_id(
dataset_service=dataset_service, dataset_id=dataset_id, job_name=job_name
)
return dataset
if __name__ == "__main__":
maybe_start_dataset_service()