Allow disabling legacy sublattice path

covalent/_workflow/electron.py

@@ -19,6 +19,7 @@
 import inspect
 import json
 import operator
+import os
 import tempfile
 from builtins import list
 from copy import deepcopy
@@ -67,6 +68,9 @@
 log_stack_info = logger.log_stack_info
 
 
+DISABLE_LEGACY_SUBLATTICES = os.environ.get("COVALENT_DISABLE_LEGACY_SUBLATTICES") == "1"
+
+
 class Electron:
     """
     An electron (or task) object that is a modular component of a
@@ -863,8 +867,6 @@ def to_decoded_electron_collection(**x):
 
 # Copied from runner.py
 def _build_sublattice_graph(sub: Lattice, json_parent_metadata: str, *args, **kwargs):
-    import os
-
     parent_metadata = json.loads(json_parent_metadata)
     for k in sub.metadata.keys():
         if not sub.metadata[k] and k != "triggers":
@@ -899,5 +901,7 @@ def _build_sublattice_graph(sub: Lattice, json_parent_metadata: str, *args, **kw
 
     except Exception as ex:
         # Fall back to legacy sublattice handling
+        if DISABLE_LEGACY_SUBLATTICES:
+            raise
         print("Falling back to legacy sublattice handling")
         return sub.serialize_to_json()