hatchet-sdk 0.47.0__py3-none-any.whl → 0.47.1__py3-none-any.whl

hatchet_sdk/clients/rest_client.py

@@ -9,6 +9,7 @@ from pydantic import StrictInt
 from hatchet_sdk.clients.rest.api.event_api import EventApi
 from hatchet_sdk.clients.rest.api.log_api import LogApi
 from hatchet_sdk.clients.rest.api.step_run_api import StepRunApi
+from hatchet_sdk.clients.rest.api.worker_api import WorkerApi
 from hatchet_sdk.clients.rest.api.workflow_api import WorkflowApi
 from hatchet_sdk.clients.rest.api.workflow_run_api import WorkflowRunApi
 from hatchet_sdk.clients.rest.api.workflow_runs_api import WorkflowRunsApi
@@ -83,6 +84,7 @@ class AsyncRestApi:
         self._step_run_api = None
         self._event_api = None
         self._log_api = None
+        self._worker_api: WorkerApi | None = None
 
     @property
     def api_client(self):
@@ -102,6 +104,13 @@ class AsyncRestApi:
             self._workflow_run_api = WorkflowRunApi(self.api_client)
         return self._workflow_run_api
 
+    @property
+    def worker_api(self):
+        if self._worker_api is None:
+            self._worker_api = WorkerApi(self.api_client)
+
+        return self._worker_api
+
     @property
     def step_run_api(self):
         if self._step_run_api is None:

hatchet_sdk/worker/action_listener_process.py

@@ -8,12 +8,14 @@ from typing import Any, List, Mapping, Optional
 
 import grpc
 
+from hatchet_sdk.client import Client, new_client_raw
 from hatchet_sdk.clients.dispatcher.action_listener import Action
 from hatchet_sdk.clients.dispatcher.dispatcher import (
     ActionListener,
     GetActionListenerRequest,
     new_dispatcher,
 )
