Repository: AgnostiqHQ/covalent (python)

Issue:
A `TypeError` is raised if a dictionary with a non-string key is passed as the input to an `@electron`
### Environment

- **Covalent version**: 0.234.0-rc.0
- **Python version**: 3.11
- **Operating system**: Ubuntu


### What is happening?

When passing a dictionary whose key is not a string as a keyword argument in an `Electron`, a `TypeError` is raised despite the fact that this is valid Python.

### How can we reproduce the issue?

This is obviously a contrived, toy example not representative of an actual use case. But it gets the point across.

```python
import covalent as ct

@ct.electron
def job(d: dict):
    return d

@ct.lattice
def workflow():
    return job({0: "hello"})

dispatch_id = ct.dispatch(workflow)()
ct.get_result(dispatch_id, wait=True).result
```

or

```python
import covalent as ct
from pathlib import Path

@ct.electron
def job2(d: dict):
    return d

@ct.lattice
def workflow():
    return job2({Path.cwd(): "hello"})

dispatch_id = ct.dispatch(workflow)()
ct.get_result(dispatch_id, wait=True).result
```

### What should happen?

I should not get the following traceback since non-string dictionary keys are valid Python:

```python
/home/rosen/software/miniconda/envs/covalent/lib/python3.11/site-packages/covalent/_workflow/lattice.py:235: UserWarning: Please make sure you are not manipulating an object inside the lattice.
  warnings.warn(
Traceback (most recent call last):
  File "/home/rosen/test.py", line 18, in <module>
    dispatch_id = ct.dispatch(workflow)()
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rosen/software/miniconda/envs/covalent/lib/python3.11/site-packages/covalent/_dispatcher_plugins/local.py", line 155, in wrapper
    dispatch_id = LocalDispatcher.register(orig_lattice, dispatcher_addr)(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rosen/software/miniconda/envs/covalent/lib/python3.11/site-packages/covalent/_dispatcher_plugins/local.py", line 217, in wrapper
    lattice.build_graph(*args, **kwargs)
  File "/home/rosen/software/miniconda/envs/covalent/lib/python3.11/site-packages/covalent/_workflow/lattice.py", line 233, in build_graph
    retval = workflow_function(*new_args, **new_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rosen/test.py", line 15, in workflow
    job2({job1()["directory"]: "hello"})
  File "/home/rosen/software/miniconda/envs/covalent/lib/python3.11/site-packages/covalent/_workflow/electron.py", line 840, in wrapper
    return electron_object(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rosen/software/miniconda/envs/covalent/lib/python3.11/site-packages/covalent/_workflow/electron.py", line 478, in __call__
    self.connect_node_with_others(
  File "/home/rosen/software/miniconda/envs/covalent/lib/python3.11/site-packages/covalent/_workflow/electron.py", line 583, in connect_node_with_others
    bound_electron = dict_electron(**param_value)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: keywords must be strings
```

### Any suggestions?

Nope, at least not until/unless I have time to look at the source code some more.

<details>
As an aside, once the above works, we should also confirm that the slightly more complex example below works too:

```python
import covalent as ct
from pathlib import Path

@ct.electron
def job1():
    return {"directory": Path.cwd()}

@ct.electron
def job2(d: dict[Path, str]):
    for k, v in d.items():
        Path(k, v).touch()

@ct.lattice
def workflow():
    output1 = job1()
    job2({output1["directory"]: "hello"})
    return "I <3 covalent"

dispatch_id = ct.dispatch(workflow)()
ct.get_result(dispatch_id, wait=True).result
```
</details>

Pull request description:
Do not require dictionary keys to be `str`
The collector electron now assembles the dictionary from two lists -- one list of keys, one list of corresponding values.

<!--
⚠️ Your pull request title should be short detailed and understandable for all.
⚠️ Add a note to /CHANGELOG.md summarizing the changes.
⚠️ If your pull request fixes an open issue, please link to the issue.
⚠️ Ensure that your branch is up-to-date with the base branch.
-->

- [x] I have added the tests to cover my changes.
- [x] I have updated the documentation and CHANGELOG accordingly.
- [ ] I have read the CONTRIBUTING document.

Closes #1937 .

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 17344263..76c2a972 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Sublattice electron function strings are now parsed correctly
+- The keys of dictionary inputs to electrons no longer need be strings.
+- Fixed inaccuracies in task packing exposed by no longer uploading null attributes upon dispatch.
 
 ### Operations
 
