Repository: run-llama/llama_deploy

Issue:
Timeout for nested network workflows fixed at 120 seconds
The NetworkWorkflow which is used to allow workflows to use other services/workflows in Llama deploy uses the AsyncLlamaDeployClient this client has a default timeout of 120 seconds which is used both as a request timeout to the llama deploy control plane, and as the timeout for session.run(...).

There is currently no way to override this so even when talking to a downstream service with a longer timeout it caps the request at 120 seconds.

When issuing a task to llama deploy perhaps the Client should pick up the timeout from the service i.e. if WorkflowA's configured timeout is 60 seconds, when issuing a request to that service the default wait_for should be 60 seconds?

Alternatively (and I think this might influence the Workflow API) workflow.run(...) should take a timeout as an override for the default timeout in the workflow. That way the session.run(..) API can also take a timeout and remain compatible. So then different calls to nested workflows at different points in a parent workflow could specify different timeouts.




Discussion:
I think re-using the same timeout in NetworkWorkflow to also set the timeout for the inner client is a fair assumption to make, opening a PR with this approach.



Patch:
diff --git a/llama_deploy/client/base.py b/llama_deploy/client/base.py
index b5ac08ee..ebe7fe3d 100644
--- a/llama_deploy/client/base.py
+++ b/llama_deploy/client/base.py
@@ -16,7 +16,7 @@ class _BaseClient(BaseSettings):
     api_server_url: str = "http://localhost:4501"
     control_plane_url: str = "http://localhost:8000"
     disable_ssl: bool = False
