feat(py): global rate limits (#328)

* chore: regen protos * feat: admin put rate limit * fix: client type * feat: rate limit example * feat: workflow config * feat: step config * release: py 0.21.0 * fix: RateLimitDuration in init --------- Co-authored-by: gabriel ruttner <gabe@hatchet.run>
2026-01-06 08:49:53 -06:00 · 2024-04-02 10:01:22 -07:00
parent 0da379e423
commit 91b0dda46e
11 changed files with 190 additions and 27 deletions
--- a/python-sdk/examples/rate_limit/event_test.py
+++ b/python-sdk/examples/rate_limit/event_test.py
@@ -0,0 +1,25 @@
+from dotenv import load_dotenv
+from hatchet_sdk.hatchet import Hatchet
+
+load_dotenv()
+
+hatchet = Hatchet(debug=True)
+
+hatchet.client.event.push(
+    "rate_limit:create",
+    {
+        "test": "1"
+    }
+)
+hatchet.client.event.push(
+    "rate_limit:create",
+    {
+        "test": "2"
+    }
+)
+hatchet.client.event.push(
+    "rate_limit:create",
+    {
+        "test": "3"
+    }
+)
--- a/python-sdk/examples/rate_limit/worker.py
+++ b/python-sdk/examples/rate_limit/worker.py
@@ -0,0 +1,24 @@
+from hatchet_sdk import Hatchet, Context
+from dotenv import load_dotenv
+from hatchet_sdk.rate_limit import RateLimit, RateLimitDuration
+
+load_dotenv()
+
+hatchet = Hatchet(debug=True)
+
+@hatchet.workflow(on_events=["rate_limit:create"])
+class RateLimitWorkflow:
+    def __init__(self):
+        self.my_value = "test"
+
+    @hatchet.step(rate_limits = [RateLimit(key='test-limit', units=1)])
+    def step1(self, context: Context):
+        print("executed step1")
+        pass
+
+hatchet.client.admin.put_rate_limit('test-limit', 2, RateLimitDuration.MINUTE)
+
+worker = hatchet.worker('test-worker', max_runs=4)
+worker.register_workflow(RateLimitWorkflow())
+
+worker.start()
--- a/python-sdk/hatchet_sdk/init.py
+++ b/python-sdk/hatchet_sdk/init.py
@@ -4,6 +4,7 @@ from .worker import Worker
 from .client import new_client
 from .context import Context
 from .workflows_pb2 import ConcurrencyLimitStrategy
+from .workflows_pb2 import RateLimitDuration

 # import models into sdk package
 from hatchet_sdk.clients.rest.models.api_error import APIError
--- a/python-sdk/hatchet_sdk/clients/admin.py
+++ b/python-sdk/hatchet_sdk/clients/admin.py
@@ -3,7 +3,7 @@ from typing import List, Union
 import grpc
 from google.protobuf import timestamp_pb2
 from ..workflows_pb2_grpc import WorkflowServiceStub
-from ..workflows_pb2 import CreateWorkflowVersionOpts, ScheduleWorkflowRequest, TriggerWorkflowRequest, PutWorkflowRequest, TriggerWorkflowResponse
+from ..workflows_pb2 import CreateWorkflowVersionOpts, PutRateLimitRequest, RateLimitDuration, ScheduleWorkflowRequest, TriggerWorkflowRequest, PutWorkflowRequest, TriggerWorkflowResponse
 from ..loader import ClientConfig
 from ..metadata import get_metadata
 import json
@@ -36,6 +36,19 @@ class AdminClientImpl:
            )
        except grpc.RpcError as e:
            raise ValueError(f"Could not put workflow: {e}")
