Improved Formatting

pytorch · Mar 12, 2024 · b43498f · b43498f
1 parent d6ccf07
commit b43498f
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 22 deletions.
diff --git a/torchx/schedulers/kueue_scheduler.py b/torchx/schedulers/kueue_scheduler.py
@@ -63,6 +63,7 @@
 from torchx.workspace.docker_workspace import DockerWorkspaceMixin
 from typing_extensions import TypedDict
 from torchx.util.role_to_pod import role_to_pod
+
 if TYPE_CHECKING:
     from docker import DockerClient
     from kubernetes.client import ApiClient, BatchV1Api, CoreV1Api, CustomObjectsApi
@@ -134,6 +135,7 @@ def sanitize_for_serialization(obj: object) -> object:
     api = client.ApiClient()
     return api.sanitize_for_serialization(obj)
 
+
 def app_to_resource(
     app: AppDef,
     service_account: Optional[str],

diff --git a/torchx/schedulers/test/kueue_scheduler_test.py b/torchx/schedulers/test/kueue_scheduler_test.py
@@ -100,9 +100,7 @@ def test_create_scheduler(self) -> None:
     def test_app_to_resource_resolved_macros(self) -> None:
         app = _test_app()
         unique_app_name = "app-name-42"
-        with patch(
-            "torchx.schedulers.kueue_scheduler.make_unique"
-        ) as make_unique_ctx:
+        with patch("torchx.schedulers.kueue_scheduler.make_unique") as make_unique_ctx:
             make_unique_ctx.return_value = unique_app_name
             resource = app_to_resource(
                 app, service_account=None, local_queue="default-kueue"
@@ -238,14 +236,10 @@ def test_role_to_pod(self) -> None:
         )
 
     def test_submit_dryrun(self) -> None:
-        cfg = KueueOpts(
-            {"namespace": "testnamespace", "local_queue": "default-kueue"}
-        )
+        cfg = KueueOpts({"namespace": "testnamespace", "local_queue": "default-kueue"})
         scheduler = create_scheduler("test")
         app = _test_app()
-        with patch(
-            "torchx.schedulers.kueue_scheduler.make_unique"
-        ) as make_unique_ctx:
+        with patch("torchx.schedulers.kueue_scheduler.make_unique") as make_unique_ctx:
             make_unique_ctx.return_value = "app-name-42"
             info = scheduler.submit_dryrun(app, cfg)
 
@@ -477,12 +471,8 @@ def test_rank0_env(self) -> None:
 
         scheduler = create_scheduler("test")
         app = _test_app(num_replicas=2)
-        cfg = KueueOpts(
-            {"namespace": "testnamespace", "local_queue": "default-kueue"}
-        )
-        with patch(
-            "torchx.schedulers.kueue_scheduler.make_unique"
-        ) as make_unique_ctx:
+        cfg = KueueOpts({"namespace": "testnamespace", "local_queue": "default-kueue"})
+        with patch("torchx.schedulers.kueue_scheduler.make_unique") as make_unique_ctx:
             make_unique_ctx.return_value = "app-name-42"
             info = scheduler.submit_dryrun(app, cfg)
 
@@ -499,9 +489,7 @@ def test_submit_dryrun_patch(self) -> None:
         cfg = KueueOpts(
             {"image_repo": "example.com/some/repo", "local_queue": "default-kueue"}
         )
-        with patch(
-            "torchx.schedulers.kueue_scheduler.make_unique"
-        ) as make_unique_ctx:
+        with patch("torchx.schedulers.kueue_scheduler.make_unique") as make_unique_ctx:
             make_unique_ctx.return_value = "app-name-42"
             info = scheduler.submit_dryrun(app, cfg)
 
@@ -1013,9 +1001,7 @@ def test_describe(self) -> None:
     def test_dryrun(self) -> None:
         scheduler = kueue_scheduler.create_scheduler("foo")
         app = _test_app()
-        cfg = KueueOpts(
-            {"namespace": "testnamespace", "local_queue": "default-kueue"}
-        )
+        cfg = KueueOpts({"namespace": "testnamespace", "local_queue": "default-kueue"})
 
         with self.assertRaises(ModuleNotFoundError):
             scheduler.submit_dryrun(app, cfg)
diff --git a/torchx/util/role_to_pod.py b/torchx/util/role_to_pod.py
@@ -18,6 +18,7 @@
 LABEL_INSTANCE_TYPE = "node.kubernetes.io/instance-type"
 ANNOTATION_ISTIO_SIDECAR = "sidecar.istio.io/inject"
 
+
 def role_to_pod(name: str, role: Role, service_account: Optional[str]) -> "V1Pod":
     from kubernetes.client.models import (  # noqa: F811 redefinition of unused
         V1Container,
@@ -179,4 +180,4 @@ def role_to_pod(name: str, role: Role, service_account: Optional[str]) -> "V1Pod
             },
             labels={},
         ),
-    )
+    )