Repository: run-llama/llama_deploy

Issue:
Allow re-creating a deployment that already exists
At the moment trying to deploy the same `deployment.yml` results in an error. 

We could leave the default behaviour like it is, but allow users to force a reload of the same deployment  via a `--force` or `--recreate` flag to the CLI that can be replicated in the corresponding apiserver API with a simple query parameter.

Discussion:
coveralls: 
[![Coverage Status](https://coveralls.io/builds/71563702/badge)](https://coveralls.io/builds/71563702)

coverage: 73.826% (+0.5%) from 73.324%
when pulling **8fa42061124739358fcd2b555e3e7467ad4a244c on massi/322**
into **4c5b832cc262dd35be28086ca6fba0d52d9f8ee5 on main**.


Patch:
diff --git a/llama_deploy/apiserver/deployment.py b/llama_deploy/apiserver/deployment.py
index 0a50607..7f72bb3 100644
--- a/llama_deploy/apiserver/deployment.py
+++ b/llama_deploy/apiserver/deployment.py
@@ -5,6 +5,7 @@ import subprocess
 import sys
 from multiprocessing.pool import ThreadPool
 from pathlib import Path
+from shutil import rmtree
 from typing import Any
 
 from dotenv import dotenv_values
@@ -68,6 +69,9 @@ class Deployment:
         self._workflow_services: list[WorkflowService] = self._load_services(config)
         self._client = Client(control_plane_url=config.control_plane.url)
         self._default_service = config.default_service
+        self._running = False
+        self._service_tasks: list[asyncio.Task] = []
+        self._service_startup_complete = asyncio.Event()
 
     @property
     def default_service(self) -> str | None:
@@ -95,6 +99,7 @@ class Deployment:
         All the tasks are gathered before returning.
         """
         tasks = []
+        self._running = True
 
         # Spawn SimpleMessageQueue if needed
         if self._simple_message_queue_server:
@@ -112,17 +117,51 @@ class Deployment:
         tasks.append(asyncio.create_task(cp_consumer_fn()))
 
         # Services
-        for wfs in self._workflow_services:
-            service_task = asyncio.create_task(wfs.launch_server())
-            tasks.append(service_task)
-            consumer_fn = await wfs.register_to_message_queue()
-            control_plane_url = f"http://{self._control_plane_config.host}:{self._control_plane_config.port}"
-            await wfs.register_to_control_plane(control_plane_url)
-            consumer_task = asyncio.create_task(consumer_fn())
-            tasks.append(consumer_task)
+        tasks.append(asyncio.create_task(self._run_services()))
 
         # Run allthethings
         await asyncio.gather(*tasks)
+        self._running = False
+
+    async def reload(self, config: Config) -> None:
+        """Reload this deployment by restarting its services.
+
+        The reload process consists in cancelling the services tasks
+        and rely on the fact that _run_services() will restart them
+        with the new configuration. This function won't return until
+        _run_services will trigger the _service_startup_complete signal.
+        """
+        self._workflow_services = self._load_services(config)
+        self._default_service = config.default_service
+
+        for t in self._service_tasks:
+            # t is awaited in _run_services(), we don't need to await here
+            t.cancel()
+
+        # Hold until _run_services() has restarted all the tasks
+        await self._service_startup_complete.wait()
+
+    async def _run_services(self) -> None:
+        """Start an asyncio task for each service and gather them.
+
+        For the time self._running holds true, the tasks will be restarted
+        if they are all cancelled. This is to support the reload process
+        (see reload() for more details).
+        """
+        while self._running:
+            self._service_tasks = []
+            # If this is a reload, self._workflow_services contains the updated configurations
+            for wfs in self._workflow_services:
+                service_task = asyncio.create_task(wfs.launch_server())
+                self._service_tasks.append(service_task)
+                consumer_fn = await wfs.register_to_message_queue()
+                control_plane_url = f"http://{self._control_plane_config.host}:{self._control_plane_config.port}"
+                await wfs.register_to_control_plane(control_plane_url)
+                consumer_task = asyncio.create_task(consumer_fn())
+                self._service_tasks.append(consumer_task)
+            # If this is a reload, unblock the reload() function signalling that tasks are up and running
+            self._service_startup_complete.set()
+            await asyncio.gather(*self._service_tasks)
 
     def _load_services(self, config: Config) -> list[WorkflowService]:
         """Creates WorkflowService instances according to the configuration object."""
@@ -151,9 +190,19 @@ class Deployment:
                 raise ValueError(msg)
 
             # Sync the service source
-            destination = self._path / service_id
+            destination = (self._path / service_id).resolve()
+
+            if destination.exists():
+                # FIXME: this could be managed at the source manager level, so that
+                # each implementation can decide what to do with existing data. For
+                # example, the git source manager might decide to perform a git pull
+                # instead of a brand new git clone. Leaving these optimnizations for
+                # later, for the time being having an empty data folder works smoothly
+                # for any source manager currently supported.
+                rmtree(str(destination))
+
             source_manager = SOURCE_MANAGERS[source.type]
-            source_manager.sync(source.name, str(destination.resolve()))
+            source_manager.sync(source.name, str(destination))
 
             # Install dependencies
             self._install_dependencies(service_config)
@@ -278,7 +327,7 @@ class Manager:
         self._deployments_path = deployments_path
         self._max_deployments = max_deployments
         self._pool = ThreadPool(processes=max_deployments)
-        self._control_plane_port = 8002
+        self._last_control_plane_port = 8002
 
     @property
     def deployment_names(self) -> list[str]:
@@ -297,34 +346,46 @@ class Manager:
         except asyncio.CancelledError:
             pass
 
-    def deploy(self, config: Config) -> None:
+    async def deploy(self, config: Config, reload: bool = False) -> None:
         """Creates a Deployment instance and starts the relative runtime.
 
         Args:
             config: The deployment configuration.
+            reload: Reload an existing deployment instead of raising an error.
 
         Raises:
             ValueError: If a deployment with the same name already exists or the maximum number of deployment exceeded.
             DeploymentError: If it wasn't possible to create a deployment.
         """
-        if config.name in self._deployments:
-            msg = f"Deployment already exists: {config.name}"
-            raise ValueError(msg)
+        if not reload:
+            # Raise an error if deployment already exists
+            if config.name in self._deployments:
+                msg = f"Deployment already exists: {config.name}"
+                raise ValueError(msg)
 
-        if len(self._deployments) == self._max_deployments:
-            msg = "Reached the maximum number of deployments, cannot schedule more"
-            raise ValueError(msg)
+            # Raise an error if we can't create any new deployment
+            if len(self._deployments) == self._max_deployments:
+                msg = "Reached the maximum number of deployments, cannot schedule more"
+                raise ValueError(msg)
+
+            # Set the control plane TCP port in the config where not specified
+            self._assign_control_plane_address(config)
 
-        self._assign_control_plane_address(config)
+            deployment = Deployment(config=config, root_path=self._deployments_path)
+            self._deployments[config.name] = deployment
+            self._pool.apply_async(func=asyncio.run, args=(deployment.start(),))
+        else:
+            if config.name not in self._deployments:
+                msg = f"Cannot find deployment to reload: {config.name}"
+                raise ValueError(msg)
 
-        deployment = Deployment(config=config, root_path=self._deployments_path)
-        self._deployments[config.name] = deployment
-        self._pool.apply_async(func=asyncio.run, args=(deployment.start(),))
+            deployment = self._deployments[config.name]
+            await deployment.reload(config)
 
     def _assign_control_plane_address(self, config: Config) -> None:
         for service in config.services.values():
             if not service.port:
-                service.port = self._control_plane_port
-                self._control_plane_port += 1
+                service.port = self._last_control_plane_port
+                self._last_control_plane_port += 1
             if not service.host:
                 service.host = "localhost"
diff --git a/llama_deploy/apiserver/routers/deployments.py b/llama_deploy/apiserver/routers/deployments.py
index 5f81e0c..aa50b44 100644
--- a/llama_deploy/apiserver/routers/deployments.py
+++ b/llama_deploy/apiserver/routers/deployments.py
@@ -8,13 +8,12 @@ from llama_deploy.apiserver.config_parser import Config
 from llama_deploy.apiserver.server import manager
 from llama_deploy.types import (
     DeploymentDefinition,
+    EventDefinition,
     SessionDefinition,
     TaskDefinition,
-    EventDefinition,
 )
 from llama_deploy.types.core import TaskResult
 
-
 deployments_router = APIRouter(
     prefix="/deployments",
 )
@@ -37,11 +36,11 @@ async def read_deployment(deployment_name: str) -> DeploymentDefinition:
 
 @deployments_router.post("/create")
 async def create_deployment(
-    config_file: UploadFile = File(...),
+    config_file: UploadFile = File(...), reload: bool = False
 ) -> DeploymentDefinition:
     """Creates a new deployment by uploading a configuration file."""
     config = Config.from_yaml_bytes(await config_file.read())
-    manager.deploy(config)
+    await manager.deploy(config, reload)
 
     return DeploymentDefinition(name=config.name)
 
diff --git a/llama_deploy/cli/deploy.py b/llama_deploy/cli/deploy.py
index 025d05c..9d13620 100644
--- a/llama_deploy/cli/deploy.py
+++ b/llama_deploy/cli/deploy.py
@@ -7,13 +7,16 @@ from llama_deploy import Client
 
 @click.command()
 @click.pass_obj  # global_config
+@click.option("--reload", is_flag=True)
 @click.argument("deployment_config_file", type=click.File("rb"))
-def deploy(global_config: tuple, deployment_config_file: IO) -> None:
+def deploy(global_config: tuple, reload: bool, deployment_config_file: IO) -> None:
     server_url, disable_ssl, timeout = global_config
     client = Client(api_server_url=server_url, disable_ssl=disable_ssl, timeout=timeout)
 
     try:
-        deployment = client.sync.apiserver.deployments.create(deployment_config_file)
+        deployment = client.sync.apiserver.deployments.create(
+            deployment_config_file, reload
+        )
     except Exception as e:
         raise click.ClickException(str(e))
 
diff --git a/llama_deploy/client/models/apiserver.py b/llama_deploy/client/models/apiserver.py
index 90339af..c63debe 100644
--- a/llama_deploy/client/models/apiserver.py
+++ b/llama_deploy/client/models/apiserver.py
@@ -3,9 +3,8 @@ import json
 from typing import Any, AsyncGenerator, TextIO
 
 import httpx
-
-from llama_index.core.workflow.events import Event
 from llama_index.core.workflow.context_serializers import JsonSerializer
+from llama_index.core.workflow.events import Event
 
 from llama_deploy.types.apiserver import Status, StatusEnum
 from llama_deploy.types.core import (
@@ -227,9 +226,12 @@ class Deployment(Model):
 class DeploymentCollection(Collection):
     """A model representing a collection of deployments currently active."""
 
-    async def create(self, config: TextIO) -> Deployment:
+    async def create(self, config: TextIO, reload: bool = False) -> Deployment:
         """Creates a new deployment from a deployment file.
 
+        If `reload` is true, an existing deployment will be reloaded, otherwise
+        an error will be raised.
+
         Example:
             ```
             with open("deployment.yml") as f:
@@ -243,6 +245,7 @@ class DeploymentCollection(Collection):
             "POST",
             create_url,
             files=files,
+            params={"reload": reload},
             verify=not self.client.disable_ssl,
             timeout=self.client.timeout,
         )
diff --git a/llama_deploy/services/workflow.py b/llama_deploy/services/workflow.py
index ce52602..5e43610 100644
--- a/llama_deploy/services/workflow.py
+++ b/llama_deploy/services/workflow.py
@@ -284,6 +284,9 @@ class WorkflowService(BaseService):
             current_call (WorkflowState):
                 The state of the current task, including run_kwargs and other session state.
         """
+        # create send_event background task
+        close_send_events = asyncio.Event()
+
         try:
             # load the state
             ctx = await self.get_workflow_state(current_call)
@@ -291,9 +294,6 @@ class WorkflowService(BaseService):
             # run the workflow
             handler = self.workflow.run(ctx=ctx, **current_call.run_kwargs)
 
-            # create send_event background task
-            close_send_events = asyncio.Event()
-
             async def send_events(
                 handler: WorkflowHandler, close_event: asyncio.Event
             ) -> None:
@@ -512,4 +512,8 @@ class WorkflowService(BaseService):
 
         cfg = uvicorn.Config(self._app, host=host, port=port)
         server = CustomServer(cfg)
-        await server.serve()
+
+        try:
+            await server.serve()
+        except asyncio.CancelledError:
+            await asyncio.gather(server.shutdown(), return_exceptions=True)


Test patch:
diff --git a/e2e_tests/apiserver/deployments/deployment1.yml b/e2e_tests/apiserver/deployments/deployment1.yml
index e63acff..a8097d8 100644
--- a/e2e_tests/apiserver/deployments/deployment1.yml
+++ b/e2e_tests/apiserver/deployments/deployment1.yml
@@ -11,5 +11,5 @@ services:
     host: localhost
     source:
       type: git
-      name: https://github.com/run-llama/llama_deploy.git
-    path: tests/apiserver/data/workflow:my_workflow
+      name: https://github.com/run-llama/llama_deploy.git@massi/322
+    path: e2e_tests/apiserver/deployments/src:my_workflow
diff --git a/e2e_tests/apiserver/deployments/deployment2.yml b/e2e_tests/apiserver/deployments/deployment2.yml
index 1699d78..0cfe96c 100644
--- a/e2e_tests/apiserver/deployments/deployment2.yml
+++ b/e2e_tests/apiserver/deployments/deployment2.yml
@@ -11,5 +11,5 @@ services:
     host: localhost
     source:
       type: git
-      name: https://github.com/run-llama/llama_deploy.git
-    path: tests/apiserver/data/workflow:my_workflow
+      name: https://github.com/run-llama/llama_deploy.git@massi/322
+    path: e2e_tests/apiserver/deployments/src:my_workflow
diff --git a/e2e_tests/apiserver/deployments/deployment_reload1.yml b/e2e_tests/apiserver/deployments/deployment_reload1.yml
new file mode 100644
index 0000000..9520d4e
--- /dev/null
+++ b/e2e_tests/apiserver/deployments/deployment_reload1.yml
@@ -0,0 +1,15 @@
+name: ReloadMe
+
+control-plane: {}
+
+default-service: test-workflow
+
+services:
+  test-workflow:
+    name: Test Workflow
+    port: 8002
+    host: localhost
+    source:
+      type: git
+      name: https://github.com/run-llama/llama_deploy.git@massi/322
+    path: e2e_tests/apiserver/deployments/src:echo_workflow_en
diff --git a/e2e_tests/apiserver/deployments/deployment_reload2.yml b/e2e_tests/apiserver/deployments/deployment_reload2.yml
new file mode 100644
index 0000000..67ff4b6
--- /dev/null
+++ b/e2e_tests/apiserver/deployments/deployment_reload2.yml
@@ -0,0 +1,15 @@
+name: ReloadMe
+
+control-plane: {}
+
+default-service: test-workflow
+
+services:
+  test-workflow:
+    name: Test Workflow
+    port: 8002
+    host: localhost
+    source:
+      type: git
+      name: https://github.com/run-llama/llama_deploy.git@massi/322
+    path: e2e_tests/apiserver/deployments/src:echo_workflow_it
diff --git a/e2e_tests/apiserver/deployments/src/__init__.py b/e2e_tests/apiserver/deployments/src/__init__.py
new file mode 100644
index 0000000..79b7b26
--- /dev/null
+++ b/e2e_tests/apiserver/deployments/src/__init__.py
@@ -0,0 +1,6 @@
+from .workflow import EchoWorkflow
+from .workflow_reload import EchoWithPrompt
+
+my_workflow = EchoWorkflow()
+echo_workflow_en = EchoWithPrompt(prompt_msg="I have received:")
+echo_workflow_it = EchoWithPrompt(prompt_msg="Ho ricevuto:")
diff --git a/e2e_tests/apiserver/deployments/src/workflow_reload.py b/e2e_tests/apiserver/deployments/src/workflow_reload.py
new file mode 100644
index 0000000..4e2b3f5
--- /dev/null
+++ b/e2e_tests/apiserver/deployments/src/workflow_reload.py
@@ -0,0 +1,11 @@
+from llama_index.core.workflow import Context, StartEvent, StopEvent, Workflow, step
+
+
+class EchoWithPrompt(Workflow):
+    def __init__(self, prompt_msg):
+        super().__init__()
+        self._prompt_msg = prompt_msg
+
+    @step
+    def do_something(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        return StopEvent(result=f"{self._prompt_msg}{ev.data}")
diff --git a/e2e_tests/apiserver/test_reload.py b/e2e_tests/apiserver/test_reload.py
new file mode 100644
index 0000000..1e82ebf
--- /dev/null
+++ b/e2e_tests/apiserver/test_reload.py
@@ -0,0 +1,26 @@
+import asyncio
+from pathlib import Path
+
+import pytest
+
+from llama_deploy.types import TaskDefinition
+
+
+@pytest.mark.asyncio
+async def test_reload(apiserver, client):
+    here = Path(__file__).parent
+    with open(here / "deployments" / "deployment_reload1.yml") as f:
+        deployment = await client.apiserver.deployments.create(f)
+        await asyncio.sleep(3)
+
+    tasks = deployment.tasks
+    res = await tasks.run(TaskDefinition(input='{"data": "bar"}'))
+    assert res == "I have received:bar"
+
+    with open(here / "deployments" / "deployment_reload2.yml") as f:
+        deployment = await client.apiserver.deployments.create(f, reload=True)
+        await asyncio.sleep(3)
+
+    tasks = deployment.tasks
+    res = await tasks.run(TaskDefinition(input='{"data": "bar"}'))
+    assert res == "Ho ricevuto:bar"
diff --git a/tests/apiserver/routers/test_deployments.py b/tests/apiserver/routers/test_deployments.py
index 7a22d08..bec6162 100644
--- a/tests/apiserver/routers/test_deployments.py
+++ b/tests/apiserver/routers/test_deployments.py
@@ -4,12 +4,12 @@ from unittest import mock
 
 import pytest
 from fastapi.testclient import TestClient
-from llama_index.core.workflow.events import Event, HumanResponseEvent
 from llama_index.core.workflow.context_serializers import JsonSerializer
+from llama_index.core.workflow.events import Event, HumanResponseEvent
 
 from llama_deploy.apiserver import Config
 from llama_deploy.types import TaskResult
-from llama_deploy.types.core import TaskDefinition, EventDefinition
+from llama_deploy.types.core import EventDefinition, TaskDefinition
 
 
 def test_read_deployments(http_client: TestClient) -> None:
@@ -37,6 +37,7 @@ def test_create_deployment(http_client: TestClient, data_path: Path) -> None:
     with mock.patch(
         "llama_deploy.apiserver.routers.deployments.manager"
     ) as mocked_manager:
+        mocked_manager.deploy = mock.AsyncMock()
         config_file = data_path / "git_service.yaml"
 
         with open(config_file, "rb") as f:
@@ -47,7 +48,7 @@ def test_create_deployment(http_client: TestClient, data_path: Path) -> None:
             )
 
         assert response.status_code == 200
-        mocked_manager.deploy.assert_called_with(actual_config)
+        mocked_manager.deploy.assert_awaited_with(actual_config, False)
 
 
 def test_create_deployment_task_not_found(
@@ -64,6 +65,32 @@ def test_create_deployment_task_not_found(
         assert response.status_code == 404
 
 
+def test_run_task_not_found(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = None
+        response = http_client.post(
+            "/deployments/test-deployment/tasks/run/",
+            json={"input": "{}"},
+        )
+        assert response.status_code == 404
+
+
+def test_run_task_no_default_service(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = mock.MagicMock(
+            default_service=None
+        )
+        response = http_client.post(
+            "/deployments/test-deployment/tasks/run/",
+            json={"input": "{}"},
+        )
+        assert response.status_code == 400
+
+
 def test_create_deployment_task_missing_service(
     http_client: TestClient, data_path: Path
 ) -> None:
@@ -90,10 +117,15 @@ def test_run_deployment_task(http_client: TestClient, data_path: Path) -> None:
     ) as mocked_manager:
         deployment = mock.AsyncMock()
         deployment.default_service = "TestService"
-        session = mock.AsyncMock()
+
+        session = mock.AsyncMock(id="42")
         deployment.client.core.sessions.create.return_value = session
         session.run.return_value = {"result": "test_result"}
-        session.id = "42"
+
+        session_from_get = mock.AsyncMock(id="84")
+        deployment.client.core.sessions.get.return_value = session_from_get
+        session_from_get.run.return_value = {"result": "test_result_from_existing"}
+
         mocked_manager.get_deployment.return_value = deployment
         response = http_client.post(
             "/deployments/test-deployment/tasks/run/",
@@ -102,6 +134,15 @@ def test_run_deployment_task(http_client: TestClient, data_path: Path) -> None:
         assert response.status_code == 200
         deployment.client.core.sessions.delete.assert_called_with("42")
 
+        deployment.reset_mock()
+        response = http_client.post(
+            "/deployments/test-deployment/tasks/run/",
+            json={"input": "{}"},
+            params={"session_id": 84},
+        )
+        assert response.status_code == 200
+        deployment.client.core.sessions.delete.assert_not_called()
+
 
 def test_create_deployment_task(http_client: TestClient, data_path: Path) -> None:
     with mock.patch(
@@ -109,10 +150,15 @@ def test_create_deployment_task(http_client: TestClient, data_path: Path) -> Non
     ) as mocked_manager:
         deployment = mock.AsyncMock()
         deployment.default_service = "TestService"
-        session = mock.AsyncMock()
+
+        session = mock.AsyncMock(id="42")
         deployment.client.core.sessions.create.return_value = session
-        session.id = "42"
         session.run_nowait.return_value = "test_task_id"
+
+        session_from_get = mock.AsyncMock(id="84")
+        deployment.client.core.sessions.get.return_value = session_from_get
+        session_from_get.run_nowait.return_value = "another_test_task_id"
+
         mocked_manager.get_deployment.return_value = deployment
         response = http_client.post(
             "/deployments/test-deployment/tasks/create/",
@@ -122,6 +168,28 @@ def test_create_deployment_task(http_client: TestClient, data_path: Path) -> Non
         td = TaskDefinition(**response.json())
         assert td.task_id == "test_task_id"
 
+        deployment.reset_mock()
+        response = http_client.post(
+            "/deployments/test-deployment/tasks/create/",
+            json={"input": "{}"},
+            params={"session_id": 84},
+        )
+        assert response.status_code == 200
+        deployment.client.core.sessions.delete.assert_not_called()
+
+
+def test_send_event_not_found(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = None
+        response = http_client.post(
+            "/deployments/test-deployment/tasks/test_task_id/events",
+            json=EventDefinition(agent_id="foo", event_obj_str="bar").model_dump(),
+            params={"session_id": 42},
+        )
+        assert response.status_code == 404
+
 
 def test_send_event(http_client: TestClient, data_path: Path) -> None:
     with mock.patch(
@@ -151,6 +219,18 @@ def test_send_event(http_client: TestClient, data_path: Path) -> None:
         assert ev_def.event_obj_str == event_def.event_obj_str
 
 
+def test_get_event_not_found(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = None
+        response = http_client.get(
+            "/deployments/test-deployment/tasks/test_task_id/events",
+            params={"session_id": "42", "task_id": "84"},
+        )
+        assert response.status_code == 404
+
+
 @pytest.mark.asyncio
 async def test_get_event_stream(http_client: TestClient, data_path: Path) -> None:
     mock_events = [
@@ -186,6 +266,46 @@ async def test_get_event_stream(http_client: TestClient, data_path: Path) -> Non
         session.get_task_result_stream.assert_called_with("test_task_id")
 
 
+def test_get_task_result_not_found(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = None
+        response = http_client.get(
+            "/deployments/test-deployment/tasks/test_task_id/results/?session_id=42",
+        )
+        assert response.status_code == 404
+
+
+def test_get_tasks_not_found(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = None
+        response = http_client.get(
+            "/deployments/test-deployment/tasks",
+        )
+        assert response.status_code == 404
+
+
+def test_get_tasks(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        deployment = mock.AsyncMock()
+        mocked_manager.get_deployment.return_value = deployment
+        session = mock.AsyncMock()
+        session.get_tasks.return_value = [TaskDefinition(input="foo")]
+        deployment.client.core.sessions.list.return_value = [session]
+
+        response = http_client.get(
+            "/deployments/test-deployment/tasks",
+        )
+        assert response.status_code == 200
+        assert len(response.json()) == 1
+        assert response.json()[0]["input"] == "foo"
+
+
 def test_get_task_result(http_client: TestClient, data_path: Path) -> None:
     with mock.patch(
         "llama_deploy.apiserver.routers.deployments.manager"
@@ -208,6 +328,17 @@ def test_get_task_result(http_client: TestClient, data_path: Path) -> None:
         deployment.client.core.sessions.get.assert_called_with("42")
 
 
+def test_get_sessions_not_found(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = None
+        response = http_client.get(
+            "/deployments/test-deployment/sessions/",
+        )
+        assert response.status_code == 404
+
+
 def test_get_sessions(http_client: TestClient, data_path: Path) -> None:
     with mock.patch(
         "llama_deploy.apiserver.routers.deployments.manager"
@@ -224,6 +355,18 @@ def test_get_sessions(http_client: TestClient, data_path: Path) -> None:
         assert response.json() == []
 
 
+def test_delete_session_not_found(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = None
+        response = http_client.post(
+            "/deployments/test-deployment/sessions/delete/?session_id=42",
+        )
+        assert response.status_code == 404
+        assert response.json() == {"detail": "Deployment not found"}
+
+
 def test_delete_session(http_client: TestClient, data_path: Path) -> None:
     with mock.patch(
         "llama_deploy.apiserver.routers.deployments.manager"
@@ -237,3 +380,64 @@ def test_delete_session(http_client: TestClient, data_path: Path) -> None:
         )
         assert response.status_code == 200
         deployment.client.core.sessions.delete.assert_called_with("42")
+
+
+def test_get_session_not_found(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = None
+        response = http_client.get(
+            "/deployments/test-deployment/sessions/foo",
+        )
+        assert response.status_code == 404
+
+
+def test_get_session(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        deployment = mock.AsyncMock()
+        mocked_manager.get_deployment.return_value = deployment
+        session = mock.AsyncMock(id="foo")
+        deployment.client.core.sessions.get.return_value = session
+        response = http_client.get("/deployments/test-deployment/sessions/foo")
+        assert response.status_code == 200
+        assert response.json() == {"session_id": "foo", "state": {}, "task_ids": []}
+
+
+def test_create_session_not_found(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        mocked_manager.get_deployment.return_value = None
+        response = http_client.post(
+            "/deployments/test-deployment/sessions/create",
+        )
+
+        assert response.status_code == 404
+
+
+def test_create_session(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        deployment = mock.AsyncMock()
+        session = mock.AsyncMock(id="test-session-id")
+        deployment.client.core.sessions.create.return_value = session
+        mocked_manager.get_deployment.return_value = deployment
+
+        response = http_client.post(
+            "/deployments/test-deployment/sessions/create",
+        )
+
+        assert response.status_code == 200
+        assert response.json() == {
+            "session_id": "test-session-id",
+            "state": {},
+            "task_ids": [],
+        }
+
+        # Verify the mocked calls
+        mocked_manager.get_deployment.assert_called_once_with("test-deployment")
+        deployment.client.core.sessions.create.assert_called_once()
diff --git a/tests/apiserver/test_deployment.py b/tests/apiserver/test_deployment.py
index 4c69949..84b571b 100644
--- a/tests/apiserver/test_deployment.py
+++ b/tests/apiserver/test_deployment.py
@@ -178,17 +178,19 @@ def test_manager_ctor() -> None:
     assert m.get_deployment("foo") is None
 
 
-def test_manager_deploy_duplicate(data_path: Path) -> None:
+@pytest.mark.asyncio
+async def test_manager_deploy_duplicate(data_path: Path) -> None:
     config = Config.from_yaml(data_path / "git_service.yaml")
 
     m = Manager()
     m._deployments["TestDeployment"] = mock.MagicMock()
 
     with pytest.raises(ValueError, match="Deployment already exists: TestDeployment"):
-        m.deploy(config)
+        await m.deploy(config)
 
 
-def test_manager_deploy_maximum_reached(data_path: Path) -> None:
+@pytest.mark.asyncio
+async def test_manager_deploy_maximum_reached(data_path: Path) -> None:
     config = Config.from_yaml(data_path / "git_service.yaml")
 
     m = Manager(max_deployments=1)
@@ -198,16 +200,17 @@ def test_manager_deploy_maximum_reached(data_path: Path) -> None:
         ValueError,
         match="Reached the maximum number of deployments, cannot schedule more",
     ):
-        m.deploy(config)
+        await m.deploy(config)
 
 
-def test_manager_deploy(data_path: Path) -> None:
+@pytest.mark.asyncio
+async def test_manager_deploy(data_path: Path) -> None:
     config = Config.from_yaml(data_path / "git_service.yaml")
     with mock.patch(
         "llama_deploy.apiserver.deployment.Deployment"
     ) as mocked_deployment:
         m = Manager()
-        m.deploy(config)
+        await m.deploy(config)
         mocked_deployment.assert_called_once()
         assert m.deployment_names == ["TestDeployment"]
         assert m.get_deployment("TestDeployment") is not None
diff --git a/tests/client/models/test_apiserver.py b/tests/client/models/test_apiserver.py
index 0efddef..deab02b 100644
--- a/tests/client/models/test_apiserver.py
+++ b/tests/client/models/test_apiserver.py
@@ -218,6 +218,7 @@ async def test_task_deployment_collection_create(client: Any) -> None:
         "POST",
         "http://localhost:4501/deployments/create",
         files={"config_file": "some config"},
+        params={"reload": False},
         verify=True,
         timeout=120.0,
     )