chore: import upstream snapshot with attribution

python/ray/train/tensorflow/__init__.py

@@ -0,0 +1,30 @@
+# isort: off
+try:
+    import tensorflow as tf  # noqa: F401
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "TensorFlow isn't installed. To install TensorFlow, run 'pip install "
+        "tensorflow'."
+    )
+# isort: on
+
+from ray.train.tensorflow.config import TensorflowConfig
+from ray.train.tensorflow.tensorflow_checkpoint import TensorflowCheckpoint
+from ray.train.tensorflow.tensorflow_trainer import TensorflowTrainer
+from ray.train.tensorflow.train_loop_utils import prepare_dataset_shard
+from ray.train.v2._internal.constants import is_v2_enabled
+
+if is_v2_enabled():
+    from ray.train.v2.tensorflow.tensorflow_trainer import (  # noqa: F811
+        TensorflowTrainer,
+    )
+
+__all__ = [
+    "TensorflowCheckpoint",
+    "TensorflowConfig",
+    "prepare_dataset_shard",
+    "TensorflowTrainer",
+]
+
+
+# DO NOT ADD ANYTHING AFTER THIS LINE.

python/ray/train/tensorflow/config.py

@@ -0,0 +1,70 @@
+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)

python/ray/train/tensorflow/keras.py

