Fix local executor tests

Also demo file transfers in local executor

Author: cjao

covalent/executor/executor_plugins/local.py

@@ -162,24 +162,26 @@ def _send(
         resources: ResourceMap,
         task_group_metadata: dict,
     ):
+        os.makedirs(self.workdir, exist_ok=True)
         dispatch_id = task_group_metadata["dispatch_id"]
         task_ids = task_group_metadata["node_ids"]
         gid = task_group_metadata["task_group_id"]
         output_uris = []
         for node_id in task_ids:
-            result_uri = os.path.join(self.cache_dir, f"result_{dispatch_id}-{node_id}.pkl")
-            stdout_uri = os.path.join(self.cache_dir, f"stdout_{dispatch_id}-{node_id}.txt")
-            stderr_uri = os.path.join(self.cache_dir, f"stderr_{dispatch_id}-{node_id}.txt")
+            result_uri = os.path.join(self.workdir, f"result_{dispatch_id}-{node_id}.pkl")
+            stdout_uri = os.path.join(self.workdir, f"stdout_{dispatch_id}-{node_id}.txt")
+            stderr_uri = os.path.join(self.workdir, f"stderr_{dispatch_id}-{node_id}.txt")
             output_uris.append((result_uri, stdout_uri, stderr_uri))
 
         server_url = format_server_url()
 
         app_log.debug(f"Running task group {dispatch_id}:{task_ids}")
+        app_log.debug(f"Generated artifacts will be saved at: {output_uris}")
         future = proc_pool.submit(
             run_task_group,
             list(map(lambda t: t.model_dump(), task_specs)),
             output_uris,
-            self.cache_dir,
+            self.workdir,
             task_group_metadata,
             server_url,
         )
@@ -190,6 +192,9 @@ def handle_cancelled(fut):
             if ex is not None:
                 tb = "".join(traceback.TracebackException.from_exception(ex).format())
                 app_log.debug(tb)
+                for task_id in task_ids:
+                    url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/{task_id}/job"
+                    requests.put(url, json={"status": "FAILED"})
             if fut.cancelled():
                 for task_id in task_ids:
                     url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/{task_id}/job"

covalent/executor/utils/wrappers.py

@@ -29,6 +29,7 @@
 
 import requests
 
+from covalent._file_transfer import FileTransfer
 from covalent._workflow.depsbash import DepsBash
 from covalent._workflow.depscall import RESERVED_RETVAL_KEY__FILES, DepsCall
 from covalent._workflow.depspip import DepsPip
