fix: workspace from dump may not match the one from yaml (#1756)

jina/jaml/parsers/executor/legacy.py

@@ -18,7 +18,7 @@ def _get_dump_path_from_config(meta_config: Dict):
             if work_dir:
                 # then try to see if it can be loaded from its regular expected workspace (ref_indexer)
                 dump_path = BaseExecutor.get_shard_workspace(work_dir, name, pea_id)
-                bin_dump_path = os.path.join(dump_path, f'{name}.{"bin"}')
+                bin_dump_path = os.path.join(dump_path, f'{name}.bin')
                 if os.path.exists(bin_dump_path):
                     return bin_dump_path
 
@@ -52,6 +52,12 @@ def parse(self, cls: Type['BaseExecutor'], data: Dict) -> 'BaseExecutor':
         if dump_path:
             obj = cls.load(dump_path)
             obj.logger.success(f'restore {cls.__name__} from {dump_path}')
+            # consider the case where `dump_path` is not based on `obj.workspace`. This is needed
+            # for
+            workspace_loaded_from = data.get('metas', {})['workspace']
+            workspace_in_dump = getattr(obj, 'workspace', None)
+            if workspace_in_dump != workspace_loaded_from:
+                obj.workspace = workspace_loaded_from
             load_from_dump = True
         else:
             cls._init_from_yaml = True

tests/unit/test_workspace.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 
 import numpy as np
 import pytest
@@ -165,6 +166,46 @@ def test_compound_indexer_ref_indexer(test_workspace, pea_id):
         assert indexer.num_dim == 512
 
 
+@pytest.fixture()
+def test_workspace_move(tmpdir):
+    os.environ['JINA_TEST_WORKSPACE'] = os.path.join(str(tmpdir), 'host')
+    yield
+    del os.environ['JINA_TEST_WORKSPACE']
+
+
+# This test tries to simulate the situation where an executor workspace is mapped to a docker container, and therefore
+# its workspace has changed.
+@pytest.mark.parametrize('pea_id', [-1, 0, 1, 2, 3])
+def test_simple_indexer_workspace_move_to_docker(test_workspace_move, tmpdir, pea_id):
+    keys = [0, 1]
+    content = [b'a', b'b']
+    old_tmpdir = os.environ['JINA_TEST_WORKSPACE']
+    docker_tmpdir = os.path.join(tmpdir, 'docker')
+
+    with BaseExecutor.load_config(os.path.join(cur_dir, 'yaml/test-kvindexer-workspace.yml'),
+                                  pea_id=pea_id) as indexer:
+        indexer.add(keys, content)
+
+    if pea_id > 0:
+        assert os.path.exists(
+            os.path.join(old_tmpdir, f'{indexer.name}-{indexer.pea_id}', f'{indexer.name}.bin'))
+    else:
+        assert os.path.exists(os.path.join(old_tmpdir, f'{indexer.name}.bin'))
+
+    shutil.copytree(os.environ['JINA_TEST_WORKSPACE'],
+                    docker_tmpdir)
+
+    shutil.rmtree(os.environ['JINA_TEST_WORKSPACE'])
+
+    os.environ['JINA_TEST_WORKSPACE'] = str(docker_tmpdir)
+
+    with BaseExecutor.load_config(os.path.join(cur_dir, 'yaml/test-kvindexer-workspace.yml'),
+                                  pea_id=pea_id) as indexer:
+        assert indexer.query(keys[0]) == content[0]
+        assert indexer.query(keys[1]) == content[1]
+        assert indexer.workspace == docker_tmpdir
+
+
 def test_compound_indexer_rw(test_workspace):
     all_vecs = np.random.random([6, 5])
     for j in range(3):

tests/unit/yaml/test-kvindexer-workspace.yml

@@ -0,0 +1,6 @@
+!BinaryPbIndexer
+with:
+  index_filename: 'inner_indexer.gz'
+metas:
+  name: kv_indexer
+  workspace: $JINA_TEST_WORKSPACE