+        
+    def put_rate_limit(self, key: str, limit: int, duration: RateLimitDuration = RateLimitDuration.SECOND):
+        try:
+            self.client.PutRateLimit(
+                PutRateLimitRequest(
+                    key=key,
+                    limit=limit,
+                    duration=duration,
+                ),
+                metadata=get_metadata(self.token),
+            )
+        except grpc.RpcError as e:
+            raise ValueError(f"Could not put rate limit: {e}")

    def schedule_workflow(self, name: str, schedules: List[Union[datetime, timestamp_pb2.Timestamp]], input={}, options: TriggerWorkflowParentOptions = None):
        timestamp_schedules = []
--- a/python-sdk/hatchet_sdk/hatchet.py
+++ b/python-sdk/hatchet_sdk/hatchet.py
@@ -1,13 +1,17 @@
-from .client import new_client 
+from hatchet_sdk.rate_limit import RateLimit
+from .client import ClientImpl, new_client 
 from typing import List
 import asyncio
 from functools import wraps
 from .workflow import WorkflowMeta
 from .worker import Worker
 from .logger import logger
-from .workflows_pb2 import ConcurrencyLimitStrategy
+from .workflows_pb2 import ConcurrencyLimitStrategy, CreateStepRateLimit
+

 class Hatchet:
+    client: ClientImpl
+
    def __init__(self, debug=False):
        # initialize a client
        self.client = new_client()
@@ -40,7 +44,11 @@ class Hatchet:
        
        return inner

-    def step(self, name: str='', timeout: str='', parents: List[str] = [], retries: int = 0):
+    def step(self, name: str='',
+              timeout: str='', 
+              parents: List[str] = [], 
+              retries: int = 0, 
+              rate_limits: List[RateLimit] | None = None):
        def inner(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
@@ -49,10 +57,15 @@ class Hatchet:
                else:
                    return func(*args, **kwargs)

+            limits = None
+            if rate_limits:
+                limits = [CreateStepRateLimit(key=rate_limit.key, units=rate_limit.units) for rate_limit in rate_limits or []]
+
            wrapper._step_name = name or func.__name__
            wrapper._step_parents = parents
            wrapper._step_timeout = timeout
            wrapper._step_retries = retries
+            wrapper._step_rate_limits = limits
            return wrapper

        return inner
--- a/python-sdk/hatchet_sdk/rate_limit.py
+++ b/python-sdk/hatchet_sdk/rate_limit.py
@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class RateLimit:
+    key: str
+    units: int
+
+class RateLimitDuration:
+    SECOND='SECOND'
+    MINUTE='MINUTE'
+    HOUR='HOUR'
--- a/python-sdk/hatchet_sdk/workflow.py
+++ b/python-sdk/hatchet_sdk/workflow.py
@@ -51,6 +51,7 @@ class WorkflowMeta(type):
                inputs='{}',
                parents=[x for x in func._step_parents],
                retries=func._step_retries,
+                rate_limits=func._step_rate_limits,
            )
            for func_name, func in attrs.items() if hasattr(func, '_step_name')
        ]
--- a/python-sdk/hatchet_sdk/workflows_pb2.py
+++ b/python-sdk/hatchet_sdk/workflows_pb2.py
@@ -15,7 +15,7 @@ _sym_db = _symbol_database.Default()
 from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2