+from hatchet_sdk.clients.rest.models.update_worker_request import UpdateWorkerRequest
 from hatchet_sdk.contracts.dispatcher_pb2 import (
     GROUP_KEY_EVENT_TYPE_STARTED,
     STEP_EVENT_TYPE_STARTED,
@@ -41,10 +43,6 @@ BLOCKED_THREAD_WARNING = (
 )
 
 
-def noop_handler():
-    pass
-
-
 @dataclass
 class WorkerActionListenerProcess:
     name: str
@@ -70,9 +68,15 @@ class WorkerActionListenerProcess:
         if self.debug:
             logger.setLevel(logging.DEBUG)
 
+        self.client = new_client_raw(self.config, self.debug)
+
         loop = asyncio.get_event_loop()
-        loop.add_signal_handler(signal.SIGINT, noop_handler)
-        loop.add_signal_handler(signal.SIGTERM, noop_handler)
+        loop.add_signal_handler(
+            signal.SIGINT, lambda: asyncio.create_task(self.pause_task_assignment())
+        )
+        loop.add_signal_handler(
+            signal.SIGTERM, lambda: asyncio.create_task(self.pause_task_assignment())
+        )
         loop.add_signal_handler(
             signal.SIGQUIT, lambda: asyncio.create_task(self.exit_gracefully())
         )
@@ -249,7 +253,15 @@ class WorkerActionListenerProcess:
 
         self.event_queue.put(STOP_LOOP)
 
+    async def pause_task_assignment(self) -> None:
+        await self.client.rest.aio.worker_api.worker_update(
+            worker=self.listener.worker_id,
+            update_worker_request=UpdateWorkerRequest(isPaused=True),
+        )
+
     async def exit_gracefully(self, skip_unregister=False):
+        await self.pause_task_assignment()
+
         if self.killing:
             return
 

{hatchet_sdk-0.47.0.dist-info → hatchet_sdk-0.47.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hatchet-sdk
-Version: 0.47.0
+Version: 0.47.1
 Summary: 
 Author: Alexander Belanger
 Author-email: alexander@hatchet.run

{hatchet_sdk-0.47.0.dist-info → hatchet_sdk-0.47.1.dist-info}/RECORD

@@ -186,7 +186,7 @@ hatchet_sdk/clients/rest/models/workflow_version_meta.py,sha256=TW4R7bAuYAg_LraN
 hatchet_sdk/clients/rest/models/workflow_workers_count.py,sha256=qhzqfvjjIDyARkiiLGluMIqEmqO-diHTsjlu0Doi0yg,2875
 hatchet_sdk/clients/rest/rest.py,sha256=G83F1k4g_ePzvXW95rApzvaRDQPcaxrj-JmZyq1LvGw,6606
 hatchet_sdk/clients/rest/tenacity_utils.py,sha256=gy500kHXQ-4ZrZH1biHGmavhfh0NarCJcaIfW-A9Qd0,1051
-hatchet_sdk/clients/rest_client.py,sha256=jEu9KNTaGgy64bqIIMV6r2st3EqAsyr5Vz77NipJtjU,21867
+hatchet_sdk/clients/rest_client.py,sha256=Jp4qRRsQdXzrWm3Fca6X2k6nMvECTQ87ux2jGyMgqfQ,22148
 hatchet_sdk/clients/run_event_listener.py,sha256=51WTg52_aISgYPOFPHJ21rb4IO6aEd7Ugp7FCf4HnfM,10184
 hatchet_sdk/clients/workflow_listener.py,sha256=HG7qwPaeUmsChYMhvj95IBZ6cuBDt3BDvsSi1h5PXcg,9596
 hatchet_sdk/connection.py,sha256=593aUGAj7Ouf00lcVwx_pmhdQ9NOC5ANT1Jrf8nwkHs,2165
@@ -223,7 +223,7 @@ hatchet_sdk/v2/callable.py,sha256=DqBc3VS7WAk5qU9Ef1HnyAzdgxUq1hDzj-CwQYQrvcA,69
 hatchet_sdk/v2/concurrency.py,sha256=mRsbCj3G2kwkAJjHOg67y226mYdXVmd0uka3ypdRKUQ,1434
 hatchet_sdk/v2/hatchet.py,sha256=mQZoloWuRCIEwDwXeTaJL2kBjgs9YLHy7UPWxBD5070,7282
 hatchet_sdk/worker/__init__.py,sha256=1Ze1seDuXx5yD1IfHmqGFgK5qrRazVW4ZcDVGl-Pddw,61
-hatchet_sdk/worker/action_listener_process.py,sha256=tmlzDgyHWxGl8fJWE9NKqjvhqpGi9SMmOh5dFyiVL-Q,9979
+hatchet_sdk/worker/action_listener_process.py,sha256=QFGxxsSGdAp4PzcENcGRHNmhDHgfuP3omt9-kyolEyk,10566
 hatchet_sdk/worker/runner/run_loop_manager.py,sha256=nV7fhNxJKCcrBm0ci118aszF_7AxenBkOTIe1UsBEt4,3490
 hatchet_sdk/worker/runner/runner.py,sha256=2aT5rumcisfzliI9ghKelA7LuZir-4tLzYk_N8a6S9c,16723
 hatchet_sdk/worker/runner/utils/capture_logs.py,sha256=s_BGxeykelVbusx6u31EPx3vv9c2BHkuBnYcaLW680E,2381
@@ -231,7 +231,7 @@ hatchet_sdk/worker/runner/utils/error_with_traceback.py,sha256=Iih_s8JNqrinXETFJ
 hatchet_sdk/worker/worker.py,sha256=7UPm3qTzNYSSm9QTNX6zBBMJqVA6nKFeCbAdqLLjUBs,13007
 hatchet_sdk/workflow.py,sha256=XRj5jcCQSvPQMXxBipf-ZlARua2E8Z9igRzGcQ5alkI,9375
 hatchet_sdk/workflow_run.py,sha256=-RUhI48NOr1Z8ha9entmVu3yVlHBbr0LWZKVHUfPVzY,1743
-hatchet_sdk-0.47.0.dist-info/METADATA,sha256=9xwNmXBdLQocIrJYm505zY4kMuQ2Y1bM1hAIgxEbqKY,1826
-hatchet_sdk-0.47.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-hatchet_sdk-0.47.0.dist-info/entry_points.txt,sha256=lyWn37JwpYymrqp6Awq-Var43vhDG6rfauPS6KQa8JY,1124
-hatchet_sdk-0.47.0.dist-info/RECORD,,
+hatchet_sdk-0.47.1.dist-info/METADATA,sha256=47a2WMiEGAa1OL7_iL0s2PiOwP2uxMahlkZ0KRLfBv0,1826
+hatchet_sdk-0.47.1.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+hatchet_sdk-0.47.1.dist-info/entry_points.txt,sha256=lyWn37JwpYymrqp6Awq-Var43vhDG6rfauPS6KQa8JY,1124
+hatchet_sdk-0.47.1.dist-info/RECORD,,