From efb038b6f9f268f4a4b34f08427ba39a8691fd74 Mon Sep 17 00:00:00 2001
From: Wei Ouyang <oeway007@gmail.com>
Date: Mon, 25 Nov 2024 23:41:05 -0800
Subject: [PATCH] Improve redis event bus (#718)

* small fix

* Improve the speed

* Fix tests

* Fix teardown

* Fix file count

* small fix
---
 CHANGELOG.md                                  |  6 ++
 docs/getting-started.md                       |  2 +-
 docs/migration-guide.md                       | 10 +--
 docs/service-type-annotation.md               |  2 +-
 helm-charts/aks-hypha.md                      |  2 +-
 helm-charts/hypha-server/Chart.yaml           |  2 +-
 helm-charts/hypha-server/values.yaml          |  2 +-
 hypha/VERSION                                 |  2 +-
 hypha/artifact.py                             | 79 ++++++++++---------
 hypha/core/__init__.py                        | 74 ++++++++++-------
 hypha/core/store.py                           |  7 ++
 .../hypha-core-app/hypha-app-webpython.js     |  2 +-
 hypha/templates/ws/index.html                 |  4 +-
 hypha/websocket.py                            | 25 ++++--
 pytest.ini                                    |  5 +-
 requirements.txt                              |  2 +-
 setup.py                                      |  2 +-
 tests/test_artifact.py                        | 19 +++--
 18 files changed, 148 insertions(+), 99 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2c4c6669..494ca152 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Hypha Change Log
 
+### 0.20.40
+ 
+ - Add vector store service to support vector search and retrieval.
+ - Fix zenodo file upload issue
+ - Speed up server by removing the `asyncio.sleep(0.01)` throttling and support concurrent handling of events in the redis event bus.
+
 ### 0.20.39
 
  - Revise artifact manager to use artifact id as the primary key, remove `prefix` based keys.
diff --git a/docs/getting-started.md b/docs/getting-started.md
index eb1fdaed..6db5031c 100644
--- a/docs/getting-started.md
+++ b/docs/getting-started.md
@@ -220,7 +220,7 @@ svc = await get_remote_service("http://localhost:9527/ws-user-scintillating-lawy
 Include the following script in your HTML file to load the `hypha-rpc` client:
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.39/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.40/dist/hypha-rpc-websocket.min.js"></script>
 ```
 
 Use the following code in JavaScript to connect to the server and access an existing service:
diff --git a/docs/migration-guide.md b/docs/migration-guide.md
index 7ad14308..5be913e4 100644
--- a/docs/migration-guide.md
+++ b/docs/migration-guide.md
@@ -15,7 +15,7 @@ To connect to the server, instead of installing the `imjoy-rpc` module, you will
 pip install -U hypha-rpc # new install
 ```
 
-We also changed our versioning strategy, we use the same version number for the server and client, so it's easier to match the client and server versions. For example, `hypha-rpc` version `0.20.39` is compatible with Hypha server version `0.20.39`.
+We also changed our versioning strategy, we use the same version number for the server and client, so it's easier to match the client and server versions. For example, `hypha-rpc` version `0.20.40` is compatible with Hypha server version `0.20.40`.
 
 #### 2. Change the imports to use `hypha-rpc`
 
@@ -128,10 +128,10 @@ loop.run_forever()
 To connect to the server, instead of using the `imjoy-rpc` module, you will need to use the `hypha-rpc` module. The `hypha-rpc` module is a standalone module that provides the RPC connection to the Hypha server. You can include it in your HTML using a script tag:
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.39/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.40/dist/hypha-rpc-websocket.min.js"></script>
 ```
 
-We also changed our versioning strategy, we use the same version number for the server and client, so it's easier to match the client and server versions. For example, `hypha-rpc` version `0.20.39` is compatible with Hypha server version `0.20.39`.
+We also changed our versioning strategy, we use the same version number for the server and client, so it's easier to match the client and server versions. For example, `hypha-rpc` version `0.20.40` is compatible with Hypha server version `0.20.40`.
 
 #### 2. Change the connection method and use camelCase for service function names
 
@@ -149,7 +149,7 @@ Here is a suggested list of search and replace operations to update your code:
 Here is an example of how the updated code might look:
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.39/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.40/dist/hypha-rpc-websocket.min.js"></script>
 <script>
 async function main(){
     const server = await hyphaWebsocketClient.connectToServer({"server_url": "https://hypha.amun.ai"});
@@ -197,7 +197,7 @@ We created a tutorial to introduce this new feature: [service type annotation](.
 Here is a quick example in JavaScript:
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.39/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.40/dist/hypha-rpc-websocket.min.js"></script>
 
 <script>
 async function main(){
diff --git a/docs/service-type-annotation.md b/docs/service-type-annotation.md
index a94e1193..b7ba9510 100644
--- a/docs/service-type-annotation.md
+++ b/docs/service-type-annotation.md
@@ -134,7 +134,7 @@ if __name__ == "__main__":
 **JavaScript Client: Service Usage**
 
 ```html
-<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.39/dist/hypha-rpc-websocket.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/hypha-rpc@0.20.40/dist/hypha-rpc-websocket.min.js"></script>
 <script>
 async function main() {
     const server = await hyphaWebsocketClient.connectToServer({"server_url": "https://hypha.amun.ai"});
diff --git a/helm-charts/aks-hypha.md b/helm-charts/aks-hypha.md
index bed2891e..ed737f92 100644
--- a/helm-charts/aks-hypha.md
+++ b/helm-charts/aks-hypha.md
@@ -146,7 +146,7 @@ replicaCount: 1
 image:
   repository: ghcr.io/amun-ai/hypha
   pullPolicy: IfNotPresent
-  tag: "0.20.39"
+  tag: "0.20.40"
 
 serviceAccount:
   create: true
diff --git a/helm-charts/hypha-server/Chart.yaml b/helm-charts/hypha-server/Chart.yaml
index 62b1e20d..7c46c6fa 100644
--- a/helm-charts/hypha-server/Chart.yaml
+++ b/helm-charts/hypha-server/Chart.yaml
@@ -15,7 +15,7 @@ type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
 # Versions are expected to follow Semantic Versioning (https://semver.org/)
-version: 0.20.39
+version: 0.20.40
 
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application. Versions are not expected to
diff --git a/helm-charts/hypha-server/values.yaml b/helm-charts/hypha-server/values.yaml
index 945e3443..d36a7770 100644
--- a/helm-charts/hypha-server/values.yaml
+++ b/helm-charts/hypha-server/values.yaml
@@ -8,7 +8,7 @@ image:
   repository: ghcr.io/amun-ai/hypha
   pullPolicy: IfNotPresent
   # Overrides the image tag whose default is the chart appVersion.
-  tag: "0.20.39"
+  tag: "0.20.40"
 
 imagePullSecrets: []
 nameOverride: ""
diff --git a/hypha/VERSION b/hypha/VERSION
index 99b962fd..0a3d03b9 100644
--- a/hypha/VERSION
+++ b/hypha/VERSION
@@ -1,3 +1,3 @@
 {
-    "version": "0.20.39.post20"
+    "version": "0.20.40"
 }
diff --git a/hypha/artifact.py b/hypha/artifact.py
index 731e95b9..b155896c 100644
--- a/hypha/artifact.py
+++ b/hypha/artifact.py
@@ -1798,46 +1798,47 @@ async def commit(
 
                 versions = artifact.versions or []
                 artifact.config = artifact.config or {}
-                if artifact.staging:
-                    s3_config = self._get_s3_config(artifact, parent_artifact)
-                    async with self._create_client_async(
-                        s3_config,
-                    ) as s3_client:
-                        download_weights = {}
-                        for file_info in artifact.staging:
-                            file_key = safe_join(
-                                s3_config["prefix"],
-                                artifact.workspace,
-                                f"{self._artifacts_dir}/{artifact.id}/v{version_index}/{file_info['path']}",
-                            )
-                            try:
-                                await s3_client.head_object(
-                                    Bucket=s3_config["bucket"], Key=file_key
-                                )
-                            except ClientError:
-                                raise FileNotFoundError(
-                                    f"File '{file_info['path']}' does not exist in the artifact."
-                                )
-                            if (
-                                file_info.get("download_weight") is not None
-                                and file_info["download_weight"] > 0
-                            ):
-                                download_weights[file_info["path"]] = file_info[
-                                    "download_weight"
-                                ]
-                        if download_weights:
-                            artifact.config["download_weights"] = download_weights
-                            flag_modified(artifact, "config")
-
-                        artifact.file_count = await self._count_files_in_prefix(
-                            s3_client,
-                            s3_config["bucket"],
-                            safe_join(
-                                s3_config["prefix"],
-                                artifact.workspace,
-                                f"{self._artifacts_dir}/{artifact.id}/v{version_index}",
-                            ),
+
+                s3_config = self._get_s3_config(artifact, parent_artifact)
+                async with self._create_client_async(
+                    s3_config,
+                ) as s3_client:
+                    download_weights = {}
+                    for file_info in artifact.staging or []:
+                        file_key = safe_join(
+                            s3_config["prefix"],
+                            artifact.workspace,
+                            f"{self._artifacts_dir}/{artifact.id}/v{version_index}/{file_info['path']}",
                         )
+                        try:
+                            await s3_client.head_object(
+                                Bucket=s3_config["bucket"], Key=file_key
+                            )
+                        except ClientError:
+                            raise FileNotFoundError(
+                                f"File '{file_info['path']}' does not exist in the artifact."
+                            )
+                        if (
+                            file_info.get("download_weight") is not None
+                            and file_info["download_weight"] > 0
+                        ):
+                            download_weights[file_info["path"]] = file_info[
+                                "download_weight"
+                            ]
+
+                    if download_weights:
+                        artifact.config["download_weights"] = download_weights
+                        flag_modified(artifact, "config")
+
+                    artifact.file_count = await self._count_files_in_prefix(
+                        s3_client,
+                        s3_config["bucket"],
+                        safe_join(
+                            s3_config["prefix"],
+                            artifact.workspace,
+                            f"{self._artifacts_dir}/{artifact.id}/v{version_index}",
+                        ),
+                    )
 
                 parent_artifact_config = (
                     parent_artifact.config if parent_artifact else {}
diff --git a/hypha/core/__init__.py b/hypha/core/__init__.py
index 50757292..b1e367d9 100644
--- a/hypha/core/__init__.py
+++ b/hypha/core/__init__.py
@@ -5,6 +5,7 @@
 import json
 import logging
 import sys
+import os
 from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from pydantic import BaseModel, Field, field_validator
@@ -547,7 +548,11 @@ async def init(self):
         loop = asyncio.get_running_loop()
         self._ready = loop.create_future()
         self._loop = loop
+
+        # Start the Redis subscription task
         loop.create_task(self._subscribe_redis())
+
+        # Wait for readiness signal
         await self._ready
 
     def on(self, event_name, func):
@@ -652,39 +657,50 @@ def stop(self):
         self._stop = True
 
     async def _subscribe_redis(self):
+        cpu_count = os.cpu_count() or 1
+        concurrent_tasks = cpu_count * 10
         pubsub = self._redis.pubsub()
         self._stop = False
-        try:
-            await pubsub.psubscribe("event:*")
-            self._ready.set_result(True)
-            while self._stop is False:
-                msg = await pubsub.get_message(ignore_subscribe_messages=True)
+        semaphore = asyncio.Semaphore(concurrent_tasks)  # Limit concurrent tasks
+
+        async def process_message(msg):
+            """Process a single message while respecting the semaphore."""
+            async with semaphore:  # Acquire semaphore
                 try:
-                    if msg:
-                        channel = msg["channel"].decode("utf-8")
-                        RedisEventBus._counter.labels(event="*").inc()
-                        if channel.startswith("event:b:"):
-                            event_type = channel[8:]
-                            data = msg["data"]
-                            await self._redis_event_bus.emit(event_type, data)
-                            if ":" not in event_type:
-                                RedisEventBus._counter.labels(event=event_type).inc()
-                        elif channel.startswith("event:d:"):
-                            event_type = channel[8:]
-                            data = json.loads(msg["data"])
-                            await self._redis_event_bus.emit(event_type, data)
-                            if ":" not in event_type:
-                                RedisEventBus._counter.labels(event=event_type).inc()
-                        elif channel.startswith("event:s:"):
-                            event_type = channel[8:]
-                            data = msg["data"].decode("utf-8")
-                            await self._redis_event_bus.emit(event_type, data)
-                            if ":" not in event_type:
-                                RedisEventBus._counter.labels(event=event_type).inc()
-                        else:
-                            logger.info("Unknown channel: %s", channel)
+                    channel = msg["channel"].decode("utf-8")
+                    RedisEventBus._counter.labels(event="*").inc()
+
+                    if channel.startswith("event:b:"):
+                        event_type = channel[8:]
+                        data = msg["data"]
+                        await self._redis_event_bus.emit(event_type, data)
+                        if ":" not in event_type:
+                            RedisEventBus._counter.labels(event=event_type).inc()
+                    elif channel.startswith("event:d:"):
+                        event_type = channel[8:]
+                        data = json.loads(msg["data"])
+                        await self._redis_event_bus.emit(event_type, data)
+                        if ":" not in event_type:
+                            RedisEventBus._counter.labels(event=event_type).inc()
+                    elif channel.startswith("event:s:"):
+                        event_type = channel[8:]
+                        data = msg["data"].decode("utf-8")
+                        await self._redis_event_bus.emit(event_type, data)
+                        if ":" not in event_type:
+                            RedisEventBus._counter.labels(event=event_type).inc()
+                    else:
+                        logger.info("Unknown channel: %s", channel)
                 except Exception as exp:
                     logger.exception(f"Error processing message: {exp}")
-                await asyncio.sleep(0)
+
+        try:
+            await pubsub.psubscribe("event:*")
+            self._ready.set_result(True)
+            while not self._stop:
+                msg = await pubsub.get_message(
+                    ignore_subscribe_messages=True, timeout=0.05
+                )
+                if msg:
+                    asyncio.create_task(process_message(msg))  # Add task to pool
         except Exception as exp:
             self._ready.set_exception(exp)
diff --git a/hypha/core/store.py b/hypha/core/store.py
index 1f2f3c65..d9c15018 100644
--- a/hypha/core/store.py
+++ b/hypha/core/store.py
@@ -907,4 +907,11 @@ async def teardown(self):
         await self.remove_client(
             client_id, self._root_user.get_workspace(), self._root_user, unload=True
         )
+        if self._websocket_server:
+            websockets = self._websocket_server.get_websockets()
+            for ws in websockets.values():
+                try:
+                    await ws.close()
+                except GeneratorExit:
+                    pass
         logger.info("Teardown complete")
diff --git a/hypha/templates/hypha-core-app/hypha-app-webpython.js b/hypha/templates/hypha-core-app/hypha-app-webpython.js
index 25489962..041dc0fb 100644
--- a/hypha/templates/hypha-core-app/hypha-app-webpython.js
+++ b/hypha/templates/hypha-core-app/hypha-app-webpython.js
@@ -27,7 +27,7 @@ loadPyodide().then(async (pyodide) => {
     pyodide.setStderr({ batched: (msg) => console.error(msg) });
     await pyodide.loadPackage("micropip");
     const micropip = pyodide.pyimport("micropip");
-    await micropip.install('hypha-rpc==0.20.39');
+    await micropip.install('hypha-rpc==0.20.40');
     const isWindow = typeof window !== "undefined";
     
     setTimeout(() => {
diff --git a/hypha/templates/ws/index.html b/hypha/templates/ws/index.html
index 30c11994..4464aece 100644
--- a/hypha/templates/ws/index.html
+++ b/hypha/templates/ws/index.html
@@ -160,7 +160,7 @@
 <body class="bg-black text-white font-poppins">
     <div id="app"></div>
     <script type="module">
-        import { HyphaCore } from "https://cdn.jsdelivr.net/npm/hypha-core@0.20.39/dist/hypha-core.mjs";
+        import { HyphaCore } from "https://cdn.jsdelivr.net/npm/hypha-core@0.20.40/dist/hypha-core.mjs";
 
         const defaultService = {
             getServerConfig(context){
@@ -765,7 +765,7 @@ <h2 className="text-2xl font-bold text-white">Artifact Details</h2>
                                     <h3 className="text-lg font-bold">Files:</h3>
                                     <ul className="list-disc list-inside text-gray-300 mt-2">
                                         {files.map((file) => (
-                                            <li key={file.name} className="break-all flex items-center justify-between">
+                                            <li key={file.name} className="break-all flex mb-1 items-center justify-between">
                                                 <span><i className={`fas ${file.type==='directory' ? 'fa-folder' : 'fa-file-alt'} mr-2`}></i>{file.name}</span>
                                                 {file.type === "file" && (
                                                     <div className="flex space-x-2">
diff --git a/hypha/websocket.py b/hypha/websocket.py
index b69a204a..aef22fb8 100644
--- a/hypha/websocket.py
+++ b/hypha/websocket.py
@@ -66,10 +66,13 @@ async def websocket_endpoint(
                     return
             else:
                 logger.warning("Rejecting legacy imjoy-rpc client (%s)", client_id)
-                await websocket.close(
-                    code=status.WS_1008_POLICY_VIOLATION,
-                    reason="Connection rejected, please  upgrade to `hypha-rpc` which pass the authentication information in the first message",
-                )
+                try:
+                    await websocket.close(
+                        code=status.WS_1008_POLICY_VIOLATION,
+                        reason="Connection rejected, please  upgrade to `hypha-rpc` which pass the authentication information in the first message",
+                    )
+                except GeneratorExit:
+                    pass
                 return
 
             try:
@@ -164,11 +167,18 @@ async def websocket_endpoint(
                     e,
                 )
 
+    def get_websockets(self):
+        """Get the active websockets."""
+        return self._websockets
+
     async def force_disconnect(self, workspace, client_id, code, reason):
         """Force disconnect a client."""
         assert f"{workspace}/{client_id}" in self._websockets, "Client not connected"
         websocket = self._websockets[f"{workspace}/{client_id}"]
-        await websocket.close(code=code, reason=reason)
+        try:
+            await websocket.close(code=code, reason=reason)
+        except GeneratorExit:
+            pass  # Suppress GeneratorExit to avoid RuntimeError
 
     async def check_client(self, client_id, workspace, user_info):
         """Check if the client is already connected."""
@@ -367,7 +377,10 @@ async def disconnect(self, websocket, reason, code=status.WS_1000_NORMAL_CLOSURE
             await websocket.send_text(json.dumps({"type": "error", "message": reason}))
         try:
             if websocket.client_state.name not in ["CLOSED", "CLOSING", "DISCONNECTED"]:
-                await websocket.close(code=code, reason=reason)
+                try:
+                    await websocket.close(code=code, reason=reason)
+                except GeneratorExit:
+                    pass  # Suppress GeneratorExit to avoid RuntimeError
         except Exception as e:
             logger.error(f"Error disconnecting websocket: {str(e)}")
 
diff --git a/pytest.ini b/pytest.ini
index d0f9fcc3..c2b8faf1 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -1,2 +1,5 @@
 [pytest]
-asyncio_mode = strict
\ No newline at end of file
+asyncio_mode = strict
+filterwarnings =
+    ignore::DeprecationWarning
+    ignore::UserWarning
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index e4cac9c2..294ae15b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,7 +3,7 @@ aiofiles==23.2.1
 websockets==13.1
 base58==2.1.1
 fastapi==0.106.0
-hypha-rpc==0.20.39
+hypha-rpc==0.20.40
 jinja2==3.1.4
 lxml==4.9.3
 msgpack==1.0.8
diff --git a/setup.py b/setup.py
index e3cc2841..40ded274 100644
--- a/setup.py
+++ b/setup.py
@@ -12,7 +12,7 @@
 REQUIREMENTS = [
     "aiofiles",
     "fastapi>=0.70.0,<=0.106.0",
-    "hypha-rpc>=0.20.39",
+    "hypha-rpc>=0.20.40",
     "msgpack>=1.0.2",
     "numpy",
     "pydantic[email]>=2.6.1",
diff --git a/tests/test_artifact.py b/tests/test_artifact.py
index 46502593..f211c01d 100644
--- a/tests/test_artifact.py
+++ b/tests/test_artifact.py
@@ -1079,7 +1079,7 @@ async def test_edit_existing_artifact(minio_server, fastapi_server, test_user_to
 
     # Retrieve and verify the uploaded file
     get_url = await artifact_manager.get_file(
-        artifact_id=dataset.id, path="example.txt"
+        artifact_id=dataset.id, file_path="example.txt"
     )
     response = requests.get(get_url)
     assert response.status_code == 200
@@ -1170,7 +1170,8 @@ async def test_artifact_schema_validation(
     assert response.ok
 
     # Commit the valid dataset artifact (should succeed as it conforms to the schema)
-    await artifact_manager.commit(artifact_id=valid_dataset.id)
+    artifact = await artifact_manager.commit(artifact_id=valid_dataset.id)
+    assert artifact.file_count == 1
 
     # Confirm the dataset is listed in the collection
     items = await artifact_manager.list(parent_id=collection.id)
@@ -1280,7 +1281,9 @@ async def test_artifact_manager_with_collection(
     assert dataset_item["view_count"] == 1
 
     # Retrieve the file in the dataset via HTTP to verify contents
-    get_url = await artifact_manager.get_file(artifact_id=dataset.id, path="test.txt")
+    get_url = await artifact_manager.get_file(
+        artifact_id=dataset.id, file_path="test.txt"
+    )
     response = requests.get(get_url)
     assert response.ok
     assert response.text == file_content
@@ -1674,7 +1677,7 @@ async def test_download_count(minio_server, fastapi_server, test_user_token):
     # Increment the download count of the artifact by download the primary file
     get_url = await artifact_manager.get_file(
         artifact_id=artifact.id,
-        path="example.txt",
+        file_path="example.txt",
     )
     response = requests.get(get_url)
     assert response.ok
@@ -1689,7 +1692,7 @@ async def test_download_count(minio_server, fastapi_server, test_user_token):
     # download twice will increment the download count by 1
     get_url = await artifact_manager.get_file(
         artifact_id=artifact.id,
-        path="example.txt",
+        file_path="example.txt",
     )
     response = requests.get(get_url)
     assert response.ok
@@ -1704,14 +1707,14 @@ async def test_download_count(minio_server, fastapi_server, test_user_token):
     # If we get the example file in silent mode, the download count won't increment
     get_url = await artifact_manager.get_file(
         artifact_id=artifact.id,
-        path="example.txt",
+        file_path="example.txt",
         silent=True,
     )
     response = requests.get(get_url)
     assert response.ok
     get_url = await artifact_manager.get_file(
         artifact_id=artifact.id,
-        path="example.txt",
+        file_path="example.txt",
         silent=True,
     )
     response = requests.get(get_url)
@@ -1726,7 +1729,7 @@ async def test_download_count(minio_server, fastapi_server, test_user_token):
     # download example 2 won't increment the download count
     get_url = await artifact_manager.get_file(
         artifact_id=artifact.id,
-        path="example2.txt",
+        file_path="example2.txt",
     )
     response = requests.get(get_url)
     assert response.ok