kestrapy 2.1.0.dev528__tar.gz → 2.1.0.dev568__tar.gz

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kestrapy
-Version: 2.1.0.dev528
+Version: 2.1.0.dev568
 Summary: Kestra EE
 Home-page: https://github.com/kestra-io/client-sdk
 Author: Kestra

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/__init__.py

@@ -9,7 +9,7 @@
 """  # noqa: E501
 
 
-__version__ = "2.1.0.dev528"
+__version__ = "2.1.0.dev568"
 
 # Define package exports
 __all__ = [
@@ -310,6 +310,7 @@ __all__ = [
     "FlowUsageTaskUsage",
     "FlowUsageTriggerUsage",
     "FlowWithSource",
+    "FlowWorkerSelector",
     "FollowLogEvent",
     "ForwardSupportTicketRequest",
     "GateOutcome",
@@ -1012,6 +1013,7 @@ from kestrapy.models.flow_usage import FlowUsage as FlowUsage
 from kestrapy.models.flow_usage_task_usage import FlowUsageTaskUsage as FlowUsageTaskUsage
 from kestrapy.models.flow_usage_trigger_usage import FlowUsageTriggerUsage as FlowUsageTriggerUsage
 from kestrapy.models.flow_with_source import FlowWithSource as FlowWithSource
+from kestrapy.models.flow_worker_selector import FlowWorkerSelector as FlowWorkerSelector
 from kestrapy.models.follow_log_event import FollowLogEvent as FollowLogEvent
 from kestrapy.models.forward_support_ticket_request import ForwardSupportTicketRequest as ForwardSupportTicketRequest
 from kestrapy.models.gate_outcome import GateOutcome as GateOutcome

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/api/flows_api.py

@@ -342,9 +342,30 @@ class FlowsApi(BaseApi):
     # Concurrency
     # ========================================================================
 
-    def search_concurrency_limits(self, tenant: str) -> PagedResultsConcurrencyLimit:
+    def search_concurrency_limits(
+        self,
+        tenant: str,
+        page: Optional[int] = None,
+        size: Optional[int] = None,
+    ) -> PagedResultsConcurrencyLimit:
+        """Search for the flow, namespace and tenant concurrency limits of a tenant.
+
+        Each result carries the number of ``running`` and ``queued`` executions. A
+        namespace limit has an empty ``flow_id``, the tenant limit an empty
+        ``namespace`` and ``flow_id``.
+        """
         path = self._tenant_path(tenant, "concurrency-limit", "search")
-        return self._json_request("GET", path, PagedResultsConcurrencyLimit)
+        params = self._build_query_params(page=page, size=size)
+        return self._json_request("GET", path, PagedResultsConcurrencyLimit, params=params)
+
+    def concurrency_limit(self, namespace: str, flow_id: str, tenant: str) -> ConcurrencyLimit:
+        """Get the concurrency limit of a flow.
+
+        The server answers 404 until the counter exists, which happens once the
+        flow (declaring a concurrency limit) has run.
+        """
+        path = self._tenant_path(tenant, "concurrency-limit", namespace, flow_id)
+        return self._json_request("GET", path, ConcurrencyLimit)
 
     def update_concurrency_limit(
         self,
@@ -353,8 +374,17 @@ class FlowsApi(BaseApi):
         tenant: str,
         concurrency_limit: ConcurrencyLimit,
     ) -> ConcurrencyLimit:
-        path = self._tenant_path(tenant, "concurrency-limit", namespace, flow_id)
-        return self._json_request("PUT", path, ConcurrencyLimit, body=concurrency_limit)
+        """Set the running-execution counter of a concurrency limit to ``concurrency_limit.running``.
+
+        Pass an empty ``flow_id`` to target a namespace limit, and an empty
+        ``namespace`` and ``flow_id`` for the tenant limit. Only ``running`` is read
+        from ``concurrency_limit``. The limit must already exist, otherwise the
+        server answers 404.
+        """
+        # Since Kestra 2.1 the namespace and flow are sent in the body, not the path.
+        path = self._tenant_path(tenant, "concurrency-limit")
+        body = {"namespace": namespace, "flowId": flow_id, "running": concurrency_limit.running}
+        return self._json_request("PUT", path, ConcurrencyLimit, body=body)
 
     # ========================================================================
     # Validation

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/configuration.py

@@ -530,7 +530,7 @@ conf = kestrapy.Configuration(
                "OS: {env}\n"\
                "Python Version: {pyversion}\n"\
                "Version of the API: 2.0.0-SNAPSHOT\n"\
-               "SDK Package Version: 2.1.0.dev528".\
+               "SDK Package Version: 2.1.0.dev568".\
                format(env=sys.platform, pyversion=sys.version)
 
     def get_host_settings(self) -> List[HostSetting]:

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/models/__init__.py

@@ -285,6 +285,7 @@ from kestrapy.models.flow_usage import FlowUsage
 from kestrapy.models.flow_usage_task_usage import FlowUsageTaskUsage
 from kestrapy.models.flow_usage_trigger_usage import FlowUsageTriggerUsage
 from kestrapy.models.flow_with_source import FlowWithSource
+from kestrapy.models.flow_worker_selector import FlowWorkerSelector
 from kestrapy.models.follow_log_event import FollowLogEvent
 from kestrapy.models.forward_support_ticket_request import ForwardSupportTicketRequest
 from kestrapy.models.gate_outcome import GateOutcome

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/models/abstract_flow.py

@@ -19,7 +19,7 @@ from typing_extensions import Annotated
 from kestrapy.models.input_object import InputObject
 from kestrapy.models.label import Label
 from kestrapy.models.output import Output
-from kestrapy.models.worker_selector import WorkerSelector
+from kestrapy.models.flow_worker_selector import FlowWorkerSelector
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -38,7 +38,7 @@ class AbstractFlow(BaseModel):
     draft: StrictBool = Field(description="Whether this flow revision is a draft. Draft revisions are skipped when an execution starts without an explicit revision (webhooks, schedules, subflows, manual triggers). Executions can still target a draft by passing the revision explicitly.")
     labels: Optional[List[Label]] = None
     variables: Optional[Dict[str, Any]] = None
-    worker_selector: Optional[WorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this flow.", alias="workerSelector")
+    worker_selector: Optional[FlowWorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this flow.", alias="workerSelector")
     deleted: StrictBool
     additional_properties: Dict[str, Any] = {}
     __properties: ClassVar[List[str]] = ["id", "namespace", "revision", "updated", "description", "inputs", "outputs", "disabled", "draft", "labels", "variables", "workerSelector", "deleted"]
@@ -150,7 +150,7 @@ class AbstractFlow(BaseModel):
             "draft": obj.get("draft"),
             "labels": [Label.from_dict(_item) for _item in obj["labels"]] if obj.get("labels") is not None else None,
             "variables": obj.get("variables"),
-            "workerSelector": WorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
+            "workerSelector": FlowWorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
             "deleted": obj.get("deleted")
         })
         # store additional fields in additional_properties

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/models/abstract_trigger.py

@@ -19,7 +19,7 @@ from kestrapy.models.assets_declaration import AssetsDeclaration
 from kestrapy.models.level import Level
 from kestrapy.models.state_type import StateType
 from kestrapy.models.the_labels_to_pass_to_the_execution_created import TheLabelsToPassToTheExecutionCreated
-from kestrapy.models.worker_selector import WorkerSelector
+from kestrapy.models.flow_worker_selector import FlowWorkerSelector
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -36,7 +36,7 @@ class AbstractTrigger(BaseModel):
     # every real trigger payload (dropping plugin props like `cron` downstream).
     when: Optional[StrictStr] = Field(default=None, description="A Pebble expression evaluated at trigger time. The trigger fires only when the expression evaluates to a truthy value (`true`, a non-empty string, a non-zero number). Use this to gate trigger execution on dynamic runtime values such as execution labels, flow variables, or environment conditions.")
     disabled: Optional[StrictBool] = False
-    worker_selector: Optional[WorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this trigger.", alias="workerSelector")
+    worker_selector: Optional[FlowWorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this trigger.", alias="workerSelector")
     policy_refs: Optional[List[StrictStr]] = Field(default=None, description="Identifiers of `enforcement: REFERENCE` governance policies to attach to this trigger and everything nested under it (Enterprise Edition; ignored in the open-source edition).", alias="policyRefs")
     log_level: Optional[Level] = Field(default=None, alias="logLevel")
     labels: Optional[TheLabelsToPassToTheExecutionCreated] = None
@@ -135,7 +135,7 @@ class AbstractTrigger(BaseModel):
             "description": obj.get("description"),
             "when": obj.get("when"),
             "disabled": obj.get("disabled") if obj.get("disabled") is not None else False,
-            "workerSelector": WorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
+            "workerSelector": FlowWorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
             "policyRefs": obj.get("policyRefs"),
             "logLevel": obj.get("logLevel"),
             "labels": TheLabelsToPassToTheExecutionCreated.from_dict(obj["labels"]) if obj.get("labels") is not None else None,

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/models/concurrency_limit.py

@@ -25,8 +25,9 @@ class ConcurrencyLimit(BaseModel):
     namespace: StrictStr
     flow_id: StrictStr = Field(alias="flowId")
     running: Optional[StrictInt] = None
+    queued: Optional[StrictInt] = Field(default=None, description="Number of queued executions; only set on the results of search_concurrency_limits.")
     additional_properties: Dict[str, Any] = {}
-    __properties: ClassVar[List[str]] = ["tenantId", "namespace", "flowId", "running"]
+    __properties: ClassVar[List[str]] = ["tenantId", "namespace", "flowId", "running", "queued"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -89,7 +90,8 @@ class ConcurrencyLimit(BaseModel):
             "tenantId": obj.get("tenantId"),
             "namespace": obj.get("namespace"),
             "flowId": obj.get("flowId"),
-            "running": obj.get("running")
+            "running": obj.get("running"),
+            "queued": obj.get("queued")
         })
         # store additional fields in additional_properties
         for _key in obj.keys():

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/models/flow.py

@@ -25,7 +25,7 @@ from kestrapy.models.output import Output
 from kestrapy.models.quota import Quota
 from kestrapy.models.sla import SLA
 from kestrapy.models.task import Task
-from kestrapy.models.worker_selector import WorkerSelector
+from kestrapy.models.flow_worker_selector import FlowWorkerSelector
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -44,7 +44,7 @@ class Flow(BaseModel):
     draft: StrictBool = Field(description="Whether this flow revision is a draft. Draft revisions are skipped when an execution starts without an explicit revision (webhooks, schedules, subflows, manual triggers). Executions can still target a draft by passing the revision explicitly.")
     labels: Optional[List[Label]] = None
     variables: Optional[Dict[str, Any]] = None
-    worker_selector: Optional[WorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this flow.", alias="workerSelector")
+    worker_selector: Optional[FlowWorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this flow.", alias="workerSelector")
     deleted: StrictBool
     var_finally: Optional[List[Task]] = Field(default=None, alias="finally")
     tasks: Annotated[List[Task], Field(min_length=1)]
@@ -226,7 +226,7 @@ class Flow(BaseModel):
             "draft": obj.get("draft"),
             "labels": [Label.from_dict(_item) for _item in obj["labels"]] if obj.get("labels") is not None else None,
             "variables": obj.get("variables"),
-            "workerSelector": WorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
+            "workerSelector": FlowWorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
             "deleted": obj.get("deleted"),
             "finally": [Task.from_dict(_item) for _item in obj["finally"]] if obj.get("finally") is not None else None,
             "tasks": [Task.from_dict(_item) for _item in obj["tasks"]] if obj.get("tasks") is not None else None,

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/models/flow_for_execution.py

@@ -21,7 +21,7 @@ from kestrapy.models.input_object import InputObject
 from kestrapy.models.map_object_object import MapObjectObject
 from kestrapy.models.output import Output
 from kestrapy.models.task_for_execution import TaskForExecution
-from kestrapy.models.worker_selector import WorkerSelector
+from kestrapy.models.flow_worker_selector import FlowWorkerSelector
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -40,7 +40,7 @@ class FlowForExecution(BaseModel):
     draft: StrictBool = Field(description="Whether this flow revision is a draft. Draft revisions are skipped when an execution starts without an explicit revision (webhooks, schedules, subflows, manual triggers). Executions can still target a draft by passing the revision explicitly.")
     labels: Optional[MapObjectObject] = Field(default=None, description="Labels as a list of Label (key/value pairs) or as a map of string to string.")
     variables: Optional[Dict[str, Any]] = None
-    worker_selector: Optional[WorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this flow.", alias="workerSelector")
+    worker_selector: Optional[FlowWorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this flow.", alias="workerSelector")
     deleted: StrictBool
     tasks: Annotated[List[TaskForExecution], Field(min_length=1)]
     errors: Optional[List[TaskForExecution]] = None
@@ -188,7 +188,7 @@ class FlowForExecution(BaseModel):
             "draft": obj.get("draft"),
             "labels": MapObjectObject.from_dict(obj["labels"]) if obj.get("labels") is not None else None,
             "variables": obj.get("variables"),
-            "workerSelector": WorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
+            "workerSelector": FlowWorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
             "deleted": obj.get("deleted"),
             "tasks": [TaskForExecution.from_dict(_item) for _item in obj["tasks"]] if obj.get("tasks") is not None else None,
             "errors": [TaskForExecution.from_dict(_item) for _item in obj["errors"]] if obj.get("errors") is not None else None,

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/models/flow_interface.py

@@ -20,7 +20,7 @@ from kestrapy.models.label import Label
 from kestrapy.models.output import Output
 from kestrapy.models.quota import Quota
 from kestrapy.models.sla import SLA
-from kestrapy.models.worker_selector import WorkerSelector
+from kestrapy.models.flow_worker_selector import FlowWorkerSelector
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -40,7 +40,7 @@ class FlowInterface(BaseModel):
     inputs: Optional[List[InputObject]] = None
     outputs: Optional[List[Output]] = None
     variables: Optional[Dict[str, Dict[str, Any]]] = None
-    worker_selector: Optional[WorkerSelector] = Field(default=None, alias="workerSelector")
+    worker_selector: Optional[FlowWorkerSelector] = Field(default=None, alias="workerSelector")
     concurrency: Optional[Concurrency] = None
     quotas: Optional[List[Quota]] = None
     sla: Optional[List[SLA]] = None
@@ -159,7 +159,7 @@ class FlowInterface(BaseModel):
             "inputs": [InputObject.from_dict(_item) for _item in obj["inputs"]] if obj.get("inputs") is not None else None,
             "outputs": [Output.from_dict(_item) for _item in obj["outputs"]] if obj.get("outputs") is not None else None,
             "variables": obj.get("variables"),
-            "workerSelector": WorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
+            "workerSelector": FlowWorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
             "concurrency": Concurrency.from_dict(obj["concurrency"]) if obj.get("concurrency") is not None else None,
             "quotas": [Quota.from_dict(_item) for _item in obj["quotas"]] if obj.get("quotas") is not None else None,
             "sla": [SLA.from_dict(_item) for _item in obj["sla"]] if obj.get("sla") is not None else None,

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/models/flow_with_source.py

@@ -25,7 +25,7 @@ from kestrapy.models.output import Output
 from kestrapy.models.quota import Quota
 from kestrapy.models.sla import SLA
 from kestrapy.models.task import Task
-from kestrapy.models.worker_selector import WorkerSelector
+from kestrapy.models.flow_worker_selector import FlowWorkerSelector
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -44,7 +44,7 @@ class FlowWithSource(BaseModel):
     draft: StrictBool = Field(description="Whether this flow revision is a draft. Draft revisions are skipped when an execution starts without an explicit revision (webhooks, schedules, subflows, manual triggers). Executions can still target a draft by passing the revision explicitly.")
     labels: Optional[List[Label]] = None
     variables: Optional[Dict[str, Any]] = None
-    worker_selector: Optional[WorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this flow.", alias="workerSelector")
+    worker_selector: Optional[FlowWorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this flow.", alias="workerSelector")
     deleted: StrictBool
     var_finally: Optional[List[Task]] = Field(default=None, alias="finally")
     tasks: Annotated[List[Task], Field(min_length=1)]
@@ -227,7 +227,7 @@ class FlowWithSource(BaseModel):
             "draft": obj.get("draft"),
             "labels": [Label.from_dict(_item) for _item in obj["labels"]] if obj.get("labels") is not None else None,
             "variables": obj.get("variables"),
-            "workerSelector": WorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
+            "workerSelector": FlowWorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
             "deleted": obj.get("deleted"),
             "finally": [Task.from_dict(_item) for _item in obj["finally"]] if obj.get("finally") is not None else None,
             "tasks": [Task.from_dict(_item) for _item in obj["tasks"]] if obj.get("tasks") is not None else None,

kestrapy-2.1.0.dev568/kestrapy/models/flow_worker_selector.py

@@ -0,0 +1,104 @@
+# coding: utf-8
+
+"""
+    Kestra EE
+
+    All API operations, except for Instance-owner-only endpoints, require a tenant identifier in the HTTP path.<br/> Endpoints designated as Instance-owner-only are not tenant-scoped.
+"""  # noqa: E501
+
+
+from __future__ import annotations
+import pprint
+import regex as re
+import json
+
+from pydantic import BaseModel, ConfigDict, Field, StrictStr
+from typing import Any, ClassVar, Dict, List, Optional
+from kestrapy.models.worker_queue_fallback import WorkerQueueFallback
+from kestrapy.models.worker_selector_match import WorkerSelectorMatch
+from typing import Optional, Set
+from typing_extensions import Self
+
+class FlowWorkerSelector(BaseModel):
+    """
+    Routing requirements for a Flow, Task or Trigger. Unlike WorkerSelector (namespace and
+    tenant defaults), ``tags`` is a Kestra property: a JSON array literal such as ``["gpu"]``
+    or a Pebble expression rendered when the job is routed. The server returns it as a
+    string in both cases.
+    """ # noqa: E501
+    tags: Optional[StrictStr] = Field(default=None, description="Required tags used to route to a matching Worker Queue, as a JSON array literal (e.g. [\"gpu\"]) or a Pebble expression rendering to the tags. Each tag is an RFC 1123 label, max 63 chars per tag, max 20 tags; a tag rendering to an empty value is ignored.")
+    match: Optional[WorkerSelectorMatch] = Field(default=None, description="How selector tags are matched against a Worker Queue's tag set. ALL: queue tags must be a superset of the selector tags. ANY: queue tags must intersect the selector tags. Defaults to ALL when null.")
+    fallback: Optional[WorkerQueueFallback] = Field(default=None, description="Strategy when no worker is available. Defaults to FAIL when null.")
+    additional_properties: Dict[str, Any] = {}
+    __properties: ClassVar[List[str]] = ["tags", "match", "fallback"]
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+        validate_assignment=True,
+        protected_namespaces=(),
+    )
+
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.model_dump(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        # TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> Optional[Self]:
+        """Create an instance of FlowWorkerSelector from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Return the dictionary representation of the model using alias.
+
+        This has the following differences from calling pydantic's
+        `self.model_dump(by_alias=True)`:
+
+        * `None` is only added to the output dict for nullable fields that
+          were set at model initialization. Other fields with value `None`
+          are ignored.
+        * Fields in `self.additional_properties` are added to the output dict.
+        """
+        excluded_fields: Set[str] = set([
+            "additional_properties",
+        ])
+
+        _dict = self.model_dump(
+            by_alias=True,
+            exclude=excluded_fields,
+            exclude_none=True,
+        )
+        # puts key-value pairs in additional_properties in the top level
+        if self.additional_properties is not None:
+            for _key, _value in self.additional_properties.items():
+                _dict[_key] = _value
+
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
+        """Create an instance of FlowWorkerSelector from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return cls.model_validate(obj)
+
+        _obj = cls.model_validate({
+            "tags": obj.get("tags"),
+            "match": obj.get("match"),
+            "fallback": obj.get("fallback")
+        })
+        # store additional fields in additional_properties
+        for _key in obj.keys():
+            if _key not in cls.__properties:
+                _obj.additional_properties[_key] = obj.get(_key)
+
+        return _obj
+
+

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy/models/task.py

