Repository: jina-ai/serve
Pull request: feat: add support for multistage dockerfile
- [x] add support for multistag dockerfile
- [x] refactoring the way of adding labels to the docker images as suggested in #1732

Problem statement:
Enable the `jina hub build` to use Dockerfile with multi-stage building
**Describe the feature**
Enable the `jina hub build` to use Dockerfile with multi-stage building


**Your proposal**
<!-- copy past your code/pull request link -->

---
<!-- Optional, but really help us locate the problem faster -->

**Environment**
<!-- Run `jina --version-full` and copy paste the output here -->

**Screenshots**
<!-- If applicable, add screenshots to help explain your problem. -->

Hints from issue discussion:
This issue is stale because it has been open 20 days with no activity. Remove stale label or comment or this will be closed in 4 days

Commit message: feat: add missing files

Patch to review:
diff --git a/jina/docker/hubio.py b/jina/docker/hubio.py
index 53c21942e8f68..ed2f94fa2596e 100644
--- a/jina/docker/hubio.py
+++ b/jina/docker/hubio.py
@@ -309,15 +309,18 @@ def build(self) -> Dict:
 
             with TimeContext(f'building {colored(self.args.path, "green")}', self.logger) as tc:
                 try:
-                    self._check_completeness()
+                    _check_result = self._check_completeness()
                     self._freeze_jina_version()
 
+                    _dockerfile_path = _check_result['Dockerfile']
+                    _labels = {_label_prefix + k: str(v) for k, v in self.manifest.items()}
                     streamer = self._raw_client.build(
                         decode=True,
                         path=self.args.path,
                         tag=self.tag,
                         pull=self.args.pull,
-                        dockerfile=self.dockerfile_path_revised,
+                        dockerfile=_dockerfile_path,
+                        labels=_labels,
                         rm=True
                     )
 
@@ -344,7 +347,8 @@ def build(self) -> Dict:
                 image, log = self._client.images.build(path=self.args.path,
                                                        tag=self.tag,
                                                        pull=self.args.pull,
-                                                       dockerfile=self.dockerfile_path_revised,
+                                                       dockerfile=_dockerfile_path,
+                                                       labels=_labels,
                                                        rm=True)
 
                 # success
@@ -594,7 +598,6 @@ def _check_completeness(self) -> Dict:
 
         self.manifest = self._read_manifest(manifest_path)
         self.manifest['jina_version'] = jina_version
-        self.dockerfile_path_revised = self._get_revised_dockerfile(dockerfile_path, self.manifest)
         self.executor_name = safe_url_name(
             f'{self.args.repository}/' + f'{self.manifest["type"]}.{self.manifest["kind"]}.{self.manifest["name"]}')
         self.tag = self.executor_name + f':{self.manifest["version"]}-{jina_version}'
@@ -647,25 +650,6 @@ def _validate_manifest(self, manifest: Dict) -> None:
         for k, v in manifest.items():
             self.logger.debug(f'{k}: {v}')
 
-    def _get_revised_dockerfile(self, dockerfile_path: str, manifest: Dict) -> str:
-        # modify dockerfile
-        revised_dockerfile = []
-        with open(dockerfile_path) as fp:
-            for l in fp:
-                revised_dockerfile.append(l)
-                if l.startswith('FROM'):
-                    revised_dockerfile.append('LABEL ')
-                    revised_dockerfile.append(
-                        ' \\      \n'.join(f'{_label_prefix}{k}="{v}"' for k, v in manifest.items()))
-
-        f = tempfile.NamedTemporaryFile('w', delete=False).name
-        with open(f, 'w', encoding='utf8') as fp:
-            fp.writelines(revised_dockerfile)
-
-        for k in revised_dockerfile:
-            self.logger.debug(k)
-        return f
-
     def _write_slack_message(self, *args):
         def _expand_fn(v):
             if isinstance(v, str):
