flyte[agent] extra dependency (#3158)

* add httpx dependency for webhook agent in default flytekit Signed-off-by: Future-Outlier <[email protected]> * update Signed-off-by: Future-Outlier <[email protected]> * update Signed-off-by: Future-Outlier <[email protected]> * lint Signed-off-by: Future-Outlier <[email protected]> * update Signed-off-by: Future-Outlier <[email protected]> * update Signed-off-by: Future-Outlier <[email protected]> * update Signed-off-by: Future-Outlier <[email protected]> * update Signed-off-by: Future-Outlier <[email protected]> --------- Signed-off-by: Future-Outlier <[email protected]>
flyteorg · Feb 27, 2025 · 43a940f · 43a940f
1 parent 2e12f43
commit 43a940f
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/Dockerfile.agent b/Dockerfile.agent
@@ -8,15 +8,13 @@ ARG VERSION
 RUN apt-get update && apt-get install build-essential -y \
     && pip install uv
 
-RUN uv pip install --system --no-cache-dir -U flytekit==$VERSION \
+RUN uv pip install --system --no-cache-dir -U flytekit[agent]==$VERSION \
   flytekitplugins-airflow==$VERSION \
   flytekitplugins-bigquery==$VERSION \
   flytekitplugins-k8sdataservice==$VERSION \
   flytekitplugins-openai==$VERSION \
   flytekitplugins-snowflake==$VERSION \
   flytekitplugins-awssagemaker==$VERSION \
-  prometheus-client \
-  grpcio-health-checking \
   && apt-get clean autoclean \
   && apt-get autoremove --yes \
   && rm -rf /var/lib/{apt,dpkg,cache,log}/ \

diff --git a/flytekit/clis/sdk_in_container/serve.py b/flytekit/clis/sdk_in_container/serve.py
@@ -68,8 +68,6 @@ def agent(_: click.Context, port, prometheus_port, worker, timeout, modules):
     """
     import asyncio
 
-    from flytekit.extras.webhook import WebhookTask  # noqa: F401
-
     working_dir = os.getcwd()
     if all(os.path.realpath(path) != working_dir for path in sys.path):
         sys.path.append(working_dir)
@@ -80,7 +78,13 @@ def agent(_: click.Context, port, prometheus_port, worker, timeout, modules):
 
 
 async def _start_grpc_server(port: int, prometheus_port: int, worker: int, timeout: int):
-    from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService, SyncAgentService
+    try:
+        from flytekit.extend.backend.agent_service import AgentMetadataService, AsyncAgentService, SyncAgentService
+        from flytekit.extras.webhook import WebhookAgent  # noqa: F401 Webhook Agent Registration
+    except ImportError as e:
+        raise ImportError(
+            "Flyte agent dependencies are not installed. Please install it using `pip install flytekit[agent]`"
+        ) from e
 
     click.secho("🚀 Starting the agent service...")
     _start_http_server(prometheus_port)

diff --git a/pyproject.toml b/pyproject.toml
@@ -147,3 +147,10 @@ extend-exclude = ["tests/", "**/tests/**"]
 [tool.codespell]
 ignore-words-list = "ot,te,raison,fo,lits,assertIn"
 skip = "./docs/build,./.git,*.txt"
+
+[project.optional-dependencies]
+agent = [
+    "grpcio-health-checking<=1.68.0",
+    "httpx",
+    "prometheus-client",
+]