Reformat with ruff

Signed-off-by: thomasjhuang <thomashuang63@gmail.com>

Author: thomasjhuang

flytekit/__init__.py

@@ -146,7 +146,7 @@
 from flytekit.image_spec import ImageSpec
 from flytekit.loggers import LOGGING_RICH_FMT_ENV_VAR, logger
 from flytekit.models.common import Annotations, AuthRole, Labels
-from flytekit.models.concurrency import ConcurrencyPolicy, ConcurrencyLimitBehavior
+from flytekit.models.concurrency import ConcurrencyLimitBehavior, ConcurrencyPolicy
 from flytekit.models.core.execution import WorkflowExecutionPhase
 from flytekit.models.core.types import BlobType
 from flytekit.models.documentation import Description, Documentation, SourceCode

flytekit/core/launch_plan.py

@@ -4,7 +4,6 @@
 from typing import Any, Callable, Dict, List, Optional, Type
 
 from flytekit.core import workflow as _annotated_workflow
-from flytekit.models.concurrency import ConcurrencyPolicy
 from flytekit.core.context_manager import FlyteContext, FlyteContextManager, FlyteEntities
 from flytekit.core.interface import Interface, transform_function_to_interface, transform_inputs_to_parameters
 from flytekit.core.promise import create_and_link_node, translate_inputs_to_literals
@@ -15,6 +14,7 @@
 from flytekit.models import literals as _literal_models
 from flytekit.models import schedule as _schedule_model
 from flytekit.models import security
+from flytekit.models.concurrency import ConcurrencyPolicy
 from flytekit.models.core import workflow as _workflow_model
 
 

flytekit/models/concurrency.py

@@ -22,6 +22,7 @@ class ConcurrencyPolicy(_common.FlyteIdlEntity):
     """
     Defines the concurrency policy for a launch plan.
     """
+
     def __init__(self, max_concurrency: int, behavior: ConcurrencyLimitBehavior = None):
         self._max_concurrency = max_concurrency
         self._behavior = behavior if behavior is not None else ConcurrencyLimitBehavior.SKIP

flytekit/models/launch_plan.py

@@ -8,8 +8,8 @@
 from flytekit.models import literals as _literals
 from flytekit.models import schedule as _schedule
 from flytekit.models import security
-from flytekit.models.core import identifier as _identifier
 from flytekit.models.concurrency import ConcurrencyPolicy
+from flytekit.models.core import identifier as _identifier
 
 
 class LaunchPlanMetadata(_common.FlyteIdlEntity):
@@ -310,7 +310,8 @@ def from_flyte_idl(cls, pb2):
             else None,
             overwrite_cache=pb2.overwrite_cache if pb2.overwrite_cache else None,
             concurrency_policy=ConcurrencyPolicy.from_flyte_idl(pb2.concurrency_policy)
-            if pb2.HasField("concurrency_policy") else None,
+            if pb2.HasField("concurrency_policy")
+            else None,
         )
 
 

flytekit/tools/translator.py

@@ -32,14 +32,14 @@
 from flytekit.models import launch_plan as _launch_plan_models
 from flytekit.models.admin import workflow as admin_workflow_models
 from flytekit.models.admin.workflow import WorkflowSpec
+from flytekit.models.concurrency import ConcurrencyPolicy
 from flytekit.models.core import identifier as _identifier_model
 from flytekit.models.core import workflow as _core_wf
 from flytekit.models.core import workflow as workflow_model
 from flytekit.models.core.workflow import ApproveCondition, GateNode, SignalCondition, SleepCondition, TaskNodeOverrides
 from flytekit.models.core.workflow import ArrayNode as ArrayNodeModel
 from flytekit.models.core.workflow import BranchNode as BranchNodeModel
 from flytekit.models.task import TaskSpec, TaskTemplate
-from flytekit.models.concurrency import ConcurrencyPolicy
 
 FlyteLocalEntity = Union[
     PythonTask,
@@ -362,8 +362,7 @@ def get_serializable_launch_plan(
     concurrency_policy = None
     if entity.concurrency is not None:
         concurrency_policy = ConcurrencyPolicy(
-            max_concurrency=entity.concurrency.max_concurrency,
-            behavior=entity.concurrency.behavior
+            max_concurrency=entity.concurrency.max_concurrency, behavior=entity.concurrency.behavior
         )
 
     lps = _launch_plan_models.LaunchPlanSpec(

tests/flytekit/unit/models/test_concurrency.py

@@ -5,24 +5,24 @@
 
 def test_concurrency_limit_behavior():
     assert ConcurrencyLimitBehavior.SKIP == _launch_plan_idl.ConcurrencyLimitBehavior.SKIP
-    
+
     # Test enum to string conversion
     assert ConcurrencyLimitBehavior.enum_to_string(ConcurrencyLimitBehavior.SKIP) == "SKIP"
     assert ConcurrencyLimitBehavior.enum_to_string(999) == "<UNKNOWN>"
 
 
 def test_concurrency_policy():
     policy = ConcurrencyPolicy(max_concurrency=1, behavior=ConcurrencyLimitBehavior.SKIP)
-    
+
     assert policy.max_concurrency == 1
     assert policy.behavior == ConcurrencyLimitBehavior.SKIP
-    
+
     # Test serialization to protobuf
     pb = policy.to_flyte_idl()
     assert isinstance(pb, _launch_plan_idl.ConcurrencyPolicy)
     assert pb.max == 1
     assert pb.behavior == _launch_plan_idl.ConcurrencyLimitBehavior.SKIP
-    
+
     # Test deserialization from protobuf
     policy2 = ConcurrencyPolicy.from_flyte_idl(pb)
     assert policy2.max_concurrency == 1
@@ -33,9 +33,9 @@ def test_concurrency_policy_with_different_max():
     # Test with a higher max value
     policy = ConcurrencyPolicy(max_concurrency=5, behavior=ConcurrencyLimitBehavior.SKIP)
     assert policy.max_concurrency == 5
-    
+
     pb = policy.to_flyte_idl()
     assert pb.max == 5
-    
+
     policy2 = ConcurrencyPolicy.from_flyte_idl(pb)
-    assert policy2.max_concurrency == 5
+    assert policy2.max_concurrency == 5

tests/flytekit/unit/tools/test_translator_concurrency.py

@@ -35,14 +35,14 @@ def my_wf(a: int) -> int:
         default_inputs={"a": 10},
         concurrency=concurrency_policy
     )
-    
+
     entity_mapping = OrderedDict()
     serialized_lp = get_serializable(entity_mapping, serialization_settings, lp)
-    
+
     assert serialized_lp.spec.concurrency_policy is not None
     assert serialized_lp.spec.concurrency_policy.max_concurrency == 5
     assert serialized_lp.spec.concurrency_policy.behavior == ConcurrencyLimitBehavior.SKIP
-    
+
     # Create a launch plan without concurrency policy
     lp_no_concurrency = launch_plan.LaunchPlan.get_or_create(
         workflow=my_wf,
@@ -51,6 +51,6 @@ def my_wf(a: int) -> int:
     )
     entity_mapping_2 = OrderedDict()
     serialized_lp_2 = get_serializable(entity_mapping_2, serialization_settings, lp_no_concurrency)
-    
+
     # Verify concurrency policy is not set
     assert serialized_lp_2.spec.concurrency_policy is None