diff --git a/covalent/_workflow/electron.py b/covalent/_workflow/electron.py
index a8a9055b..0e80f0a2 100644
--- a/covalent/_workflow/electron.py
+++ b/covalent/_workflow/electron.py
@@ -572,8 +572,8 @@ class Electron:
 
         elif isinstance(param_value, dict):
 
-            def _auto_dict_node(*args, **kwargs):
-                return dict(kwargs)
+            def _auto_dict_node(keys, values):
+                return {keys[i]: values[i] for i in range(len(keys))}
 
             dict_electron = Electron(
                 function=_auto_dict_node,
@@ -581,7 +581,7 @@ class Electron:
                 task_group_id=self.task_group_id,
                 packing_tasks=True and active_lattice.task_packing,
             )  # Group the auto-generated node with the main node.
-            bound_electron = dict_electron(**param_value)
+            bound_electron = dict_electron(list(param_value.keys()), list(param_value.values()))
             transport_graph.set_node_value(bound_electron.node_id, "name", electron_dict_prefix)
             transport_graph.add_edge(
                 dict_electron.node_id,
diff --git a/covalent_dispatcher/_core/dispatcher.py b/covalent_dispatcher/_core/dispatcher.py
index e1754796..b0ecd27b 100644
--- a/covalent_dispatcher/_core/dispatcher.py
+++ b/covalent_dispatcher/_core/dispatcher.py
@@ -183,6 +183,7 @@ async def _submit_task_group(dispatch_id: str, sorted_nodes: List[int], task_gro
         app_log.debug("8A: Update node success (run_planned_workflow).")
 
     else:
+        # Nodes whose values have already been resolved
         known_nodes = []
 
         # Skip the group if all task outputs can be reused from a
@@ -196,6 +197,8 @@ async def _submit_task_group(dispatch_id: str, sorted_nodes: List[int], task_gro
             # Gather inputs for each task and send the task spec sequence to the runner
             task_specs = []
 
+            sorted_nodes_set = set(sorted_nodes)
+
             for node_id in sorted_nodes:
                 app_log.debug(f"Gathering inputs for task {node_id} (run_planned_workflow).")
 
@@ -214,8 +217,16 @@ async def _submit_task_group(dispatch_id: str, sorted_nodes: List[int], task_gro
                     "args_ids": abs_task_input["args"],
                     "kwargs_ids": abs_task_input["kwargs"],
                 }
-                known_nodes += abs_task_input["args"]
-                known_nodes += list(abs_task_input["kwargs"].values())
+                # Task inputs that don't belong to the task group have already beeen resolved
+                external_task_args = filter(
+                    lambda x: x not in sorted_nodes_set, abs_task_input["args"]
+                )
+                known_nodes.extend(external_task_args)
+                external_task_kwargs = filter(
+                    lambda x: x not in sorted_nodes_set, abs_task_input["kwargs"].values()
+                )
+                known_nodes.extend(external_task_kwargs)
+
                 task_specs.append(task_spec)
 
             app_log.debug(
diff --git a/covalent_dispatcher/_dal/importers/result.py b/covalent_dispatcher/_dal/importers/result.py
index 395516b8..7e4bd36f 100644
--- a/covalent_dispatcher/_dal/importers/result.py
+++ b/covalent_dispatcher/_dal/importers/result.py
@@ -72,7 +72,6 @@ def import_result(
     # Main case: insert new lattice, electron, edge, and job records
 
     storage_path = os.path.join(base_path, dispatch_id)
-    os.makedirs(storage_path)
 
     lattice_record_kwargs = _get_result_meta(res, storage_path, electron_id)
     lattice_record_kwargs.update(_get_lattice_meta(res.lattice, storage_path))
@@ -143,6 +142,7 @@ def _connect_result_to_electron(
         fields={"id", "cancel_requested"},
         equality_filters={"id": parent_electron_record.job_id},
         membership_filters={},
+        for_update=True,
     )[0]
     cancel_requested = parent_job_record.cancel_requested
 
diff --git a/covalent_dispatcher/_dal/importers/tg.py b/covalent_dispatcher/_dal/importers/tg.py
index 468abdad..c67cc34b 100644
--- a/covalent_dispatcher/_dal/importers/tg.py
+++ b/covalent_dispatcher/_dal/importers/tg.py
@@ -51,7 +51,9 @@ def import_transport_graph(
     # Propagate parent electron id's `cancel_requested` property to the sublattice electrons
     if electron_id is not None:
         parent_e_record = Electron.meta_type.get_by_primary_key(session, electron_id)
-        job_record = Job.get_by_primary_key(session=session, primary_key=parent_e_record.job_id)
+        job_record = Job.get_by_primary_key(
+            session=session, primary_key=parent_e_record.job_id, for_update=True
+        )
         cancel_requested = job_record.cancel_requested
     else:
         cancel_requested = False


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 58189ba9..1349a02c 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -217,9 +217,10 @@ jobs:
         if: env.BUILD_AND_RUN_ALL
         id: covalent_start
         run: |
+          export COVALENT_ENABLE_TASK_PACKING=1
           covalent db migrate
           if [ "${{ matrix.backend }}" = 'dask' ] ; then
-            COVALENT_ENABLE_TASK_PACKING=1 covalent start -d
+            covalent start -d
           elif [ "${{ matrix.backend }}" = 'local' ] ; then
             covalent start --no-cluster -d
           else
diff --git a/tests/covalent_dispatcher_tests/_core/dispatcher_db_integration_test.py b/tests/covalent_dispatcher_tests/_core/dispatcher_db_integration_test.py
index 53444a7b..1544b753 100644
--- a/tests/covalent_dispatcher_tests/_core/dispatcher_db_integration_test.py
+++ b/tests/covalent_dispatcher_tests/_core/dispatcher_db_integration_test.py
@@ -104,7 +104,7 @@ async def test_get_abstract_task_inputs(mocker, test_db):
 
     @ct.lattice
     def dict_workflow(arg):
-        return dict_task(arg)
+        return dict_task(arg=arg)
 
     #    1   2
     #     \   \
@@ -159,7 +159,7 @@ async def test_get_abstract_task_inputs(mocker, test_db):
 
     # dict-type inputs
 
-    # Nodes 0=task, 1=:electron_dict:, 2=1, 3=2
+    # Nodes 0=task, 1=:electron_dict:, 2=["a" (3), "b" (4)], 5=[1 (6), 2 (7)]
     dict_workflow.build_graph({"a": 1, "b": 2})
     abstract_args = {"a": 2, "b": 3}
     tg = dict_workflow.transport_graph
@@ -172,10 +172,31 @@ async def test_get_abstract_task_inputs(mocker, test_db):
         mock_get_incoming_edges,
     )
 
+    task_inputs = await _get_abstract_task_inputs(
+        result_object.dispatch_id, 0, tg.get_node_value(0, "name")
+    )
+    expected_inputs = {"args": [], "kwargs": {"arg": 1}}
+
+    assert task_inputs == expected_inputs
+
     task_inputs = await _get_abstract_task_inputs(
         result_object.dispatch_id, 1, tg.get_node_value(1, "name")
     )
-    expected_inputs = {"args": [], "kwargs": abstract_args}
+    expected_inputs = {"args": [2, 5], "kwargs": {}}
+
+    assert task_inputs == expected_inputs
+
+    task_inputs = await _get_abstract_task_inputs(
+        result_object.dispatch_id, 2, tg.get_node_value(2, "name")
+    )
+    expected_inputs = {"args": [3, 4], "kwargs": {}}
+
+    assert task_inputs == expected_inputs
+
+    task_inputs = await _get_abstract_task_inputs(
+        result_object.dispatch_id, 5, tg.get_node_value(5, "name")
+    )
+    expected_inputs = {"args": [6, 7], "kwargs": {}}
 
     assert task_inputs == expected_inputs
 
diff --git a/tests/covalent_dispatcher_tests/_core/execution_test.py b/tests/covalent_dispatcher_tests/_core/execution_test.py
index 6d521691..4e2c20ac 100644
--- a/tests/covalent_dispatcher_tests/_core/execution_test.py
+++ b/tests/covalent_dispatcher_tests/_core/execution_test.py
@@ -116,7 +116,7 @@ async def test_get_task_inputs(mocker, test_db):
 
     @ct.lattice
     def dict_workflow(arg):
-        return dict_task(arg)
+        return dict_task(arg=arg)
 
     #    1   2
     #     \   \
@@ -167,20 +167,36 @@ async def test_get_task_inputs(mocker, test_db):
     # dict-type inputs
 
     dict_workflow.build_graph({"a": 1, "b": 2})
-    serialized_args = {"a": ct.TransportableObject(1), "b": ct.TransportableObject(2)}
 
     # Nodes 0=task, 1=:electron_dict:, 2=1, 3=2
+    # Nodes 0=task, 1=:electron_dict:, 2=["a" (3), "b" (4)], 5=[1 (6), 2 (7)]
+
     sdkres = Result(lattice=dict_workflow, dispatch_id="asdf_dict_workflow")
     result_object = get_mock_srvresult(sdkres, test_db)
     tg = result_object.lattice.transport_graph
-    tg.set_node_value(2, "output", ct.TransportableObject(1))
-    tg.set_node_value(3, "output", ct.TransportableObject(2))
+
+    tg.set_node_value(1, "output", ct.TransportableObject("node_1_output"))
+    tg.set_node_value(3, "output", ct.TransportableObject("a"))
+    tg.set_node_value(4, "output", ct.TransportableObject("b"))
+    tg.set_node_value(6, "output", ct.TransportableObject(1))
+    tg.set_node_value(7, "output", ct.TransportableObject(2))
 
     mock_get_result = mocker.patch(
         "covalent_dispatcher._core.runner.datasvc.get_result_object", return_value=result_object
     )
-    task_inputs = await _get_task_inputs(1, tg.get_node_value(1, "name"), result_object)
-    expected_inputs = {"args": [], "kwargs": serialized_args}
+    serialized_kwargs = {"arg": ct.TransportableObject("node_1_output")}
+    task_inputs = await _get_task_inputs(0, tg.get_node_value(0, "name"), result_object)
+    expected_inputs = {"args": [], "kwargs": serialized_kwargs}
+
+    serialized_args = [ct.TransportableObject("a"), ct.TransportableObject("b")]
+    task_inputs = await _get_task_inputs(2, tg.get_node_value(2, "name"), result_object)
+    expected_inputs = {"args": serialized_args, "kwargs": {}}
+
+    assert task_inputs == expected_inputs
+
+    serialized_args = [ct.TransportableObject(1), ct.TransportableObject(2)]
+    task_inputs = await _get_task_inputs(5, tg.get_node_value(5, "name"), result_object)
+    expected_inputs = {"args": serialized_args, "kwargs": {}}
 
     assert task_inputs == expected_inputs
 
diff --git a/tests/covalent_dispatcher_tests/_dal/importers/result_import_test.py b/tests/covalent_dispatcher_tests/_dal/importers/result_import_test.py
index 440742cb..819f88bc 100644
--- a/tests/covalent_dispatcher_tests/_dal/importers/result_import_test.py
+++ b/tests/covalent_dispatcher_tests/_dal/importers/result_import_test.py
@@ -27,6 +27,7 @@ from covalent._serialize.result import serialize_result
 from covalent._shared_files.schemas.result import AssetSchema, ResultSchema
 from covalent._shared_files.util_classes import RESULT_STATUS
 from covalent_dispatcher._dal.importers.result import SERVER_URL, handle_redispatch, import_result
+from covalent_dispatcher._dal.job import Job
 from covalent_dispatcher._dal.result import get_result_object
 from covalent_dispatcher._db.datastore import DataStore
 
@@ -140,6 +141,7 @@ def test_import_previously_imported_result(mocker, test_db):
         prefix="covalent-"
     ) as srv_dir:
         sub_res = get_mock_result(sub_dispatch_id, sdk_dir)
+        sub_res.metadata.root_dispatch_id = dispatch_id
         import_result(sub_res, srv_dir, None)
         srv_res = get_result_object(dispatch_id, bare=True)
         parent_node = srv_res.lattice.transport_graph.get_node(0)
@@ -152,6 +154,49 @@ def test_import_previously_imported_result(mocker, test_db):
     assert sub_srv_res._electron_id == parent_node._electron_id
 
 
+def test_import_subdispatch_cancel_req(mocker, test_db):
+    """Test that Job.cancel_requested is propagated to sublattices"""
+
+    dispatch_id = "test_propagate_cancel_requested"
+    sub_dispatch_id = "test_propagate_cancel_requested_sub"
+
+    mocker.patch("covalent_dispatcher._dal.base.workflow_db", test_db)
+
+    mock_filter_uris = mocker.patch(
+        "covalent_dispatcher._dal.importers.result._filter_remote_uris"
+    )
+
+    with tempfile.TemporaryDirectory(prefix="covalent-") as sdk_dir, tempfile.TemporaryDirectory(
+        prefix="covalent-"
+    ) as srv_dir:
+        res = get_mock_result(dispatch_id, sdk_dir)
+        import_result(res, srv_dir, None)
+
+    with test_db.Session() as session:
+        Job.update_bulk(
+            session, values={"cancel_requested": True}, equality_filters={}, membership_filters={}
+        )
+        session.commit()
+
+    with tempfile.TemporaryDirectory(prefix="covalent-") as sdk_dir, tempfile.TemporaryDirectory(
+        prefix="covalent-"
+    ) as srv_dir:
+        sub_res = get_mock_result(sub_dispatch_id, sdk_dir)
+        sub_res.metadata.root_dispatch_id = dispatch_id
+        srv_res = get_result_object(dispatch_id, bare=True)
+        parent_node = srv_res.lattice.transport_graph.get_node(0)
+        import_result(sub_res, srv_dir, parent_node._electron_id)
+
+    with tempfile.TemporaryDirectory(prefix="covalent-") as srv_dir:
+        import_result(sub_res, srv_dir, parent_node._electron_id)
+
+    with test_db.Session() as session:
+        uncancelled = Job.get(
+            session, fields=[], equality_filters={"cancel_requested": False}, membership_filters={}
+        )
+        assert len(uncancelled) == 0
+
+
 @pytest.mark.parametrize(
     "parent_status,new_status",
     [
diff --git a/tests/covalent_tests/workflow/electron_test.py b/tests/covalent_tests/workflow/electron_test.py
index d7a3e319..2d5936ed 100644
--- a/tests/covalent_tests/workflow/electron_test.py
+++ b/tests/covalent_tests/workflow/electron_test.py
@@ -377,18 +377,30 @@ def test_autogen_dict_electrons():
     g = workflow.transport_graph._graph
 
     # Account for postprocessing node
-    assert list(g.nodes) == [0, 1, 2, 3, 4]
+    assert list(g.nodes) == [0, 1, 2, 3, 4, 5, 6, 7, 8]
     fn = g.nodes[1]["function"].get_deserialized()
-    assert fn(x=2, y=5, z=7) == {"x": 2, "y": 5, "z": 7}
-    assert g.nodes[2]["value"].get_deserialized() == 5
-    assert g.nodes[3]["value"].get_deserialized() == 7
+    assert fn(["x", "y", "z"], [2, 5, 7]) == {"x": 2, "y": 5, "z": 7}
+    fn = g.nodes[2]["function"].get_deserialized()
+    assert fn("x", "y") == ["x", "y"]
+    keys = [g.nodes[3]["value"].get_deserialized(), g.nodes[4]["value"].get_deserialized()]
+    fn = g.nodes[5]["function"].get_deserialized()
+    assert fn(2, 3) == [2, 3]
+    vals = [g.nodes[6]["value"].get_deserialized(), g.nodes[7]["value"].get_deserialized()]
+    assert keys == ["x", "y"]
+    assert vals == [5, 7]
     assert set(g.edges) == {
         (1, 0, 0),
         (2, 1, 0),
-        (3, 1, 0),
-        (0, 4, 0),
-        (0, 4, 1),
-        (1, 4, 0),
+        (3, 2, 0),
+        (4, 2, 0),
+        (5, 1, 0),
+        (6, 5, 0),
+        (7, 5, 0),
+        (0, 8, 0),
+        (0, 8, 1),
+        (1, 8, 0),
+        (2, 8, 0),
+        (5, 8, 0),
     }
 
 
diff --git a/tests/functional_tests/workflow_stack_test.py b/tests/functional_tests/workflow_stack_test.py
index 47e1578a..f20b3a0f 100644
--- a/tests/functional_tests/workflow_stack_test.py
+++ b/tests/functional_tests/workflow_stack_test.py
@@ -800,7 +800,8 @@ def test_workflows_with_dict_nodes():
         res_1 = sum_values(x)
         return square(res_1)
 
-    dispatch_id = ct.dispatch(workflow)({"x": 1, "y": 2, "z": 3})
+    # Check that non-string keys are allowed
+    dispatch_id = ct.dispatch(workflow)({"x": 1, "y": 2, 3: 3})
 
     res_obj = rm.get_result(dispatch_id, wait=True)