feat: volumes can be mapped to arbitrary path in container (#1596)

* feat: volumes can be mapped to arbitrary path in container

jina/parsers/peapods/runtimes/container.py

@@ -16,5 +16,8 @@ def mixin_container_runtime_parser(parser):
                     help='pull the latest image before running')
     gp.add_argument('--volumes', type=str, nargs='*', metavar='DIR',
                     help='the path on the host to be mounted inside the container. '
-                         'they will be mounted to the root path, i.e. /user/test/my-workspace will be mounted to '
-                         '/my-workspace inside the container. all volumes are mounted with read-write mode.')
+                         'if separated by ":" the first part will be considered as the local host path and the second '
+                         'part is the path in the container system. '
+                         'If no split provided, then the basename of that directory will be mounted into container\'s root path, e.g. --volumes="/user/test/my-workspace" '
+                         'will be mounted into "/my-workspace" inside the container. all volumes are mounted with '
+                         'read-write mode.')

jina/peapods/runtimes/container/__init__.py

@@ -29,10 +29,13 @@ def setup(self):
     def teardown(self):
         self._container.stop()
 
-    def run_forever(self):
+    def _stream_logs(self):
         for line in self._container.logs(stream=True):
             self.logger.info(line.strip().decode())
 
+    def run_forever(self):
+        self._stream_logs()
+
     def _set_network_for_dind_linux(self):
         import docker
         # recompute the control_addr, do not assign client, since this would be an expensive object to
@@ -96,9 +99,14 @@ def _docker_run(self, replay: bool = False):
                     f'"uses_internal" {self.args.uses_internal} is not like a path, please check it')
         if self.args.volumes:
             for p in self.args.volumes:
-                Path(os.path.abspath(p)).mkdir(parents=True, exist_ok=True)
-                _p = '/' + os.path.basename(p)
-                _volumes[os.path.abspath(p)] = {'bind': _p, 'mode': 'rw'}
+                paths = p.split(':')
+                local_path = paths[0]
+                Path(os.path.abspath(local_path)).mkdir(parents=True, exist_ok=True)
+                if len(paths) == 2:
+                    container_path = paths[1]
+                else:
+                    container_path = '/' + os.path.basename(p)
+                _volumes[os.path.abspath(local_path)] = {'bind': container_path, 'mode': 'rw'}
 
         _expose_port = [self.args.port_ctrl]
         if self.args.socket_in.is_bind:
@@ -121,8 +129,7 @@ def _docker_run(self, replay: bool = False):
         if replay:
             # when replay is on, it means last time it fails to start
             # therefore we know the loop below wont block the main process
-            for line in self._container.logs(stream=True):
-                self.logger.info(line.strip().decode())
+            self._stream_logs()
 
         client.close()
 

tests/unit/mwu-encoder/mwu_encoder_volume_change.yml

@@ -0,0 +1,7 @@
+!MWUUpdater
+with:
+  greetings: hello im from external yaml!
+metas:
+  name: ext-mwu-encoder
+  py_modules: mwu_encoder.py
+  workspace: '/mapped/here/abc'

tests/unit/peapods/runtimes/container/test_container_runtime.py

@@ -10,7 +10,6 @@
 from jina.parsers import set_pea_parser
 from jina.parsers.ping import set_ping_parser
 from jina.peapods import Pea
-from jina.peapods.runtimes.container import ContainerRuntime
 from tests import random_docs
 
 cur_dir = os.path.dirname(os.path.abspath(__file__))
@@ -131,6 +130,18 @@ def test_container_volume(docker_image_built, tmpdir):
     assert os.path.exists(os.path.join(abc_path, 'ext-mwu-encoder.bin'))
 
 
+def test_container_volume_arbitrary(docker_image_built, tmpdir):
+    abc_path = os.path.join(tmpdir, 'abc')
+    f = (Flow()
+         .add(name=random_name(), uses=f'docker://{img_name}', volumes=abc_path + ':' + '/mapped/here/abc',
+              uses_internal=os.path.join(cur_dir, '../../../mwu-encoder/mwu_encoder_volume_change.yml')))
+
+    with f:
+        f.index(random_docs(10))
+
+    assert os.path.exists(os.path.join(abc_path, 'ext-mwu-encoder.bin'))
+
+
 def test_container_ping(docker_image_built):
     a4 = set_pea_parser().parse_args(['--uses', f'docker://{img_name}'])
     a5 = set_ping_parser().parse_args(['0.0.0.0', str(a4.port_ctrl), '--print-response'])