-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0fworkflows.proto\x1a\x1fgoogle/protobuf/timestamp.proto\">\n\x12PutWorkflowRequest\x12(\n\x04opts\x18\x01 \x01(\x0b\x32\x1a.CreateWorkflowVersionOpts\"\xbf\x02\n\x19\x43reateWorkflowVersionOpts\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x02 \x01(\t\x12\x0f\n\x07version\x18\x03 \x01(\t\x12\x16\n\x0e\x65vent_triggers\x18\x04 \x03(\t\x12\x15\n\rcron_triggers\x18\x05 \x03(\t\x12\x36\n\x12scheduled_triggers\x18\x06 \x03(\x0b\x32\x1a.google.protobuf.Timestamp\x12$\n\x04jobs\x18\x07 \x03(\x0b\x32\x16.CreateWorkflowJobOpts\x12-\n\x0b\x63oncurrency\x18\x08 \x01(\x0b\x32\x18.WorkflowConcurrencyOpts\x12\x1d\n\x10schedule_timeout\x18\t \x01(\tH\x00\x88\x01\x01\x42\x13\n\x11_schedule_timeout\"n\n\x17WorkflowConcurrencyOpts\x12\x0e\n\x06\x61\x63tion\x18\x01 \x01(\t\x12\x10\n\x08max_runs\x18\x02 \x01(\x05\x12\x31\n\x0elimit_strategy\x18\x03 \x01(\x0e\x32\x19.ConcurrencyLimitStrategy\"s\n\x15\x43reateWorkflowJobOpts\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x02 \x01(\t\x12\x0f\n\x07timeout\x18\x03 \x01(\t\x12&\n\x05steps\x18\x04 \x03(\x0b\x32\x17.CreateWorkflowStepOpts\"\x93\x01\n\x16\x43reateWorkflowStepOpts\x12\x13\n\x0breadable_id\x18\x01 \x01(\t\x12\x0e\n\x06\x61\x63tion\x18\x02 \x01(\t\x12\x0f\n\x07timeout\x18\x03 \x01(\t\x12\x0e\n\x06inputs\x18\x04 \x01(\t\x12\x0f\n\x07parents\x18\x05 \x03(\t\x12\x11\n\tuser_data\x18\x06 \x01(\t\x12\x0f\n\x07retries\x18\x07 \x01(\x05\"\x16\n\x14ListWorkflowsRequest\"\x93\x02\n\x17ScheduleWorkflowRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12-\n\tschedules\x18\x02 \x03(\x0b\x32\x1a.google.protobuf.Timestamp\x12\r\n\x05input\x18\x03 \x01(\t\x12\x16\n\tparent_id\x18\x04 \x01(\tH\x00\x88\x01\x01\x12\x1f\n\x12parent_step_run_id\x18\x05 \x01(\tH\x01\x88\x01\x01\x12\x18\n\x0b\x63hild_index\x18\x06 \x01(\x05H\x02\x88\x01\x01\x12\x16\n\tchild_key\x18\x07 \x01(\tH\x03\x88\x01\x01\x42\x0c\n\n_parent_idB\x15\n\x13_parent_step_run_idB\x0e\n\x0c_child_indexB\x0c\n\n_child_key\"\xb2\x01\n\x0fWorkflowVersion\x12\n\n\x02id\x18\x01 \x01(\t\x12.\n\ncreated_at\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12.\n\nupdated_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07version\x18\x05 \x01(\t\x12\r\n\x05order\x18\x06 \x01(\x05\x12\x13\n\x0bworkflow_id\x18\x07 \x01(\t\"?\n\x17WorkflowTriggerEventRef\x12\x11\n\tparent_id\x18\x01 \x01(\t\x12\x11\n\tevent_key\x18\x02 \x01(\t\"9\n\x16WorkflowTriggerCronRef\x12\x11\n\tparent_id\x18\x01 \x01(\t\x12\x0c\n\x04\x63ron\x18\x02 \x01(\t\"\xe3\x01\n\x16TriggerWorkflowRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05input\x18\x02 \x01(\t\x12\x16\n\tparent_id\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x1f\n\x12parent_step_run_id\x18\x04 \x01(\tH\x01\x88\x01\x01\x12\x18\n\x0b\x63hild_index\x18\x05 \x01(\x05H\x02\x88\x01\x01\x12\x16\n\tchild_key\x18\x06 \x01(\tH\x03\x88\x01\x01\x42\x0c\n\n_parent_idB\x15\n\x13_parent_step_run_idB\x0e\n\x0c_child_indexB\x0c\n\n_child_key\"2\n\x17TriggerWorkflowResponse\x12\x17\n\x0fworkflow_run_id\x18\x01 \x01(\t*l\n\x18\x43oncurrencyLimitStrategy\x12\x16\n\x12\x43\x41NCEL_IN_PROGRESS\x10\x00\x12\x0f\n\x0b\x44ROP_NEWEST\x10\x01\x12\x10\n\x0cQUEUE_NEWEST\x10\x02\x12\x15\n\x11GROUP_ROUND_ROBIN\x10\x03\x32\xcd\x01\n\x0fWorkflowService\x12\x34\n\x0bPutWorkflow\x12\x13.PutWorkflowRequest\x1a\x10.WorkflowVersion\x12>\n\x10ScheduleWorkflow\x12\x18.ScheduleWorkflowRequest\x1a\x10.WorkflowVersion\x12\x44\n\x0fTriggerWorkflow\x12\x17.TriggerWorkflowRequest\x1a\x18.TriggerWorkflowResponseBBZ@github.com/hatchet-dev/hatchet/internal/services/admin/contractsb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0fworkflows.proto\x1a\x1fgoogle/protobuf/timestamp.proto\">\n\x12PutWorkflowRequest\x12(\n\x04opts\x18\x01 \x01(\x0b\x32\x1a.CreateWorkflowVersionOpts\"\xbf\x02\n\x19\x43reateWorkflowVersionOpts\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x02 \x01(\t\x12\x0f\n\x07version\x18\x03 \x01(\t\x12\x16\n\x0e\x65vent_triggers\x18\x04 \x03(\t\x12\x15\n\rcron_triggers\x18\x05 \x03(\t\x12\x36\n\x12scheduled_triggers\x18\x06 \x03(\x0b\x32\x1a.google.protobuf.Timestamp\x12$\n\x04jobs\x18\x07 \x03(\x0b\x32\x16.CreateWorkflowJobOpts\x12-\n\x0b\x63oncurrency\x18\x08 \x01(\x0b\x32\x18.WorkflowConcurrencyOpts\x12\x1d\n\x10schedule_timeout\x18\t \x01(\tH\x00\x88\x01\x01\x42\x13\n\x11_schedule_timeout\"n\n\x17WorkflowConcurrencyOpts\x12\x0e\n\x06\x61\x63tion\x18\x01 \x01(\t\x12\x10\n\x08max_runs\x18\x02 \x01(\x05\x12\x31\n\x0elimit_strategy\x18\x03 \x01(\x0e\x32\x19.ConcurrencyLimitStrategy\"s\n\x15\x43reateWorkflowJobOpts\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x02 \x01(\t\x12\x0f\n\x07timeout\x18\x03 \x01(\t\x12&\n\x05steps\x18\x04 \x03(\x0b\x32\x17.CreateWorkflowStepOpts\"\xbe\x01\n\x16\x43reateWorkflowStepOpts\x12\x13\n\x0breadable_id\x18\x01 \x01(\t\x12\x0e\n\x06\x61\x63tion\x18\x02 \x01(\t\x12\x0f\n\x07timeout\x18\x03 \x01(\t\x12\x0e\n\x06inputs\x18\x04 \x01(\t\x12\x0f\n\x07parents\x18\x05 \x03(\t\x12\x11\n\tuser_data\x18\x06 \x01(\t\x12\x0f\n\x07retries\x18\x07 \x01(\x05\x12)\n\x0brate_limits\x18\x08 \x03(\x0b\x32\x14.CreateStepRateLimit\"1\n\x13\x43reateStepRateLimit\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05units\x18\x02 \x01(\x05\"\x16\n\x14ListWorkflowsRequest\"\x93\x02\n\x17ScheduleWorkflowRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12-\n\tschedules\x18\x02 \x03(\x0b\x32\x1a.google.protobuf.Timestamp\x12\r\n\x05input\x18\x03 \x01(\t\x12\x16\n\tparent_id\x18\x04 \x01(\tH\x00\x88\x01\x01\x12\x1f\n\x12parent_step_run_id\x18\x05 \x01(\tH\x01\x88\x01\x01\x12\x18\n\x0b\x63hild_index\x18\x06 \x01(\x05H\x02\x88\x01\x01\x12\x16\n\tchild_key\x18\x07 \x01(\tH\x03\x88\x01\x01\x42\x0c\n\n_parent_idB\x15\n\x13_parent_step_run_idB\x0e\n\x0c_child_indexB\x0c\n\n_child_key\"\xb2\x01\n\x0fWorkflowVersion\x12\n\n\x02id\x18\x01 \x01(\t\x12.\n\ncreated_at\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12.\n\nupdated_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07version\x18\x05 \x01(\t\x12\r\n\x05order\x18\x06 \x01(\x05\x12\x13\n\x0bworkflow_id\x18\x07 \x01(\t\"?\n\x17WorkflowTriggerEventRef\x12\x11\n\tparent_id\x18\x01 \x01(\t\x12\x11\n\tevent_key\x18\x02 \x01(\t\"9\n\x16WorkflowTriggerCronRef\x12\x11\n\tparent_id\x18\x01 \x01(\t\x12\x0c\n\x04\x63ron\x18\x02 \x01(\t\"\xe3\x01\n\x16TriggerWorkflowRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05input\x18\x02 \x01(\t\x12\x16\n\tparent_id\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x1f\n\x12parent_step_run_id\x18\x04 \x01(\tH\x01\x88\x01\x01\x12\x18\n\x0b\x63hild_index\x18\x05 \x01(\x05H\x02\x88\x01\x01\x12\x16\n\tchild_key\x18\x06 \x01(\tH\x03\x88\x01\x01\x42\x0c\n\n_parent_idB\x15\n\x13_parent_step_run_idB\x0e\n\x0c_child_indexB\x0c\n\n_child_key\"2\n\x17TriggerWorkflowResponse\x12\x17\n\x0fworkflow_run_id\x18\x01 \x01(\t\"W\n\x13PutRateLimitRequest\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05limit\x18\x02 \x01(\x05\x12$\n\x08\x64uration\x18\x03 \x01(\x0e\x32\x12.RateLimitDuration\"\x16\n\x14PutRateLimitResponse*l\n\x18\x43oncurrencyLimitStrategy\x12\x16\n\x12\x43\x41NCEL_IN_PROGRESS\x10\x00\x12\x0f\n\x0b\x44ROP_NEWEST\x10\x01\x12\x10\n\x0cQUEUE_NEWEST\x10\x02\x12\x15\n\x11GROUP_ROUND_ROBIN\x10\x03*5\n\x11RateLimitDuration\x12\n\n\x06SECOND\x10\x00\x12\n\n\x06MINUTE\x10\x01\x12\x08\n\x04HOUR\x10\x02\x32\x8a\x02\n\x0fWorkflowService\x12\x34\n\x0bPutWorkflow\x12\x13.PutWorkflowRequest\x1a\x10.WorkflowVersion\x12>\n\x10ScheduleWorkflow\x12\x18.ScheduleWorkflowRequest\x1a\x10.WorkflowVersion\x12\x44\n\x0fTriggerWorkflow\x12\x17.TriggerWorkflowRequest\x1a\x18.TriggerWorkflowResponse\x12;\n\x0cPutRateLimit\x12\x14.PutRateLimitRequest\x1a\x15.PutRateLimitResponseBBZ@github.com/hatchet-dev/hatchet/internal/services/admin/contractsb\x06proto3')

 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -23,8 +23,10 @@ _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'workflows_pb2', _globals)
 if _descriptor._USE_C_DESCRIPTORS == False:
  _globals['DESCRIPTOR']._options = None
  _globals['DESCRIPTOR']._serialized_options = b'Z@github.com/hatchet-dev/hatchet/internal/services/admin/contracts'