diff --git a/tests/unit/docker/hub-mwu-multistage/Dockerfile b/tests/unit/docker/hub-mwu-multistage/Dockerfile
new file mode 100644
index 0000000000000..b4df1c0161180
--- /dev/null
+++ b/tests/unit/docker/hub-mwu-multistage/Dockerfile
@@ -0,0 +1,11 @@
+FROM jinaai/jina:test-pip as base
+
+ADD *.py mwu_encoder.yml ./
+
+FROM base
+
+ADD README.md ./
+
+FROM base
+
+ENTRYPOINT ["jina", "pod", "--uses", "mwu_encoder.yml"]
\ No newline at end of file
diff --git a/tests/unit/docker/hub-mwu-multistage/README.md b/tests/unit/docker/hub-mwu-multistage/README.md
new file mode 100644
index 0000000000000..2826920f84173
--- /dev/null
+++ b/tests/unit/docker/hub-mwu-multistage/README.md
@@ -0,0 +1,3 @@
+# MWUEncoder
+
+A dummy encoder encodes any input into a random 3-dimensional vector. This encoder has a dummy parameter `greetings` which prints a greeting message on start and on every encode. 
\ No newline at end of file
diff --git a/tests/unit/docker/hub-mwu-multistage/__init__.py b/tests/unit/docker/hub-mwu-multistage/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/tests/unit/docker/hub-mwu-multistage/manifest.yml b/tests/unit/docker/hub-mwu-multistage/manifest.yml
new file mode 100644
index 0000000000000..1fddc6ecbb12c
--- /dev/null
+++ b/tests/unit/docker/hub-mwu-multistage/manifest.yml
@@ -0,0 +1,16 @@
+manifest_version: 1
+name: Dummy MWU Encoder
+description: a minimum working unit of a containerized encoder, used for tutorial only
+author: Jina AI Dev-Team (dev-team@jina.ai)
+url: https://jina.ai
+documentation: https://github.com/jina-ai/jina-hub
+version: 0.0.6
+vendor: Jina AI Limited
+license: apache-2.0
+update: nightly
+keywords: 
+  - example
+  - demo
+  - toy-example
+kind: encoder
+type: pod
diff --git a/tests/unit/docker/hub-mwu-multistage/mwu_encoder.py b/tests/unit/docker/hub-mwu-multistage/mwu_encoder.py
new file mode 100644
index 0000000000000..027002551a732
--- /dev/null
+++ b/tests/unit/docker/hub-mwu-multistage/mwu_encoder.py
@@ -0,0 +1,17 @@
+from typing import Any
+
+import numpy as np
+
+from jina.executors.encoders import BaseEncoder
+
+
+class MWUEncoder(BaseEncoder):
+
+    def __init__(self, greetings: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._greetings = greetings
+        self.logger.success(f'look at me! {greetings}')
+
+    def encode(self, data: Any, *args, **kwargs) -> Any:
+        self.logger.info(f'{self._greetings} {data}')
+        return np.random.random([data.shape[0], 3])
diff --git a/tests/unit/docker/hub-mwu-multistage/mwu_encoder.yml b/tests/unit/docker/hub-mwu-multistage/mwu_encoder.yml
new file mode 100644
index 0000000000000..23b9927d74473
--- /dev/null
+++ b/tests/unit/docker/hub-mwu-multistage/mwu_encoder.yml
@@ -0,0 +1,7 @@
+!MWUEncoder
+with:
+  greetings: im from internal yaml!
+metas:
+  name: my-mwu-encoder
+  py_modules: mwu_encoder.py
+  workspace: ./
\ No newline at end of file
diff --git a/tests/unit/docker/test_hub_build.py b/tests/unit/docker/test_hub_build.py
index 9111cc60234f4..eea13a1fc5b80 100644
--- a/tests/unit/docker/test_hub_build.py
+++ b/tests/unit/docker/test_hub_build.py
@@ -86,3 +86,32 @@ def test_jina_version_freeze_no_jina_dependency(requirements, tmpdir):
     with open(requirements_file, 'r') as fp:
         requirements = pkg_resources.parse_requirements(fp)
         assert len(list(filter(lambda x: 'jina' in str(x), requirements))) == 0
+
+
+def test_labels():
+    class MockContainers:
+        def __init__(self):
+            pass
+
+        def build(self, *args, **kwargs):
+            labels = kwargs['labels']
+            assert all([isinstance(v, str) for k, v in labels.items()])
+            assert 'ai.jina.hub.version' in labels
+            assert 'ai.jina.hub.name' in labels
+
+            raise BaseException('labels all good')
+
+    args = set_hub_build_parser().parse_args(
+        [os.path.join(cur_dir, 'hub-mwu'), '--test-uses', '--raise-error'])
+    hubio = HubIO(args)
+    hubio._raw_client = MockContainers()
+
+    with pytest.raises(BaseException, match='labels all good'):
+        hubio.build()
+
+
+@pytest.mark.timeout(360)
+def test_hub_build_multistage():
+    args = set_hub_build_parser().parse_args(
+        [os.path.join(cur_dir, 'hub-mwu-multistage'), '--test-uses', '--raise-error'])
+    HubIO(args).build()
\ No newline at end of file


Code review comments:
diff_hunk: @@ -309,15 +309,18 @@ def build(self) -> Dict:
 
             with TimeContext(f'building {colored(self.args.path, "green")}', self.logger) as tc:
                 try:
-                    self._check_completeness()
+                    _check_result = self._check_completeness()
                     self._freeze_jina_version()
 
+                    _dockerfile_path = _check_result['Dockerfile']
path: jina/docker/hubio.py
text: @user1:
So I think here it should be `_dockerfile_path = 'Dockerfile'`, or alternatively something like `_check_result['Dockerfile'].split('/')[-1]` (to be better prepared for different dockerfile names, such as the one for GPU). This could also be changed in `_check_completeness`

This is because in the call to the client below, `path=self.args.path` is passed, which means that the `dockerfile` location will be relative to that `path`. However in `_check_completness` the dockerfile path is obtained by joining `self.args.path` and `Dockerfile`. I think this was the reason that #1732 broke the hub builds.

@user1:
Also, back then the issue was that CI didn't catch this error for some reason. I hope now it will

@author:
hmm, the tests actually work locally. I've put them into the integration tests and let's see how ci says.

@user1:
Aaaah I figured it out! 

The reason test worked and hub build broke was because in tests `args.path` is specified with an absolute path - for example `hub_build_test` uses `os.path.join(cur_dir, 'hub-mwu')`, where `cur_dir` is an absolute path, while building in Jina Hub is done using relative paths, e.g. `encoders/audio/MFCCTimbreEncoder/`. 

Tests worked absolutely fine, but they were only done for the absolute path - where both `path` and `dockerfile` are absolute paths. If both are relative, then docker SDK tries to join them together, which doesn't work unless `dockerfile` is really just the name of the file inside `args.path`.

So yes, extracting the path using `os.path.basename` as you are doing now is the right way (btw you have to do it in another location in the file, that's why the CI is failing). All that should be added is a test where `args.path` is passed as a relative path, and we are good to go :D

@author:
Great catch! Added another test on the relative pathes.