[train] enable isort (ray-project#40172)

Enable `isort` for `python/ray/train` so that it will be formatted with `./scripts/format.sh`.

---------

Signed-off-by: Matthew Deng <matt@anyscale.com>

.isort.cfg

@@ -19,7 +19,7 @@ filter_files=True
 # python/ray/setup-dev.py
 # For the rest we will gradually remove them from the blacklist as we
 # reformat the code to follow the style guide.
-skip_glob=doc/*,python/ray/__init__.py,python/ray/setup-dev.py,python/build/*,python/ray/cloudpickle/*,python/ray/thirdparty_files/*,python/ray/_private/thirdparty/*,python/ray/_private/runtime_env/agent/thirdparty_files/*,python/ray/dag/*.py,ci/*,python/ray/_private/*,python/ray/air/*,dashboard/*,python/ray/includes/*,python/ray/internal/*,python/ray/ray_operator/*,python/ray/scripts/*,python/ray/serve/generated/serve_pb2.py,python/ray/sgd/*,python/ray/streaming/*,python/ray/tests/*,python/ray/tests/*,python/ray/train/*,python/ray/tune/*,python/ray/util/*,python/ray/workers/*,python/ray/workflow/*,rllib/*,release/*,
+skip_glob=doc/*,python/ray/__init__.py,python/ray/setup-dev.py,python/build/*,python/ray/cloudpickle/*,python/ray/thirdparty_files/*,python/ray/_private/thirdparty/*,python/ray/_private/runtime_env/agent/thirdparty_files/*,python/ray/dag/*.py,ci/*,python/ray/_private/*,python/ray/air/*,dashboard/*,python/ray/includes/*,python/ray/internal/*,python/ray/ray_operator/*,python/ray/scripts/*,python/ray/serve/generated/serve_pb2.py,python/ray/sgd/*,python/ray/streaming/*,python/ray/tests/*,python/ray/tests/*,python/ray/tune/*,python/ray/util/*,python/ray/workers/*,python/ray/workflow/*,rllib/*,release/*,
 
 known_local_folder=ray
 known_afterray=psutil,setproctitle

python/ray/train/__init__.py

@@ -1,17 +1,21 @@
 # Try import ray[train] core requirements (defined in setup.py)
+# isort: off
 try:
+    import fsspec  # noqa: F401
     import pandas  # noqa: F401
-    import requests  # noqa: F401
     import pyarrow  # noqa: F401
-    import fsspec  # noqa: F401
+    import requests  # noqa: F401
 except ImportError as exc:
     raise ImportError(
         "Can't import ray.train as some dependencies are missing. "
         'Run `pip install "ray[train]"` to fix.'
     ) from exc
+# isort: on
 
 
 from ray._private.usage import usage_lib
+from ray.air.config import CheckpointConfig, FailureConfig, RunConfig, ScalingConfig
+from ray.air.result import Result
 
 # Import this first so it can be used in other modules
 from ray.train._checkpoint import Checkpoint
@@ -23,9 +27,6 @@
 from ray.train.context import get_context
 from ray.train.trainer import TrainingIterator
 
-from ray.air.config import CheckpointConfig, FailureConfig, RunConfig, ScalingConfig
-from ray.air.result import Result
-
 usage_lib.record_library_usage("train")
 
 Checkpoint.__module__ = "ray.train"

python/ray/train/_checkpoint.py

@@ -7,8 +7,8 @@
 import shutil
 import tempfile
 import traceback
-from typing import Any, Dict, Iterator, List, Optional, Union
 import uuid
+from typing import Any, Dict, Iterator, List, Optional, Union
 
 import pyarrow.fs
 

python/ray/train/_internal/backend_executor.py

@@ -2,17 +2,17 @@
 import os
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar, Any
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar
 
 import ray
 import ray._private.ray_constants as ray_constants
-from ray.data import Dataset
 from ray._private.ray_constants import env_integer
+from ray.data import Dataset
 from ray.exceptions import RayActorError
 from ray.train import Checkpoint, DataConfig
 from ray.train._internal.session import (
-    _TrainingResult,
     TrialInfo,
+    _TrainingResult,
     get_session,
     init_session,
     shutdown_session,
@@ -24,9 +24,9 @@
 from ray.train.constants import (
     ENABLE_DETAILED_AUTOFILLED_METRICS_ENV,
     ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
+    ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
     TRAIN_ENABLE_WORKER_SPREAD_ENV,
     TRAIN_PLACEMENT_GROUP_TIMEOUT_S_ENV,
-    ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
 )
 from ray.util.placement_group import get_current_placement_group, remove_placement_group
 

python/ray/train/_internal/checkpoint_manager.py

@@ -3,12 +3,11 @@
 from typing import Any, Callable, List, Optional, Tuple
 
 from ray._private.dict import flatten_dict
-from ray.air.config import MAX
 from ray.air._internal.util import is_nan
+from ray.air.config import MAX
 from ray.train import CheckpointConfig
-from ray.train._internal.storage import _delete_fs_path
 from ray.train._internal.session import _TrainingResult
-
+from ray.train._internal.storage import _delete_fs_path
 
 logger = logging.getLogger(__name__)
 

python/ray/train/_internal/data_config.py

@@ -1,20 +1,14 @@
-from typing import Optional, Union, Dict, List
+import sys
+from typing import Dict, List, Optional, Union
 
 import ray
 from ray.actor import ActorHandle
+from ray.data import DataIterator, Dataset, ExecutionOptions, NodeIdStr
+from ray.data.preprocessor import Preprocessor
 
 # TODO(justinvyu): Fix the circular import error
 from ray.train.constants import TRAIN_DATASET_KEY  # noqa
-from ray.util.annotations import PublicAPI, DeveloperAPI
-from ray.data import (
-    Dataset,
-    DataIterator,
-    ExecutionOptions,
-    NodeIdStr,
-)
-from ray.data.preprocessor import Preprocessor
-
-import sys
+from ray.util.annotations import DeveloperAPI, PublicAPI
 
 if sys.version_info >= (3, 8):
     from typing import Literal

python/ray/train/_internal/dataset_spec.py

@@ -2,7 +2,6 @@
 from typing import Callable, Dict, List, Optional, Union
 
 from ray.actor import ActorHandle
-
 from ray.data import Dataset, DatasetPipeline
 
 RayDataset = Union["Dataset", "DatasetPipeline"]

python/ray/train/_internal/session.py

@@ -1,38 +1,38 @@
-import os
+import functools
 import logging
+import os
 import platform
 import queue
 import sys
 import threading
 import time
+import warnings
 from dataclasses import dataclass
 from datetime import datetime
-import functools
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Set, Type, Union
-import warnings
 
 import ray
 from ray.air._internal.session import _get_session
-from ray.air._internal.util import StartTraceback, RunnerThread
+from ray.air._internal.util import RunnerThread, StartTraceback
 from ray.air.constants import (
-    _RESULT_FETCH_TIMEOUT,
     _ERROR_FETCH_TIMEOUT,
+    _RESULT_FETCH_TIMEOUT,
     SESSION_MISUSE_LOG_ONCE_KEY,
-    TIMESTAMP,
     TIME_THIS_ITER_S,
+    TIMESTAMP,
 )
 from ray.data import Dataset, DatasetPipeline
 from ray.train import Checkpoint
 from ray.train._internal.accelerator import Accelerator
-from ray.train._internal.storage import _use_storage_context, StorageContext
+from ray.train._internal.storage import StorageContext, _use_storage_context
 from ray.train.constants import (
+    CHECKPOINT_DIR_NAME,
     DETAILED_AUTOFILLED_KEYS,
+    RAY_CHDIR_TO_TRIAL_DIR,
+    TIME_TOTAL_S,
     WORKER_HOSTNAME,
     WORKER_NODE_IP,
     WORKER_PID,
-    TIME_TOTAL_S,
-    RAY_CHDIR_TO_TRIAL_DIR,
-    CHECKPOINT_DIR_NAME,
 )
 from ray.train.error import SessionMisuseError
 from ray.util.annotations import DeveloperAPI, PublicAPI
@@ -43,7 +43,6 @@
     SchedulingStrategyT,
 )
 
-
 if TYPE_CHECKING:
     from ray.data import DataIterator
     from ray.tune.execution.placement_groups import PlacementGroupFactory

python/ray/train/_internal/storage.py

@@ -1,11 +1,5 @@
-import dataclasses
-import fnmatch
-import logging
-import os
-from pathlib import Path
-import shutil
-from typing import Callable, Dict, List, Optional, Tuple, Type, Union, TYPE_CHECKING
-
+# Try import ray[train] core requirements (defined in setup.py)
+# isort: off
 try:
     import fsspec  # noqa
     from fsspec.implementations.local import LocalFileSystem
@@ -25,11 +19,19 @@
         "pyarrow is a required dependency of Ray Train and Ray Tune. "
         "Please install with: `pip install pyarrow`"
     ) from e
+# isort: on
 
+import dataclasses
+import fnmatch
+import logging
+import os
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Type, Union
 
 from ray._private.storage import _get_storage_uri
 from ray.air._internal.filelock import TempFileLock
-from ray.train._internal.syncer import Syncer, SyncConfig, _BackgroundSyncer
+from ray.train._internal.syncer import SyncConfig, Syncer, _BackgroundSyncer
 from ray.train.constants import _get_defaults_results_dir
 
 if TYPE_CHECKING:

python/ray/train/_internal/syncer.py

@@ -1,36 +1,25 @@
 import abc
+import logging
 import threading
+import time
 import traceback
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    Union,
-    Optional,
-    Tuple,
-)
 import warnings
-
-import logging
-import time
 from dataclasses import dataclass
-
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from ray._private.thirdparty.tabulate.tabulate import tabulate
 from ray.air._internal.remote_storage import (
-    fs_hint,
-    upload_to_uri,
-    download_from_uri,
     delete_at_uri,
+    download_from_uri,
+    fs_hint,
     is_non_local_path_uri,
+    upload_to_uri,
 )
 from ray.train.constants import _DEPRECATED_VALUE
 from ray.util import log_once
-from ray.util.annotations import PublicAPI, DeveloperAPI
+from ray.util.annotations import DeveloperAPI, PublicAPI
 from ray.widgets import Template
 
-
 logger = logging.getLogger(__name__)
 
 # Syncing period for syncing checkpoints between nodes or to cloud.

python/ray/train/_internal/utils.py

@@ -1,28 +1,17 @@
 import abc
 import functools
 import inspect
-import os
 import logging
+import os
 from pathlib import Path
-
-from typing import (
-    Tuple,
-    Dict,
-    List,
-    Any,
-    Union,
-    Callable,
-    TypeVar,
-    Optional,
-)
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
 
 import ray
-from ray.air._internal.util import find_free_port, StartTraceback
 from ray.actor import ActorHandle
+from ray.air._internal.util import StartTraceback, find_free_port
 from ray.exceptions import RayActorError
 from ray.types import ObjectRef
 
-
 T = TypeVar("T")
 
 logger = logging.getLogger(__name__)

python/ray/train/_internal/worker_group.py

@@ -3,11 +3,11 @@
 import socket
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Callable, List, TypeVar, Optional, Dict, Type, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
 
 import ray
 from ray.actor import ActorHandle
-from ray.air._internal.util import skip_exceptions, exception_cause
+from ray.air._internal.util import exception_cause, skip_exceptions
 from ray.types import ObjectRef
 from ray.util.placement_group import PlacementGroup
 

python/ray/train/base_trainer.py

@@ -4,33 +4,30 @@
 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 warnings
 
 import pyarrow.fs
 
 import ray
 import ray.cloudpickle as pickle
-from ray.air._internal.config import ensure_only_allowed_dataclass_keys_updated
+from ray._private.dict import merge_dicts
 from ray.air._internal import usage as air_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 import Checkpoint
 from ray.train.constants import TRAIN_DATASET_KEY
 from ray.util import PublicAPI
 from ray.util.annotations import DeveloperAPI
-from ray._private.dict import merge_dicts
-
 
 if TYPE_CHECKING:
     from ray.data import Dataset
     from ray.data.preprocessor import Preprocessor
-
     from ray.tune import Trainable
 
 _TRAINER_PKL = "trainer.pkl"
@@ -578,8 +575,8 @@ def fit(self) -> Result:
             TrainingFailedError: If any failures during the execution of
             ``self.as_trainable()``, or during the Tune execution loop.
         """
-        from ray.tune.tuner import Tuner, TunerInternal
         from ray.tune import TuneError
+        from ray.tune.tuner import Tuner, TunerInternal
 
         trainable = self.as_trainable()
         param_space = self._extract_fields_for_tuner_param_space()

python/ray/train/constants.py

@@ -2,11 +2,11 @@
 from pathlib import Path
 
 from ray.air.constants import (  # noqa: F401
+    COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV,
     EVALUATION_DATASET_KEY,
     MODEL_KEY,
     PREPROCESSOR_KEY,
     TRAIN_DATASET_KEY,
-    COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV,
 )
 
 

python/ray/train/context.py

@@ -1,11 +1,10 @@
 import threading
-from typing import TYPE_CHECKING, Optional, Dict, Any
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from ray.train._internal import session
 from ray.train._internal.storage import StorageContext
 from ray.util.annotations import DeveloperAPI, PublicAPI
 
-
 if TYPE_CHECKING:
     from ray.tune.execution.placement_groups import PlacementGroupFactory