@@ -0,0 +1,210 @@
+import shutil
+from abc import abstractmethod
+from typing import Dict, List, Optional, Union
+
+from tensorflow.keras.callbacks import Callback as KerasCallback
+
+import ray
+from ray.train.tensorflow import TensorflowCheckpoint
+from ray.util.annotations import PublicAPI
+
+
+class _Callback(KerasCallback):
+    """Base class for Ray Train's Keras callbacks."""
+
+    _allowed = [
+        "epoch_begin",
+        "epoch_end",
+        "train_batch_begin",
+        "train_batch_end",
+        "test_batch_begin",
+        "test_batch_end",
+        "predict_batch_begin",
+        "predict_batch_end",
+        "train_begin",
+        "train_end",
+        "test_begin",
+        "test_end",
+        "predict_begin",
+        "predict_end",
+    ]
+
+    def __init__(self, on: Union[str, List[str]] = "validation_end"):
+        super(_Callback, self).__init__()
+
+        if not isinstance(on, list):
+            on = [on]
+        if any(w not in self._allowed for w in on):
+            raise ValueError(
+                "Invalid trigger time selected: {}. Must be one of {}".format(
+                    on, self._allowed
+                )
+            )
+        self._on = on
+
+    def _handle(self, logs: Dict, when: str):
+        raise NotImplementedError
+
+    def on_epoch_begin(self, epoch, logs=None):
+        if "epoch_begin" in self._on:
+            self._handle(logs, "epoch_begin")
+
+    def on_epoch_end(self, epoch, logs=None):
+        if "epoch_end" in self._on:
+            self._handle(logs, "epoch_end")
+
+    def on_train_batch_begin(self, batch, logs=None):
+        if "train_batch_begin" in self._on:
+            self._handle(logs, "train_batch_begin")
+
+    def on_train_batch_end(self, batch, logs=None):
+        if "train_batch_end" in self._on:
+            self._handle(logs, "train_batch_end")
+
+    def on_test_batch_begin(self, batch, logs=None):
+        if "test_batch_begin" in self._on:
+            self._handle(logs, "test_batch_begin")
+
+    def on_test_batch_end(self, batch, logs=None):
+        if "test_batch_end" in self._on:
+            self._handle(logs, "test_batch_end")
+
+    def on_predict_batch_begin(self, batch, logs=None):
+        if "predict_batch_begin" in self._on:
+            self._handle(logs, "predict_batch_begin")
+
+    def on_predict_batch_end(self, batch, logs=None):
+        if "predict_batch_end" in self._on:
+            self._handle(logs, "predict_batch_end")
+
+    def on_train_begin(self, logs=None):
+        if "train_begin" in self._on:
+            self._handle(logs, "train_begin")
+
+    def on_train_end(self, logs=None):
+        if "train_end" in self._on:
+            self._handle(logs, "train_end")
+
+    def on_test_begin(self, logs=None):
+        if "test_begin" in self._on:
+            self._handle(logs, "test_begin")
+
+    def on_test_end(self, logs=None):
+        if "test_end" in self._on:
+            self._handle(logs, "test_end")
+
+    def on_predict_begin(self, logs=None):
+        if "predict_begin" in self._on:
+            self._handle(logs, "predict_begin")
+
+    def on_predict_end(self, logs=None):
+        if "predict_end" in self._on:
+            self._handle(logs, "predict_end")
+
+
+class RayReportCallback(_Callback):
+    def __init__(
+        self,
+        checkpoint_on: Union[str, List[str]] = "epoch_end",
+        report_metrics_on: Union[str, List[str]] = "epoch_end",
+        metrics: Optional[Union[str, List[str], Dict[str, str]]] = None,
+    ):
+        if isinstance(checkpoint_on, str):
+            checkpoint_on = [checkpoint_on]
+        if isinstance(report_metrics_on, str):
+            report_metrics_on = [report_metrics_on]
+
+        on = list(set(checkpoint_on + report_metrics_on))
+        super().__init__(on=on)
+
+        self._checkpoint_on: List[str] = checkpoint_on
+        self._report_metrics_on: List[str] = report_metrics_on
+        self._metrics = metrics
+
+    def _get_reported_metrics(self, logs: Dict) -> Dict:
+        assert isinstance(self._metrics, (type(None), str, list, dict))
+
+        if self._metrics is None:
+            reported_metrics = logs
+        elif isinstance(self._metrics, str):
+            reported_metrics = {self._metrics: logs[self._metrics]}
+        elif isinstance(self._metrics, list):
+            reported_metrics = {metric: logs[metric] for metric in self._metrics}
+        elif isinstance(self._metrics, dict):
+            reported_metrics = {
+                key: logs[metric] for key, metric in self._metrics.items()
+            }
+
+        assert isinstance(reported_metrics, dict)
+        return reported_metrics
+
+    @abstractmethod
+    def _save_and_report_checkpoint(
+        self, metrics: Dict, checkpoint: TensorflowCheckpoint
+    ):
+        """Save checkpoint and report metrics corresonding to this checkpoint."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _report_metrics(self, metrics: Dict):
+        """Report metrics."""
+        raise NotImplementedError
+
+    def _handle(self, logs: Dict, when: str):
+        assert when in self._checkpoint_on or when in self._report_metrics_on
+
+        metrics = self._get_reported_metrics(logs)
+
+        should_checkpoint = when in self._checkpoint_on
+        if should_checkpoint:
+            checkpoint = TensorflowCheckpoint.from_model(self.model)
+            self._save_and_report_checkpoint(metrics, checkpoint)
+            # Clean up temporary checkpoint
+            shutil.rmtree(checkpoint.path, ignore_errors=True)
+        else:
+            self._report_metrics(metrics)
+
+
+@PublicAPI(stability="alpha")
+class ReportCheckpointCallback(RayReportCallback):
+    """Keras callback for Ray Train reporting and checkpointing.
+
+    .. note::
+        Metrics are always reported with checkpoints, even if the event isn't specified
+        in ``report_metrics_on``.
+
+    Example:
+        .. testcode:: python
+
+            ############# Using it in TrainSession ###############
+            from ray.air.integrations.keras import ReportCheckpointCallback
+            def train_loop_per_worker():
+                strategy = tf.distribute.MultiWorkerMirroredStrategy()
+                with strategy.scope():
+                    model = build_model()
+
+                model.fit(dataset_shard, callbacks=[ReportCheckpointCallback()])
+
+    Args:
+        metrics: Metrics to report. If this is a list, each item describes
+            the metric key reported to Keras, and it's reported under the
+            same name. If this is a dict, each key is the name reported
+            and the respective value is the metric key reported to Keras.
+            If this is None, all Keras logs are reported.
+        report_metrics_on: When to report metrics. Must be one of
+            the Keras event hooks (less the ``on_``), e.g.
+            "train_start" or "predict_end". Defaults to "epoch_end".
+        checkpoint_on: When to save checkpoints. Must be one of the Keras event hooks
+            (less the ``on_``), e.g. "train_start" or "predict_end". Defaults to
+            "epoch_end".
+    """
+
+    def _save_and_report_checkpoint(
+        self, metrics: Dict, checkpoint: TensorflowCheckpoint
+    ):
+        """Save checkpoint and report metrics corresonding to this checkpoint."""
+        ray.train.report(metrics, checkpoint=checkpoint)
+
+    def _report_metrics(self, metrics: Dict):
+        """Report metrics."""
+        ray.train.report(metrics, checkpoint=None)

python/ray/train/tensorflow/tensorflow_checkpoint.py

@@ -0,0 +1,155 @@
+import os
+import shutil
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+import tensorflow as tf
+from tensorflow import keras
+
+from ray.train._internal.framework_checkpoint import FrameworkCheckpoint
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.preprocessor import Preprocessor
+
+
+@PublicAPI(stability="beta")
+class TensorflowCheckpoint(FrameworkCheckpoint):
+    """A :py:class:`~ray.train.Checkpoint` with TensorFlow-specific functionality."""
+
+    MODEL_FILENAME_KEY = "_model_filename"
+
+    @classmethod
+    def from_model(
+        cls,
+        model: keras.Model,
+        *,
+        preprocessor: Optional["Preprocessor"] = None,
+    ) -> "TensorflowCheckpoint":
+        """Create a :py:class:`~ray.train.Checkpoint` that stores a Keras model.
+
+        The checkpoint created with this method needs to be paired with
+        `model` when used.
+
+        Args:
+            model: The Keras model, whose weights are stored in the checkpoint.
+            preprocessor: A fitted preprocessor to be applied before inference.
+
+        Returns:
+            A :py:class:`TensorflowCheckpoint` containing the specified model.
+
+        Examples:
+
+            .. testcode::
+
+                from ray.train.tensorflow import TensorflowCheckpoint
+                import tensorflow as tf
+
+                model = tf.keras.applications.resnet.ResNet101()
+                checkpoint = TensorflowCheckpoint.from_model(model)
+
+            .. testoutput::
+                :options: +MOCK
+                :hide:
+
+                ...  # Model may or may not be downloaded
+
+        """
+        tempdir = tempfile.mkdtemp()
+        filename = "model.keras"
+        model.save(Path(tempdir, filename).as_posix())
+
+        checkpoint = cls.from_directory(tempdir)
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+        checkpoint.update_metadata({cls.MODEL_FILENAME_KEY: filename})
+        return checkpoint
+
+    @classmethod
+    def from_h5(
+        cls, file_path: str, *, preprocessor: Optional["Preprocessor"] = None
+    ) -> "TensorflowCheckpoint":
+        """Create a :py:class:`~ray.train.Checkpoint` that stores a Keras
+        model from H5 format.
+
+        The checkpoint generated by this method contains all the information needed.
+        Thus no `model` is needed to be supplied when using this checkpoint.
+
+        Args:
+            file_path: The path to the .h5 file to load model from. This is the
+                same path that is used for ``model.save(path)``.
+            preprocessor: A fitted preprocessor to be applied before inference.
+
+        Returns:
+            A :py:class:`TensorflowCheckpoint` converted from h5 format.
+
+        """
+        if not os.path.isfile(file_path) or not file_path.endswith(".h5"):
+            raise ValueError(
+                "Please supply a h5 file path to `TensorflowCheckpoint.from_h5()`."
+            )
+        tempdir = tempfile.mkdtemp()
+        filename = os.path.basename(file_path)
+        new_checkpoint_file = Path(tempdir, filename).as_posix()
+        shutil.copy(file_path, new_checkpoint_file)
+
+        checkpoint = cls.from_directory(tempdir)
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+        checkpoint.update_metadata({cls.MODEL_FILENAME_KEY: filename})
+        return checkpoint
+
+    @classmethod
+    def from_saved_model(
+        cls, dir_path: str, *, preprocessor: Optional["Preprocessor"] = None
+    ) -> "TensorflowCheckpoint":
+        """Create a :py:class:`~ray.train.Checkpoint` that stores a Keras
+        model from SavedModel format.
+
+        The checkpoint generated by this method contains all the information needed.
+        Thus no `model` is needed to be supplied when using this checkpoint.
+
+        Args:
+            dir_path: The directory containing the saved model. This is the same
+                directory as used by ``model.save(dir_path)``.
+            preprocessor: A fitted preprocessor to be applied before inference.
+
+        Returns:
+            A :py:class:`TensorflowCheckpoint` converted from SavedModel format.
+
+        """
+        if not os.path.isdir(dir_path):
+            raise ValueError(
+                "Please supply a directory to `TensorflowCheckpoint.from_saved_model`"
+            )
+        tempdir = tempfile.mkdtemp()
+        # TODO(ml-team): Replace this with copytree()
+        os.rmdir(tempdir)
+        shutil.copytree(dir_path, tempdir)
+
+        checkpoint = cls.from_directory(tempdir)
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+        # NOTE: The entire directory is the checkpoint.
+        checkpoint.update_metadata({cls.MODEL_FILENAME_KEY: "."})
+        return checkpoint
+
+    def get_model(
+        self,
+    ) -> tf.keras.Model:
+        """Retrieve the model stored in this checkpoint.
+
+        Returns:
+            The Tensorflow Keras model stored in the checkpoint.
+        """
+        metadata = self.get_metadata()
+        if self.MODEL_FILENAME_KEY not in metadata:
+            raise ValueError(
+                "`TensorflowCheckpoint` cannot retrieve the model if you override the "
+                "checkpoint metadata. Please use `Checkpoint.update_metadata` instead."
+            )
+        model_filename = metadata[self.MODEL_FILENAME_KEY]
+        with self.as_directory() as checkpoint_dir:
+            model_path = Path(checkpoint_dir, model_filename).as_posix()
+            return keras.models.load_model(model_path)

python/ray/train/tensorflow/tensorflow_trainer.py

@@ -0,0 +1,191 @@
+from typing import Any, Callable, Dict, Optional, Union
+
+from ray.train import Checkpoint, DataConfig, RunConfig, ScalingConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.tensorflow.config import TensorflowConfig
+from ray.train.trainer import GenDataset
+from ray.util import PublicAPI
+
+
+@PublicAPI(stability="beta")
+class TensorflowTrainer(DataParallelTrainer):
+    """A Trainer for data parallel Tensorflow training.
+
+    This Trainer runs the function ``train_loop_per_worker`` on multiple Ray
+    Actors. These actors already have the necessary TensorFlow process group already
+    configured for distributed TensorFlow training.
+
+    The ``train_loop_per_worker`` function is expected to take in either 0 or 1
+    arguments:
+
+    .. testcode::
+
+        def train_loop_per_worker():
+            ...
+
+    .. testcode::
+
+        def train_loop_per_worker(config: Dict):
+            ...
+
+    If ``train_loop_per_worker`` accepts an argument, then
+    ``train_loop_config`` will be passed in as the argument. This is useful if you
+    want to tune the values in ``train_loop_config`` as hyperparameters.
+
+    If the ``datasets`` dict contains a training dataset (denoted by
+    the "train" key), then it will be split into multiple dataset
+    shards that can then be accessed by ``ray.train.get_dataset_shard("train")`` inside
+    ``train_loop_per_worker``. All the other datasets will not be split and
+    ``ray.train.get_dataset_shard(...)`` will return the entire Dataset.
+
+    Inside the ``train_loop_per_worker`` function, you can use any of the
+    :ref:`Ray Train loop methods <train-loop-api>`.
+
+    .. warning::
+        Ray will not automatically set any environment variables or configuration
+        related to local parallelism / threading
+        :ref:`aside from "OMP_NUM_THREADS" <omp-num-thread-note>`.
+        If you desire greater control over TensorFlow threading, use
+        the ``tf.config.threading`` module (eg.
+        ``tf.config.threading.set_inter_op_parallelism_threads(num_cpus)``)
+        at the beginning of your ``train_loop_per_worker`` function.
+
+
+    .. testcode::
+
+        from ray import train
+
+        def train_loop_per_worker():
+            # Report intermediate results for callbacks or logging and
+            # checkpoint data.
+            train.report(...)
+
+            # Returns dict of last saved checkpoint.
+            train.get_checkpoint()
+
+            # Returns the Dataset shard for the given key.
+            train.get_dataset_shard("my_dataset")
+
+            # Returns the total number of workers executing training.
+            train.get_context().get_world_size()
+
+            # Returns the rank of this worker.
+            train.get_context().get_world_rank()
+
+            # Returns the rank of the worker on the current node.
+            train.get_context().get_local_rank()
+
+    Any returns from the ``train_loop_per_worker`` will be discarded and not
+    used or persisted anywhere.
+
+    Example:
+
+    .. testcode::
+
+        import os
+        import tempfile
+        import tensorflow as tf
+        import numpy as np
+
+        import ray
+        from ray import train
+        from ray.train import Checkpoint, ScalingConfig
+        from ray.train.tensorflow import TensorflowTrainer
+
+        def build_model():
+            # toy neural network : 1-layer
+            return tf.keras.Sequential(
+                [tf.keras.layers.Dense(
+                    1, activation="linear", input_shape=(1,))]
+            )
+
+        def train_loop_per_worker(config):
+            dataset_shard = train.get_dataset_shard("train")
+            strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+            with strategy.scope():
+                model = build_model()
+                model.compile(
+                    optimizer="Adam", loss="mean_squared_error", metrics=["mse"])
+
+            tf_dataset = dataset_shard.to_tf(
+                feature_columns="x",
+                label_columns="y",
+                batch_size=1
+            )
+            for epoch in range(config["num_epochs"]):
+                model.fit(tf_dataset)
+
+                # Create checkpoint.
+                checkpoint_dir = tempfile.mkdtemp()
+                model.save_weights(
+                    os.path.join(checkpoint_dir, "my_checkpoint")
+                )
+                checkpoint = Checkpoint.from_directory(checkpoint_dir)
+
+                train.report(
+                    {},
+                    checkpoint=checkpoint,
+                )
+
+        train_dataset = ray.data.from_items([{"x": np.array([x], dtype=np.float32), "y": x + 1} for x in range(32)])
+        trainer = TensorflowTrainer(
+            train_loop_per_worker=train_loop_per_worker,
+            scaling_config=ScalingConfig(num_workers=3, use_gpu=True),
+            datasets={"train": train_dataset},
+            train_loop_config={"num_epochs": 2},
+        )
+        result = trainer.fit()
+
+    .. testoutput::
+        :options:+ELLIPSIS
+        :hide:
+
+        ...
+
+    Args:
+        train_loop_per_worker: The training function to execute.
+            This can either take in no arguments or a ``config`` dict.
+        train_loop_config: Configurations to pass into
+            ``train_loop_per_worker`` if it accepts an argument.
+        tensorflow_config: Configuration for setting up the TensorFlow backend.
+            If set to None, use the default configuration. This replaces the
+            ``backend_config`` arg of ``DataParallelTrainer``.
+        scaling_config: Configuration for how to scale data parallel training.
+        dataset_config: Configuration for dataset ingest.
+        run_config: Configuration for the execution of the training run.
+        datasets: Any Datasets to use for training. Use
+            the key "train" to denote which dataset is the training
+            dataset.
+        metadata: Dict that should be made available via
+            `ray.train.get_context().get_metadata()` and in `checkpoint.get_metadata()`
+            for checkpoints saved from this Trainer. Must be JSON-serializable.
+        resume_from_checkpoint: A checkpoint to resume training from.
+    """
+
+    def __init__(
+        self,
+        train_loop_per_worker: Union[Callable[[], None], Callable[[Dict], None]],
+        *,
+        train_loop_config: Optional[Dict] = None,
+        tensorflow_config: Optional[TensorflowConfig] = None,
+        scaling_config: Optional[ScalingConfig] = None,
+        dataset_config: Optional[DataConfig] = None,
+        run_config: Optional[RunConfig] = None,
+        datasets: Optional[Dict[str, GenDataset]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        resume_from_checkpoint: Optional[Checkpoint] = None,
+    ):
+        if not tensorflow_config:
+            tensorflow_config = TensorflowConfig()
+
+        super(TensorflowTrainer, self).__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            train_loop_config=train_loop_config,
+            backend_config=tensorflow_config,
+            scaling_config=scaling_config,
+            dataset_config=dataset_config,
+            run_config=run_config,
+            datasets=datasets,
+            resume_from_checkpoint=resume_from_checkpoint,
+            metadata=metadata,
+        )

python/ray/train/tensorflow/train_loop_utils.py

@@ -0,0 +1,27 @@
+import tensorflow as tf
+
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="beta")
+def prepare_dataset_shard(tf_dataset_shard: tf.data.Dataset):
+    """A utility function that overrides default config for Tensorflow Dataset.
+
+    This should be used on a TensorFlow ``Dataset`` created by calling
+    ``iter_tf_batches()`` on a ``ray.data.Dataset`` returned by
+    ``ray.train.get_dataset_shard()`` since the dataset has already
+    been sharded across the workers.
+
+    Args:
+        tf_dataset_shard: A TensorFlow Dataset.
+
+    Returns:
+        A TensorFlow Dataset with:
+            - autosharding turned off
+            - prefetching turned on with autotune enabled
+    """
+    options = tf.data.Options()
+    options.experimental_distribute.auto_shard_policy = (
+        tf.data.experimental.AutoShardPolicy.OFF
+    )
+    return tf_dataset_shard.with_options(options).prefetch(tf.data.AUTOTUNE)