chore: Add proper ruff configuration

.pre-commit-config.yaml

@@ -7,12 +7,11 @@ repos:
       - id: end-of-file-fixer
       - id: trailing-whitespace
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.4
+    # Pin to the latest release as of today
+    rev: v0.12.10
     hooks:
-      - id: ruff
-exclude: |
-  (?x)^(
-    kubeflow/trainer/__init__.py|
-    kubeflow/trainer/api/__init__.py|
-    kubeflow/trainer/models/.*|
-  )$
+      # Lint + auto-fix (must run before format)
+      - id: ruff-check
+        args: [ --fix ]
+      # Format after fixes
+      - id: ruff-format

Makefile

@@ -48,12 +48,13 @@ uv: ## Install UV
 
 .PHONY: ruff
 ruff: ## Install Ruff
-	@uvx ruff --help &> /dev/null || uv tool install ruff
+	@uv run ruff --help &> /dev/null || uv tool install ruff
 
 .PHONY: verify
 verify: install-dev  ## install all required tools
 	@uv lock --check
-	@uvx ruff check --show-fixes
+	@uv run ruff check --show-fixes --output-format=github .
+	@uv run ruff format --check kubeflow
 
 .PHONY: uv-venv
 uv-venv:

kubeflow/trainer/__init__.py

@@ -13,11 +13,13 @@
 # limitations under the License.
 
 
-from __future__ import absolute_import
-
 # Import the Kubeflow Trainer client.
 from kubeflow.trainer.api.trainer_client import TrainerClient  # noqa: F401
 
+# import backends and its associated configs
+from kubeflow.trainer.backends.kubernetes.types import KubernetesBackendConfig
+from kubeflow.trainer.backends.localprocess.types import LocalProcessBackendConfig
+
 # Import the Kubeflow Trainer constants.
 from kubeflow.trainer.constants.constants import DATASET_PATH, MODEL_PATH  # noqa: F401
 
@@ -32,17 +34,12 @@
     Initializer,
     Loss,
     Runtime,
+    RuntimeTrainer,
     TorchTuneConfig,
     TorchTuneInstructDataset,
-    RuntimeTrainer,
     TrainerType,
 )
 