-  _globals['_CONCURRENCYLIMITSTRATEGY']._serialized_start=1706
-  _globals['_CONCURRENCYLIMITSTRATEGY']._serialized_end=1814
+  _globals['_CONCURRENCYLIMITSTRATEGY']._serialized_start=1913
+  _globals['_CONCURRENCYLIMITSTRATEGY']._serialized_end=2021
+  _globals['_RATELIMITDURATION']._serialized_start=2023
+  _globals['_RATELIMITDURATION']._serialized_end=2076
  _globals['_PUTWORKFLOWREQUEST']._serialized_start=52
  _globals['_PUTWORKFLOWREQUEST']._serialized_end=114
  _globals['_CREATEWORKFLOWVERSIONOPTS']._serialized_start=117
@@ -34,21 +36,27 @@ if _descriptor._USE_C_DESCRIPTORS == False:
  _globals['_CREATEWORKFLOWJOBOPTS']._serialized_start=550
  _globals['_CREATEWORKFLOWJOBOPTS']._serialized_end=665
  _globals['_CREATEWORKFLOWSTEPOPTS']._serialized_start=668
-  _globals['_CREATEWORKFLOWSTEPOPTS']._serialized_end=815
-  _globals['_LISTWORKFLOWSREQUEST']._serialized_start=817
-  _globals['_LISTWORKFLOWSREQUEST']._serialized_end=839
-  _globals['_SCHEDULEWORKFLOWREQUEST']._serialized_start=842
-  _globals['_SCHEDULEWORKFLOWREQUEST']._serialized_end=1117
-  _globals['_WORKFLOWVERSION']._serialized_start=1120
-  _globals['_WORKFLOWVERSION']._serialized_end=1298
-  _globals['_WORKFLOWTRIGGEREVENTREF']._serialized_start=1300
-  _globals['_WORKFLOWTRIGGEREVENTREF']._serialized_end=1363
-  _globals['_WORKFLOWTRIGGERCRONREF']._serialized_start=1365
-  _globals['_WORKFLOWTRIGGERCRONREF']._serialized_end=1422
-  _globals['_TRIGGERWORKFLOWREQUEST']._serialized_start=1425
-  _globals['_TRIGGERWORKFLOWREQUEST']._serialized_end=1652
-  _globals['_TRIGGERWORKFLOWRESPONSE']._serialized_start=1654
-  _globals['_TRIGGERWORKFLOWRESPONSE']._serialized_end=1704
-  _globals['_WORKFLOWSERVICE']._serialized_start=1817
-  _globals['_WORKFLOWSERVICE']._serialized_end=2022
+  _globals['_CREATEWORKFLOWSTEPOPTS']._serialized_end=858
+  _globals['_CREATESTEPRATELIMIT']._serialized_start=860
+  _globals['_CREATESTEPRATELIMIT']._serialized_end=909
+  _globals['_LISTWORKFLOWSREQUEST']._serialized_start=911
+  _globals['_LISTWORKFLOWSREQUEST']._serialized_end=933
+  _globals['_SCHEDULEWORKFLOWREQUEST']._serialized_start=936
+  _globals['_SCHEDULEWORKFLOWREQUEST']._serialized_end=1211
+  _globals['_WORKFLOWVERSION']._serialized_start=1214
+  _globals['_WORKFLOWVERSION']._serialized_end=1392
+  _globals['_WORKFLOWTRIGGEREVENTREF']._serialized_start=1394
+  _globals['_WORKFLOWTRIGGEREVENTREF']._serialized_end=1457
+  _globals['_WORKFLOWTRIGGERCRONREF']._serialized_start=1459
+  _globals['_WORKFLOWTRIGGERCRONREF']._serialized_end=1516
+  _globals['_TRIGGERWORKFLOWREQUEST']._serialized_start=1519
+  _globals['_TRIGGERWORKFLOWREQUEST']._serialized_end=1746
+  _globals['_TRIGGERWORKFLOWRESPONSE']._serialized_start=1748
+  _globals['_TRIGGERWORKFLOWRESPONSE']._serialized_end=1798
+  _globals['_PUTRATELIMITREQUEST']._serialized_start=1800
+  _globals['_PUTRATELIMITREQUEST']._serialized_end=1887
+  _globals['_PUTRATELIMITRESPONSE']._serialized_start=1889
+  _globals['_PUTRATELIMITRESPONSE']._serialized_end=1911
+  _globals['_WORKFLOWSERVICE']._serialized_start=2079
+  _globals['_WORKFLOWSERVICE']._serialized_end=2345
 # @@protoc_insertion_point(module_scope)
