class_docstr: Defines interface for distributed training on Ray.
  
  Note: The base ``BaseTrainer`` class cannot be instantiated directly. Only
  one of its subclasses can be used.
  
  Note to developers: If a new trainer is added, please update
  `air/_internal/usage.py`.
  
  **How does a trainer work?**
  
  - First, initialize the Trainer. The initialization runs locally,
    so heavyweight setup should not be done in ``__init__``.
  - Then, when you call ``trainer.fit()``, the Trainer is serialized
    and copied to a remote Ray actor. The following methods are then
    called in sequence on the remote actor.
  - ``trainer.setup()``: Any heavyweight Trainer setup should be
    specified here.
  - ``trainer.preprocess_datasets()``: The datasets passed to the Trainer will be
    setup here.
  - ``trainer.train_loop()``: Executes the main training logic.
  - Calling ``trainer.fit()`` will return a ``ray.result.Result``
    object where you can access metrics from your training run, as well
    as any checkpoints that may have been saved.
  
  **How do I create a new Trainer?**
  
  Subclass ``ray.train.trainer.BaseTrainer``, and override the ``training_loop``
  method, and optionally ``setup``.
  
  .. testcode::
  
      import torch
  
      from ray.train.trainer import BaseTrainer
      from ray import train, tune
  
  
      class MyPytorchTrainer(BaseTrainer):
          def setup(self):
              self.model = torch.nn.Linear(1, 1)
              self.optimizer = torch.optim.SGD(
                  self.model.parameters(), lr=0.1)
  
          def training_loop(self):
              # You can access any Trainer attributes directly in this method.
              # self.datasets["train"] has already been
              dataset = self.datasets["train"]
  
              torch_ds = dataset.iter_torch_batches(dtypes=torch.float)
              loss_fn = torch.nn.MSELoss()
  
              for epoch_idx in range(10):
                  loss = 0
                  num_batches = 0
                  torch_ds = dataset.iter_torch_batches(
                      dtypes=torch.float, batch_size=2
                  )
                  for batch in torch_ds:
                      X = torch.unsqueeze(batch["x"], 1)
                      y = torch.unsqueeze(batch["y"], 1)
                      # Compute prediction error
                      pred = self.model(X)
                      batch_loss = loss_fn(pred, y)
  
                      # Backpropagation
                      self.optimizer.zero_grad()
                      batch_loss.backward()
                      self.optimizer.step()
  
                      loss += batch_loss.item()
                      num_batches += 1
                  loss /= num_batches
  
                  # Use Tune functions to report intermediate
                  # results.
                  train.report({"loss": loss, "epoch": epoch_idx})
  
  
      # Initialize the Trainer, and call Trainer.fit()
      import ray
      train_dataset = ray.data.from_items(
          [{"x": i, "y": i} for i in range(10)])
      my_trainer = MyPytorchTrainer(datasets={"train": train_dataset})
      result = my_trainer.fit()
  
  .. testoutput::
          :hide:
  
          ...
  
  Args:
      scaling_config: Configuration for how to scale training.
      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
          `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.
class_docstr_tok: ["Defines","interface","for","distributed","training","on","Ray",".","Note",":","The","base","``","BaseTrainer","``","class","can","not","be","instantiated","directly",".","Only","one","of","its","subclasses","can","be","used",".","Note","to","developers",":","If","a","new","trainer","is","added",",","please","update","`","air\/_internal\/usage.py","`",".","*","*","How","does","a","trainer","work","?","*","*","-","First",",","initialize","the","Trainer",".","The","initialization","runs","locally",",","so","heavyweight","setup","should","not","be","done","in","``","__init__","``",".","-","Then",",","when","you","call","``","trainer.fit","(",")","``",",","the","Trainer","is","serialized","and","copied","to","a","remote","Ray","actor",".","The","following","methods","are","then","called","in","sequence","on","the","remote","actor",".","-","``","trainer.setup","(",")","``",":","Any","heavyweight","Trainer","setup","should","be","specified","here",".","-","``","trainer.preprocess_datasets","(",")","``",":","The","datasets","passed","to","the","Trainer","will","be","setup","here",".","-","``","trainer.train_loop","(",")","``",":","Executes","the","main","training","logic",".","-","Calling","``","trainer.fit","(",")","``","will","return","a","``","ray.result.Result","``","object","where","you","can","access","metrics","from","your","training","run",",","as","well","as","any","checkpoints","that","may","have","been","saved",".","*","*","How","do","I","create","a","new","Trainer","?","*","*","Subclass","``","ray.train.trainer.BaseTrainer","``",",","and","override","the","``","training_loop","``","method",",","and","optionally","``","setup","``",".","..","testcode",":",":","import","torch","from","ray.train.trainer","import","BaseTrainer","from","ray","import","train",",","tune","class","MyPytorchTrainer","(","BaseTrainer",")",":","def","setup","(","self",")",":","self.model","=","torch.nn.Linear","(","1",",","1",")","self.optimizer","=","torch.optim.SGD","(","self.model.parameters","(",")",",","lr=0.1",")","def","training_loop","(","self",")",":","#","You","can","access","any","Trainer","attributes","directly","in","this","method",".","#","self.datasets","[","``","train","''","]","has","already","been","dataset","=","self.datasets","[","``","train","''","]","torch_ds","=","dataset.iter_torch_batches","(","dtypes=torch.float",")","loss_fn","=","torch.nn.MSELoss","(",")","for","epoch_idx","in","range","(","10",")",":","loss","=","0","num_batches","=","0","torch_ds","=","dataset.iter_torch_batches","(","dtypes=torch.float",",","batch_size=2",")","for","batch","in","torch_ds",":","X","=","torch.unsqueeze","(","batch","[","``","x","''","]",",","1",")","y","=","torch.unsqueeze","(","batch","[","``","y","''","]",",","1",")","#","Compute","prediction","error","pred","=","self.model","(","X",")","batch_loss","=","loss_fn","(","pred",",","y",")","#","Backpropagation","self.optimizer.zero_grad","(",")","batch_loss.backward","(",")","self.optimizer.step","(",")","loss","+=","batch_loss.item","(",")","num_batches","+=","1","loss","\/=","num_batches","#","Use","Tune","functions","to","report","intermediate","#","results",".","train.report","(","{","``","loss","''",":","loss",",","``","epoch","''",":","epoch_idx","}",")","#","Initialize","the","Trainer",",","and","call","Trainer.fit","(",")","import","ray","train_dataset","=","ray.data.from_items","(","[","{","``","x","''",":","i",",","``","y","''",":","i","}","for","i","in","range","(","10",")","]",")","my_trainer","=","MyPytorchTrainer","(","datasets=","{","``","train","''",":","train_dataset","}",")","result","=","my_trainer.fit","(",")","..","testoutput",":",":",":","hide",":","...","Args",":","scaling_config",":","Configuration","for","how","to","scale","training",".","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","`","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","."]
docstr: Restores a Train experiment from a previously interrupted/failed run.
  
  Restore should be used for experiment-level fault tolerance in the event
  that the head node crashes (e.g., OOM or some other runtime error) or the
  entire cluster goes down (e.g., network error affecting all nodes).
  
  A run that has already completed successfully will not be resumed from this API.
  To continue training from a successful run, launch a new run with the
  ``<Framework>Trainer(resume_from_checkpoint)`` API instead, passing in a
  checkpoint from the previous run to start with.
  
  .. note::
  
      Restoring an experiment from a path that's pointing to a *different*
      location than the original experiment path is supported. However, Ray Train
      assumes that the full experiment directory is available
      (including checkpoints) so that it's possible to resume trials from their
      latest state.
  
      For example, if the original experiment path was run locally, then the
      results are uploaded to cloud storage, Ray Train expects the full contents
      to be available in cloud storage if attempting to resume
      via ``<Framework>Trainer.restore("s3://...")``. The restored run will
      continue writing results to the same cloud storage location.
  
  The following example can be paired with implementing job retry using
  :ref:`Ray Jobs <jobs-overview>` to produce a Train experiment that will
  attempt to resume on both experiment-level and trial-level failures:
  
  .. testcode::
  
      import os
      import ray
      from ray import train
      from ray.train.trainer import BaseTrainer
  
      experiment_name = "unique_experiment_name"
      storage_path = os.path.expanduser("~/ray_results")
      experiment_dir = os.path.join(storage_path, experiment_name)
  
      # Define some dummy inputs for demonstration purposes
      datasets = {"train": ray.data.from_items([{"a": i} for i in range(10)])}
  
      class CustomTrainer(BaseTrainer):
          def training_loop(self):
              pass
  
      if CustomTrainer.can_restore(experiment_dir):
          trainer = CustomTrainer.restore(
              experiment_dir, datasets=datasets
          )
      else:
          trainer = CustomTrainer(
              datasets=datasets,
              run_config=train.RunConfig(
                  name=experiment_name,
                  storage_path=storage_path,
                  # Tip: You can also enable retries on failure for
                  # worker-level fault tolerance
                  failure_config=train.FailureConfig(max_failures=3),
              ),
          )
  
      result = trainer.fit()
  
  .. testoutput::
      :hide:
  
      ...
  
  Args:
      path: The path to the experiment directory of the training run to restore.
          This can be a local path or a remote URI if the experiment was
          uploaded to the cloud.
      storage_filesystem: Custom ``pyarrow.fs.FileSystem``
          corresponding to the ``path``. This may be necessary if the original
          experiment passed in a custom filesystem.
      datasets: Re-specified datasets used in the original training run.
          This must include all the datasets that were passed in the
          original trainer constructor.
      scaling_config: Optionally re-specified scaling config. This can be
          modified to be different from the original spec.
      **kwargs: Other optionally re-specified arguments, passed in by subclasses.
  
  Raises:
      ValueError: If all datasets were not re-supplied on restore.
  
  Returns:
      BaseTrainer: A restored instance of the class that is calling this method.
