Memory improvements (1/3): Introduce new data access layer and schemas

covalent/_file_transfer/strategies/shutil_strategy.py

@@ -34,6 +34,17 @@ def __init__(
 
     # return callable to copy files in the local file system
     def cp(self, from_file: File, to_file: File = File()) -> None:
+        """
+        Get a callable that copies a file from one location to another locally
+
+        Args:
+            from_file: File to copy from
+            to_file: File to copy to. Defaults to File().
+
+        Returns:
+            A callable that copies a file from one location to another locally
+        """
+
         def callable():
             shutil.copyfile(from_file.filepath, to_file.filepath)
 

covalent/_results_manager/result.py

@@ -63,7 +63,9 @@
     """
 
     NEW_OBJ = RESULT_STATUS.NEW_OBJECT
-    PENDING_REUSE = RESULT_STATUS.PENDING_REUSE
+    PENDING_REUSE = (
+        RESULT_STATUS.PENDING_REUSE
+    )  # Facilitates reuse of previous electrons in the new dispatcher design
     COMPLETED = RESULT_STATUS.COMPLETED
     POSTPROCESSING = RESULT_STATUS.POSTPROCESSING
     PENDING_POSTPROCESSING = RESULT_STATUS.PENDING_POSTPROCESSING
@@ -106,14 +108,14 @@
             pattern = re.compile(regex)
             m = pattern.match(input_string)
             if m:
-                arg_str_repr = m.group(1).rstrip(",")
-                kwarg_str_repr = m.group(2)
+                arg_str_repr = m[1].rstrip(",")
+                kwarg_str_repr = m[2]
             else:
                 arg_str_repr = str(None)
                 kwarg_str_repr = str(None)
         else:
             arg_str_repr = str(None)
             kwarg_str_repr = str(None)
 
         show_result_str = f"""
 Lattice Result
@@ -212,7 +214,7 @@
         if self._result is not None:
             return self._result.get_deserialized()
         else:
             return None
 
     @property
     def inputs(self) -> dict:

covalent/_serialize/common.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+""" Serialization/Deserialization methods for Assets """
 
 import hashlib
 import json
@@ -26,17 +27,41 @@
 from .._shared_files.schemas.asset import AssetSchema
 from .._workflow.transportable_object import TransportableObject
 
+__all__ = [
+    "AssetType",
+    "save_asset",
+    "load_asset",
+]
+
+
 CHECKSUM_ALGORITHM = "sha"
 
 
 class AssetType(Enum):
-    OBJECT = 0
-    TRANSPORTABLE = 1
+    """
+    Enum for the type of Asset data
+
+    """
+
+    OBJECT = 0  # Fallback to cloudpickling
+    TRANSPORTABLE = 1  # Custom TO serialization
     JSONABLE = 2
-    TEXT = 3
+    TEXT = 3  # Mainly for stdout, stderr, docstrings, etc.
 
 
 def serialize_asset(data: Any, data_type: AssetType) -> bytes:
+    """
+    Serialize the asset data
+
+    Args:
+        data: Data to serialize
+        data_type: Type of the Asset data to serialize
+
+    Returns:
+        Serialized data as bytes
+
+    """
+
     if data_type == AssetType.OBJECT:
         return cloudpickle.dumps(data)
     elif data_type == AssetType.TRANSPORTABLE:
@@ -46,10 +71,22 @@
     elif data_type == AssetType.TEXT:
         return data.encode("utf-8")
     else:
         raise TypeError(f"Unsupported data type {type(data)}")
 
 
 def deserialize_asset(data: bytes, data_type: AssetType) -> Any:
+    """
+    Deserialize the asset data
+
+    Args:
+        data: Data to deserialize
+        data_type: Type of the Asset data to deserialize
+
+    Returns:
+        Deserialized data
+
+    """
+
     if data_type == AssetType.OBJECT:
         return cloudpickle.loads(data)
     elif data_type == AssetType.TRANSPORTABLE:
@@ -59,14 +96,39 @@
     elif data_type == AssetType.TEXT:
         return data.decode("utf-8")
     else:
         raise TypeError("Unsupported data type")
 
 
 def _sha1_asset(data: bytes) -> str:
+    """
+    Compute the sha1 checksum of the asset data
+
+    Args:
+        data: Data to compute checksum for
+
+    Returns:
+        sha1 checksum of the data
+
+    """
+
     return hashlib.sha1(data).hexdigest()
 
 
 def save_asset(data: Any, data_type: AssetType, storage_path: str, filename: str) -> AssetSchema:
+    """
+    Save the asset data to the storage path
+
+    Args:
+        data: Data to save
+        data_type: Type of the Asset data to save
+        storage_path: Path to save the data to
+        filename: Name of the file to save the data to
+
+    Returns:
+        AssetSchema object containing metadata about the saved data
+
+    """
+
     scheme = "file"
 
     serialized = serialize_asset(data, data_type)
@@ -80,16 +142,26 @@
 
 
 def load_asset(asset_meta: AssetSchema, data_type: AssetType) -> Any:
+    """
+    Load the asset data from the storage path
+
+    Args:
+        asset_meta: Metadata about the asset to load
+        data_type: Type of the Asset data to load
+
+    Returns:
+        Asset data
+
+    """
+
     scheme_prefix = "file://"
     uri = asset_meta.uri
 
     if not uri:
         return None
 
-    if uri.startswith(scheme_prefix):
-        path = uri[len(scheme_prefix) :]
-    else:
-        path = uri
+    path = uri[len(scheme_prefix) :] if uri.startswith(scheme_prefix) else uri
+
     with open(path, "rb") as f:
         data = f.read()
     return deserialize_asset(data, data_type)

covalent/_serialize/electron.py

@@ -29,6 +29,12 @@
 from .._workflow.transportable_object import TransportableObject
 from .common import AssetType, load_asset, save_asset
 
+__all__ = [
+    "serialize_node",
+    "deserialize_node",
+]
+
+
 ASSET_TYPES = {
     "function": AssetType.TRANSPORTABLE,
     "function_string": AssetType.TEXT,
@@ -52,13 +58,13 @@
     # Optional
     status = node_attrs.get("status", RESULT_STATUS.NEW_OBJECT)
 
-    start_time = node_attrs.get("start_time", None)
+    start_time = node_attrs.get("start_time")
     if start_time:
         start_time = start_time.isoformat()
 
-    end_time = node_attrs.get("end_time", None)
+    end_time = node_attrs.get("end_time")
     if end_time:
         end_time = end_time.isoformat()
 
     return ElectronMetadata(
         task_group_id=task_group_id,

covalent/_serialize/lattice.py

@@ -29,6 +29,12 @@
 from .common import AssetType, load_asset, save_asset
 from .transport_graph import deserialize_transport_graph, serialize_transport_graph
 
+__all__ = [
+    "serialize_lattice",
+    "deserialize_lattice",
+]
+
+
 ASSET_TYPES = {
     "workflow_function": AssetType.TRANSPORTABLE,
     "workflow_function_string": AssetType.TEXT,
@@ -87,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"],
@@ -218,7 +224,7 @@
 
 def deserialize_lattice(model: LatticeSchema) -> Lattice:
     def dummy_function(x):
         return x
 
     lat = Lattice(dummy_function)
 

covalent/_serialize/result.py

@@ -30,6 +30,15 @@
 from .common import AssetType, load_asset, save_asset
 from .lattice import deserialize_lattice, serialize_lattice
 
+__all__ = [
+    "serialize_result",
+    "deserialize_result",
+    "strip_local_uris",
+    "merge_response_manifest",
+    "extract_assets",
+]
+
+
 ASSET_TYPES = {
     "error": AssetType.TEXT,
     "result": AssetType.TRANSPORTABLE,
@@ -101,28 +110,28 @@
 
 def strip_local_uris(res: ResultSchema) -> ResultSchema:
     # Create a copy with the local uris removed for submission
     manifest = res.copy(deep=True).dict()
 
     # Strip workflow asset uris:
     dispatch_assets = manifest["assets"]
     for _, asset in dispatch_assets.items():
         asset["uri"] = ""
 
     lattice = manifest["lattice"]
     lattice_assets = lattice["assets"]
     for _, asset in lattice_assets.items():
         asset["uri"] = ""
 
     # Node assets
     tg = lattice["transport_graph"]
 
     nodes = tg["nodes"]
     for node in nodes:
         node_assets = node["assets"]
         for _, asset in node_assets.items():
             asset["uri"] = ""
 
     return ResultSchema.parse_obj(manifest)
 
 
 # Functions to postprocess response from dispatcher
@@ -136,6 +145,7 @@
         response: The manifest returned from `/register`.
     Returns:
         A combined manifest with asset `remote_uri`s populated.
+
     """
 
     manifest.metadata.dispatch_id = response.metadata.dispatch_id
@@ -170,32 +180,28 @@
 
 
 def extract_assets(manifest: ResultSchema) -> List[AssetSchema]:
-    """Extract all of the asset metadata from a manifest dictionary.
+    """
+    Extract all of the asset metadata from a manifest dictionary.
 
     Args:
         manifest: A result manifest
 
     Returns:
         A list of assets
-    """
 
-    assets = []
+    """
 
     # workflow-level assets
     dispatch_assets = manifest.assets
-    for key, asset in dispatch_assets:
-        assets.append(asset)
-
+    assets = [asset for key, asset in dispatch_assets]
     lattice = manifest.lattice
     lattice_assets = lattice.assets
-    for key, asset in lattice_assets:
-        assets.append(asset)
+    assets.extend(asset for key, asset in lattice_assets)
 
     # Node assets
     tg = lattice.transport_graph
     nodes = tg.nodes
     for node in nodes:
         node_assets = node.assets
-        for key, asset in node_assets:
-            assets.append(asset)
+        assets.extend(asset for key, asset in node_assets)
     return assets