use the entity cache

fgcz · Oct 18, 2024 · 5488b2b · 5488b2b
1 parent 0439fea
commit 5488b2b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/src/bfabric/experimental/app_interface/cli/app.py b/src/bfabric/experimental/app_interface/cli/app.py
@@ -9,6 +9,7 @@
 from bfabric.cli_formatting import setup_script_logging
 from bfabric.experimental.app_interface.app_runner._spec import AppSpec
 from bfabric.experimental.app_interface.app_runner.runner import run_app, Runner
+from bfabric.experimental.entity_lookup_cache import EntityLookupCache
 
 app_app = cyclopts.App("app", help="Run an app.")
 
@@ -27,14 +28,15 @@ def run(
     setup_script_logging()
     client = Bfabric.from_config()
     app_spec_parsed = AppSpec.model_validate(yaml.safe_load(app_spec.read_text()))
-    run_app(
-        app_spec=app_spec_parsed,
-        workunit_ref=workunit_ref,
-        work_dir=work_dir,
-        client=client,
-        ssh_user=ssh_user,
-        read_only=read_only,
-    )
+    with EntityLookupCache.enable():
+        run_app(
+            app_spec=app_spec_parsed,
+            workunit_ref=workunit_ref,
+            work_dir=work_dir,
+            client=client,
+            ssh_user=ssh_user,
+            read_only=read_only,
+        )
 
 
 @app_app.command()

diff --git a/src/bfabric/experimental/app_interface/cli/chunk.py b/src/bfabric/experimental/app_interface/cli/chunk.py
@@ -9,6 +9,7 @@
 from bfabric.cli_formatting import setup_script_logging
 from bfabric.experimental.app_interface.app_runner._spec import AppSpec
 from bfabric.experimental.app_interface.app_runner.runner import run_app, Runner
+from bfabric.experimental.entity_lookup_cache import EntityLookupCache
 
 app_chunk = cyclopts.App("chunk", help="Run an app on a chunk. You can create the chunks with `app dispatch`.")
 
@@ -58,8 +59,9 @@ def process(app_spec: Path, chunk_dir: Path) -> None:
     chunk_dir = chunk_dir.resolve()
     app_spec_parsed = AppSpec.model_validate(yaml.safe_load(app_spec.read_text()))
 
-    runner = Runner(spec=app_spec_parsed, client=client, ssh_user=None)
-    runner.run_process(chunk_dir=chunk_dir)
+    with EntityLookupCache.enable():
+        runner = Runner(spec=app_spec_parsed, client=client, ssh_user=None)
+        runner.run_process(chunk_dir=chunk_dir)
 
 
 @app_chunk.command()