chore: import upstream snapshot with attribution

python/ray/train/trainer.py

@@ -0,0 +1,196 @@
+import logging
+import traceback
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, TypeVar, Union
+
+from ray.air._internal.util import (
+    StartTraceback,
+    StartTracebackWithWorkerRank,
+    skip_exceptions,
+)
+from ray.train import Checkpoint, DataConfig
+from ray.train._internal.backend_executor import (
+    BackendExecutor,
+    InactiveWorkerGroupError,
+    TrainBackendError,
+    TrainingWorkerError,
+)
+from ray.train._internal.session import _TrainingResult, _TrainSession, get_session
+from ray.train._internal.utils import ActorWrapper
+from ray.train.backend import BackendConfig
+from ray.train.base_trainer import (  # noqa: F401
+    BaseTrainer,
+    GenDataset,
+    TrainingFailedError,
+)
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.data import Dataset
+
+T = TypeVar("T")
+S = TypeVar("S")
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class TrainingIterator:
+    """An iterator over Train results. Returned by ``trainer.run_iterator``."""
+
+    def __init__(
+        self,
+        backend_executor: Union[BackendExecutor, ActorWrapper],
+        backend_config: BackendConfig,
+        train_func: Union[Callable[[], T], Callable[[Dict[str, Any]], T]],
+        datasets: Dict[str, "Dataset"],
+        metadata: Dict[str, Any],
+        data_config: DataConfig,
+        checkpoint: Optional[Union[Dict, str, Path, Checkpoint]],
+    ):
+        self._backend_executor = backend_executor
+        self._backend = backend_config.backend_cls()
+        self._train_func = train_func
+        self._datasets = datasets
+        self._metadata = metadata
+        self._data_config = data_config
+
+        self._start_training(
+            train_func=train_func,
+            datasets=self._datasets,
+            metadata=self._metadata,
+            data_config=self._data_config,
+            checkpoint=checkpoint,
+        )
+
+        self._finished_training = False
+
+    def __iter__(self):
+        return self
+
+    def _start_training(
+        self,
+        train_func,
+        datasets,
+        metadata,
+        data_config,
+        checkpoint: Optional[Checkpoint] = None,
+    ):
+        tune_session: _TrainSession = get_session()
+        assert tune_session, "`_start_training` should only be called from within Tune"
+        storage = tune_session.storage
+
+        self._run_with_error_handling(
+            lambda: self._backend_executor.start_training(
+                train_func=train_func,
+                datasets=datasets,
+                metadata=metadata,
+                data_config=data_config,
+                storage=storage,
+                checkpoint=checkpoint,
+            )
+        )
+
+    def _run_with_error_handling(self, func: Callable):
+        try:
+            return func()
+        except TrainingWorkerError:
+            # TODO(ml-team): This Train fault-tolerance code doesn't get used
+            # since max_retries=0
+            # Workers have already been restarted.
+            logger.info(
+                "Workers have been successfully restarted. Resuming "
+                "training from latest checkpoint."
+            )
+            self._start_training(
+                self._train_func,
+                self._datasets,
+                self._metadata,
+                self._data_config,
+            )
+            return self._run_with_error_handling(func)
+        except InactiveWorkerGroupError:
+            raise RuntimeError(
+                "This Trainer is not active. It is either shutdown "
+                "already or never started in the first place. "
+                "Either create a new Trainer or start this one."
+            ) from None
+        except TrainBackendError:
+            raise RuntimeError(
+                "Training failed. You should not be seeing "
+                "this error and this is a bug. Please create "
+                "a new issue at "
+                "https://github.com/ray-project/ray."
+            ) from None
+
+    def __next__(self):
+        if self.is_finished():
+            self._backend_executor.report_final_run_status(errored=False)
+            raise StopIteration
+        try:
+            next_results = self._run_with_error_handling(self._fetch_next_result)
+            if next_results is None:
+                self._backend_executor.report_final_run_status(errored=False)
+                self._run_with_error_handling(self._finish_training)
+                self._finished_training = True
+                raise StopIteration
+            else:
+                return next_results
+        except StartTraceback as e:
+            # If this is a StartTraceback, then this is a user error.
+            # We raise it directly
+            if isinstance(e, StartTracebackWithWorkerRank):
+                failed_rank = e.worker_rank
+            else:
+                failed_rank = None
+
+            # Extract the stack trace from the exception
+            e = skip_exceptions(e)
+            stack_trace = "".join(
+                traceback.format_exception(type(e), e, e.__traceback__)
+            )
+
+            self._backend_executor.report_final_run_status(
+                errored=True, stack_trace=stack_trace, failed_rank=failed_rank
+            )
+            try:
+                # Exception raised in at least one training worker. Immediately raise
+                # this error to the user and do not attempt to terminate gracefully.
+                self._backend_executor.shutdown(graceful_termination=False)
+                self._finished_training = True
+            except Exception:
+                pass
+            raise
+
+    def _fetch_next_result(self) -> Optional[List[Dict]]:
+        """Fetch next results produced by ``session.report()`` from each worker.
+
+        Assumes ``start_training`` has already been called.
+
+        Returns:
+            A list of dictionaries of values passed to ``session.report()`` from
+                each worker. Each item corresponds to an intermediate result
+                a single worker. If there are no more items to fetch,
+                returns None.
+        """
+        results = self._backend_executor.get_next_results()
+        if results is None:
+            return None
+        assert all(isinstance(result, _TrainingResult) for result in results)
+        return results
+
+    def _finish_training(self):
+        """Finish training and return final results. Propagate any exceptions.
+
+        Blocks until training is finished on all workers.
+
+        Assumes `start_training` has already been called.
+
+        Returns:
+            A list of return values from calling ``train_func`` on each worker.
+                Each item corresponds to the return value from a single worker.
+        """
+        return self._backend_executor.finish_training()
+
+    def is_finished(self) -> bool:
+        return self._finished_training