@@ -169,14 +170,10 @@ def run_task_group(
 
     """
 
-    prefix = "file://"
-    prefix_len = len(prefix)
-
     outputs = {}
     results = []
     dispatch_id = task_group_metadata["dispatch_id"]
     task_ids = task_group_metadata["node_ids"]
-    gid = task_group_metadata["task_group_id"]
 
     os.environ["COVALENT_DISPATCH_ID"] = dispatch_id
     os.environ["COVALENT_DISPATCHER_URL"] = server_url
@@ -288,9 +285,8 @@ def run_task_group(
                         uri_resp = requests.post(upload_url, headers=headers)
                         uri_resp.raise_for_status()
                         remote_uri = uri_resp.json()["remote_uri"]
-
-                        with open(result_uri, "rb") as f:
-                            requests.put(remote_uri, data=f)
+                        _, cp = FileTransfer(f"file://{result_uri}", remote_uri).cp()
+                        cp()
 
                     sys.stdout.flush()
                     if stdout_uri:
@@ -299,10 +295,8 @@ def run_task_group(
                         uri_resp = requests.post(upload_url, headers=headers)
                         uri_resp.raise_for_status()
                         remote_uri = uri_resp.json()["remote_uri"]
-
-                        with open(stdout_uri, "rb") as f:
-
-                            requests.put(remote_uri, data=f)
+                        _, cp = FileTransfer(f"file://{stdout_uri}", remote_uri).cp()
+                        cp()
 
                     sys.stderr.flush()
                     if stderr_uri:
@@ -311,10 +305,8 @@ def run_task_group(
                         uri_resp = requests.post(upload_url, headers=headers)
                         uri_resp.raise_for_status()
                         remote_uri = uri_resp.json()["remote_uri"]
-
-                        with open(stderr_uri, "rb") as f:
-                            headers = {"Content-Length": os.path.getsize(stderr_uri)}
-                            requests.put(remote_uri, data=f)
+                        _, cp = FileTransfer(f"file://{stderr_uri}", remote_uri).cp()
+                        cp()
 
                     result_path = os.path.join(results_dir, f"result-{dispatch_id}:{task_id}.json")
 

tests/covalent_tests/executor/executor_plugins/local_test.py

@@ -281,47 +281,80 @@ def task(x, y):
     node_0_function_url = (
         f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/function"
     )
+    node_0_function_file_url = f"{server_url}/files/node_0_function"
+    node_0_output_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/output"
+    node_0_output_file_url = f"{server_url}/files/node_0_output"
+    node_0_stdout_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/stdout"
+    node_0_stdout_file_url = f"{server_url}/files/node_0_stdout"
+    node_0_stderr_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/stderr"
+    node_0_stderr_file_url = f"{server_url}/files/node_0_stderr"
 
     hooks_file = tempfile.NamedTemporaryFile("wb")
     hooks_file.write(ser_hooks)
     hooks_file.flush()
     hooks_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/hooks"
+    hooks_file_url = f"{server_url}/files/hooks"
 
     node_1_file = tempfile.NamedTemporaryFile("wb")
     node_1_file.write(ser_x)
     node_1_file.flush()
     node_1_output_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/1/assets/output"
+    node_1_output_file_url = f"{server_url}/node_1_output"
 
     node_2_file = tempfile.NamedTemporaryFile("wb")
     node_2_file.write(ser_y)
     node_2_file.flush()
     node_2_output_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/2/assets/output"
+    node_2_output_file_url = f"{server_url}/node_2_output"
 
     task_spec = TaskSpec(
         electron_id=0,
         args=[1, 2],
         kwargs={},
     )
 
+    # GET/POST URLs
+    url_map = {
+        node_0_function_url: node_0_function_file_url,
+        node_0_output_url: node_0_output_file_url,
+        node_0_stdout_url: node_0_stdout_file_url,
+        node_0_stderr_url: node_0_stderr_file_url,
+        hooks_url: hooks_file_url,
+        node_1_output_url: node_1_output_file_url,
+        node_2_output_url: node_2_output_file_url,
+    }
+
+    # GET/PUT files
     resources = {
-        node_0_function_url: ser_task,
-        node_1_output_url: ser_x,
-        node_2_output_url: ser_y,
-        hooks_url: ser_hooks,
+        node_0_function_file_url: ser_task,
+        node_1_output_file_url: ser_x,
+        node_2_output_file_url: ser_y,
+        hooks_file_url: ser_hooks,
     }
 
-    def mock_req_get(url, stream):
+    def mock_req_get(url, **kwargs):
         mock_resp = MagicMock()
         mock_resp.status_code = 200
-        mock_resp.content = resources[url]
+        if url in url_map:
+            mock_resp.json.return_value = {"remote_uri": url_map[url]}
+        else:
+            mock_resp.content = resources[url]
         return mock_resp
 
-    def mock_req_post(url, files):
-        resources[url] = files["asset_file"].read()
+    def mock_req_post(url, **kwargs):
+        mock_resp = MagicMock()
+        mock_resp.status_code = 200
+        mock_resp.json.return_value = {"remote_uri": url_map[url]}
+        return mock_resp
+
+    def mock_req_put(url, data=None, headers={}, json={}):
+        if data is not None:
+            resources[url] = data if isinstance(data, bytes) else data.read()
+        return MagicMock()
 
     mocker.patch("requests.get", mock_req_get)
     mocker.patch("requests.post", mock_req_post)
-    mock_put = mocker.patch("requests.put")
+    mocker.patch("requests.put", mock_req_put)
     task_group_metadata = {
         "dispatch_id": dispatch_id,
         "node_ids": [node_id],
@@ -342,8 +375,7 @@ def mock_req_post(url, files):
         server_url=server_url,
     )
 
-    with open(result_file.name, "rb") as f:
-        output = TransportableObject.deserialize(f.read())
+    output = TransportableObject.deserialize(resources[node_0_output_file_url])
     assert output.get_deserialized() == 3
 
     with open(cb_tmpfile.name, "r") as f:
@@ -352,8 +384,6 @@ def mock_req_post(url, files):
     with open(ca_tmpfile.name, "r") as f:
         assert f.read() == "Bye\n"
 
-    mock_put.assert_called()
-
 
 def test_run_task_group_exception(mocker):
     """Test the wrapper submitted to local"""
@@ -398,47 +428,80 @@ def task(x, y):
     node_0_function_url = (
         f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/function"
     )
+    node_0_function_file_url = f"{server_url}/files/node_0_function"
+    node_0_output_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/output"
+    node_0_output_file_url = f"{server_url}/files/node_0_output"
+    node_0_stdout_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/stdout"
+    node_0_stdout_file_url = f"{server_url}/files/node_0_stdout"
+    node_0_stderr_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/stderr"
+    node_0_stderr_file_url = f"{server_url}/files/node_0_stderr"
 
     hooks_file = tempfile.NamedTemporaryFile("wb")
     hooks_file.write(ser_hooks)
     hooks_file.flush()
     hooks_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/0/assets/hooks"
+    hooks_file_url = f"{server_url}/files/hooks"
 
     node_1_file = tempfile.NamedTemporaryFile("wb")
     node_1_file.write(ser_x)
     node_1_file.flush()
     node_1_output_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/1/assets/output"
+    node_1_output_file_url = f"{server_url}/node_1_output"
 
     node_2_file = tempfile.NamedTemporaryFile("wb")
     node_2_file.write(ser_y)
     node_2_file.flush()
     node_2_output_url = f"{server_url}/api/v2/dispatches/{dispatch_id}/electrons/2/assets/output"
+    node_2_output_file_url = f"{server_url}/node_2_output"
 
     task_spec = TaskSpec(
         electron_id=0,
         args=[1],
         kwargs={"y": 2},
     )
 
+    # GET/POST URLs
+    url_map = {
+        node_0_function_url: node_0_function_file_url,
+        node_0_output_url: node_0_output_file_url,
+        node_0_stdout_url: node_0_stdout_file_url,
+        node_0_stderr_url: node_0_stderr_file_url,
+        hooks_url: hooks_file_url,
+        node_1_output_url: node_1_output_file_url,
+        node_2_output_url: node_2_output_file_url,
+    }
+
+    # GET/PUT files
     resources = {
-        node_0_function_url: ser_task,
-        node_1_output_url: ser_x,
-        node_2_output_url: ser_y,
-        hooks_url: ser_hooks,
+        node_0_function_file_url: ser_task,
+        node_1_output_file_url: ser_x,
+        node_2_output_file_url: ser_y,
+        hooks_file_url: ser_hooks,
     }
 
-    def mock_req_get(url, stream):
+    def mock_req_get(url, **kwargs):
         mock_resp = MagicMock()
         mock_resp.status_code = 200
-        mock_resp.content = resources[url]
+        if url in url_map:
+            mock_resp.json.return_value = {"remote_uri": url_map[url]}
+        else:
+            mock_resp.content = resources[url]
         return mock_resp
 
-    def mock_req_post(url, files):
-        resources[url] = files["asset_file"].read()
+    def mock_req_post(url, **kwargs):
+        mock_resp = MagicMock()
+        mock_resp.status_code = 200
+        mock_resp.json.return_value = {"remote_uri": url_map[url]}
+        return mock_resp
+
+    def mock_req_put(url, data=None, headers={}, json={}):
+        if data is not None:
+            resources[url] = data if isinstance(data, bytes) else data.read()
+        return MagicMock()
 
     mocker.patch("requests.get", mock_req_get)
     mocker.patch("requests.post", mock_req_post)
-    mocker.patch("requests.put")
+    mocker.patch("requests.put", mock_req_put)
     task_group_metadata = {
         "dispatch_id": dispatch_id,
         "node_ids": [node_id],
@@ -459,6 +522,9 @@ def mock_req_post(url, files):
         server_url=server_url,
     )
 
+    stderr = resources[node_0_stderr_file_url].decode("utf-8")
+    assert "AssertionError" in stderr
+
     summary_file_path = f"{results_dir.name}/result-{dispatch_id}:{node_id}.json"
 
     with open(summary_file_path, "r") as f:
@@ -571,7 +637,7 @@ def test_send_internal(
         run_task_group,
         list(map(lambda t: t.dict(), test_case["task_specs"])),
         test_case["expected_output_uris"],
-        "mock_cache_dir",
+        local_exec.workdir,
         test_case["task_group_metadata"],
         test_case["expected_server_url"],
     )