--- a/python-sdk/hatchet_sdk/workflows_pb2.pyi
+++ b/python-sdk/hatchet_sdk/workflows_pb2.pyi
@@ -13,10 +13,19 @@ class ConcurrencyLimitStrategy(int, metaclass=_enum_type_wrapper.EnumTypeWrapper
    DROP_NEWEST: _ClassVar[ConcurrencyLimitStrategy]
    QUEUE_NEWEST: _ClassVar[ConcurrencyLimitStrategy]
    GROUP_ROUND_ROBIN: _ClassVar[ConcurrencyLimitStrategy]
+
+class RateLimitDuration(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = ()
+    SECOND: _ClassVar[RateLimitDuration]
+    MINUTE: _ClassVar[RateLimitDuration]
+    HOUR: _ClassVar[RateLimitDuration]
 CANCEL_IN_PROGRESS: ConcurrencyLimitStrategy
 DROP_NEWEST: ConcurrencyLimitStrategy
 QUEUE_NEWEST: ConcurrencyLimitStrategy
 GROUP_ROUND_ROBIN: ConcurrencyLimitStrategy
+SECOND: RateLimitDuration
+MINUTE: RateLimitDuration
+HOUR: RateLimitDuration

 class PutWorkflowRequest(_message.Message):
    __slots__ = ("opts",)
@@ -69,7 +78,7 @@ class CreateWorkflowJobOpts(_message.Message):
    def __init__(self, name: _Optional[str] = ..., description: _Optional[str] = ..., timeout: _Optional[str] = ..., steps: _Optional[_Iterable[_Union[CreateWorkflowStepOpts, _Mapping]]] = ...) -> None: ...

 class CreateWorkflowStepOpts(_message.Message):
-    __slots__ = ("readable_id", "action", "timeout", "inputs", "parents", "user_data", "retries")
+    __slots__ = ("readable_id", "action", "timeout", "inputs", "parents", "user_data", "retries", "rate_limits")
    READABLE_ID_FIELD_NUMBER: _ClassVar[int]
    ACTION_FIELD_NUMBER: _ClassVar[int]
    TIMEOUT_FIELD_NUMBER: _ClassVar[int]
@@ -77,6 +86,7 @@ class CreateWorkflowStepOpts(_message.Message):
    PARENTS_FIELD_NUMBER: _ClassVar[int]
    USER_DATA_FIELD_NUMBER: _ClassVar[int]
    RETRIES_FIELD_NUMBER: _ClassVar[int]
+    RATE_LIMITS_FIELD_NUMBER: _ClassVar[int]
    readable_id: str
    action: str
    timeout: str
@@ -84,7 +94,16 @@ class CreateWorkflowStepOpts(_message.Message):
    parents: _containers.RepeatedScalarFieldContainer[str]
    user_data: str
    retries: int
-    def __init__(self, readable_id: _Optional[str] = ..., action: _Optional[str] = ..., timeout: _Optional[str] = ..., inputs: _Optional[str] = ..., parents: _Optional[_Iterable[str]] = ..., user_data: _Optional[str] = ..., retries: _Optional[int] = ...) -> None: ...
+    rate_limits: _containers.RepeatedCompositeFieldContainer[CreateStepRateLimit]
+    def __init__(self, readable_id: _Optional[str] = ..., action: _Optional[str] = ..., timeout: _Optional[str] = ..., inputs: _Optional[str] = ..., parents: _Optional[_Iterable[str]] = ..., user_data: _Optional[str] = ..., retries: _Optional[int] = ..., rate_limits: _Optional[_Iterable[_Union[CreateStepRateLimit, _Mapping]]] = ...) -> None: ...
+
+class CreateStepRateLimit(_message.Message):
+    __slots__ = ("key", "units")
+    KEY_FIELD_NUMBER: _ClassVar[int]
+    UNITS_FIELD_NUMBER: _ClassVar[int]
+    key: str
+    units: int
+    def __init__(self, key: _Optional[str] = ..., units: _Optional[int] = ...) -> None: ...

 class ListWorkflowsRequest(_message.Message):
    __slots__ = ()
@@ -161,3 +180,17 @@ class TriggerWorkflowResponse(_message.Message):
    WORKFLOW_RUN_ID_FIELD_NUMBER: _ClassVar[int]
    workflow_run_id: str
    def __init__(self, workflow_run_id: _Optional[str] = ...) -> None: ...
+
+class PutRateLimitRequest(_message.Message):
+    __slots__ = ("key", "limit", "duration")
+    KEY_FIELD_NUMBER: _ClassVar[int]
+    LIMIT_FIELD_NUMBER: _ClassVar[int]
+    DURATION_FIELD_NUMBER: _ClassVar[int]
+    key: str
+    limit: int
+    duration: RateLimitDuration
+    def __init__(self, key: _Optional[str] = ..., limit: _Optional[int] = ..., duration: _Optional[_Union[RateLimitDuration, str]] = ...) -> None: ...
+
+class PutRateLimitResponse(_message.Message):
+    __slots__ = ()
+    def __init__(self) -> None: ...
--- a/python-sdk/hatchet_sdk/workflows_pb2_grpc.py
+++ b/python-sdk/hatchet_sdk/workflows_pb2_grpc.py
@@ -30,6 +30,11 @@ class WorkflowServiceStub(object):
                request_serializer=workflows__pb2.TriggerWorkflowRequest.SerializeToString,
                response_deserializer=workflows__pb2.TriggerWorkflowResponse.FromString,
                )
+        self.PutRateLimit = channel.unary_unary(
+                '/WorkflowService/PutRateLimit',
+                request_serializer=workflows__pb2.PutRateLimitRequest.SerializeToString,
+                response_deserializer=workflows__pb2.PutRateLimitResponse.FromString,
+                )


 class WorkflowServiceServicer(object):
@@ -54,6 +59,12 @@ class WorkflowServiceServicer(object):
        context.set_details('Method not implemented!')
        raise NotImplementedError('Method not implemented!')

+    def PutRateLimit(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+

 def add_WorkflowServiceServicer_to_server(servicer, server):
    rpc_method_handlers = {
@@ -72,6 +83,11 @@ def add_WorkflowServiceServicer_to_server(servicer, server):
                    request_deserializer=workflows__pb2.TriggerWorkflowRequest.FromString,
                    response_serializer=workflows__pb2.TriggerWorkflowResponse.SerializeToString,
            ),
+            'PutRateLimit': grpc.unary_unary_rpc_method_handler(
+                    servicer.PutRateLimit,
+                    request_deserializer=workflows__pb2.PutRateLimitRequest.FromString,
+                    response_serializer=workflows__pb2.PutRateLimitResponse.SerializeToString,
+            ),
    }
    generic_handler = grpc.method_handlers_generic_handler(
            'WorkflowService', rpc_method_handlers)
@@ -133,3 +149,20 @@ class WorkflowService(object):
            workflows__pb2.TriggerWorkflowResponse.FromString,
            options, channel_credentials,
            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
+    @staticmethod
+    def PutRateLimit(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/WorkflowService/PutRateLimit',
+            workflows__pb2.PutRateLimitRequest.SerializeToString,
+            workflows__pb2.PutRateLimitResponse.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
--- a/python-sdk/pyproject.toml
+++ b/python-sdk/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "hatchet-sdk"
-version = "0.20.0"
+version = "0.21.0"
 description = ""
 authors = ["Alexander Belanger <alexander@hatchet.run>"]
 readme = "README.md"