Memory improvements (3/3): Introduce new runner and executor API

.github/workflows/tests.yml

@@ -211,16 +211,25 @@ jobs:
         run: |
           covalent db migrate
           if [ "${{ matrix.backend }}" = 'dask' ] ; then
-            covalent start -d
+            COVALENT_ENABLE_TASK_PACKING=1 covalent start -d
           elif [ "${{ matrix.backend }}" = 'local' ] ; then
             covalent start --no-cluster -d
           else
             echo "Invalid backend specified in test matrix."
             exit 1
           fi
+          cat $HOME/.config/covalent/covalent.conf
         env:
           COVALENT_EXECUTOR_DIR: doc/source/how_to/execution/custom_executors
 
+      - name: Print Covalent status
+        if: env.BUILD_AND_RUN_ALL
+        id: covalent_status
+        run: |
+          covalent status
+          covalent cluster --info
+          covalent cluster --logs
+
       - name: Run functional tests and measure coverage
         id: functional-tests
         if: env.BUILD_AND_RUN_ALL

CHANGELOG.md

@@ -38,11 +38,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Documentation and test cases for database triggers.
 - Added the `__pow__` method to the `Electron` class
+- New Runner and executor API to bypass server-side memory when running tasks.
 
 ### Docs
 
 - Added federated learning showcase code
-- Updated tutorial for redispatching workflows with Streamlit 
+- Updated tutorial for redispatching workflows with Streamlit
 
 ### Tests
 

covalent/_serialize/lattice.py

@@ -42,8 +42,8 @@
     "inputs": AssetType.TRANSPORTABLE,
     "named_args": AssetType.TRANSPORTABLE,
     "named_kwargs": AssetType.TRANSPORTABLE,
-    "cova_imports": AssetType.OBJECT,
-    "lattice_imports": AssetType.OBJECT,
+    "cova_imports": AssetType.JSONABLE,
+    "lattice_imports": AssetType.TEXT,
     "deps": AssetType.JSONABLE,
     "call_before": AssetType.JSONABLE,
     "call_after": AssetType.JSONABLE,
@@ -93,8 +93,8 @@
 
     try:
         workflow_func_string = lat.workflow_function_string
     except AttributeError:
         workflow_func_string = ""
     workflow_func_str_asset = save_asset(
         workflow_func_string,
         ASSET_TYPES["workflow_function_string"],
@@ -224,7 +224,7 @@
 
 def deserialize_lattice(model: LatticeSchema) -> Lattice:
     def dummy_function(x):
         return x
 
     lat = Lattice(dummy_function)
 

covalent/_shared_files/schemas/lattice.py

@@ -62,8 +62,8 @@
 LATTICE_DEPS_FILENAME = "deps.json"
 LATTICE_CALL_BEFORE_FILENAME = "call_before.json"
 LATTICE_CALL_AFTER_FILENAME = "call_after.json"
-LATTICE_COVA_IMPORTS_FILENAME = "cova_imports.pkl"
-LATTICE_LATTICE_IMPORTS_FILENAME = "lattice_imports.pkl"
+LATTICE_COVA_IMPORTS_FILENAME = "cova_imports.json"
+LATTICE_LATTICE_IMPORTS_FILENAME = "lattice_imports.txt"
 LATTICE_STORAGE_TYPE = "file"
 
 
@@ -120,5 +120,5 @@
         if v is not None:
             for key in v:
                 if key in ASSET_FILENAME_MAP:
                     raise ValueError(f"Asset {key} conflicts with built-in key")
         return v

covalent/_shared_files/utils.py

@@ -21,7 +21,7 @@
 import shutil
 import socket
 from datetime import timedelta
-from typing import Any, Callable, Dict, Set, Tuple
+from typing import Any, Callable, Dict, List, Tuple
 
 import cloudpickle
 from pennylane._device import Device
@@ -37,9 +37,6 @@
 DEFAULT_UI_PORT = get_config("user_interface.port")
 
 
-# Dictionary to map Dask clients to their scheduler addresses
-_address_client_mapper = {}
-
 _IMPORT_PATH_SEPARATOR = ":"
 
 
@@ -141,7 +138,7 @@
     return function_str + "\n\n"
 
 
-def get_imports(func: Callable) -> Tuple[str, Set[str]]:
+def get_imports(func: Callable) -> Tuple[str, List[str]]:
     """
     Given an input workflow function, find the imports that were used, and determine
         which ones are Covalent-related.
@@ -155,7 +152,7 @@
     """
 
     imports_str = ""
-    cova_imports = set()
+    cova_imports = []
     for i, j in func.__globals__.items():
         if inspect.ismodule(j) or (
             inspect.isfunction(j) and j.__name__ in ["lattice", "electron"]
@@ -167,7 +164,7 @@
 
             if j.__name__ in ["covalent", "lattice", "electron"]:
                 import_line = f"# {import_line}"
-                cova_imports.add(i)
+                cova_imports.append(i)
 
             imports_str += import_line
 
@@ -257,12 +254,12 @@
     if src_uri.startswith(scheme_prefix):
         src_path = src_uri[len(scheme_prefix) :]
     else:
         raise TypeError(f"{src_uri} is not a valid URI")
         # src_path = src_uri
     if dest_uri.startswith(scheme_prefix):
         dest_path = dest_uri[len(scheme_prefix) :]
     else:
         raise TypeError(f"{dest_uri} is not a valid URI")
 
     shutil.copyfile(src_path, dest_path)
 

covalent/_workflow/lattice.py

@@ -82,7 +82,6 @@ def __init__(
         self.named_kwargs = None
         self.electron_outputs = {}
         self.lattice_imports, self.cova_imports = get_imports(self.workflow_function)
-        self.cova_imports.update({"electron"})
 
         self.workflow_function = TransportableObject.make_transportable(self.workflow_function)
 
@@ -110,13 +109,11 @@ def serialize_to_json(self) -> str:
         for node_name, output in self.electron_outputs.items():
             attributes["electron_outputs"][node_name] = output.to_dict()
 
-        attributes["cova_imports"] = list(self.cova_imports)
         return json.dumps(attributes)
 
     @staticmethod
     def deserialize_from_json(json_data: str) -> None:
         attributes = json.loads(json_data)
-        attributes["cova_imports"] = set(attributes["cova_imports"])
 
         for node_name, object_dict in attributes["electron_outputs"].items():
             attributes["electron_outputs"][node_name] = TransportableObject.from_dict(object_dict)
@@ -211,6 +208,7 @@ def build_graph(self, *args, **kwargs) -> None:
         self.inputs = TransportableObject({"args": args, "kwargs": kwargs})
         self.named_args = TransportableObject(named_args)
         self.named_kwargs = TransportableObject(named_kwargs)
+        self.lattice_imports, self.cova_imports = get_imports(workflow_function)
 
         # Set any lattice metadata not explicitly set by the user
         constraint_names = {"executor", "workflow_executor", "deps", "call_before", "call_after"}

covalent/executor/__init__.py

@@ -32,7 +32,7 @@
 from .._shared_files import logger
 from .._shared_files.config import get_config, update_config
 from ..quantum import QCluster, Simulator
-from .base import BaseExecutor, wrapper_fn
+from .base import BaseExecutor
 
 app_log = logger.app_log
 log_stack_info = logger.log_stack_info