-# import backends and its associated configs
-from kubeflow.trainer.backends.kubernetes.types import KubernetesBackendConfig
-from kubeflow.trainer.backends.localprocess.types import LocalProcessBackendConfig
-
-
 __all__ = [
     "BuiltinTrainer",
     "CustomTrainer",

kubeflow/trainer/api/__init__.py

@@ -1,4 +1,3 @@
 # ruff: noqa
 
 # import apis into api package
-

kubeflow/trainer/api/trainer_client.py

@@ -12,26 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections.abc import Iterator
 import logging
-from typing import Optional, Union, Iterator
+from typing import Optional, Union
 
-from kubeflow.trainer.constants import constants
-from kubeflow.trainer.types import types
 from kubeflow.trainer.backends.kubernetes.backend import KubernetesBackend
 from kubeflow.trainer.backends.kubernetes.types import KubernetesBackendConfig
-from kubeflow.trainer.backends.localprocess.backend import LocalProcessBackend
-from kubeflow.trainer.backends.localprocess.backend import LocalProcessBackendConfig
-
+from kubeflow.trainer.backends.localprocess.backend import (
+    LocalProcessBackend,
+    LocalProcessBackendConfig,
+)
+from kubeflow.trainer.constants import constants
+from kubeflow.trainer.types import types
 
 logger = logging.getLogger(__name__)
 
 
 class TrainerClient:
     def __init__(
         self,
-        backend_config: Union[
-            KubernetesBackendConfig, LocalProcessBackendConfig
-        ] = KubernetesBackendConfig(),
+        backend_config: Union[KubernetesBackendConfig, LocalProcessBackendConfig] = None,
     ):
         """Initialize a Kubeflow Trainer client.
 
@@ -45,12 +45,15 @@ def __init__(
 
         """
         # initialize training backend
+        if not backend_config:
+            backend_config = KubernetesBackendConfig()
+
         if isinstance(backend_config, KubernetesBackendConfig):
             self.backend = KubernetesBackend(backend_config)
         elif isinstance(backend_config, LocalProcessBackendConfig):
             self.backend = LocalProcessBackend(backend_config)
         else:
-            raise ValueError("Invalid backend config '{}'".format(backend_config))
+            raise ValueError(f"Invalid backend config '{backend_config}'")
 
     def list_runtimes(self) -> list[types.Runtime]:
         """List of the available runtimes.

kubeflow/trainer/backends/base.py

@@ -13,22 +13,27 @@
 # limitations under the License.
 
 import abc
+from collections.abc import Iterator
+from typing import Optional, Union
 
-from typing import Optional, Union, Iterator
 from kubeflow.trainer.constants import constants
 from kubeflow.trainer.types import types
 
 
 class ExecutionBackend(abc.ABC):
+    @abc.abstractmethod
     def list_runtimes(self) -> list[types.Runtime]:
         raise NotImplementedError()
 
+    @abc.abstractmethod
     def get_runtime(self, name: str) -> types.Runtime:
         raise NotImplementedError()
 
+    @abc.abstractmethod
     def get_runtime_packages(self, runtime: types.Runtime):
         raise NotImplementedError()
 
+    @abc.abstractmethod
     def train(
         self,
         runtime: Optional[types.Runtime] = None,
@@ -37,12 +42,15 @@ def train(
     ) -> str:
         raise NotImplementedError()
 
+    @abc.abstractmethod
     def list_jobs(self, runtime: Optional[types.Runtime] = None) -> list[types.TrainJob]:
         raise NotImplementedError()
 
+    @abc.abstractmethod
     def get_job(self, name: str) -> types.TrainJob:
         raise NotImplementedError()
 
+    @abc.abstractmethod
     def get_job_logs(
         self,
         name: str,
@@ -51,6 +59,7 @@ def get_job_logs(
     ) -> Iterator[str]:
         raise NotImplementedError()
 
+    @abc.abstractmethod
     def wait_for_job_status(
         self,
         name: str,
@@ -60,5 +69,6 @@ def wait_for_job_status(
     ) -> types.TrainJob:
         raise NotImplementedError()
 
+    @abc.abstractmethod
     def delete_job(self, name: str):
         raise NotImplementedError()

kubeflow/trainer/backends/kubernetes/backend.py

@@ -12,23 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections.abc import Iterator
 import copy
 import logging
 import multiprocessing
 import random
+import re
 import string
 import time
+from typing import Optional, Union
 import uuid
-from typing import Optional, Union, Iterator
-import re
 
-from kubeflow.trainer.constants import constants
-from kubeflow.trainer.types import types
-from kubeflow.trainer.utils import utils
 from kubeflow_trainer_api import models
 from kubernetes import client, config, watch
+
 from kubeflow.trainer.backends.base import ExecutionBackend
 from kubeflow.trainer.backends.kubernetes import types as k8s_types
+from kubeflow.trainer.constants import constants
+from kubeflow.trainer.types import types
+from kubeflow.trainer.utils import utils
 
 logger = logging.getLogger(__name__)
 
@@ -141,8 +143,8 @@ def get_runtime_packages(self, runtime: types.Runtime):
             runtime_copy.trainer.set_command(tuple(mpi_command))
 
         def print_packages():
-            import subprocess
             import shutil
+            import subprocess
             import sys
 
             # Print Python version.
@@ -353,17 +355,15 @@ def get_job_logs(
                 )
 
                 # Stream logs incrementally.
-                for logline in log_stream:
-                    yield logline  # type:ignore
+                yield from log_stream
             else:
                 logs = self.core_api.read_namespaced_pod_log(
                     name=pod_name,
                     namespace=self.namespace,
                     container=container_name,
                 )
 
-                for line in logs.splitlines():
-                    yield line
+                yield from logs.splitlines()
 
         except Exception as e:
             raise RuntimeError(
@@ -554,9 +554,12 @@ def __get_trainjob_from_crd(
         # Update the TrainJob status from its conditions.
         if trainjob_crd.status and trainjob_crd.status.conditions:
             for c in trainjob_crd.status.conditions:
-                if c.type == constants.TRAINJOB_COMPLETE and c.status == "True":
-                    trainjob.status = c.type
-                elif c.type == constants.TRAINJOB_FAILED and c.status == "True":
+                if (
+                    c.type == constants.TRAINJOB_COMPLETE
+                    and c.status == "True"
+                    or c.type == constants.TRAINJOB_FAILED
+                    and c.status == "True"
+                ):
                     trainjob.status = c.type
         else:
             # The TrainJob running status is defined when all training node (e.g. Pods) are

kubeflow/trainer/backends/kubernetes/backend_test.py

@@ -19,31 +19,31 @@
 It tests KubernetesBackend's behavior across job listing, resource creation etc
 """
 
+from dataclasses import asdict
 import datetime
 import multiprocessing
 import random
 import string
-import uuid
-from dataclasses import asdict
 from typing import Optional
 from unittest.mock import Mock, patch
+import uuid
 
-import pytest
 from kubeflow_trainer_api import models
+import pytest
 
-from kubeflow.trainer.constants import constants
-from kubeflow.trainer.types import types
-from kubeflow.trainer.utils import utils
 from kubeflow.trainer.backends.kubernetes.backend import KubernetesBackend
 from kubeflow.trainer.backends.kubernetes.types import KubernetesBackendConfig
-from kubeflow.trainer.test.common import TestCase
+from kubeflow.trainer.constants import constants
 from kubeflow.trainer.test.common import (
-    SUCCESS,
-    FAILED,
     DEFAULT_NAMESPACE,
-    TIMEOUT,
+    FAILED,
     RUNTIME,
+    SUCCESS,
+    TIMEOUT,
+    TestCase,
 )
+from kubeflow.trainer.types import types
+from kubeflow.trainer.utils import utils
 
 # In all tests runtime name is equal to the framework name.
 TORCH_RUNTIME = "torch"
@@ -788,7 +788,6 @@ def test_get_runtime_packages(kubernetes_backend, test_case):
             },
             expected_error=ValueError,
         ),
-
     ],
 )
 def test_train(kubernetes_backend, test_case):

kubeflow/trainer/backends/kubernetes/types.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from typing import Optional
+
 from kubernetes import client
 from pydantic import BaseModel