code: def restore(
      cls: Type["BaseTrainer"],
      path: Union[str, os.PathLike],
      storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
      datasets: Optional[Dict[str, GenDataset]] = None,
      preprocessor: Optional["Preprocessor"] = None,
      scaling_config: Optional[ScalingConfig] = None,
      **kwargs,
  ) -> "BaseTrainer":
      """Restores a Train experiment from a previously interrupted/failed run.
  
      Restore should be used for experiment-level fault tolerance in the event
      that the head node crashes (e.g., OOM or some other runtime error) or the
      entire cluster goes down (e.g., network error affecting all nodes).
  
      A run that has already completed successfully will not be resumed from this API.
      To continue training from a successful run, launch a new run with the
      ``<Framework>Trainer(resume_from_checkpoint)`` API instead, passing in a
      checkpoint from the previous run to start with.
  
      .. note::
  
          Restoring an experiment from a path that's pointing to a *different*
          location than the original experiment path is supported. However, Ray Train
          assumes that the full experiment directory is available
          (including checkpoints) so that it's possible to resume trials from their
          latest state.
  
          For example, if the original experiment path was run locally, then the
          results are uploaded to cloud storage, Ray Train expects the full contents
          to be available in cloud storage if attempting to resume
          via ``<Framework>Trainer.restore("s3://...")``. The restored run will
          continue writing results to the same cloud storage location.
  
      The following example can be paired with implementing job retry using
      :ref:`Ray Jobs <jobs-overview>` to produce a Train experiment that will
      attempt to resume on both experiment-level and trial-level failures:
  
      .. testcode::
  
          import os
          import ray
          from ray import train
          from ray.train.trainer import BaseTrainer
  
          experiment_name = "unique_experiment_name"
          storage_path = os.path.expanduser("~/ray_results")
          experiment_dir = os.path.join(storage_path, experiment_name)
  
          # Define some dummy inputs for demonstration purposes
          datasets = {"train": ray.data.from_items([{"a": i} for i in range(10)])}
  
          class CustomTrainer(BaseTrainer):
              def training_loop(self):
                  pass
  
          if CustomTrainer.can_restore(experiment_dir):
              trainer = CustomTrainer.restore(
                  experiment_dir, datasets=datasets
              )
          else:
              trainer = CustomTrainer(
                  datasets=datasets,
                  run_config=train.RunConfig(
                      name=experiment_name,
                      storage_path=storage_path,
                      # Tip: You can also enable retries on failure for
                      # worker-level fault tolerance
                      failure_config=train.FailureConfig(max_failures=3),
                  ),
              )
  
          result = trainer.fit()
  
      .. testoutput::
          :hide:
  
          ...
  
      Args:
          path: The path to the experiment directory of the training run to restore.
              This can be a local path or a remote URI if the experiment was
              uploaded to the cloud.
          storage_filesystem: Custom ``pyarrow.fs.FileSystem``
              corresponding to the ``path``. This may be necessary if the original
              experiment passed in a custom filesystem.
          datasets: Re-specified datasets used in the original training run.
              This must include all the datasets that were passed in the
              original trainer constructor.
          scaling_config: Optionally re-specified scaling config. This can be
              modified to be different from the original spec.
          **kwargs: Other optionally re-specified arguments, passed in by subclasses.
  
      Raises:
          ValueError: If all datasets were not re-supplied on restore.
  
      Returns:
          BaseTrainer: A restored instance of the class that is calling this method.
      """
      if not cls.can_restore(path, storage_filesystem):
          raise ValueError(
              f"Invalid restore path: {path}. Make sure that this path exists and "
              "is the experiment directory that results from a call to "
              "`trainer.fit()`."
          )
      fs, fs_path = get_fs_and_path(path, storage_filesystem)
      with fs.open_input_file(os.path.join(fs_path, _TRAINER_PKL)) as f:
          trainer_cls, param_dict = pickle.loads(f.readall())
  
      if trainer_cls is not cls:
          warnings.warn(
              f"Invalid trainer type. You are attempting to restore a trainer of type"
              f" {trainer_cls} with `{cls.__name__}.restore`, "
              "which will most likely fail. "
              f"Use `{trainer_cls.__name__}.restore` instead."
          )
  
      original_datasets = param_dict.pop("datasets", {})
      if original_datasets and not datasets:
          raise ValueError(
              "The following datasets need to be provided again on restore: "
              f"{list(original_datasets.keys())}\n"
              f"Use {cls.__name__}.restore(..., datasets=datasets) "
              "with the datasets that were provided to the original trainer."
          )
      datasets = datasets or {}
      if set(original_datasets) != set(datasets):
          raise ValueError(
              "The provided datasets don't match the original dataset keys.\n"
              f"  Expected datasets for the keys: {list(original_datasets.keys())}\n"
              f"  Actual datasets provided: {list(datasets.keys())}"
          )
      param_dict["datasets"] = datasets
  
      # If no preprocessor is re-specified, then it will be set to None
      # here and loaded from the latest checkpoint
      param_dict["preprocessor"] = preprocessor
  
      if scaling_config:
          param_dict["scaling_config"] = scaling_config
  
      for param_name, val in kwargs.items():
          # Overwrite the old value if something is passed into restore
          if val is not None:
              param_dict[param_name] = val
  
      try:
          trainer = cls(**param_dict)
      except Exception as e:
          raise ValueError(
              "Trainer restoration failed (see above for the stack trace). "
              "Make sure that you use the right trainer class to restore: "
              f"`{cls.__name__}.restore`\n"
          ) from e
      trainer._restore_path = path
      trainer._restore_storage_filesystem = storage_filesystem
      return trainer
  
imports: import abc
  import copy
  import inspect
  import json
  import logging
  import os
  import warnings
  from pathlib import Path
  from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
  import pyarrow.fs
  import ray
  import ray.cloudpickle
  from ray._private.dict import merge_dicts
  from ray.air._internal import usage
  from ray.air._internal.config import ensure_only_allowed_dataclass_keys_updated
  from ray.air._internal.usage import AirEntrypoint
  from ray.air.config import RunConfig, ScalingConfig
  from ray.air.result import Result
  from ray.train import Checkpoint
  from ray.train._internal.session import _get_session
  from ray.train._internal.storage import _exists_at_fs_path, get_fs_and_path
  from ray.train.constants import TRAIN_DATASET_KEY
  from ray.util import PublicAPI
  from ray.util.annotations import DeveloperAPI
label_desc: Use image node_id 2 for calling the BaseTrainer obj's underlying member method code with example usage: obj.restore(cls, path, storage_filesystem, datasets, preprocessor, scaling_config) and returns: trainer