fix: Asynchronously refresh registry for the feast ui command (#2672)

* fix: Allow Feast UI to asynchronously refresh the registry (and pass host + port values in CLI)

Signed-off-by: Danny Chiao <[email protected]>

* remove lint

Signed-off-by: Danny Chiao <[email protected]>

* revert change

Signed-off-by: Danny Chiao <[email protected]>

* increment version

Signed-off-by: Danny Chiao <[email protected]>

Author: adchia

sdk/python/feast/cli.py

@@ -110,16 +110,42 @@ def version():
 
 
 @cli.command()
+@click.option(
+    "--host",
+    "-h",
+    type=click.STRING,
+    default="0.0.0.0",
+    help="Specify a host for the server [default: 0.0.0.0]",
+)
+@click.option(
+    "--port",
+    "-p",
+    type=click.INT,
+    default=8888,
+    help="Specify a port for the server [default: 8888]",
+)
+@click.option(
+    "--registry_ttl_sec",
+    "-r",
+    help="Number of seconds after which the registry is refreshed. Default is 5 seconds.",
+    type=int,
+    default=5,
+)
 @click.pass_context
-def ui(ctx: click.Context):
+def ui(ctx: click.Context, host: str, port: int, registry_ttl_sec: int):
     """
     Shows the Feast UI over the current directory
     """
     repo = ctx.obj["CHDIR"]
     cli_check_repo(repo)
     store = FeatureStore(repo_path=str(repo))
-    repo_config = load_repo_config(repo)
-    store.serve_ui(registry_dump(repo_config, repo_path=repo))
+    # Pass in the registry_dump method to get around a circular dependency
+    store.serve_ui(
+        host=host,
+        port=port,
+        get_registry_dump=registry_dump,
+        registry_ttl_sec=registry_ttl_sec,
+    )
 
 
 @cli.command()

sdk/python/feast/feature_store.py

@@ -21,6 +21,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     Dict,
     Iterable,
     List,
@@ -1995,14 +1996,23 @@ def get_feature_server_endpoint(self) -> Optional[str]:
         return self._provider.get_feature_server_endpoint()
 
     @log_exceptions_and_usage
-    def serve_ui(self, registry_dump: str) -> None:
+    def serve_ui(
+        self, host: str, port: int, get_registry_dump: Callable, registry_ttl_sec: int
+    ) -> None:
         """Start the UI server locally"""
         warnings.warn(
             "The Feast UI is an experimental feature. "
             "We do not guarantee that future changes will maintain backward compatibility.",
             RuntimeWarning,
         )
-        ui_server.start_server(self, registry_dump, self.config.project)
+        ui_server.start_server(
+            self,
+            host=host,
+            port=port,
+            get_registry_dump=get_registry_dump,
+            project_id=self.config.project,
+            registry_ttl_sec=registry_ttl_sec,
+        )
 
     @log_exceptions_and_usage
     def serve_transformations(self, port: int) -> None:

sdk/python/feast/ui/package.json

@@ -6,7 +6,7 @@
     "@elastic/datemath": "^5.0.3",
     "@elastic/eui": "^57.0.0",
     "@emotion/react": "^11.9.0",
-    "@feast-dev/feast-ui": "^0.20.2",
+    "@feast-dev/feast-ui": "^0.20.3",
     "@testing-library/jest-dom": "^5.16.4",
     "@testing-library/react": "^13.2.0",
     "@testing-library/user-event": "^13.5.0",

sdk/python/feast/ui_server.py

@@ -1,4 +1,6 @@
 import json
+import threading
+from typing import Callable, Optional
 
 import pkg_resources
 import uvicorn
@@ -9,7 +11,12 @@
 import feast
 
 
-def get_app(store: "feast.FeatureStore", registry_dump: str, project_id: str):
+def get_app(
+    store: "feast.FeatureStore",
+    get_registry_dump: Callable,
+    project_id: str,
+    registry_ttl_secs: int,
+):
     ui_dir = pkg_resources.resource_filename(__name__, "ui/build/")
 
     app = FastAPI()
@@ -22,9 +29,33 @@ def get_app(store: "feast.FeatureStore", registry_dump: str, project_id: str):
         allow_headers=["*"],
     )
 
+    # Asynchronously refresh registry, notifying shutdown and canceling the active timer if the app is shutting down
+    registry_json = ""
+    shutting_down = False
+    active_timer: Optional[threading.Timer] = None
+
+    def async_refresh():
+        store.refresh_registry()
+        nonlocal registry_json
+        registry_json = get_registry_dump(store.config, store.repo_path)
+        if shutting_down:
+            return
+        nonlocal active_timer
+        active_timer = threading.Timer(registry_ttl_secs, async_refresh)
+        active_timer.start()
+
+    @app.on_event("shutdown")
+    def shutdown_event():
+        nonlocal shutting_down
+        shutting_down = True
+        if active_timer:
+            active_timer.cancel()
+
+    async_refresh()
+
     @app.get("/registry")
     def read_registry():
-        return json.loads(registry_dump)
+        return json.loads(registry_json)
 
     # Generate projects-list json that points to the current repo's project
     # TODO(adchia): Enable users to also add project name + description fields in feature_store.yaml
@@ -59,6 +90,13 @@ def catch_all():
     return app
 
 
-def start_server(store: "feast.FeatureStore", registry_dump: str, project_id: str):
-    app = get_app(store, registry_dump, project_id)
-    uvicorn.run(app, host="0.0.0.0", port=8888)
+def start_server(
+    store: "feast.FeatureStore",
+    host: str,
+    port: int,
+    get_registry_dump: Callable,
+    project_id: str,
+    registry_ttl_sec: int,
+):
+    app = get_app(store, get_registry_dump, project_id, registry_ttl_sec)
+    uvicorn.run(app, host=host, port=port)

ui/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@feast-dev/feast-ui",
-  "version": "0.20.2",
+  "version": "0.20.3",
   "private": false,
   "files": [
     "dist"

ui/src/index.tsx

@@ -75,13 +75,6 @@ ReactDOM.render(
       reactQueryClient={queryClient}
       feastUIConfigs={{
         tabsRegistry: tabsRegistry,
-        projectListPromise: fetch("http://0.0.0.0:8888/projects-list", {
-          headers: {
-            "Content-Type": "application/json",
-          },
-        }).then((res) => {
-          return res.json();
-        })
       }}
     />
   </React.StrictMode>,