[Train] Make get_checkpoint top level function in ray.train (ray-proj…

…ect#37906)

This is bringing the API up-to-date with ray-project/enhancements#36

python/ray/train/__init__.py

@@ -1,6 +1,6 @@
 from ray._private.usage import usage_lib
 from ray.train._internal.data_config import DataConfig
-from ray.train._internal.session import get_dataset_shard, report
+from ray.train._internal.session import get_checkpoint, get_dataset_shard, report
 from ray.train.backend import BackendConfig
 from ray.train.constants import TRAIN_DATASET_KEY
 from ray.train.context import get_context
@@ -13,6 +13,7 @@
 usage_lib.record_library_usage("train")
 
 __all__ = [
+ "get_checkpoint",
  "get_context",
  "get_dataset_shard",
  "report",

python/ray/train/context.py

@@ -1,7 +1,6 @@
 import threading
 from typing import TYPE_CHECKING, Optional
 
-from ray.air import Checkpoint
 from ray.train._internal import session
 from ray.util.annotations import PublicAPI
 
@@ -27,10 +26,6 @@ def wrapped(func):
 class TrainContext:
  """Context for Ray training executions."""
 
- @_copy_doc(session.get_checkpoint)
- def get_checkpoint(self) -> Optional[Checkpoint]:
- return session.get_checkpoint()
-
  @_copy_doc(session.get_experiment_name)
  def get_experiment_name(self) -> str:
  return session.get_experiment_name()