feat: api normalizer support parameter build_args_envs (#48)

* feat: api normalizer support parameter build_args_envs

* fix: remove assert in func insert_build_args_envs

* fix: add optional dry_run  in normalize func

* test: optimization test

* test: replace the test token value

* test: improve test core

* test: improve test core

* test: improve  multi-stage case

* refactor: remove func add_build_args_envs, optimize code formatting and remove unused reference packages

* test: improve test core

* test: remove the useless package

* fix: support multi-stage build, environment variable setting

* test: improve test docker

.gitignore

@@ -147,6 +147,7 @@ tests/cases/executor_3/config.yml
 tests/cases/executor_3/Dockerfile
 tests/cases/executor_4/__jina__.Dockerfile
 tests/cases/executor_4/Dockerfile
+tests/cases/executor_7/__jina__.Dockerfile
 tests/cases/nested/__jina__.Dockerfile
 tests/cases/nested/config.yml
 tests/cases/nested/Dockerfile

normalizer/core.py

@@ -368,6 +368,7 @@ def normalize(
     work_path: 'pathlib.Path',
     meta: Dict = {'jina': '2'},
     env: Dict = {},
+    build_args_envs: Dict = {},
     dry_run: bool = False,
     **kwargs,
 ) -> ExecutorModel:
@@ -586,7 +587,8 @@ def normalize(
             docker_file=dockerfile_path,
             build_args={'JINA_VERSION': f'{jina_version}'},
         )
-
+        if len(build_args_envs.keys()):
+            dockerfile.insert_build_args_envs(build_args_envs)
         # if dockerfile.is_multistage():
         #     # Don't support multi-stage Dockerfie Optimization
         #     return
@@ -597,13 +599,17 @@ def normalize(
         #         f'RUN pip install jina=={jina_version}', at_start=True
         #     )
         #     dockerfile.dump(work_path / 'Dockerfile.normed')
+        if not dry_run:
+            dockerfile.dump(dockerfile_path)
     else:
         logger.debug('=> generating Dockerfile ...')
         dockerfile = ExecutorDockerfile(build_args={'JINA_VERSION': jina_image_tag})
 
         # if len(base_images) > 0:
         #     logger.debug(f'=> use base image: {base_images}')
         #     dockerfile.baseimage = base_images.pop()
+        if len(build_args_envs.keys()):
+            dockerfile.insert_build_args_envs(build_args_envs)
 
         dockerfile.add_work_dir()
         # dockerfile._parser.add_lines(f'RUN pip install jina=={jina_version}')

normalizer/docker/parser.py

@@ -5,6 +5,8 @@
 from posixpath import basename
 from textwrap import dedent
 from typing import Dict, List
+import re
+FROM_VAR_RE = re.compile(r"(?P<var>(?P<name>FROM .*))")
 
 from dockerfile_parse import DockerfileParser
 
@@ -41,6 +43,47 @@ def __init__(
 
     def __str__(self):
         return self.content
+
+    def get_build_args_envs(self, build_args_envs): 
+        build_args = '';
+        build_args_line_num = 0;
+        for index, item in enumerate(build_args_envs):
+            build_args_line_num += 1
+            build_args += dedent(
+            f"""\
+            ARG {item}
+            """
+        )
+        build_envs = '';
+        build_envs_line_num = 0
+        for index, item in enumerate(build_args_envs):
+            build_envs_line_num += 1
+            build_envs += dedent(
+            f"""\
+            ENV {item}=${item}
+            """
+        )
+        return (build_args, build_envs, build_args_line_num, build_envs_line_num)
+
+    def insert_build_args_envs(self, build_args_envs):
+        build_args, build_envs, build_args_line_num, build_envs_line_num = self.get_build_args_envs(build_args_envs)
+        build_args_envs_line_num = 0;
+        build_args_envs_str = '';
+        if len(build_args): 
+            build_args_envs_str += build_args
+            build_args_envs_line_num += build_args_line_num
+
+        if len(build_envs): 
+            build_args_envs_str += build_envs
+            build_args_envs_line_num += build_envs_line_num
+
+        if len(build_args_envs_str):
+            insert_num = 0
+            for index, line in enumerate(self._parser.lines):
+                if (FROM_VAR_RE.match(line) is not None):
+                    insert_line_num = index + insert_num * build_args_envs_line_num + 1
+                    self._parser.add_lines_at(insert_line_num, build_args_envs_str)
+                    insert_num += 1
 
     def add_apt_installs(self, tools):
         instruction_template = dedent(

normalizer/models.py

@@ -38,6 +38,7 @@ class PackagePayload(BaseModel):
     package_path: Path
     meta: Optional[Dict] = {'jina': 'master'}
     env: Optional[Dict] = {}
+    build_args_envs: Optional[Dict] = {}
 
 
 class NormalizeResult(BaseModel):

server/routes/normalizer.py

@@ -32,6 +32,7 @@ def normalize(
             block_data.package_path,
             meta=block_data.meta,
             env=block_data.env,
+            build_args_envs=block_data.build_args_envs
         )
 
     except Exception as ex:

tests/cases/executor_1/Dockerfile.expect

@@ -0,0 +1,14 @@
+# This file is automatically generated by Jina executor normalizer plugin.
+# It is not intended for manual editing.
+
+FROM jinaai/jina:2-py37-perf
+ARG AUTH_TOKEN
+ARG TOKEN
+ENV AUTH_TOKEN=$AUTH_TOKEN
+ENV TOKEN=$TOKEN
+
+# setup the workspace
+COPY . /workspace
+WORKDIR /workspace
+
+ENTRYPOINT ["jina", "executor", "--uses", "config.yml"]

tests/cases/executor_7/Dockerfile

@@ -0,0 +1,16 @@
+# This file is automatically generated by Jina executor normalizer plugin.
+# It is not intended for manual editing.
+
+FROM jinaai/jina:3.6.10-py39-perf        
+
+# install the third-party requirements
+RUN  "the first stage" 
+ENTRYPOINT ["jina", "executor", "--uses", "config.yml"]
+
+
+FROM jinaai/jina:3.6.10-py39-perf     
+
+# install the third-party requirements
+RUN  "the second stage" 
+
+ENTRYPOINT ["jina", "executor", "--uses", "config.yml"]

tests/cases/executor_7/Dockerfile.expect

@@ -0,0 +1,24 @@
+# This file is automatically generated by Jina executor normalizer plugin.
+# It is not intended for manual editing.
+
+FROM jinaai/jina:3.6.10-py39-perf        
+ARG AUTH_TOKEN
+ARG TOKEN
+ENV AUTH_TOKEN=$AUTH_TOKEN
+ENV TOKEN=$TOKEN
+
+# install the third-party requirements
+RUN  "the first stage" 
+ENTRYPOINT ["jina", "executor", "--uses", "config.yml"]
+
+
+FROM jinaai/jina:3.6.10-py39-perf     
+ARG AUTH_TOKEN
+ARG TOKEN
+ENV AUTH_TOKEN=$AUTH_TOKEN
+ENV TOKEN=$TOKEN
+
+# install the third-party requirements
+RUN  "the second stage" 
+
+ENTRYPOINT ["jina", "executor", "--uses", "config.yml"]

tests/cases/executor_7/README.md

@@ -0,0 +1,2 @@
+# executor_7
+

tests/cases/executor_7/config.yml

@@ -0,0 +1,4 @@
+jtype: executor_7
+metas:
+  py_modules:
+    - executor.py

tests/cases/executor_7/executor.py

@@ -0,0 +1,7 @@
+from jina import Executor, DocumentArray, requests
+
+
+class executor_7(Executor):
+    @requests
+    def foo(self, docs: DocumentArray, **kwargs):
+        pass

tests/cases/executor_7/requirements.txt

@@ -0,0 +1 @@
+

tests/docker_cases/Dockerfile.case1.expect

@@ -0,0 +1,12 @@
+FROM jinaai/jina:2.0-perf
+ARG AUTH_TOKEN
+ARG TOKEN
+ENV AUTH_TOKEN=$AUTH_TOKEN
+ENV TOKEN=$TOKEN
+
+COPY . /workspace
+WORKDIR /workspace
+
+RUN pip install .
+
+ENTRYPOINT ["jina", "executor", "--uses", "config.yml"]

tests/test_core.py

@@ -1,6 +1,7 @@
 import json
 from pathlib import Path
 import pytest
+import os
 
 from normalizer import deps, core
 from normalizer.models import ExecutorModel
@@ -40,47 +41,47 @@ def test_inspect_dummy_execs():
     [
         (
             Path(__file__).parent / 'cases' / 'executor_1',
-            Path(__file__).parent / 'cases' / 'executor_1.json',
+            Path(__file__).parent / 'cases' / 'executor_1.json'
         ),
         (
             Path(__file__).parent / 'cases' / 'executor_2',
-            Path(__file__).parent / 'cases' / 'executor_2.json',
+            Path(__file__).parent / 'cases' / 'executor_2.json'
         ),
         (
             Path(__file__).parent / 'cases' / 'executor_3',
-            Path(__file__).parent / 'cases' / 'executor_3.json',
+            Path(__file__).parent / 'cases' / 'executor_3.json'
         ),
         (
             Path(__file__).parent / 'cases' / 'executor_4',
-            Path(__file__).parent / 'cases' / 'executor_4.json',
+            Path(__file__).parent / 'cases' / 'executor_4.json'
         ),
         (
             Path(__file__).parent / 'cases' / 'executor_5',
-            None,
+            None
         ),
         (
             Path(__file__).parent / 'cases' / 'executor_6',
-            None,
+            None
         ),
         (
             Path(__file__).parent / 'cases' / 'nested',
-            Path(__file__).parent / 'cases' / 'nested.json',
+            Path(__file__).parent / 'cases' / 'nested.json'
         ),
         (
             Path(__file__).parent / 'cases' / 'nested_2',
-            None,
+            None
         ),
         (
             Path(__file__).parent / 'cases' / 'nested_3',
-            None,
+            None
         ),
         (
             Path(__file__).parent / 'cases' / 'nested_4',
-            None,
+            None
         ),
         (
             Path(__file__).parent / 'cases' / 'nested_5',
-            None,
+            None
         ),
     ],
 )
@@ -94,6 +95,57 @@ def test_get_executor_args(package_path, expected_path):
             assert executor == expected_executor
     else:
         core.normalize(package_path, dry_run=True)
+
+
+@pytest.mark.parametrize(
+    'package_path, build_args_envs',
+    [   
+        (
+            Path(__file__).parent / 'cases' / 'executor_1',
+            { 
+                'AUTH_TOKEN': "AUTH_TOKEN",
+                'TOKEN': 'ghp_Nwh9o70GDSzs'
+            }
+        ),
+        (
+            Path(__file__).parent / 'cases' / 'executor_7',
+            { 
+                'AUTH_TOKEN': "AUTH_TOKEN",
+                'TOKEN': 'ghp_Nwh9o70GDSzs'
+            }
+        ),
+    ],
+)
+def test_compare_dockerfile_env_vars(package_path, build_args_envs):
+
+    dockerfile_path = Path(package_path / 'Dockerfile') 
+    dockerfile_expected_path = Path(package_path / 'Dockerfile.expect') 
+
+    originDockerfileStr = None;
+    if dockerfile_path.exists():
+        with open(dockerfile_path, 'r') as fp:
+            originDockerfileStr = str(fp.read())
+
+    core.normalize(package_path, build_args_envs=build_args_envs, dry_run=False)
+    assert dockerfile_path.exists() == True;
+
+    dockerfileStr = None
+    with open(dockerfile_path, 'r') as fp:
+        dockerfileStr = str(fp.read())
+
+    dockerfileExpectedStr = ''
+    with open(dockerfile_expected_path, 'r') as fp:
+        dockerfileExpectedStr = str(fp.read())
+
+    if originDockerfileStr:
+        with open(dockerfile_path, 'w') as fp:
+            fp.write(originDockerfileStr)
+    else:
+        os.remove(dockerfile_path)
+
+    assert dockerfileExpectedStr == dockerfileStr
+
+
 def test_prelude():
     imports = [
         deps.Package(name='tensorflow', version='2.5.0'),

tests/test_docker.py

@@ -34,9 +34,27 @@ def test_baseimage(exe_dockerfile):
     assert exe_dockerfile.lines[3] == 'FROM jinaai/jina:latest-perf\n'
 
 
-def test_load_dockerfile():
+@pytest.mark.parametrize(
+    'build_args_envs',
+    [   
+        (
+             { 
+                'AUTH_TOKEN': "AUTH_TOKEN",
+                'TOKEN': 'ghp_Nwh9o70GDSzs'
+            }
+        ),
+    ],
+)
+def test_load_dockerfile(build_args_envs):
     docker_file = Path(__file__).parent / 'docker_cases' / 'Dockerfile.case1'
+    docker_expect_file = Path(__file__).parent / 'docker_cases' / 'Dockerfile.case1.expect'
+
     parser = ExecutorDockerfile(docker_file=docker_file)
+    parser.insert_build_args_envs(build_args_envs)
+
+    expect_parser = ExecutorDockerfile(docker_file=docker_expect_file)
+
+    assert str(parser) == str(expect_parser)
     assert len(parser.parent_images) == 1
     assert parser.baseimage == 'jinaai/jina:2.0-perf'
     assert parser.entrypoint == '["jina", "executor", "--uses", "config.yml"]'