-    timeout: float = 120.0
+    timeout: float | None = 120.0
     poll_interval: float = 0.5
 
     async def request(
diff --git a/llama_deploy/client/models/core.py b/llama_deploy/client/models/core.py
index 4a32e0f2..90e8c208 100644
--- a/llama_deploy/client/models/core.py
+++ b/llama_deploy/client/models/core.py
@@ -146,7 +146,10 @@ async def get_task_result_stream(
             except httpx.HTTPStatusError as e:
                 if e.response.status_code != 404:
                     raise  # Re-raise if it's not a 404 error
-                if time.time() - start_time < self.client.timeout:
+                if (
+                    self.client.timeout is None  # means no timeout, always poll
+                    or time.time() - start_time < self.client.timeout
+                ):
                     await asyncio.sleep(self.client.poll_interval)
                 else:
                     raise TimeoutError(
diff --git a/llama_deploy/deploy/deploy.py b/llama_deploy/deploy/deploy.py
index 1497952c..650b6151 100644
--- a/llama_deploy/deploy/deploy.py
+++ b/llama_deploy/deploy/deploy.py
@@ -166,7 +166,7 @@ async def deploy_workflow(
 
     # override the service manager, while maintaining dict of existing services
     workflow._service_manager = NetworkServiceManager(
-        control_plane_config, workflow._service_manager._services
+        workflow._service_manager._services
     )
 
     service = WorkflowService(
diff --git a/llama_deploy/deploy/network_workflow.py b/llama_deploy/deploy/network_workflow.py
index e5bddaf3..95e65442 100644
--- a/llama_deploy/deploy/network_workflow.py
+++ b/llama_deploy/deploy/network_workflow.py
@@ -1,32 +1,36 @@
+import warnings
 from typing import Any, Dict, Optional
 
 from llama_index.core.workflow import StartEvent, StopEvent, Workflow, step
 from llama_index.core.workflow.service import ServiceManager, ServiceNotFoundError
 
-from llama_deploy.client.async_client import AsyncLlamaDeployClient
-from llama_deploy.client.sync_client import LlamaDeployClient
+from llama_deploy.client import Client
 from llama_deploy.control_plane.server import ControlPlaneConfig
 
 
 class NetworkWorkflow(Workflow):
     def __init__(
         self,
-        control_plane_config: ControlPlaneConfig,
         remote_service_name: str,
+        control_plane_config: ControlPlaneConfig | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
-        self.control_plane_config = control_plane_config
+        if control_plane_config is not None:
+            warnings.warn(
+                "The control_plane_config parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         self.remote_service_name = remote_service_name
+        # Use the same timeout for the client
+        self._client = Client(timeout=self._timeout)
 
     @step
     async def run_remote_workflow(self, ev: StartEvent) -> StopEvent:
-        client = AsyncLlamaDeployClient(self.control_plane_config)
-        kwargs = ev.dict()
-
-        session = await client.create_session()
-        result = await session.run(self.remote_service_name, **kwargs)
-        await client.delete_session(session.session_id)
+        session = await self._client.core.sessions.create()
+        result = await session.run(self.remote_service_name, **ev.dict())
+        await self._client.core.sessions.delete(session.id)
 
         return StopEvent(result=result)
 
@@ -34,13 +38,19 @@ async def run_remote_workflow(self, ev: StartEvent) -> StopEvent:
 class NetworkServiceManager(ServiceManager):
     def __init__(
         self,
-        control_plane_config: ControlPlaneConfig,
-        existing_services: Dict[str, Workflow],
+        existing_services: Dict[str, Workflow] | None = None,
+        control_plane_config: ControlPlaneConfig | None = None,
     ) -> None:
         super().__init__()
+        if control_plane_config is not None:
+            warnings.warn(
+                "The control_plane_config parameter is deprecated and will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         # override with passed in/inherited services
-        self._services = existing_services
-        self.control_plane_config = control_plane_config
+        self._services: Dict[str, Workflow] = existing_services or {}
+        self._client = Client()
 
     def get(self, name: str, default: Optional["Workflow"] = None) -> "Workflow":
         try:
@@ -48,9 +58,7 @@ def get(self, name: str, default: Optional["Workflow"] = None) -> "Workflow":
         except ServiceNotFoundError:
             local_workflow = None
 
-        # TODO: service manager does not support async
-        client = LlamaDeployClient(self.control_plane_config)
-        services = client.list_services()
+        services = self._client.sync.core.services.list()
 
         remote_service = None
         for service in services:
@@ -60,7 +68,7 @@ def get(self, name: str, default: Optional["Workflow"] = None) -> "Workflow":
 
         # If the remove service exists, swap it in
         if remote_service is not None:
-            return NetworkWorkflow(self.control_plane_config, name, timeout=None)
+            return NetworkWorkflow(name, timeout=None)
 
         # else default to the local workflow -- if it exists
         if local_workflow is None:


Test patch:
diff --git a/tests/deploy/test_network_workflow.py b/tests/deploy/test_network_workflow.py
new file mode 100644
index 00000000..485bcf4a
--- /dev/null
+++ b/tests/deploy/test_network_workflow.py
@@ -0,0 +1,99 @@
+from unittest import mock
+
+import pytest
+from llama_index.core.workflow import StartEvent
+
+from llama_deploy import ControlPlaneConfig
+from llama_deploy.deploy.network_workflow import (
+    NetworkServiceManager,
+    NetworkWorkflow,
+    ServiceNotFoundError,
+)
+
+
+def test_network_workflow_ctor() -> None:
+    nw = NetworkWorkflow(remote_service_name="test_workflow")
+    assert nw.remote_service_name == "test_workflow"
+
+    with pytest.warns(
+        DeprecationWarning, match="The control_plane_config parameter is deprecated"
+    ):
+        nw = NetworkWorkflow(
+            remote_service_name="test_workflow",
+            control_plane_config=ControlPlaneConfig(),
+        )
+    assert nw.remote_service_name == "test_workflow"
+
+
+@pytest.mark.asyncio
+async def test_network_workflow_run_remote_workflow() -> None:
+    # Setup
+    mock_client = mock.MagicMock()
+    mock_session = mock.MagicMock()
+    mock_session.run = mock.AsyncMock(return_value={"result": "success"})
+    mock_session.id = "test_session_id"
+
+    mock_client.core.sessions.create = mock.AsyncMock(return_value=mock_session)
+    mock_client.core.sessions.delete = mock.AsyncMock()
+
+    # Create workflow
+    workflow = NetworkWorkflow(remote_service_name="test_workflow")
+    workflow._client = mock_client  # Inject mock client
+
+    # Create start event with test data
+    start_event = StartEvent(param1="value1", param2="value2")  # type:ignore
+
+    # Run workflow
+    result = await workflow.run_remote_workflow(start_event)
+
+    # Assertions
+    assert result.result == {"result": "success"}
+    mock_client.core.sessions.create.assert_awaited_once()
+    mock_session.run.assert_awaited_once_with(
+        "test_workflow", param1="value1", param2="value2"
+    )
+    mock_client.core.sessions.delete.assert_awaited_once_with("test_session_id")
+
+
+def test_service_manager_ctor() -> None:
+    sm = NetworkServiceManager()
+    assert not sm._services
+
+    with pytest.warns(
+        DeprecationWarning, match="The control_plane_config parameter is deprecated"
+    ):
+        NetworkServiceManager(control_plane_config=ControlPlaneConfig())
+
+
+def test_service_manager_get() -> None:
+    mock_client = mock.MagicMock()
+    mock_service = mock.MagicMock()
+    mock_service.service_name = "test_service"
+
+    # Mock the client.sync.core.services.list() call
+    mock_client.sync.core.services.list.return_value = [mock_service]
+
+    # Create service manager with some existing services
+    existing_services = {"local_service": mock.MagicMock()}
+    sm = NetworkServiceManager(existing_services=existing_services)  # type: ignore
+    sm._client = mock_client  # Inject mock client
+
+    # Test getting existing local service
+    local_workflow = sm.get("local_service")
+    assert local_workflow == existing_services["local_service"]
+
+    # Test getting remote service
+    remote_workflow = sm.get("test_service")
+    assert isinstance(remote_workflow, NetworkWorkflow)
+    assert remote_workflow.remote_service_name == "test_service"
+    assert remote_workflow._timeout == remote_workflow._client.timeout
+    assert remote_workflow._timeout is None
+
+    # Test service not found
+    with pytest.raises(
+        ServiceNotFoundError, match="Service nonexistent_service not found"
+    ):
+        sm.get("nonexistent_service")
+
+    # Verify client method was called
+    mock_client.sync.core.services.list.assert_called()