Annotate more api (ray-project#26501)

Signed-off-by: Stefan van der Kleij <[email protected]>

python/ray/train/horovod/horovod_trainer.py

@@ -7,11 +7,13 @@
 
 from ray.train.data_parallel_trainer import DataParallelTrainer
 from ray.train.horovod.config import HorovodConfig
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 class HorovodTrainer(DataParallelTrainer):
  """A Trainer for data parallel Horovod training.
 

python/ray/train/huggingface/huggingface_predictor.py

@@ -11,11 +11,13 @@
 from ray.air.checkpoint import Checkpoint
 from ray.air.constants import TENSOR_COLUMN_NAME
 from ray.train.predictor import Predictor
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 class HuggingFacePredictor(Predictor):
  """A predictor for HuggingFace Transformers PyTorch models.
 

python/ray/train/huggingface/utils.py

@@ -13,11 +13,13 @@
 )
 from ray.air._internal.torch_utils import load_torch_model
 from ray.air.checkpoint import Checkpoint
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 def load_checkpoint(
  checkpoint: Checkpoint,
  model: Union[Type[transformers.modeling_utils.PreTrainedModel], torch.nn.Module],

python/ray/train/lightgbm/lightgbm_predictor.py

@@ -8,11 +8,13 @@
 from ray.air.constants import TENSOR_COLUMN_NAME
 from ray.train.lightgbm.utils import load_checkpoint
 from ray.train.predictor import Predictor
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 class LightGBMPredictor(Predictor):
  """A predictor for LightGBM models.
 

python/ray/train/lightgbm/utils.py

@@ -9,11 +9,13 @@
 )
 from ray.air.checkpoint import Checkpoint
 from ray.air.constants import MODEL_KEY
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 def to_air_checkpoint(
  path: str,
  booster: lightgbm.Booster,
@@ -39,6 +41,7 @@ def to_air_checkpoint(
  return checkpoint
 
 
+@PublicAPI(stability="alpha")
 def load_checkpoint(
  checkpoint: Checkpoint,
 ) -> Tuple[lightgbm.Booster, Optional["Preprocessor"]]:

python/ray/train/rl/rl_predictor.py

@@ -9,11 +9,13 @@
 from ray.rllib.utils.typing import EnvType
 from ray.train.predictor import Predictor
 from ray.train.rl.utils import load_checkpoint
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 class RLPredictor(Predictor):
  """A predictor for RLlib policies.
 

python/ray/train/rl/utils.py

@@ -8,6 +8,7 @@
 )
 from ray.rllib.policy.policy import Policy
 from ray.rllib.utils.typing import EnvType
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
@@ -16,6 +17,7 @@
 RL_CONFIG_FILE = "config.pkl"
 
 
+@PublicAPI(stability="alpha")
 def load_checkpoint(
  checkpoint: Checkpoint,
  env: Optional[EnvType] = None,

python/ray/train/sklearn/sklearn_predictor.py

@@ -10,11 +10,13 @@
 from ray.train.sklearn._sklearn_utils import _set_cpu_params
 from ray.train.sklearn.utils import load_checkpoint
 from ray.util.joblib import register_ray
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 class SklearnPredictor(Predictor):
  """A predictor for scikit-learn compatible estimators.
 

python/ray/train/sklearn/utils.py

@@ -10,11 +10,13 @@
 )
 from ray.air.checkpoint import Checkpoint
 from ray.air.constants import MODEL_KEY
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 def to_air_checkpoint(
  path: str,
  estimator: BaseEstimator,
@@ -41,6 +43,7 @@ def to_air_checkpoint(
  return checkpoint
 
 
+@PublicAPI(stability="alpha")
 def load_checkpoint(
  checkpoint: Checkpoint,
 ) -> Tuple[BaseEstimator, Optional["Preprocessor"]]:

python/ray/train/tensorflow/tensorflow_predictor.py

@@ -10,13 +10,15 @@
 from ray.air.checkpoint import Checkpoint
 from ray.train.data_parallel_trainer import _load_checkpoint
 from ray.train._internal.dl_predictor import DLPredictor
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 logger = logging.getLogger(__name__)
 
 
+@PublicAPI(stability="alpha")
 class TensorflowPredictor(DLPredictor):
  """A predictor for TensorFlow models.
 

python/ray/train/tensorflow/utils.py

@@ -6,11 +6,13 @@
 from ray.air.checkpoint import Checkpoint
 from ray.air.constants import MODEL_KEY, PREPROCESSOR_KEY
 from ray.train.data_parallel_trainer import _load_checkpoint
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 def to_air_checkpoint(
  model: keras.Model, preprocessor: Optional["Preprocessor"] = None
 ) -> Checkpoint:
@@ -29,6 +31,7 @@ def to_air_checkpoint(
  return checkpoint
 
 
+@PublicAPI(stability="alpha")
 def load_checkpoint(
  checkpoint: Checkpoint,
  model: Union[Callable[[], tf.keras.Model], Type[tf.keras.Model], tf.keras.Model],

python/ray/train/torch/torch_predictor.py

@@ -9,13 +9,15 @@
 from ray.air.checkpoint import Checkpoint
 from ray.train.torch.utils import load_checkpoint
 from ray.train._internal.dl_predictor import DLPredictor
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 logger = logging.getLogger(__name__)
 
 
+@PublicAPI(stability="alpha")
 class TorchPredictor(DLPredictor):
  """A predictor for PyTorch models.
 

python/ray/train/torch/utils.py

@@ -6,11 +6,13 @@
 from ray.air.constants import MODEL_KEY, PREPROCESSOR_KEY
 from ray.train.data_parallel_trainer import _load_checkpoint
 from ray.air._internal.torch_utils import load_torch_model
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 def to_air_checkpoint(
  model: torch.nn.Module, preprocessor: Optional["Preprocessor"] = None
 ) -> Checkpoint:
@@ -29,6 +31,7 @@ def to_air_checkpoint(
  return checkpoint
 
 
+@PublicAPI(stability="alpha")
 def load_checkpoint(
  checkpoint: Checkpoint, model: Optional[torch.nn.Module] = None
 ) -> Tuple[torch.nn.Module, Optional["Preprocessor"]]:

python/ray/train/xgboost/utils.py

@@ -9,11 +9,13 @@
 )
 from ray.air.checkpoint import Checkpoint
 from ray.air.constants import MODEL_KEY
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 def to_air_checkpoint(
  path: str,
  booster: xgboost.Booster,
@@ -39,6 +41,7 @@ def to_air_checkpoint(
  return checkpoint
 
 
+@PublicAPI(stability="alpha")
 def load_checkpoint(
  checkpoint: Checkpoint,
 ) -> Tuple[xgboost.Booster, Optional["Preprocessor"]]:

python/ray/train/xgboost/xgboost_predictor.py

@@ -7,11 +7,13 @@
 from ray.air.constants import TENSOR_COLUMN_NAME
 from ray.train.predictor import Predictor
 from ray.train.xgboost.utils import load_checkpoint
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
  from ray.data.preprocessor import Preprocessor
 
 
+@PublicAPI(stability="alpha")
 class XGBoostPredictor(Predictor):
  """A predictor for XGBoost models.