Repository: jina-ai/serve
Pull request: feat(daemon): scale endpoint for flows & pods
- [x] Separated the `rolling_update` & `scale` endpoints for Flows / Pods.
- [x] Separated `rolling_update` & `scale` methods for `JinaDClient`.
- [x] Removed `UpdateOperation` enums.
- [x] Added unittests for `scale` store, partial-store & client.
- [x] Added distributed tests for scale endpoints (Local Flow-remote Executor, Remote Flow, Remote Pod)

![image](https://user-images.githubusercontent.com/9050737/141942121-e7dc1f1b-b0c3-4705-8437-65436cf27605.png)


Problem statement:
add distribuetd test for scale method
**TODO**
In order to support `scale` endpoint for Remote Flows and Pods, JinaD needs to:

- [ ] Clarify the Update endpoint and accept different parameters depending on the `kind` of update operation
- [ ] Add a new UpdateOperation scale
- [ ] Unittest the `update` methods
- [ ] Add distributed tests for scale endpoint. (Remote Flow, and local flow with remote executors)

Hints from issue discussion:
Blocked until #3816 is merged

Commit message: feat(daemon): scale endpoint

Patch to review:
diff --git a/daemon/api/endpoints/flows.py b/daemon/api/endpoints/flows.py
index 04f3a6d37565e..7d8f31abe1e11 100644
--- a/daemon/api/endpoints/flows.py
+++ b/daemon/api/endpoints/flows.py
@@ -45,13 +45,11 @@ async def _create(flow: FlowDepends = Depends(FlowDepends)):
 
 
 @router.put(
-    path='/{id}',
-    summary='Trigger an update operation on the Flow object',
-    description='Types supported: "rolling_update"',
+    path='/rolling_update/{id}',
+    summary='Trigger a rolling_update operation on the Flow object',
 )
-async def _update(
+async def rolling_update(
     id: DaemonID,
-    kind: UpdateOperation,
     pod_name: str,
     dump_path: Optional[str] = None,
     uses_with: Optional[Dict[str, Any]] = None,
@@ -63,7 +61,18 @@ async def _update(
             else:
                 uses_with = {'dump_path': dump_path}
 
-        return await store.update(id, kind, pod_name=pod_name, uses_with=uses_with)
+        return await store.rolling_update(id=id, pod_name=pod_name, uses_with=uses_with)
+    except Exception as ex:
+        raise Runtime400Exception from ex
+
+
+@router.put(
+    path='/scale/{id}',
+    summary='Trigger a scale operation on the Flow object',
+)
+async def scale(id: DaemonID, pod_name: str, replicas: int):
+    try:
+        return await store.scale(id=id, pod_name=pod_name, replicas=replicas)
     except Exception as ex:
         raise Runtime400Exception from ex
 
diff --git a/daemon/api/endpoints/partial/flow.py b/daemon/api/endpoints/partial/flow.py
index 897e9faff45ff..a6ae3599ccb71 100644
--- a/daemon/api/endpoints/partial/flow.py
+++ b/daemon/api/endpoints/partial/flow.py
@@ -5,7 +5,6 @@
 from jina.parsers.flow import set_flow_parser
 
 from ....models import FlowModel
-from ....models.enums import UpdateOperation
 from ....models.ports import PortMappings
 from ....models.partial import PartialFlowItem
 from ....excepts import PartialDaemon400Exception
@@ -44,13 +43,11 @@ async def _create(flow: 'FlowModel', ports: Optional[PortMappings] = None):
 
 
 @router.put(
-    path='',
-    summary='Run an update operation on the Flow object',
-    description='Types supported: "rolling_update"',
+    path='/rolling_update',
+    summary='Run a rolling_update operation on the Flow object',
     response_model=PartialFlowItem,
 )
-async def _update(
-    kind: UpdateOperation,
+async def rolling_update(
     pod_name: str,
     uses_with: Optional[Dict[str, Any]] = None,
 ):
@@ -60,7 +57,24 @@ async def _update(
     .. #noqa: DAR201
     """
     try:
-        return store.update(kind, pod_name, uses_with=uses_with)
+        return store.rolling_update(pod_name=pod_name, uses_with=uses_with)
+    except ValueError as ex:
+        raise PartialDaemon400Exception from ex
+
+
+@router.put(
+    path='/scale',
+    summary='Run a rolling_update operation on the Flow object',
+    response_model=PartialFlowItem,
+)
+async def scale(pod_name: str, replicas: int):
+    """
+
+    .. #noqa: DAR101
+    .. #noqa: DAR201
+    """
+    try:
+        return store.scale(pod_name=pod_name, replicas=replicas)
     except ValueError as ex:
         raise PartialDaemon400Exception from ex
 
diff --git a/daemon/api/endpoints/partial/pod.py b/daemon/api/endpoints/partial/pod.py
index 3c01795caad92..2227aded77b3f 100644
--- a/daemon/api/endpoints/partial/pod.py
+++ b/daemon/api/endpoints/partial/pod.py
@@ -7,7 +7,6 @@
 from ....models import PodModel
 from ....models.partial import PartialStoreItem
 from ....stores import partial_store as store
-from ....models.enums import UpdateOperation
 
 router = APIRouter(prefix='/pod', tags=['pod'])
 
@@ -45,22 +44,35 @@ async def _create(pod: 'PodModel'):
 
 
 @router.put(
-    path='',
-    summary='Run an update operation on the Pod object',
-    description='Types supported: "rolling_update"',
+    path='/rolling_update',
+    summary='Run a rolling_update operation on the Pod object',
+    response_model=PartialStoreItem,
+)
+async def rolling_update(uses_with: Optional[Dict[str, Any]] = None):
+    """
+
+    .. #noqa: DAR101
+    .. #noqa: DAR201
+    """
+    try:
+        return await store.rolling_update(uses_with=uses_with)
+    except ValueError as ex:
+        raise PartialDaemon400Exception from ex
+
+
+@router.put(
+    path='/scale',
+    summary='Run a scale operation on the Pod object',
     response_model=PartialStoreItem,
 )
-async def _update(
-    kind: UpdateOperation,
-    uses_with: Optional[Dict[str, Any]] = None,
-):
+async def scale(replicas: int):
     """
 
     .. #noqa: DAR101
     .. #noqa: DAR201
     """
     try:
-        return store.update(kind, uses_with=uses_with)
+        return await store.scale(replicas=replicas)
     except ValueError as ex:
         raise PartialDaemon400Exception from ex
 
diff --git a/daemon/api/endpoints/pods.py b/daemon/api/endpoints/pods.py
index 94b168a639da6..156c29394c84c 100644
--- a/daemon/api/endpoints/pods.py
+++ b/daemon/api/endpoints/pods.py
@@ -45,13 +45,11 @@ async def _create(pod: PodDepends = Depends(PodDepends)):
 
 
 @router.put(
-    path='/{id}',
-    summary='Trigger a rolling update on this Pod',
-    description='Types supported: "rolling_update"',
+    path='/rolling_update/{id}',
+    summary='Trigger a rolling_update operation on the Pod object',
 )
-async def _update(
+async def rolling_update(
     id: DaemonID,
-    kind: UpdateOperation,
     dump_path: Optional[str] = None,
     uses_with: Optional[Dict[str, Any]] = None,
 ):
@@ -61,7 +59,19 @@ async def _update(
                 uses_with['dump_path'] = dump_path
             else:
                 uses_with = {'dump_path': dump_path}
-        return await store.update(id, kind, uses_with=uses_with)
+
+        return await store.rolling_update(id=id, uses_with=uses_with)
+    except Exception as ex:
+        raise Runtime400Exception from ex
+
+
+@router.put(
+    path='/scale/{id}',
+    summary='Trigger a scale operation on the Pod object',
+)
+async def scale(id: DaemonID, replicas: int):
+    try:
+        return await store.scale(id=id, replicas=replicas)
     except Exception as ex:
         raise Runtime400Exception from ex
 
diff --git a/daemon/models/enums.py b/daemon/models/enums.py
index c8b80191c8e08..9af25c23cfabd 100644
--- a/daemon/models/enums.py
+++ b/daemon/models/enums.py
@@ -85,4 +85,5 @@ class UpdateOperation(DaemonEnum):
     We consider these an `update` operation since they **change** the underlying state
     """
 
+    SCALE = 'scale'
     ROLLING_UPDATE = 'rolling_update'
diff --git a/daemon/stores/__init__.py b/daemon/stores/__init__.py
index 0b580bea9b991..e9a8c59ef61fa 100644
--- a/daemon/stores/__init__.py
+++ b/daemon/stores/__init__.py
@@ -1,4 +1,4 @@
-from typing import Optional, TYPE_CHECKING, Type
+from typing import Optional, TYPE_CHECKING, Type, Union
 
 from .base import BaseStore
 from .flows import FlowStore
@@ -10,7 +10,7 @@
 from ..models.enums import IDLiterals
 
 if TYPE_CHECKING:
-    from .partial import PartialStore
+    from .partial import PartialPeaStore, PartialPodStore, PartialFlowStore
 
 
 def _get_store(cls: Type[BaseStore]) -> BaseStore:
@@ -29,7 +29,9 @@ def _get_store(cls: Type[BaseStore]) -> BaseStore:
             return cls()
 
 
-def _get_partial_store() -> Optional['PartialStore']:
+def _get_partial_store() -> Optional[
+    Union['PartialPeaStore', 'PartialPodStore', 'PartialFlowStore']
+]:
     """Get partial store object
 
     :return: partial store object
diff --git a/daemon/stores/containers.py b/daemon/stores/containers.py
index b374bd4884ab1..284036b5a27ea 100644
--- a/daemon/stores/containers.py
+++ b/daemon/stores/containers.py
@@ -4,6 +4,7 @@
 from copy import deepcopy
 from platform import uname
 from http import HTTPStatus
+from abc import ABC, abstractmethod
 from typing import Dict, TYPE_CHECKING, List, Optional, Union
 
 import aiohttp
@@ -20,7 +21,6 @@
 from ..helper import if_alive, id_cleaner, error_msg_from
 from ..models import DaemonID
 from ..models.ports import PortMappings
-from ..models.enums import UpdateOperation, IDLiterals
 from ..models.containers import (
     ContainerArguments,
     ContainerItem,
@@ -32,48 +32,25 @@
     from pydantic import BaseModel
 
 
-class ContainerStore(BaseStore):
+class ContainerStore(BaseStore, ABC):
     """A Store of Containers spawned by daemon"""
 
     _kind = 'container'
     _status_model = ContainerStoreStatus
 
-    async def _add(self, uri, *args, **kwargs):
+    @abstractmethod
+    async def add_in_partial(self, uri, *args, **kwargs):
         """Implements jina object creation in `partial-daemon`
 
         .. #noqa: DAR101"""
-        raise NotImplementedError
+        ...
 
-    @if_alive
-    async def _update(
-        self, uri: str, params: Dict, json: Optional[Dict] = None, **kwargs
-    ) -> Dict:
-        """Sends `PUT` request to `partial-daemon` to execute a command on a Flow.
-
-        :param uri: uri of partial-daemon
-        :param params: query params to be sent
-        :param json: json payload to be sent as body
-        :param kwargs: keyword args
-        :raises PartialDaemon400Exception: if update fails
-        :return: response from partial-daemon
-        """
-
-        self._logger.debug(
-            f'sending PUT request to partial-daemon on {uri}/{self._kind}'
-        )
-        async with aiohttp.request(
-            method='PUT', url=f'{uri}/{self._kind}', params=params, json=json
-        ) as response:
-            response_json = await response.json()
-            if response.status != HTTPStatus.OK:
-                raise PartialDaemon400Exception(error_msg_from(response_json))
-            return response_json
-
-    async def _delete(self, uri, *args, **kwargs):
+    @abstractmethod
+    async def delete_in_partial(self, uri, *args, **kwargs):
         """Implements jina object termination in `partial-daemon`
 
         .. #noqa: DAR101"""
-        raise NotImplementedError
+        ...
 
     async def ready(self, uri) -> bool:
         """Check if the container with partial-daemon is alive
@@ -208,7 +185,7 @@ async def add(
             kwargs.update(
                 {'ports': ports.dict()} if isinstance(ports, PortMappings) else {}
             )
-            object = await self._add(uri=uri, params=params, **kwargs)
+            object = await self.add_in_partial(uri=uri, params=params, **kwargs)
         except Exception as e:
             self._logger.error(f'{self._kind} creation failed as {e}')
             if container is not None:
@@ -249,47 +226,6 @@ async def add(
             workspace_store[workspace_id].metadata.managed_objects.add(id)
             return id
 
-    @BaseStore.dump
-    async def update(
-        self,
-        id: DaemonID,
-        kind: UpdateOperation,
-        pod_name: str,
-        uses_with: Optional[Dict] = None,
-        **kwargs,
-    ) -> DaemonID:
-        """Update the container in the store
-
-        :param id: id of the container
-        :param kind: type of update command to execute (only rolling_update for now)
-        :param pod_name: pod to target with the dump request
-        :param uses_with: the uses_with arguments to update the executor in pod_name
-        :param kwargs: keyword args
-        :raises KeyError: if id doesn't exist in the store
-        :return: id of the container
-        """
-        if id not in self:
-            raise KeyError(f'{colored(id, "red")} not found in store.')
-
-        if id.jtype == IDLiterals.JFLOW:
-            params = {'kind': kind.value, 'pod_name': pod_name}
-        elif id.jtype == IDLiterals.JPOD:
-            params = {'kind': kind.value}
-        else:
-            self._logger.error(f'update not supported for {id.type} {id}')
-            return id
-
-        uri = self[id].metadata.uri
-        try:
-            object = await self._update(uri, params, json=uses_with)
-        except Exception as e:
-            self._logger.error(f'Error while updating the {self._kind.title()}: \n{e}')
-            raise
-        else:
-            self[id].arguments.object = object
-            self._logger.success(f'{colored(id, "green")} is updated successfully')
-            return id
-
     @BaseStore.dump
     async def delete(self, id: DaemonID, **kwargs) -> None:
         """Delete a container from the store
@@ -303,9 +239,9 @@ async def delete(self, id: DaemonID, **kwargs) -> None:
 
         uri = self[id].metadata.uri
         try:
-            await self._delete(uri=uri)
+            await self.delete_in_partial(uri=uri)
         except Exception as e:
-            self._logger.error(f'Error while updating the {self._kind.title()}: \n{e}')
+            self._logger.error(f'Error while deleting the {self._kind.title()}: \n{e}')
             raise
         else:
             workspace_id = self[id].workspace_id
diff --git a/daemon/stores/flows.py b/daemon/stores/flows.py
index eff41d05e87bf..28e921b75a895 100644
--- a/daemon/stores/flows.py
+++ b/daemon/stores/flows.py
@@ -1,58 +1,99 @@
-from typing import Dict
-from http import HTTPStatus
-
-import aiohttp
+from typing import TYPE_CHECKING, Dict, Optional
 
+from jina.helper import colored
+from .base import BaseStore
+from .mixin import AiohttpMixin
 from .containers import ContainerStore
-from ..excepts import PartialDaemon400Exception
-from ..helper import if_alive, error_msg_from
+
+if TYPE_CHECKING:
+    from ..models import DaemonID
 
 
-class FlowStore(ContainerStore):
+class FlowStore(ContainerStore, AiohttpMixin):
     """A Store of Flows spawned as Containers by Daemon"""
 
     _kind = 'flow'
 
-    @if_alive
-    async def _add(self, uri: str, params: Dict, **kwargs) -> Dict:
+    async def add_in_partial(self, uri: str, params: Dict, **kwargs) -> Dict:
         """Sends `POST` request to `partial-daemon` to create a Flow.
 
         :param uri: uri of partial-daemon
-        :param params: json payload to be sent
+        :param params: Flow params
         :param kwargs: keyword args
-        :raises PartialDaemon400Exception: if creation fails
-        :return: response from partial-daemon
+        :return: response from POST request
         """
-        self._logger.debug(
-            f'sending POST request to partial-daemon on {uri}/{self._kind}'
-        )
         ports = kwargs.get('ports', [])
-        async with aiohttp.request(
-            method='POST',
+        return await self.POST(
             url=f'{uri}/{self._kind}',
+            params=None,
             json={'flow': params, 'ports': ports},
-        ) as response:
-            response_json = await response.json()
-            if response.status != HTTPStatus.CREATED:
-                raise PartialDaemon400Exception(error_msg_from(response_json))
-            return response_json
+        )
 
-    @if_alive
-    async def _delete(self, uri, **kwargs) -> Dict:
-        """Sends a `DELETE` request to terminate the Flow & remove the container
+    async def delete_in_partial(self, uri: str, **kwargs) -> Dict:
+        """Sends a `DELETE` request to `partial-daemon` to terminate the Flow
+        and, remove the container.
 
         :param uri: uri of partial-daemon
         :param kwargs: keyword args
         :raises PartialDaemon400Exception: if deletion fails
-        :return: response from partial-daemon
+        :return: response from DELETE request
         """
-        self._logger.debug(
-            f'sending DELETE request to partial-daemon on {uri}/{self._kind}'
-        )
-        async with aiohttp.request(
-            method='DELETE', url=f'{uri}/{self._kind}'
-        ) as response:
-            response_json = await response.json()
-            if response.status != HTTPStatus.OK:
-                raise PartialDaemon400Exception(error_msg_from(response_json))
-            return response_json
+        return await self.DELETE(url=f'{uri}/{self._kind}')
+
+    @BaseStore.dump
+    async def rolling_update(
+        self, id: 'DaemonID', pod_name: str, uses_with: Optional[Dict] = None
+    ) -> 'DaemonID':
+        """rolling_update the Flow inside the container
+
+        :param id: id of the Flow
+        :param pod_name: Pod in the Flow to be rolling updated
+        :param uses_with: the uses_with arguments to update the executor in pod_name
+        :raises KeyError: if id doesn't exist in the store
+        :return: id of the Flow
+        """
+
+        if id not in self:
+            raise KeyError(f'{colored(id, "red")} not found in store.')
+
+        try:
+            object = await self.PUT(
+                url=f'{self[id].metadata.uri}/{self._kind}/rolling_update',
+                params={'pod_name': pod_name},
+                json=uses_with,
+            )
+        except Exception as e:
+            self._logger.error(f'Error while sending rolling_update to the Flow: \n{e}')
+            raise
+        else:
+            self[id].arguments.object = object
+            self._logger.success(f'{colored(id, "green")} is updated successfully')
+            return id
+
+    @BaseStore.dump
+    async def scale(self, id: 'DaemonID', pod_name: str, replicas: int) -> 'DaemonID':
+        """Scale the Flow inside the container
+
+        :param id: id of the Flow
+        :param pod_name: Pod in the Flow to be rolling updated
+        :param replicas: The number of replicas to scale to
+        :raises KeyError: if id doesn't exist in the store
+        :return: id of the Flow
+        """
+
+        if id not in self:
+            raise KeyError(f'{colored(id, "red")} not found in store.')
+
+        try:
+            object = await self.PUT(
+                url=f'{self[id].metadata.uri}/{self._kind}/scale',
+                params={'pod_name': pod_name, 'replicas': replicas},
+                json=None,
+            )
+        except Exception as e:
+            self._logger.error(f'Error while scaling the Flow: \n{e}')
+            raise
+        else:
+            self[id].arguments.object = object
+            self._logger.success(f'{colored(id, "green")} is scaled successfully')
+            return id
diff --git a/daemon/stores/mixin.py b/daemon/stores/mixin.py
new file mode 100644
index 0000000000000..1265a203b068e
--- /dev/null
+++ b/daemon/stores/mixin.py
@@ -0,0 +1,67 @@
+from http import HTTPStatus
+from typing import Dict, Optional
+
+import aiohttp
+from jina.helper import T
+
+from ..excepts import PartialDaemon400Exception
+from ..helper import if_alive, error_msg_from
+
+
+class AiohttpMixin:
+    @if_alive
+    async def POST(
+        self: T, *, url: str, params: Optional[Dict] = None, json: Optional[Dict] = None
+    ):
+        """Sends `POST` request to `partial-daemon` to create a Flow/Pea/Pod.
+
+        :param url: uri of partial-daemon
+        :param params: query-params to be sent
+        :param json: json payload to be sent
+        :raises PartialDaemon400Exception: if POST request fails
+        :return: response from partial-daemon
+        """
+        self._logger.debug(f'sending POST request to partial-daemon on {url}')
+        async with aiohttp.request(
+            method='POST', url=url, params=params, json=json
+        ) as response:
+            response_json = await response.json()
+            if response.status != HTTPStatus.CREATED:
+                raise PartialDaemon400Exception(error_msg_from(response_json))
+            return response_json
+
+    @if_alive
+    async def PUT(
+        self: T, *, url: str, params: Optional[Dict] = None, json: Optional[Dict] = None
+    ):
+        """Sends `PUT` request to `partial-daemon` to update a Flow/Pea/Pod.
+
+        :param url: uri of partial-daemon
+        :param params: query-params to be sent
+        :param json: json payload to be sent
+        :raises PartialDaemon400Exception: if PUT request fails
+        :return: response from partial-daemon
+        """
+        self._logger.debug(f'sending PUT request to partial-daemon on {url}')
+        async with aiohttp.request(
+            method='PUT', url=url, params=params, json=json
+        ) as response:
+            response_json = await response.json()
+            if response.status != HTTPStatus.OK:
+                raise PartialDaemon400Exception(error_msg_from(response_json))
+            return response_json
+
+    @if_alive
+    async def DELETE(self: T, *, url: str):
+        """Sends `DELETE` request to `partial-daemon` to terminate a Flow/Pea/Pod.
+
+        :param url: uri of partial-daemon
+        :raises PartialDaemon400Exception: if POST request fails
+        :return: response from partial-daemon
+        """
+        self._logger.debug(f'sending DELETE request to partial-daemon on {url}')
+        async with aiohttp.request(method='DELETE', url=url) as response:
+            response_json = await response.json()
+            if response.status != HTTPStatus.OK:
+                raise PartialDaemon400Exception(error_msg_from(response_json))
+            return response_json
diff --git a/daemon/stores/partial.py b/daemon/stores/partial.py
index a15503f9baa3e..d42942e9ebffe 100644
--- a/daemon/stores/partial.py
+++ b/daemon/stores/partial.py
@@ -1,8 +1,9 @@
 from pathlib import Path
+from abc import ABC, abstractmethod
 from argparse import Namespace
 from typing import Dict, Optional, Union
 
-from jina.helper import colored, random_port
+from jina.helper import colored, random_port, T
 from jina.peapods import Pea, Pod, CompoundPod
 from jina.peapods.peas.helper import update_runtime_cls
 from jina import Flow, __docker_host__
@@ -15,31 +16,26 @@
 from ..models.partial import PartialFlowItem, PartialStoreItem
 
 
-class PartialStore:
+class PartialStore(ABC):
     """A store spawned inside partial-daemon container"""
 
     def __init__(self):
         self._logger = JinaLogger(self.__class__.__name__, **vars(jinad_args))
         self.item = PartialStoreItem()
+        self.object: Union['Pea', 'Pod', 'Flow'] = None
 
+    @abstractmethod
     def add(self, *args, **kwargs) -> PartialStoreItem:
         """Add a new element to the store. This method needs to be overridden by the subclass
 
 
         .. #noqa: DAR101"""
-        raise NotImplementedError
-
-    def update(self, *args, **kwargs) -> PartialStoreItem:
-        """Updates the element to the store. This method needs to be overridden by the subclass
-
-
-        .. #noqa: DAR101"""
-        raise NotImplementedError
+        ...
 
     def delete(self) -> None:
         """Terminates the object in the store & stops the server"""
         try:
-            if hasattr(self, 'object'):
+            if hasattr(self.object, 'close'):
                 self.object.close()
             else:
                 self._logger.warning(f'nothing to close. exiting')
@@ -83,31 +79,38 @@ class PartialPodStore(PartialPeaStore):
 
     peapod_cls = Pod
 
-    def update(
-        self,
-        kind: UpdateOperation,
-        pod_name: str,
-        uses_with: Optional[Dict] = None,
-        **kwargs,
+    async def rolling_update(
+        self, uses_with: Optional[Dict] = None
     ) -> PartialStoreItem:
-        """Runs an update operation on the Flow.
-        :param kind: type of update command to execute (rolling_update)
-        :param pod_name: pod to target with the dump request
-        :param uses_with: the uses_with for the new executors to be rolled updated
-        :param kwargs: keyword args
+        """Perform rolling_update on current Pod
+
+        :param uses_with: a Dictionary of arguments to restart the executor with
         :return: Item describing the Flow object
         """
         try:
-            if kind == UpdateOperation.ROLLING_UPDATE:
-                self.object.rolling_update(uses_with=uses_with)
-            else:
-                self._logger.error(f'unsupported kind: {kind}, no changes done')
-                return self.item
+            await self.object.rolling_update(uses_with=uses_with)
         except Exception as e:
             self._logger.error(f'{e!r}')
             raise
         else:
             self.item.arguments = vars(self.object.args)
+            self._logger.success(f'Pod is successfully rolling_updated!')
+            return self.item
+
+    async def scale(self, replicas: int) -> PartialStoreItem:
+        """Scale the current Pod
+
+        :param replicas: number of replicas for the Pod
+        :return: Item describing the Flow object
+        """
+        try:
+            await self.object.scale(replicas=replicas)
+        except Exception as e:
+            self._logger.error(f'{e!r}')
+            raise
+        else:
+            self.item.arguments = vars(self.object.args)
+            self._logger.success(f'Pod is successfully scaled!')
             return self.item
 
 
@@ -222,30 +225,38 @@ def _set_pea_ports(self, pea_args, port_mapping, port_name):
                 ),
             )
 
-    def update(
-        self,
-        kind: UpdateOperation,
-        pod_name: str,
-        uses_with: Optional[Dict] = None,
-        **kwargs,
+    def rolling_update(
+        self, pod_name: str, uses_with: Optional[Dict] = None
     ) -> PartialFlowItem:
-        """Runs an update operation on the Flow.
-        :param kind: type of update command to execute (rolling_update)
-        :param pod_name: pod to target with the dump request
-        :param uses_with: the uses_with for the new executors to be rolled updated
-        :param kwargs: keyword args
+        """Perform rolling_update on the Pod in current Flow
+
+        :param pod_name: Pod in the Flow to be rolling updated
+        :param uses_with: a Dictionary of arguments to restart the executor with
         :return: Item describing the Flow object
         """
         try:
-            if kind == UpdateOperation.ROLLING_UPDATE:
-                self.object.rolling_update(pod_name=pod_name, uses_with=uses_with)
-            else:
-                self._logger.error(f'unsupported kind: {kind}, no changes done')
-                return self.item
+            self.object.rolling_update(pod_name=pod_name, uses_with=uses_with)
+        except Exception as e:
+            self._logger.error(f'{e!r}')
+            raise
+        else:
+            self.item.arguments = vars(self.object.args)
+            self._logger.success(f'Flow is successfully rolling_updated!')
+            return self.item
+
+    def scale(self, pod_name: str, replicas: int) -> PartialFlowItem:
+        """Scale the Pod in current Flow
+
+        :param pod_name: Pod to be scaled
+        :param replicas: number of replicas for the Pod
+        :return: Item describing the Flow object
+        """
+        try:
+            self.object.scale(pod_name=pod_name, replicas=replicas)
         except Exception as e:
             self._logger.error(f'{e!r}')
             raise
         else:
             self.item.arguments = vars(self.object.args)
-            self._logger.success(f'Flow is updated successfully!')
+            self._logger.success(f'Flow is successfully scaled!')
             return self.item
diff --git a/daemon/stores/peas.py b/daemon/stores/peas.py
index 29c688b6a9c86..84c76278c34ae 100644
--- a/daemon/stores/peas.py
+++ b/daemon/stores/peas.py
@@ -1,51 +1,35 @@
 from typing import Dict
-from http import HTTPStatus
-
-import aiohttp
 
+from .mixin import AiohttpMixin
 from .containers import ContainerStore
-from ..excepts import PartialDaemon400Exception
-from ..helper import if_alive, error_msg_from
 
 
-class PeaStore(ContainerStore):
+class PeaStore(ContainerStore, AiohttpMixin):
     """A Store of Peas spawned as Containers by Daemon"""
 
     _kind = 'pea'
 
-    @if_alive
-    async def _add(self, uri: str, params: Dict, **kwargs) -> Dict:
-        """Sends `POST` request to `mini-jinad` to create a Pea/Pod.
+    async def add_in_partial(self, uri: str, params: Dict, **kwargs) -> Dict:
+        """Sends `POST` request to `partial-daemon` to create a Pea/Pod.
 
-        :param uri: uri of mini-jinad
+        :param uri: uri of partial-daemon
         :param params: json payload to be sent
         :param kwargs: keyword args
         :raises PartialDaemon400Exception: if creation fails
         :return: response from mini-jinad
         """
-        self._logger.debug(f'sending POST request to mini-jinad on {uri}/{self._kind}')
-        async with aiohttp.request(
-            method='POST', url=f'{uri}/{self._kind}', json=params
-        ) as response:
-            response_json = await response.json()
-            if response.status != HTTPStatus.CREATED:
-                raise PartialDaemon400Exception(error_msg_from(response_json))
-            return response_json
-
-    @if_alive
-    async def _delete(self, uri, **kwargs) -> Dict:
-        """Sends a `DELETE` request to `mini-jinad` to terminate a Pea/Pod
-
-        :param uri: uri of mini-jinad
+        return await self.POST(
+            url=f'{uri}/{self._kind}',
+            params=None,
+            json=params,
+        )
+
+    async def delete_in_partial(self, uri, **kwargs) -> Dict:
+        """Sends a `DELETE` request to `partial-daemon` to terminate a Pea/Pod
+
+        :param uri: uri of partial-daemon
         :param kwargs: keyword args
         :raises PartialDaemon400Exception: if deletion fails
-        :return: response from mini-jinad
+        :return: response from partial-daemon
         """
-        self._logger.debug(
-            f'sending DELETE request to mini-jinad on {uri}/{self._kind}'
-        )
-        async with aiohttp.request('DELETE', url=f'{uri}/{self._kind}') as response:
-            response_json = await response.json()
-            if response.status != HTTPStatus.OK:
-                raise PartialDaemon400Exception(error_msg_from(response_json))
-            return response_json
+        return await self.DELETE(url=f'{uri}/{self._kind}')
diff --git a/daemon/stores/pods.py b/daemon/stores/pods.py
index f112d50ef7279..db1053ed814be 100644
--- a/daemon/stores/pods.py
+++ b/daemon/stores/pods.py
@@ -1,7 +1,70 @@
+from typing import TYPE_CHECKING, Dict, Optional
+
+from jina.helper import colored
+from .base import BaseStore
 from .peas import PeaStore
 
+if TYPE_CHECKING:
+    from ..models import DaemonID
+
 
 class PodStore(PeaStore):
     """A Store of Pods spawned as Containers by Daemon"""
 
     _kind = 'pod'
+
+    @BaseStore.dump
+    async def rolling_update(
+        self, id: 'DaemonID', uses_with: Optional[Dict] = None
+    ) -> 'DaemonID':
+        """rolling_update the Pod inside the container
+
+        :param id: id of the Pod
+        :param uses_with: the uses_with arguments to update the executor in pod_name
+        :raises KeyError: if id doesn't exist in the store
+        :return: id of the Flow
+        """
+
+        if id not in self:
+            raise KeyError(f'{colored(id, "red")} not found in store.')
+
+        try:
+            object = await self.PUT(
+                url=f'{self[id].metadata.uri}/{self._kind}/rolling_update',
+                params=None,
+                json=uses_with,
+            )
+        except Exception as e:
+            self._logger.error(f'Error while sending rolling_update to the Pod: \n{e}')
+            raise
+        else:
+            self[id].arguments.object = object
+            self._logger.success(f'{colored(id, "green")} is updated successfully')
+            return id
+
+    @BaseStore.dump
+    async def scale(self, id: 'DaemonID', replicas: int) -> 'DaemonID':
+        """Scale the Pod inside the container
+
+        :param id: id of the Pod
+        :param replicas: The number of replicas to scale to
+        :raises KeyError: if id doesn't exist in the store
+        :return: id of the Flow
+        """
+
+        if id not in self:
+            raise KeyError(f'{colored(id, "red")} not found in store.')
+
+        try:
+            object = await self.PUT(
+                url=f'{self[id].metadata.uri}/{self._kind}/scale',
+                params={'replicas': replicas},
+                json=None,
+            )
+        except Exception as e:
+            self._logger.error(f'Error while scaling the Pod: \n{e}')
+            raise
+        else:
+            self[id].arguments.object = object
+            self._logger.success(f'{colored(id, "green")} is scaled successfully')
+            return id
diff --git a/jina/peapods/runtimes/jinad/__init__.py b/jina/peapods/runtimes/jinad/__init__.py
index 26f57436f7d33..aa8ccae2c0403 100644
--- a/jina/peapods/runtimes/jinad/__init__.py
+++ b/jina/peapods/runtimes/jinad/__init__.py
@@ -32,6 +32,8 @@ def __init__(
         args: 'argparse.Namespace',
         **kwargs,
     ):
+        self.pea_id = None
+        self.logstream = None
         super().__init__(args, **kwargs)
         # Need the `proper` control address to send `activate` and `deactivate` signals, from the pea in the `main`
         # process.
@@ -39,8 +41,6 @@ def __init__(
         self.timeout_ctrl = args.timeout_ctrl
         self.host = args.host
         self.port_jinad = args.port_jinad
-        self.pea_id = None
-        self.logstream = None
 
     async def async_setup(self):
         """Create Workspace, Pea on remote JinaD server"""
@@ -101,6 +101,10 @@ async def async_run_forever(self):
             )
 
     async def async_cancel(self):
+        """Skip cancel for JinadRuntime"""
+        pass
+
+    async def async_teardown(self):
         """Cancels the logstream task, removes the remote Pea"""
         if self.logstream is not None:
             self.logstream.cancel()
diff --git a/tests/distributed/test_scale_remote/__init__.py b/tests/distributed/test_scale_remote/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/tests/distributed/test_scale_remote/executors.py b/tests/distributed/test_scale_remote/executors.py
new file mode 100644
index 0000000000000..bc74e13e581db
--- /dev/null
+++ b/tests/distributed/test_scale_remote/executors.py
@@ -0,0 +1,14 @@
+from jina import Executor, requests
+
+
+class ScalableExecutor(Executor):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.replica_id = self.runtime_args.replica_id
+        self.shard_id = self.runtime_args.shard_id
+
+    @requests
+    def foo(self, docs, **kwargs):
+        for doc in docs:
+            doc.tags['replica_id'] = self.replica_id
+            doc.tags['shard_id'] = self.shard_id
diff --git a/tests/distributed/test_scale_remote/test_local_flow_remote_executor.py b/tests/distributed/test_scale_remote/test_local_flow_remote_executor.py
new file mode 100644
index 0000000000000..5f728fef1af3e
--- /dev/null
+++ b/tests/distributed/test_scale_remote/test_local_flow_remote_executor.py
@@ -0,0 +1,180 @@
+import os
+import time
+import multiprocessing
+from functools import partial
+
+import pytest
+
+from jina import Flow, Document, DocumentArray, Client
+
+cur_dir = os.path.dirname(os.path.abspath(__file__))
+IMG_NAME = 'jina/scale-executor'
+
+NUM_CONCURRENT_CLIENTS = 20
+NUM_DOCS_SENT_BY_CLIENTS = 50
+CLOUD_HOST = 'localhost:8000'
+
+
+@pytest.fixture
+def pod_params(request):
+    num_replicas, scale_to, shards = request.param
+    return num_replicas, scale_to, shards
+
+
+@pytest.fixture(scope='function')
+def docker_image_built():
+    import docker
+
+    client = docker.from_env()
+    client.images.build(
+        path=os.path.join(cur_dir, '../../integration/scale/scale-executor'),
+        tag=IMG_NAME,
+    )
+    client.close()
+    yield
+    time.sleep(2)
+    client = docker.from_env()
+    client.containers.prune()
+
+
+@pytest.fixture
+def remote_flow_with_zed_runtime(pod_params):
+    num_replicas, scale_to, shards = pod_params
+    return Flow().add(
+        name='executor',
+        uses='ScalableExecutor',
+        replicas=num_replicas,
+        shards=shards,
+        polling='ANY',
+        host=CLOUD_HOST,
+        py_modules='executors.py',
+        upload_files=cur_dir,
+    )
+
+
+@pytest.fixture
+def remote_flow_with_container_runtime(pod_params, docker_image_built):
+    num_replicas, scale_to, shards = pod_params
+    return Flow().add(
+        name='executor',
+        uses=f'docker://{IMG_NAME}',
+        replicas=num_replicas,
+        shards=shards,
+        polling='ANY',
+        host=CLOUD_HOST,
+    )
+
+
+@pytest.fixture(
+    params=['remote_flow_with_zed_runtime', 'remote_flow_with_container_runtime']
+)
+def remote_flow_with_runtime(request):
+    return request.getfixturevalue(request.param)
+
+
+@pytest.mark.parametrize(
+    'pod_params',  # (num_replicas, scale_to, shards)
+    [
+        (2, 3, 1),  # scale up 1 replica with 1 shard
+        (2, 3, 2),  # scale up 1 replica with 2 shards
+        (3, 1, 1),  # scale down 2 replicas with 1 shard
+        (3, 1, 2),  # scale down 2 replicas with 1 shard
+    ],
+    indirect=True,
+)
+def test_scale_success(remote_flow_with_runtime: Flow, pod_params):
+    num_replicas, scale_to, shards = pod_params
+    with remote_flow_with_runtime as f:
+        ret1 = f.index(
+            inputs=DocumentArray([Document() for _ in range(200)]),
+            return_results=True,
+            request_size=10,
+        )
+        f.scale(pod_name='executor', replicas=scale_to)
+        ret2 = f.index(
+            inputs=DocumentArray([Document() for _ in range(200)]),
+            return_results=True,
+            request_size=10,
+        )
+
+        assert len(ret1) == 20
+        replica_ids = set()
+        for r in ret1:
+            assert len(r.docs) == 10
+            for replica_id in r.docs.get_attributes('tags__replica_id'):
+                replica_ids.add(replica_id)
+
+        assert replica_ids == set(range(num_replicas))
+
+        assert len(ret2) == 20
+        replica_ids = set()
+        for r in ret2:
+            assert len(r.docs) == 10
+            for replica_id in r.docs.get_attributes('tags__replica_id'):
+                replica_ids.add(replica_id)
+
+        assert replica_ids == set(range(scale_to))
+
+
+@pytest.mark.parametrize(
+    'pod_params',
+    [
+        (2, 3, 1),
+        (3, 2, 1),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize('protocol', ['grpc', 'http', 'websocket'])
+def test_scale_with_concurrent_client(
+    remote_flow_with_runtime: Flow, pod_params, protocol
+):
+    def peer_client(port, protocol, peer_hash, queue):
+        rv = Client(protocol=protocol, port=port).index(
+            [Document(text=peer_hash) for _ in range(NUM_DOCS_SENT_BY_CLIENTS)],
+            request_size=5,
+            return_results=True,
+        )
+        for r in rv:
+            for doc in r.docs:
+                # our proto objects are not fit to be sent by queues
+                queue.put(doc.text)
+
+    num_replicas, scale_to, _ = pod_params
+    queue = multiprocessing.Queue()
+
+    with remote_flow_with_runtime as f:
+        f.protocol = protocol
+        port_expose = f.port_expose
+
+        thread_pool = []
+        for peer_id in range(NUM_CONCURRENT_CLIENTS):
+            # test
+            t = multiprocessing.Process(
+                target=partial(peer_client, port_expose, protocol, str(peer_id), queue)
+            )
+            t.start()
+            thread_pool.append(t)
+
+        f.scale(pod_name='executor', replicas=scale_to)
+
+        for t in thread_pool:
+            t.join()
+
+        c = Client(protocol=protocol, port=port_expose)
+        rv = c.index(
+            [Document() for _ in range(5)], request_size=1, return_results=True
+        )
+
+    assert queue.qsize() == NUM_CONCURRENT_CLIENTS * NUM_DOCS_SENT_BY_CLIENTS
+    all_docs = []
+    while not queue.empty():
+        all_docs.append(queue.get())
+
+    assert len(all_docs) == NUM_CONCURRENT_CLIENTS * NUM_DOCS_SENT_BY_CLIENTS
+
+    assert len(rv) == 5
+
+    replicas = []
+    for r in rv:
+        assert len(r.docs) == 1
+        replicas.append(r.docs[0].tags['replica_id'])
diff --git a/tests/integration/scale/test_scale.py b/tests/integration/scale/test_scale.py
index 178cf7e586ebc..f932c928d29fe 100644
--- a/tests/integration/scale/test_scale.py
+++ b/tests/integration/scale/test_scale.py
@@ -5,9 +5,8 @@
 
 import pytest
 
-from jina import Flow, Executor, Document, DocumentArray, requests
+from jina import Flow, Executor, Document, DocumentArray, requests, Client
 from jina.excepts import RuntimeFailToStart, ScalingFails
-from jina.clients import Client
 
 cur_dir = os.path.dirname(os.path.abspath(__file__))
 IMG_NAME = 'jina/scale-executor'


Code review comments:
diff_hunk: @@ -60,7 +57,24 @@ async def _update(
     .. #noqa: DAR201
     """
     try:
-        return store.update(kind, pod_name, uses_with=uses_with)
+        return store.rolling_update(pod_name=pod_name, uses_with=uses_with)
+    except ValueError as ex:
+        raise PartialDaemon400Exception from ex
+
+
+@router.put(
+    path='/scale',
+    summary='Run a rolling_update operation on the Flow object',
path: daemon/api/endpoints/partial/flow.py
text: @user1:
wrong summary