Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement web API of get_infos #2558

Merged
merged 2 commits into from
Oct 28, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions mars/services/storage/api/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
"""
19 changes: 19 additions & 0 deletions mars/services/storage/api/web.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}

Expand Down Expand Up @@ -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:
Expand All @@ -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)
5 changes: 5 additions & 0 deletions mars/services/storage/tests/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)