Implement web API of get_infos (mars-project#2558)

wjsi · Oct 28, 2021 · a7779b8 · a7779b8
1 parent b1b033f
commit a7779b8
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 0 deletions.
diff --git a/mars/services/storage/api/core.py b/mars/services/storage/api/core.py
@@ -64,3 +64,18 @@ async def put(
         object information: ObjectInfo
             the put object information
         """
+
+    @abstractmethod
+    async def get_infos(self, data_key: str) -> List[DataInfo]:
+        """
+        Get data information items for specific data key
+
+        Parameters
+        ----------
+        data_key
+
+        Returns
+        -------
+        out
+            List of information for specified key
+        """
diff --git a/mars/services/storage/api/web.py b/mars/services/storage/api/web.py
@@ -77,6 +77,12 @@ async def put_data(self, session_id: str, data_key: str):
         )
         self.write(serialize_serializable(res))
 
+    @web_api("(?P<data_key>[^/]+)", method="get", arg_filter={"action": "get_infos"})
+    async def get_infos(self, session_id: str, data_key: str):
+        oscar_api = await self._get_storage_api_by_object_id(session_id, data_key)
+        res = await oscar_api.get_infos(data_key)
+        self.write(serialize_serializable(res))
+
 
 web_handlers = {StorageWebAPIHandler.get_root_pattern(): StorageWebAPIHandler}
 
@@ -104,6 +110,7 @@ async def get(
         )
         return deserialize_serializable(res.body)
 
+    @mo.extensible
     async def put(
         self, data_key: str, obj: object, level: StorageLevel = StorageLevel.MEMORY
     ) -> DataInfo:
@@ -117,3 +124,15 @@ async def put(
             data=serialize_serializable(obj),
         )
         return deserialize_serializable(res.body)
+
+    @mo.extensible
+    async def get_infos(self, data_key: str) -> List[DataInfo]:
+        path = f"{self._address}/api/session/{self._session_id}/storage/{data_key}"
+        params = dict(action="get_infos")
+        res = await self._request_url(
+            path=path,
+            method="GET",
+            headers={"Content-Type": "application/octet-stream"},
+            params=params,
+        )
+        return deserialize_serializable(res.body)
diff --git a/mars/services/storage/tests/test_api.py b/mars/services/storage/tests/test_api.py
@@ -188,5 +188,10 @@ async def test_web_storage_api():
         )
         np.testing.assert_array_equal(value[3:5, :], sliced_value)
 
+        infos = await web_storage_api.get_infos(t.chunks[0].key)
+        assert len(infos) == 1
+        assert infos[0].level == StorageLevel.MEMORY
+        assert infos[0].memory_size == t.chunks[0].nbytes
+
         await MockStorageAPI.cleanup(pool.external_address)
         await MockClusterAPI.cleanup(pool.external_address)