always import ray.train._checkpoint.Checkpoint

Signed-off-by: Timothy Seah <tseah@anyscale.com>

Author: TimothySeah

python/ray/train/v2/api/callback.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Optional
 
-from ray.train import Checkpoint
+from ray.train._checkpoint import Checkpoint
 from ray.train.v2._internal.execution.context import TrainRunContext
 from ray.util.annotations import DeveloperAPI
 

python/ray/train/v2/api/reported_checkpoint.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Any, Dict
 
-from ray.train import Checkpoint
+from ray.train._checkpoint import Checkpoint
 from ray.util.annotations import PublicAPI
 
 

python/ray/train/v2/api/train_fn_utils.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
-from ray.train import Checkpoint
+from ray.train._checkpoint import Checkpoint
 from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
 from ray.train.v2.api.context import TrainContext
 from ray.train.v2.api.reported_checkpoint import ReportedCheckpoint