Repository: Netflix/metaflow
Pull request: Fix issue when running on batch with local metadata
Artifacts were persisted *after* the metadata was saved to S3 to be
repatriated locally which caused it to completely miss.

This addresses this issue (fixes #878) and also addresses another issue
in load_metadata when loading non existent metadata.

Problem statement:
KeyError getting results of Batch runs in 2.4.0+
I get a KeyError when getting results of a run that uses AWS Batch starting in metaflow 2.4.0, including 2.4.7. 2.3.6 works fine. Non-Batch runs work fine.

To reproduce, use the following files:

```
# branch_flow.py
from metaflow import FlowSpec, step, resources


class BranchFlow(FlowSpec):
    @step
    def start(self):
        self.next(self.a, self.b)

    @step
    def a(self):
        self.x = 1
        self.next(self.join)

    @resources(memory=400, cpu=1)
    @step
    def b(self):
        self.x = 2
        self.next(self.join)

    @step
    def join(self, inputs):
        print("a is %s" % inputs.a.x)
        print("b is %s" % inputs.b.x)
        self.x = inputs.a.x + inputs.b.x
        self.merge_artifacts(inputs)
        print("total is %d" % sum(input.x for input in inputs))
        self.next(self.end)

    @step
    def end(self):
        pass


if __name__ == "__main__":
    BranchFlow()
```

```
# inspect_run.py
from metaflow import Step, namespace, metadata


def main() -> None:
    namespace(f"user:error-repro")
    metadata("local")

    with open("run_id.txt", "r", encoding="utf-8") as run_id_file:
        run_id = run_id_file.read()

    result = getattr(Step(f"BranchFlow/{run_id}/a").task.data, "x")
    print(f"Result (should be 1): {result}")


if __name__ == "__main__":
    main()
```

Set up the necessary environment variables:

```
export METAFLOW_BATCH_JOB_QUEUE=<Batch job queue ARN here>
export METAFLOW_ECS_S3_ACCESS_IAM_ROLE=<Role ARN here>
export METAFLOW_DEFAULT_DATASTORE=s3
export METAFLOW_DATASTORE_SYSROOT_S3=s3://YOUR-BUCKET-HERE/artifacts
export METAFLOW_DATATOOLS_SYSROOT_S3=s3://YOUR-BUCKET-HERE/datatools
export METAFLOW_USER=error-repro
export METAFLOW_DEFAULT_METADATA=local
# Configure AWS credentials and region
```

Run the flow: `python branch_flow.py --no-pylint run --run-id-file run_id.txt --with batch:memory=512,cpu=2,gpu=0`

When that completes, get the results: `python inspect_run.py`

With metaflow 2.3.6 the output of inspect_run.py is:

```
Result (should be 1): 1
```

With metaflow 2.4.0 and 2.4.7 the output is:

```
Traceback (most recent call last):
  File "inspect_run.py", line 17, in <module>
    main()
  File "inspect_run.py", line 12, in main
    result = getattr(Step(f"BranchFlow/{run_id}/a").task.data, "x")
  File "/root/.cache/pypoetry/virtualenvs/metaflow-keyerror-repro-w5vfYHO1-py3.7/lib/python3.7/site-packages/metaflow/client/core.py", line 671, in __getattr__
    return self._artifacts[name].data
KeyError: 'x'
```

With metaflow 2.3.6, I see metadata files describing the location of where the artifact data is stored:

```
root@b8d3bac88e70:/shared/metaflow_keyerror_repro# ls -lah .metaflow/BranchFlow/1640123333270245/a/2/_meta
total 30K
drwxrwxrwx 1 root root 12K Dec 21 21:53 .
drwxrwxrwx 1 root root   0 Dec 21 21:52 ..
-rwxr-xr-x 1 root root 573 Dec 21 21:52 0_artifact__current_step.json
-rwxr-xr-x 1 root root 570 Dec 21 21:52 0_artifact__exception.json
-rwxr-xr-x 1 root root 579 Dec 21 21:52 0_artifact__foreach_num_splits.json
-rwxr-xr-x 1 root root 574 Dec 21 21:52 0_artifact__foreach_stack.json
-rwxr-xr-x 1 root root 572 Dec 21 21:52 0_artifact__foreach_var.json
-rwxr-xr-x 1 root root 568 Dec 21 21:52 0_artifact__success.json
-rwxr-xr-x 1 root root 568 Dec 21 21:52 0_artifact__task_ok.json
-rwxr-xr-x 1 root root 571 Dec 21 21:52 0_artifact__transition.json
-rwxr-xr-x 1 root root 564 Dec 21 21:52 0_artifact_name.json
-rwxr-xr-x 1 root root 561 Dec 21 21:52 0_artifact_x.json
-rwxr-xr-x 1 root root 282 Dec 21 21:52 _self.json
-rwxr-xr-x 1 root root 228 Dec 21 21:52 sysmeta_attempt-done_1640123575211.json
-rwxr-xr-x 1 root root 218 Dec 21 21:52 sysmeta_attempt_1640123574574.json
-rwxr-xr-x 1 root root 254 Dec 21 21:52 sysmeta_attempt_ok_1640123574938.json
-rwxr-xr-x 1 root root 263 Dec 21 21:52 sysmeta_aws-batch-awslogs-group_1640123574933.json
-rwxr-xr-x 1 root root 260 Dec 21 21:52 sysmeta_aws-batch-awslogs-region_1640123574933.json
-rwxr-xr-x 1 root root 357 Dec 21 21:52 sysmeta_aws-batch-awslogs-stream_1640123574933.json
-rwxr-xr-x 1 root root 258 Dec 21 21:52 sysmeta_aws-batch-ce-name_1640123574933.json
-rwxr-xr-x 1 root root 260 Dec 21 21:52 sysmeta_aws-batch-execution-env_1640123574933.json
-rwxr-xr-x 1 root root 246 Dec 21 21:52 sysmeta_aws-batch-job-attempt_1640123574933.json
-rwxr-xr-x 1 root root 271 Dec 21 21:52 sysmeta_aws-batch-job-id_1640123574933.json
-rwxr-xr-x 1 root root 258 Dec 21 21:52 sysmeta_aws-batch-jq-name_1640123574933.json
-rwxr-xr-x 1 root root 425 Dec 21 21:52 sysmeta_code-package_1640123574574.json
-rwxr-xr-x 1 root root 259 Dec 21 21:52 sysmeta_ds-root_1640123574574.json
-rwxr-xr-x 1 root root 219 Dec 21 21:52 sysmeta_ds-type_1640123574574.json
-rwxr-xr-x 1 root root 233 Dec 21 21:52 sysmeta_origin-run-id_1640123574574.json
```

With metaflow 2.4.0+ the artifact metadata files are not there.

```
root@b8d3bac88e70:/shared/metaflow_keyerror_repro# ls -lah .metaflow/BranchFlow/1640124153548041/a/2/_meta
total 11K
drwxrwxrwx 1 root root 4.0K Dec 21 22:06 .
drwxrwxrwx 1 root root    0 Dec 21 22:05 ..
-rwxr-xr-x 1 root root  282 Dec 21 22:05 _self.json
-rwxr-xr-x 1 root root  232 Dec 21 22:06 sysmeta_attempt_1640124371680.json
-rwxr-xr-x 1 root root  277 Dec 21 22:06 sysmeta_aws-batch-awslogs-group_1640124372108.json
-rwxr-xr-x 1 root root  274 Dec 21 22:06 sysmeta_aws-batch-awslogs-region_1640124372108.json
-rwxr-xr-x 1 root root  371 Dec 21 22:06 sysmeta_aws-batch-awslogs-stream_1640124372108.json
-rwxr-xr-x 1 root root  272 Dec 21 22:06 sysmeta_aws-batch-ce-name_1640124372108.json
-rwxr-xr-x 1 root root  274 Dec 21 22:06 sysmeta_aws-batch-execution-env_1640124372108.json
-rwxr-xr-x 1 root root  260 Dec 21 22:06 sysmeta_aws-batch-job-attempt_1640124372108.json
-rwxr-xr-x 1 root root  285 Dec 21 22:06 sysmeta_aws-batch-job-id_1640124372108.json
-rwxr-xr-x 1 root root  272 Dec 21 22:06 sysmeta_aws-batch-jq-name_1640124372108.json
-rwxr-xr-x 1 root root  439 Dec 21 22:06 sysmeta_code-package_1640124371680.json
-rwxr-xr-x 1 root root  273 Dec 21 22:06 sysmeta_ds-root_1640124371680.json
-rwxr-xr-x 1 root root  233 Dec 21 22:06 sysmeta_ds-type_1640124371680.json
-rwxr-xr-x 1 root root  247 Dec 21 22:06 sysmeta_origin-run-id_1640124371680.json
```

Hints from issue discussion:

Commit message: Fix issue when running on batch with local metadata

Artifacts were persisted *after* the metadata was saved to S3 to be
repatriated locally which caused it to completely miss.

This addresses this issue (fixes #878) and also addresses another issue
in load_metadata when loading non existent metadata.

Patch to review:
diff --git a/metaflow/datastore/task_datastore.py b/metaflow/datastore/task_datastore.py
index 713361461ed..b6d94eefd5c 100644
--- a/metaflow/datastore/task_datastore.py
+++ b/metaflow/datastore/task_datastore.py
@@ -435,6 +435,37 @@ def save_metadata(self, contents, allow_overwrite=True, add_attempt=True):
             add_attempt,
         )
 
+    @require_mode("w")
+    def _dangerous_save_metadata_post_done(
+        self, contents, allow_overwrite=True, add_attempt=True
+    ):
+        """
+        Method identical to save_metadata BUT BYPASSES THE CHECK ON DONE
+
+        @warning This method should not be used unless you know what you are doing. This
+        will write metadata to a datastore that has been marked as done.
+
+        Currently used to re-patriate metadata in the case of a remote run and no
+        metadata service
+
+        This method requires mode 'w'
+
+        Parameters
+        ----------
+        contents : Dict[string -> JSON-ifiable objects]
+            Dictionary of metadata to store
+        allow_overwrite : boolean, optional
+            If True, allows the overwriting of the metadata, defaults to True
+        add_attempt : boolean, optional
+            If True, adds the attempt identifier to the metadata. defaults to
+            True
+        """
+        return self._save_file(
+            {k: json.dumps(v).encode("utf-8") for k, v in contents.items()},
+            allow_overwrite,
+            add_attempt,
+        )
+
     @require_mode("r")
     def load_metadata(self, names, add_attempt=True):
         """
@@ -454,7 +485,8 @@ def load_metadata(self, names, add_attempt=True):
             Results indexed by the name of the metadata loaded
         """
         return {
-            k: json.loads(v) for k, v in self._load_file(names, add_attempt).items()
+            k: json.loads(v) if v is not None else None
+            for k, v in self._load_file(names, add_attempt).items()
         }
 
     @require_mode(None)
diff --git a/metaflow/metadata/util.py b/metaflow/metadata/util.py
index af16ad86ef6..ec6ab58fcd5 100644
--- a/metaflow/metadata/util.py
+++ b/metaflow/metadata/util.py
@@ -16,7 +16,7 @@ def sync_local_metadata_to_datastore(metadata_local_dir, task_ds):
             tar.add(metadata_local_dir)
         blob = buf.getvalue()
         _, key = task_ds.parent_datastore.save_data([blob], len_hint=1)[0]
-        task_ds.save_metadata({"local_metadata": key})
+        task_ds._dangerous_save_metadata_post_done({"local_metadata": key})
 
 
 def sync_local_metadata_from_datastore(metadata_local_dir, task_ds):
diff --git a/metaflow/plugins/aws/batch/batch_decorator.py b/metaflow/plugins/aws/batch/batch_decorator.py
index 08b07ac57d1..24bc11eb49f 100644
--- a/metaflow/plugins/aws/batch/batch_decorator.py
+++ b/metaflow/plugins/aws/batch/batch_decorator.py
@@ -269,27 +269,11 @@ def task_pre_step(
         if num_parallel > 1:
             _setup_multinode_environment()
 
-    def task_post_step(
-        self, step_name, flow, graph, retry_count, max_user_code_retries
+    def task_finished(
+        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
     ):
-        # task_post_step may run locally if fallback is activated for @catch
-        # decorator.
-        if "AWS_BATCH_JOB_ID" in os.environ:
-            # If `local` metadata is configured, we would need to copy task
-            # execution metadata from the AWS Batch container to user's
-            # local file system after the user code has finished execution.
-            # This happens via datastore as a communication bridge.
-            if self.metadata.TYPE == "local":
-                # Note that the datastore is *always* Amazon S3 (see
-                # runtime_task_created function).
-                sync_local_metadata_to_datastore(
-                    DATASTORE_LOCAL_DIR, self.task_datastore
-                )
 
-    def task_exception(
-        self, exception, step_name, flow, graph, retry_count, max_user_code_retries
-    ):
-        # task_exception may run locally if fallback is activated for @catch
+        # task_finished may run locally if fallback is activated for @catch
         # decorator.
         if "AWS_BATCH_JOB_ID" in os.environ:
             # If `local` metadata is configured, we would need to copy task
@@ -303,9 +287,6 @@ def task_exception(
                     DATASTORE_LOCAL_DIR, self.task_datastore
                 )
 
-    def task_finished(
-        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
-    ):
         try:
             self._save_logs_sidecar.kill()
         except:
diff --git a/metaflow/plugins/aws/eks/kubernetes_decorator.py b/metaflow/plugins/aws/eks/kubernetes_decorator.py
index 7af90322771..d8dd7194d50 100644
--- a/metaflow/plugins/aws/eks/kubernetes_decorator.py
+++ b/metaflow/plugins/aws/eks/kubernetes_decorator.py
@@ -223,27 +223,10 @@ def task_pre_step(
             # Start MFLog sidecar to collect task logs.
             self._save_logs_sidecar = SidecarSubProcess("save_logs_periodically")
 
-    def task_post_step(
-        self, step_name, flow, graph, retry_count, max_user_code_retries
-    ):
-        # task_post_step may run locally if fallback is activated for @catch
-        # decorator.
-        if "METAFLOW_KUBERNETES_WORKLOAD" in os.environ:
-            # If `local` metadata is configured, we would need to copy task
-            # execution metadata from the AWS Batch container to user's
-            # local file system after the user code has finished execution.
-            # This happens via datastore as a communication bridge.
-            if self.metadata.TYPE == "local":
-                # Note that the datastore is *always* Amazon S3 (see
-                # runtime_task_created function).
-                sync_local_metadata_to_datastore(
-                    DATASTORE_LOCAL_DIR, self.task_datastore
-                )
-
-    def task_exception(
-        self, exception, step_name, flow, graph, retry_count, max_user_code_retries
+    def task_finished(
+        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
     ):
-        # task_exception may run locally if fallback is activated for @catch
+        # task_finished may run locally if fallback is activated for @catch
         # decorator.
         if "METAFLOW_KUBERNETES_WORKLOAD" in os.environ:
             # If `local` metadata is configured, we would need to copy task
@@ -257,9 +240,6 @@ def task_exception(
                     DATASTORE_LOCAL_DIR, self.task_datastore
                 )
 
-    def task_finished(
-        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
-    ):
         try:
             self._save_logs_sidecar.kill()
         except:


Code review comments:
diff_hunk: @@ -435,6 +435,37 @@ def save_metadata(self, contents, allow_overwrite=True, add_attempt=True):
             add_attempt,
         )
 
+    @require_mode("w")
+    def _dangerous_save_metadata_post_done(
+        self, contents, allow_overwrite=True, add_attempt=True
+    ):
+        """
+        Method identical to save_metadata BUT BYPASSES THE CHECK ON DONE
+
+        @warning This method should not be used unless you know what you are doing. This
+        will write metadata to a datastore that has been marked as done.
+
+        Currently used to re-patriate metadata in the case of a remote run and no
+        metadata service
path: metaflow/datastore/task_datastore.py
text: @user1:
Maybe there is a way to phrase this in a way that is easier to understand for someone new to the codebase, see my take below

```suggestion
        Currently only used in the case when the task is executed remotely but there is no (remote) metadata service configured, and we use datastore to share metadata between the task and the Metaflow local scheduler
```