diff --git a/examples/genai/slack_bot.py b/examples/genai/slack_bot.py index 1c7605ca4..1204e7004 100644 --- a/examples/genai/slack_bot.py +++ b/examples/genai/slack_bot.py @@ -23,7 +23,7 @@ import asyncio import os import time -from typing import Any +from typing import Any, cast from slack_sdk import WebClient from slack_sdk.errors import SlackApiError @@ -71,7 +71,7 @@ async def post_slack_message(channel: str, text: str, thread_ts: str | None = No client = get_slack_client() try: response = client.chat_postMessage(channel=channel, text=text, thread_ts=thread_ts) - return response.data + return cast(dict[str, Any], response.data) except SlackApiError as e: print(f"Error posting message: {e.response['error']}") raise @@ -92,7 +92,7 @@ def get_thread_replies(client: WebClient, channel: str, thread_ts: str) -> list[ try: response = client.conversations_replies(channel=channel, ts=thread_ts) # First message is the parent, skip it - messages = response.data.get("messages", []) + messages = cast(dict[str, Any], response.data).get("messages", []) return messages[1:] if len(messages) > 1 else [] except SlackApiError as e: print(f"Error fetching replies: {e.response['error']}") @@ -168,7 +168,7 @@ async def post_greeting_message(channel_id: str, initial_message: str) -> dict[s client = get_slack_client() try: resp = client.chat_postMessage(channel=channel_id, text=initial_message) - initial_response = resp.data + initial_response = cast(dict[str, Any], resp.data) thread_ts = initial_response["ts"] thread_url = f"https://slack.com/app_redirect?channel={channel_id}&message_ts={thread_ts}" @@ -209,8 +209,9 @@ async def slack_echo_bot( # Get bot's user ID to filter out own messages auth_response = client.auth_test() - bot_user_id = auth_response.data["user_id"] - print(f"🤖 Bot authenticated as: {auth_response.data['user']}") + auth_data = cast(dict[str, Any], auth_response.data) + bot_user_id = auth_data["user_id"] + print(f"🤖 Bot authenticated as: {auth_data['user']}") # Post initial greeting (traced - won't repost on crash/restart) greeting_info = await post_greeting_message(channel_id, initial_message) diff --git a/examples/multi_cluster/dynamic_selector.py b/examples/multi_cluster/dynamic_selector.py new file mode 100644 index 000000000..075082cdc --- /dev/null +++ b/examples/multi_cluster/dynamic_selector.py @@ -0,0 +1,61 @@ +import asyncio +import random + +import flyte +import flyte.errors + +env = flyte.TaskEnvironment( + "dynamic-selector", +) + + +queues = ["dogfood-1", "dogfood-3"] + + +@env.task +async def worker(x: int, cluster: str) -> int: + return x + + +@flyte.trace +async def next_cluster() -> str: + return random.choice(queues) + + +async def assign(x: int, max_retries: int = 3) -> int: + """ + In case of assignment fails because of timeout, we will reassign to a different cluster. + Args: + x: int + max_retries: int + Returns: result + """ + retries = 0 + while True: + cluster = await next_cluster() + try: + return await worker.override(queue=cluster)(x, cluster) + except flyte.errors.TaskTimeoutError: + retries += 1 + if retries >= max_retries: + raise + + +@env.task +async def driver(n: int) -> int: + coros = [] + for i in range(n): + coros.append(assign(i)) + results = await asyncio.gather(*coros, return_exceptions=True) + total = 0 + for r in results: + if isinstance(r, BaseException): + raise r + total += r + return total + + +if __name__ == "__main__": + flyte.init_from_config() + r = flyte.run(driver, 10) + print(r.url)