fix: cache hub pull with build env (#5061)

* test: add test for commit env in hubble fetch response

* fix: make build envs picklable and thus fix hubble fetch caching

* refactor: update jina/hubble/hubio.py

Co-authored-by: Nan Wang <nan.wang@jina.ai>

* refactor: update jina/hubble/hubio.py

Co-authored-by: Nan Wang <nan.wang@jina.ai>

* fix: adjust test to new error message

Co-authored-by: Nan Wang <nan.wang@jina.ai>

jina/hubble/hubio.py

@@ -356,14 +356,14 @@ def push(self) -> None:
                 env_list = env.split('=')
                 if len(env_list) != 2:
                     raise Exception(
-                        f'The --build-env parameter: `{env}` is wrong format. you can use: `--build-env {env}=YOUR_VALUE`.'
+                        f'The `--build-env` parameter: `{env}` is wrong format. you can use: `--build-env {env}=YOUR_VALUE`.'
                     )
                 if check_requirements_env_variable(env_list[0]) is False:
                     raise Exception(
-                        f'The --build-env parameter key:`{env_list[0]}` can only consist of uppercase letter and number and underline.'
+                        f'The `--build-env` parameter key:`{env_list[0]}` can only consist of uppercase letter and number and underline.'
                     )
                 build_env_dict[env_list[0]] = env_list[1]
-            build_env = build_env_dict if len(list(build_env_dict.keys())) > 0 else None
+            build_env = build_env_dict if build_env_dict else None
 
         requirements_file = work_path / 'requirements.txt'
 
@@ -433,7 +433,7 @@ def push(self) -> None:
 
                 if build_env:
                     form_data['buildEnv'] = json.dumps(build_env)
-                
+
                 uuid8, secret = load_secret(work_path)
                 if self.args.force_update or uuid8:
                     form_data['id'] = self.args.force_update or uuid8
@@ -712,7 +712,7 @@ def _send_request_with_retry(url, **kwargs):
             image_name=image_name,
             archive_url=resp['package']['download'],
             md5sum=resp['package']['md5'],
-            build_env=buildEnv.keys() if buildEnv else [],
+            build_env=list(buildEnv.keys()) if buildEnv else [],
         )
 
     @staticmethod

tests/unit/hubble/test_hubio.py

@@ -21,15 +21,13 @@
     disk_cache_offline,
     get_requirements_env_variables,
 )
+from jina.hubble.hubapi import get_secret_path
 from jina.hubble.hubio import HubExecutor, HubIO
 from jina.parsers.hubble import (
     set_hub_new_parser,
     set_hub_pull_parser,
     set_hub_push_parser,
 )
-from jina.hubble.hubapi import (
-    get_secret_path
-)
 
 cur_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -93,23 +91,34 @@ def iter_lines(self):
 
 
 class FetchMetaMockResponse:
-    def __init__(self, response_code: int = 200, no_image=False, fail_count=0):
+    def __init__(
+        self,
+        response_code: int = 200,
+        no_image=False,
+        fail_count=0,
+        add_build_env=False,
+    ):
         self.response_code = response_code
         self.no_image = no_image
         self._tried_count = 0
         self._fail_count = fail_count
+        self._build_env = add_build_env
 
     def json(self):
         if self._tried_count <= self._fail_count:
             return {'message': 'Internal server error'}
 
+        commit_val = {'_id': 'commit_id', 'tags': ['v0']}
+        if self._build_env:
+            commit_val['commitParams'] = {'buildEnv': {'key1': 'val1', 'key2': 'val2'}}
+
         return {
             'data': {
                 'keywords': [],
                 'id': 'dummy_mwu_encoder',
                 'name': 'alias_dummy',
                 'visibility': 'public',
-                'commit': {'_id': 'commit_id', 'tags': ['v0']},
+                'commit': commit_val,
                 'package': {
                     'containers': []
                     if self.no_image
@@ -169,7 +178,6 @@ def _mock_post(url, data, headers=None, stream=True):
 
     result = HubIO(args).push()
 
-
     exec_config_path = get_secret_path(os.stat(exec_path).st_ino)
     shutil.rmtree(exec_config_path)
 
@@ -219,13 +227,13 @@ def _mock_post(url, data, headers=None, stream=True):
 @pytest.mark.parametrize(
     'env_variable_consist_error',
     [
-        'The --build-env parameter key:`{build_env_key}` can only consist of uppercase letter and number and underline.'
+        'The `--build-env` parameter key:`{build_env_key}` can only consist of uppercase letter and number and underline.'
     ],
 )
 @pytest.mark.parametrize(
     'env_variable_format_error',
     [
-        'The --build-env parameter: `{build_env}` is wrong format. you can use: `--build-env {build_env}=YOUR_VALUE`.'
+        'The `--build-env` parameter: `{build_env}` is wrong format. you can use: `--build-env {build_env}=YOUR_VALUE`.'
     ],
 )
 @pytest.mark.parametrize('path', ['dummy_executor_fail'])
@@ -313,7 +321,7 @@ def _mock_post(url, data, headers=None, stream=True):
 
     with pytest.raises(Exception) as info:
         result = HubIO(args).push()
-        
+
     assert requirements_file_need_build_env_error.format(
         env_variables_str=','.join(requirements_file_env_variables)
     ) in str(info.value)
@@ -400,7 +408,6 @@ def _mock_post(url, data, headers=None, stream=True):
 
     with pytest.raises(Exception) as info:
         HubIO(args).push()
-
 
     assert expected_error.format(dockerfile=dockerfile, work_path=args.path) in str(
         info.value
@@ -466,6 +473,41 @@ def _mock_post(url, json, headers=None):
     assert executor.tag == 'v0.1'
 
 
+@pytest.mark.parametrize('rebuild_image', [True, False])
+def test_fetch_with_buildenv(mocker, monkeypatch, rebuild_image):
+    mock = mocker.Mock()
+
+    def _mock_post(url, json, headers=None):
+        mock(url=url, json=json)
+        return FetchMetaMockResponse(response_code=200, add_build_env=True)
+
+    monkeypatch.setattr(requests, 'post', _mock_post)
+    args = set_hub_pull_parser().parse_args(['jinahub://dummy_mwu_encoder'])
+
+    executor, _ = HubIO(args).fetch_meta(
+        'dummy_mwu_encoder', None, rebuild_image=rebuild_image, force=True
+    )
+
+    assert executor.uuid == 'dummy_mwu_encoder'
+    assert executor.name == 'alias_dummy'
+    assert executor.tag == 'v0'
+    assert executor.image_name == 'jinahub/pod.dummy_mwu_encoder'
+    assert executor.md5sum == 'ecbe3fdd9cbe25dbb85abaaf6c54ec4f'
+    assert executor.build_env == ['key1', 'key2']
+
+    _, mock_kwargs = mock.call_args_list[0]
+    assert mock_kwargs['json']['rebuildImage'] is rebuild_image
+
+    executor, _ = HubIO(args).fetch_meta('dummy_mwu_encoder', '', force=True)
+    assert executor.tag == 'v0'
+
+    _, mock_kwargs = mock.call_args_list[1]
+    assert mock_kwargs['json']['rebuildImage'] is True  # default value must be True
+
+    executor, _ = HubIO(args).fetch_meta('dummy_mwu_encoder', 'v0.1', force=True)
+    assert executor.tag == 'v0.1'
+
+
 def test_fetch_with_no_image(mocker, monkeypatch):
     mock = mocker.Mock()