@@ -19,7 +19,7 @@ from kestrapy.models.assets_declaration import AssetsDeclaration
 from kestrapy.models.cache import Cache
 from kestrapy.models.level import Level
 from kestrapy.models.property_duration import PropertyDuration
-from kestrapy.models.worker_selector import WorkerSelector
+from kestrapy.models.flow_worker_selector import FlowWorkerSelector
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -34,7 +34,7 @@ class Task(BaseModel):
     retry: Optional[Dict[str, Any]] = Field(default=None, description="Retry policy applied when the task fails.")
     timeout: Optional[PropertyDuration] = None
     disabled: Optional[StrictBool] = None
-    worker_selector: Optional[WorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this task.", alias="workerSelector")
+    worker_selector: Optional[FlowWorkerSelector] = Field(default=None, description="Routing requirements (tags + fallback) for this task.", alias="workerSelector")
     policy_refs: Optional[List[StrictStr]] = Field(default=None, description="Identifiers of `enforcement: REFERENCE` governance policies to attach to this task and everything nested under it (Enterprise Edition; ignored in the open-source edition).", alias="policyRefs")
     log_level: Optional[Level] = Field(default=None, alias="logLevel")
     allow_failure: Optional[StrictBool] = Field(default=None, alias="allowFailure")
@@ -142,7 +142,7 @@ class Task(BaseModel):
             "retry": obj.get("retry"),
             "timeout": PropertyDuration.from_dict(obj["timeout"]) if obj.get("timeout") is not None else None,
             "disabled": obj.get("disabled"),
-            "workerSelector": WorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
+            "workerSelector": FlowWorkerSelector.from_dict(obj["workerSelector"]) if obj.get("workerSelector") is not None else None,
             "policyRefs": obj.get("policyRefs"),
             "logLevel": obj.get("logLevel"),
             "allowFailure": obj.get("allowFailure"),

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kestrapy
-Version: 2.1.0.dev528
+Version: 2.1.0.dev568
 Summary: Kestra EE
 Home-page: https://github.com/kestra-io/client-sdk
 Author: Kestra

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/kestrapy.egg-info/SOURCES.txt

@@ -341,6 +341,7 @@ kestrapy/models/flow_usage.py
 kestrapy/models/flow_usage_task_usage.py
 kestrapy/models/flow_usage_trigger_usage.py
 kestrapy/models/flow_with_source.py
+kestrapy/models/flow_worker_selector.py
 kestrapy/models/follow_log_event.py
 kestrapy/models/forward_support_ticket_request.py
 kestrapy/models/gate_outcome.py
@@ -784,7 +785,9 @@ testApis/test_test_suites_api.py
 testApis/test_triggers_api.py
 testApis/test_users_api.py
 tests/test_asset_relation_requests.py
+tests/test_flow_worker_selector.py
 tests/test_flow_yaml.py
+tests/test_kv_api.py
 tests/test_model_types.py
 tests/test_query_filter_golden.py
 tests/test_replay_with_inputs_empty_body.py

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/pyproject.toml

@@ -1,6 +1,6 @@
     [project]
 name = "kestrapy"
-version = "2.1.0.dev528"
+version = "2.1.0.dev568"
 description = "Kestra EE"
     authors = [
     {name = "Kestra",email = "hello@kestra.io"},

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/setup.py

@@ -19,7 +19,7 @@ from setuptools import setup, find_packages  # noqa: H301
 # prerequisite: setuptools
 # http://pypi.python.org/pypi/setuptools
 NAME = "kestrapy"
-VERSION = "2.1.0.dev528"
+VERSION = "2.1.0.dev568"
 PYTHON_REQUIRES = ">= 3.9"
 REQUIRES = [
 "requests (>= 2.32.5)",

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/testApis/test_flows_api.py

@@ -12,6 +12,8 @@ from kestrapy import (
     QueryFilter,
     QueryFilterField,
     QueryFilterOp,
+    StateType,
+    WorkerSelectorMatch,
 )
 from test_helpers import (
     TENANT,
@@ -21,9 +23,9 @@ from test_helpers import (
     log_flow_yaml_with_description,
     log_flow_yaml_with_labels,
     two_task_flow_yaml,
-    concurrency_flow_yaml,
     invalid_task_flow_yaml,
     create_flow,
+    create_execution,
     create_log_flow,
     simple_flow_fixture,
     complete_flow_body,
@@ -52,6 +54,31 @@ class TestCreateFlow:
         assert result.deleted is False
         assert result.tasks is not None and len(result.tasks) > 0
 
+    def test_worker_selector(self, client):
+        fid = random_id()
+        ns = random_namespace()
+        body = f"""\
+id: {fid}
+namespace: {ns}
+workerSelector:
+  tags:
+    - gpu
+  match: ALL
+tasks:
+  - id: h
+    type: io.kestra.plugin.core.log.Log
+    message: hi
+"""
+        created = client.flows.create_flow(TENANT, body)
+
+        # tags is a Pebble-capable property on flows: a literal list comes back as its JSON text
+        assert created.worker_selector.tags == '["gpu"]'
+        assert created.worker_selector.match == WorkerSelectorMatch.ALL
+        assert created.draft is False
+
+        fetched = client.flows.flow(ns, fid, TENANT)
+        assert fetched.worker_selector.tags == '["gpu"]'
+
     @pytest.mark.xfail(
         reason="kestra-ee 2.0 validates Schedule-trigger inputs more strictly and "
         "rejects the shared complete flow (422 / 'Missing inputs for Schedule "
@@ -1040,38 +1067,77 @@ class TestTaskFromFlow:
 # ========================================================================
 
 
-class TestConcurrency:
-    def test_search_concurrency_limits(self, client):
-        result = client.flows.search_concurrency_limits(TENANT)
+def _sleep_concurrency_flow_yaml(flow_id, namespace):
+    return f"""\
+id: {flow_id}
+namespace: {namespace}
+concurrency:
+  behavior: QUEUE
+  limit: 1
+tasks:
+  - id: wait
+    type: io.kestra.plugin.core.flow.Sleep
+    duration: PT10S
+"""
+
+
+def _start_execution_in_state(client, namespace, flow_id, expected, timeout=30):
+    execution_id = create_execution(client, namespace, flow_id).id
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        execution = client.executions.execution(execution_id, TENANT)
+        if execution.state is not None and execution.state.current == expected:
+            return execution_id
+        time.sleep(0.2)
+    raise TimeoutError(f"Execution {execution_id} did not reach {expected} within {timeout}s")
 
-        assert result is not None
-        assert result.results is not None
 
-    @pytest.mark.skip(reason="Concurrency limit PUT returns 404 on this kestra-ee image")
-    def test_update_concurrency_limit_set(self, client):
+# A concurrency-limit counter only exists once a flow declaring `concurrency` has run,
+# so these tests start one execution (running) and, for search, a second one queued
+# behind limit 1.
+class TestConcurrency:
+    def test_search_concurrency_limits_running_and_queued(self, client):
+        fid = random_id()
+        ns = random_namespace()
+        create_flow(client, _sleep_concurrency_flow_yaml(fid, ns))
+        _start_execution_in_state(client, ns, fid, StateType.RUNNING)
+        _start_execution_in_state(client, ns, fid, StateType.QUEUED)
+
+        found = None
+        deadline = time.time() + 10
+        while found is None and time.time() < deadline:
+            result = client.flows.search_concurrency_limits(TENANT, page=1, size=1000)
+            found = next((l for l in result.results if l.namespace == ns and l.flow_id == fid), None)
+            if found is None:
+                time.sleep(0.2)
+
+        assert found is not None
+        assert found.running == 1
+        assert found.queued == 1
+
+    def test_update_concurrency_limit_resets_running(self, client):
         fid = random_id()
         ns = random_namespace()
-        create_flow(client, concurrency_flow_yaml(fid, ns, 1))
+        create_flow(client, _sleep_concurrency_flow_yaml(fid, ns))
+        _start_execution_in_state(client, ns, fid, StateType.RUNNING)
 
-        limit = ConcurrencyLimit(tenant_id=TENANT, namespace=ns, flow_id=fid, running=5)
+        assert client.flows.concurrency_limit(ns, fid, TENANT).running == 1
 
+        limit = ConcurrencyLimit(tenant_id=TENANT, namespace=ns, flow_id=fid, running=0)
         result = client.flows.update_concurrency_limit(ns, fid, TENANT, limit)
 
-        assert result is not None
         assert result.namespace == ns
         assert result.flow_id == fid
+        assert result.running == 0
 
-    @pytest.mark.skip(reason="Concurrency limit PUT returns 404 on this kestra-ee image")
-    def test_update_concurrency_limit_update(self, client):
+    def test_update_concurrency_limit_unknown_is_404(self, client):
         fid = random_id()
         ns = random_namespace()
-        create_flow(client, concurrency_flow_yaml(fid, ns, 3))
+        limit = ConcurrencyLimit(tenant_id=TENANT, namespace=ns, flow_id=fid, running=0)
 
-        limit = ConcurrencyLimit(tenant_id=TENANT, namespace=ns, flow_id=fid, running=10)
-        result = client.flows.update_concurrency_limit(ns, fid, TENANT, limit)
-
-        assert result is not None
-        assert result.flow_id == fid
+        with pytest.raises(ApiException) as exc:
+            client.flows.update_concurrency_limit(ns, fid, TENANT, limit)
+        assert exc.value.status == 404
 
 
 # ========================================================================

{kestrapy-2.1.0.dev528 → kestrapy-2.1.0.dev568}/testApis/test_instance_family_api.py

@@ -231,7 +231,8 @@ def test_kill_switch_create_search_delete_roundtrip(client):
     ks_name = "ks-" + random_id()[:8]
     with _tolerate_gating("create_kill_switch"):
         created = client.kill_switches.create_kill_switch(
-            KillSwitch(name=ks_name, startDate=start, evaluationType=EvaluationType.KILL)
+            # Kestra requires enabled (422 when null); False keeps the switch inert for other tests.
+            KillSwitch(name=ks_name, startDate=start, evaluationType=EvaluationType.KILL, enabled=False)
         )
     ks_id = created.id
     try: