From 30c8e69d2729871065ec44e67f7c84db990f7e53 Mon Sep 17 00:00:00 2001
From: Noam Bernstein <noam.bernstein@nrl.navy.mil>
Date: Tue, 22 Oct 2024 17:02:46 -0400
Subject: [PATCH] Remove _ConfigSet_loc from Atoms in do_remotely() when
 iterable is not a ConfigSet (see
 https://github.com/libAtoms/workflow/issues/344)

---
 tests/local_scripts/complete_pytest.tin       |  2 +-
 tests/test_list_with_nested_configset_info.py | 59 +++++++++++++++++++
 wfl/autoparallelize/remote.py                 |  5 ++
 3 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_list_with_nested_configset_info.py

diff --git a/tests/local_scripts/complete_pytest.tin b/tests/local_scripts/complete_pytest.tin
index c587a6b1..4d50dc7b 100755
--- a/tests/local_scripts/complete_pytest.tin
+++ b/tests/local_scripts/complete_pytest.tin
@@ -67,7 +67,7 @@ l=`egrep '^=.*(passed|failed|skipped|xfailed|error).* in ' complete_pytest.tin.o
 echo "summary line $l"
 lp=$( echo $l | sed -E -e 's/ in .*//' -e 's/\s*,\s*/\n/g' )
 
-declare -A expected_n=( ["passed"]="175" ["skipped"]="21" ["warnings"]=823 ["xfailed"]=2 ["xpassed"]=1 )
+declare -A expected_n=( ["passed"]="176" ["skipped"]="21" ["warnings"]=823 ["xfailed"]=2 ["xpassed"]=1 )
 IFS=$'\n'
 t_stat=0
 for out in $lp; do
diff --git a/tests/test_list_with_nested_configset_info.py b/tests/test_list_with_nested_configset_info.py
new file mode 100644
index 00000000..fec634f3
--- /dev/null
+++ b/tests/test_list_with_nested_configset_info.py
@@ -0,0 +1,59 @@
+from wfl.configset import ConfigSet, OutputSpec
+from wfl.calculators.generic import calculate
+from wfl.generate.md import md
+
+import ase.io
+from ase.atoms import Atoms
+from ase.calculators.emt import EMT
+
+import pytest
+
+pytestmark = pytest.mark.remote
+
+def test_list_with_nested_configset_info(tmp_path, expyre_systems, remoteinfo_env):
+    for sys_name in expyre_systems:
+        if sys_name.startswith('_'):
+            continue
+
+        do_test_list_with_nested_configset_info(tmp_path, sys_name, remoteinfo_env)
+
+def do_test_list_with_nested_configset_info(tmp_path, sys_name, remoteinfo_env):
+    ri = {'sys_name': sys_name, 'job_name': 'pytest_'+sys_name,
+          'resources': {'max_time': '1h', 'num_nodes': 1},
+          'num_inputs_per_queued_job': 1, 'check_interval': 10}
+
+    remoteinfo_env(ri)
+
+    print('RemoteInfo', ri)
+
+    cs = [Atoms('Al', cell=[3.0] * 3, pbc=[True] * 3) for _ in range(20)]
+
+    for at_i, at in enumerate(cs):
+        at.info["_ConfigSet_loc"] = f" / {at_i} / 0 / 10000"
+
+    ase.io.write(tmp_path / "tt.extxyz", cs)
+
+    configs = ConfigSet("tt.extxyz", file_root=tmp_path)
+
+    os = OutputSpec("t1.extxyz", file_root=tmp_path)
+    cc = md(configs, os, (EMT, [], {}), steps=10, dt=1.0, autopara_info={'remote_info': ri})
+
+    traj_final_configs = []
+    for traj_grp in cc.groups():
+        traj_final_configs.append([atoms for atoms in traj_grp][-1])
+
+    # ConfigSet should work
+    print("trying t2")
+    configs = ConfigSet(traj_final_configs)
+
+    os = OutputSpec("t2.extxyz", file_root=tmp_path)
+    _ = md(configs, os, (EMT, [], {}), steps=10, dt=1.0, autopara_info={'remote_info': {'num_inputs_per_queued_job': 1,
+        'sys_name': 'tin', 'job_name': 'wif_test', 'resources': {'num_cores': 16, 'max_time': '1h', 'partitions': 'n2013,n2016'}}})
+
+    # list originally failed, as in https://github.com/libAtoms/workflow/issues/344
+    print("trying t3")
+    configs = traj_final_configs
+
+    os = OutputSpec("t3.extxyz", file_root=tmp_path)
+    _ = md(configs, os, (EMT, [], {}), steps=10, dt=1.0, autopara_info={'remote_info': {'num_inputs_per_queued_job': 1,
+        'sys_name': 'tin', 'job_name': 'wif_test', 'resources': {'num_cores': 16, 'max_time': '1h', 'partitions': 'n2013,n2016'}}})
diff --git a/wfl/autoparallelize/remote.py b/wfl/autoparallelize/remote.py
index e4fca6eb..88c5d939 100644
--- a/wfl/autoparallelize/remote.py
+++ b/wfl/autoparallelize/remote.py
@@ -61,6 +61,11 @@ def do_remotely(autopara_info, iterable=None, outputspec=None, op=None, rng=None
                 # special things to do when item is Atoms
                 if 'EXPYRE_REMOTE_JOB_FAILED' in item.info:
                     del item.info['EXPYRE_REMOTE_JOB_FAILED']
+                # remote ConfigSet_loc info from things that are not ConfigSets,
+                # e.g. lists that were extracted from a ConfigSet iterator
+                if not isinstance(iterable, ConfigSet):
+                    if item.info.pop("_ConfigSet_loc", None) is not None:
+                        warnings.warn("Removed _ConfigSet_loc info field from Atoms because input is not a ConfigSet")
 
             item_list.append(item)
             item_i_list.append(item_i)