feat: support jinaai executor scheme (#5515)

jina/orchestrate/deployments/config/helper.py

@@ -1,7 +1,7 @@
 import os
 from typing import Dict
 
-from hubble.executor.helper import parse_hub_uri
+from hubble.executor.helper import is_valid_docker_uri, parse_hub_uri
 from hubble.executor.hubio import HubIO
 
 from jina import (
@@ -139,8 +139,6 @@ def validate_uses(uses: str):
         return True
 
     try:
-        scheme, _, _, _ = parse_hub_uri(uses)
-        if scheme in {'jinahub+docker', 'jinahub+sandbox'}:
-            return True
+        return is_valid_docker_uri(uses)
     except ValueError:
         return False

tests/integration/hub_usage/test_hub_usage.py

@@ -62,7 +62,10 @@ def local_hub_executor(tmpdir):
     )
 
 
-def test_use_from_local_hub_deployment_level(mocker, monkeypatch, local_hub_executor):
+@pytest.mark.parametrize('uses', ['jinahub://hello', 'jinaai://jina-ai/hello'])
+def test_use_from_local_hub_deployment_level(
+    mocker, monkeypatch, local_hub_executor, uses
+):
     from hubble.executor.hubio import HubExecutor, HubIO
 
     mock = mocker.Mock()
@@ -91,12 +94,13 @@ def _mock_fetch(
         )
 
     monkeypatch.setattr(HubIO, 'fetch_meta', _mock_fetch)
-    a = set_deployment_parser().parse_args(['--uses', 'jinahub://hello'])
+    a = set_deployment_parser().parse_args(['--uses', uses])
     with Deployment(a):
         pass
 
 
-def test_use_from_local_hub_flow_level(mocker, monkeypatch, local_hub_executor):
+@pytest.mark.parametrize('uses', ['jinahub://hello', 'jinaai://jina-ai/hello'])
+def test_use_from_local_hub_flow_level(mocker, monkeypatch, local_hub_executor, uses):
     from hubble.executor.hubio import HubExecutor, HubIO
 
     mock = mocker.Mock()
@@ -126,5 +130,5 @@ def _mock_fetch(
 
     monkeypatch.setattr(HubIO, 'fetch_meta', _mock_fetch)
 
-    with Flow().add(uses='jinahub://hello', install_requirements=True):
+    with Flow().add(uses=uses, install_requirements=True):
         pass

tests/integration/issues/github_4488/test_deployment_protocol.py

@@ -32,13 +32,17 @@ def key_pem(cert_prefix):
 
 @pytest.mark.parametrize('protocol', ['http', 'websocket', 'grpc'])
 @pytest.mark.parametrize('tls', [True, False])
-def test_deployment_protocol(protocol, tls, cert_pem, key_pem):
+@pytest.mark.parametrize(
+    'uses',
+    ['jinaai+sandbox://jina-ai/DummyHubExecutor'],
+)
+def test_deployment_protocol(protocol, tls, cert_pem, key_pem, uses):
     cert = cert_pem if tls else None
     key = key_pem if tls else None
     f = (
         Flow(protocol=protocol, ssl_certfile=cert, ssl_keyfile=key)
         .add(uses=MyExec)
-        .add(uses='jinahub+sandbox://DummyHubExecutor')
+        .add(uses=uses)
     )
     with f:
         for node, v in f._deployment_nodes.items():

tests/integration/sandbox/test_sandbox.py

@@ -4,8 +4,11 @@
 
 
 @pytest.mark.parametrize('endpoint', ['foo', 'bar'])
-def test_sandbox(endpoint):
-    with Flow().add(uses='jinahub+sandbox://Hello') as f:
+@pytest.mark.parametrize(
+    'uses', ['jinaai+sandbox://jina-ai/Hello']
+)
+def test_sandbox(endpoint, uses):
+    with Flow().add(uses=uses) as f:
         da = f.post(
             endpoint,
             [

tests/unit/orchestrate/deployments/config/test_docker_compose_pod_config.py

@@ -29,7 +29,13 @@ def namespace_equal(
 
 @pytest.mark.parametrize('shards', [1, 5])
 @pytest.mark.parametrize('replicas', [1, 5])
-@pytest.mark.parametrize('uses_before', [None, 'jinahub+docker://HubBeforeExecutor'])
+@pytest.mark.parametrize(
+    'uses_before',
+    [
+        None,
+        'jinaai+docker://jina-ai/HubBeforeExecutor',
+    ],
+)
 @pytest.mark.parametrize('uses_after', [None, 'docker://docker_after_image:latest'])
 @pytest.mark.parametrize('uses_with', ['{"paramkey": "paramvalue"}', None])
 @pytest.mark.parametrize('uses_metas', ['{"workspace": "workspacevalue"}', None])
@@ -399,10 +405,26 @@ def test_docker_compose_gateway(deployments_addresses, custom_gateway):
 @pytest.mark.parametrize('shards', [3, 1])
 @pytest.mark.parametrize('replicas', [3, 1])
 @pytest.mark.parametrize(
-    'uses', ['jinahub+docker://HubExecutor', 'docker://docker_image:latest']
+    'uses',
+    [
+        'docker://docker_image:latest',
+        'jinaai+docker://jina-ai/HubExecutor',
+    ],
+)
+@pytest.mark.parametrize(
+    'uses_before',
+    [
+        None,
+        'jinaai+docker://jina-ai/HubBeforeExecutor',
+    ],
+)
+@pytest.mark.parametrize(
+    'uses_after',
+    [
+        None,
+        'jinaai+docker://jina-ai/HubAfterExecutor',
+    ],
 )
-@pytest.mark.parametrize('uses_before', [None, 'jinahub+docker://HubBeforeExecutor'])
-@pytest.mark.parametrize('uses_after', [None, 'jinahub+docker://HubAfterExecutor'])
 @pytest.mark.parametrize('uses_with', ['{"paramkey": "paramvalue"}', None])
 @pytest.mark.parametrize('uses_metas', ['{"workspace": "workspacevalue"}', None])
 @pytest.mark.parametrize('polling', ['ANY', 'ALL'])
@@ -533,7 +555,10 @@ def _mock_fetch(
         if uses_before is not None:
             uses_before_name, uses_before_config = yaml_configs[1]
             assert uses_before_name == 'executor-uses-before'
-            assert uses_before_config['image'] == 'jinahub/HubBeforeExecutor'
+            assert uses_before_config['image'] in {
+                'jinahub/HubBeforeExecutor',
+                'jinahub/jina-ai/HubBeforeExecutor',
+            }
             assert uses_before_config['entrypoint'] == ['jina']
             uses_before_args = uses_before_config['command']
             assert uses_before_args[0] == 'executor'
@@ -554,7 +579,10 @@ def _mock_fetch(
                 yaml_configs[1] if uses_before is None else yaml_configs[2]
             )
             assert uses_after_name == 'executor-uses-after'
-            assert uses_after_config['image'] == 'jinahub/HubAfterExecutor'
+            assert uses_after_config['image'] in {
+                'jinahub/HubAfterExecutor',
+                'jinahub/jina-ai/HubAfterExecutor',
+            }
             assert uses_after_config['entrypoint'] == ['jina']
             uses_after_args = uses_after_config['command']
             assert uses_after_args[0] == 'executor'
@@ -586,11 +614,11 @@ def _mock_fetch(
                 expected_name += f'-rep-{i_replica}'
                 expected_arg_name += f'/rep-{i_replica}'
             assert replica_name == expected_name
-            assert (
-                replica_config['image'] == 'docker_image:latest'
-                if uses == 'docker_image:latest'
-                else 'jinahub/HubExecutor'
-            )
+            assert replica_config['image'] in {
+                'docker_image:latest',
+                'jinahub/HubExecutor',
+                'jinahub/jina-ai/HubExecutor',
+            }
             assert replica_config['entrypoint'] == ['jina']
             replica_args = replica_config['command']
             assert replica_args[0] == 'executor'

tests/unit/orchestrate/deployments/config/test_helper.py

@@ -42,7 +42,10 @@ def test_to_compatible_name():
     assert to_compatible_name('executor/hey-ha_HO') == 'executor-hey-ha-ho'
 
 
-def test_get_image_name(mocker, monkeypatch):
+@pytest.mark.parametrize(
+    'uses', ['jinaai://jina-ai/DummyExecutor']
+)
+def test_get_image_name(mocker, monkeypatch, uses):
     mock = mocker.Mock()
 
     def _mock_fetch(
@@ -71,11 +74,9 @@ def _mock_fetch(
 
     monkeypatch.setattr(HubIO, 'fetch_meta', _mock_fetch)
 
-    uses = 'jinahub://DummyExecutor'
-
     image_name = get_image_name(uses)
 
-    assert image_name == 'jinahub/DummyExecutor'
+    assert image_name in {'jinahub/DummyExecutor', 'jinahub/jina-ai/DummyExecutor'}
 
     _, mock_kwargs = mock.call_args_list[0]
     assert mock_kwargs['rebuild_image'] is True  # default value must be True

tests/unit/orchestrate/deployments/config/test_k8s_deployment_config.py

@@ -28,7 +28,13 @@ def namespace_equal(
 
 
 @pytest.mark.parametrize('shards', [1, 5])
-@pytest.mark.parametrize('uses_before', [None, 'jinahub+docker://HubBeforeExecutor'])
+@pytest.mark.parametrize(
+    'uses_before',
+    [
+        None,
+        'jinaai+docker://jina-ai/HubBeforeExecutor',
+    ],
+)
 @pytest.mark.parametrize('uses_after', [None, 'docker://docker_after_image:latest'])
 @pytest.mark.parametrize('uses_with', ['{"paramkey": "paramvalue"}', None])
 @pytest.mark.parametrize('uses_metas', ['{"workspace": "workspacevalue"}', None])
@@ -376,10 +382,26 @@ def assert_port_config(port_dict: Dict, name: str, port: int):
 
 @pytest.mark.parametrize('shards', [3, 1])
 @pytest.mark.parametrize(
-    'uses', ['jinahub+docker://HubExecutor', 'docker://docker_image:latest']
+    'uses',
+    [
+        'docker://docker_image:latest',
+        'jinaai+docker://jina-ai/HubExecutor',
+    ],
+)
+@pytest.mark.parametrize(
+    'uses_before',
+    [
+        None,
+        'jinaai+docker://jina-ai/HubBeforeExecutor',
+    ],
+)
+@pytest.mark.parametrize(
+    'uses_after',
+    [
+        None,
+        'jinaai+docker://jina-ai/HubAfterExecutor',
+    ],
 )
-@pytest.mark.parametrize('uses_before', [None, 'jinahub+docker://HubBeforeExecutor'])
-@pytest.mark.parametrize('uses_after', [None, 'jinahub+docker://HubAfterExecutor'])
 @pytest.mark.parametrize('uses_with', ['{"paramkey": "paramvalue"}', None])
 @pytest.mark.parametrize('uses_metas', ['{"workspace": "workspacevalue"}', None])
 @pytest.mark.parametrize('polling', ['ANY', 'ALL'])
@@ -580,7 +602,10 @@ def _mock_fetch(
         if uses_before is not None:
             uses_before_container = head_containers[1]
             assert uses_before_container['name'] == 'uses-before'
-            assert uses_before_container['image'] == 'jinahub/HubBeforeExecutor'
+            assert uses_before_container['image'] in {
+                'jinahub/HubBeforeExecutor',
+                'jinahub/jina-ai/HubBeforeExecutor',
+            }
             assert uses_before_container['imagePullPolicy'] == 'IfNotPresent'
             assert uses_before_container['command'] == ['jina']
             uses_before_runtime_container_args = uses_before_container['args']
@@ -614,7 +639,10 @@ def _mock_fetch(
         if uses_after is not None:
             uses_after_container = head_containers[-1]
             assert uses_after_container['name'] == 'uses-after'
-            assert uses_after_container['image'] == 'jinahub/HubAfterExecutor'
+            assert uses_after_container['image'] in {
+                'jinahub/HubAfterExecutor',
+                'jinahub/jina-ai/HubAfterExecutor',
+            }
             assert uses_after_container['imagePullPolicy'] == 'IfNotPresent'
             assert uses_after_container['command'] == ['jina']
             uses_after_runtime_container_args = uses_after_container['args']
@@ -708,6 +736,7 @@ def _mock_fetch(
         assert shard_container['name'] == 'executor'
         assert shard_container['image'] in {
             'jinahub/HubExecutor',
+            'jinahub/jina-ai/HubExecutor',
             'docker_image:latest',
         }
         assert shard_container['imagePullPolicy'] == 'IfNotPresent'

tests/unit/orchestrate/flow/flow-construct/test_flow_to_docker_compose_yaml.py

@@ -391,10 +391,12 @@ def test_disable_auto_volume(tmpdir):
     assert 'volumes' not in services['executor0']
 
 
-def test_flow_to_docker_compose_sandbox(tmpdir):
-    flow = Flow(name='test-flow', port=8080).add(
-        uses=f'jinahub+sandbox://DummyHubExecutor'
-    )
+@pytest.mark.parametrize(
+    'uses',
+    ['jinaai+sandbox://jina-ai/DummyHubExecutor'],
+)
+def test_flow_to_docker_compose_sandbox(tmpdir, uses):
+    flow = Flow(name='test-flow', port=8080).add(uses=uses)
 
     dump_path = os.path.join(str(tmpdir), 'test_flow_docker_compose.yml')
 

tests/unit/orchestrate/flow/flow-construct/test_flow_to_k8s_yaml.py

@@ -575,11 +575,15 @@ def test_raise_exception_invalid_executor(tmpdir):
         f.to_kubernetes_yaml(str(tmpdir))
 
 
-def test_flow_to_k8s_yaml_sandbox(tmpdir):
-
-    flow = Flow(name='test-flow', port=8080).add(
-        uses=f'jinahub+sandbox://DummyHubExecutor'
-    )
+@pytest.mark.parametrize(
+    'uses',
+    [
+        f'jinaai+sandbox://jina-ai/DummyHubExecutor',
+    ],
+)
+def test_flow_to_k8s_yaml_sandbox(tmpdir, uses):
+
+    flow = Flow(name='test-flow', port=8080).add(uses=uses)
 
     dump_path = os.path.join(str(tmpdir), 'test_flow_k8s')
 

tests/unit/orchestrate/pods/test_pod_factory.py

@@ -1,18 +1,22 @@
+import pytest
 from hubble.executor.hubio import HubIO
 
 from jina.orchestrate.pods.factory import PodFactory
 from jina.parsers import set_pod_parser
 
 
-def test_container_pod(mocker, monkeypatch):
+@pytest.mark.parametrize(
+    'uses', ['jinaai+docker://jina-ai/DummyExecutor']
+)
+def test_container_pod(mocker, monkeypatch, uses):
     mock = mocker.Mock()
 
     def _mock_pull(self):
         return 'docker://jinahub/dummy_executor'
 
     monkeypatch.setattr(HubIO, 'pull', _mock_pull)
 
-    args = set_pod_parser().parse_args(['--uses', 'jinahub+docker://DummyExecutor'])
+    args = set_pod_parser().parse_args(['--uses', uses])
     pod = PodFactory.build_pod(args)
     assert pod.args.uses == 'docker://jinahub/dummy_executor'
     assert pod.name == 'ContainerPod'

tests/unit/serve/executors/test_executor.py

@@ -74,10 +74,11 @@ def serve_exec(**kwargs):
     t.join()
 
 
-def test_executor_load_from_hub():
-    exec = Executor.from_hub(
-        'jinahub://DummyHubExecutor', uses_metas={'name': 'hello123'}
-    )
+@pytest.mark.parametrize(
+    'uses', ['jinaai://jina-ai/DummyHubExecutor']
+)
+def test_executor_load_from_hub(uses):
+    exec = Executor.from_hub(uses, uses_metas={'name': 'hello123'})
     da = DocumentArray([Document()])
     exec.foo(da)
     assert da.texts == ['hello']
@@ -438,11 +439,15 @@ def test_default_workspace(tmpdir):
     'exec_type',
     [Executor.StandaloneExecutorType.EXTERNAL, Executor.StandaloneExecutorType.SHARED],
 )
-def test_to_k8s_yaml(tmpdir, exec_type):
+@pytest.mark.parametrize(
+    'uses',
+    ['jinahub+docker://DummyHubExecutor', 'jinaai+docker://jina-ai/DummyHubExecutor'],
+)
+def test_to_k8s_yaml(tmpdir, exec_type, uses):
     Executor.to_kubernetes_yaml(
         output_base_path=tmpdir,
         port_expose=2020,
-        uses='jinahub+docker://DummyHubExecutor',
+        uses=uses,
         executor_type=exec_type,
     )
 
@@ -480,12 +485,16 @@ def test_to_k8s_yaml(tmpdir, exec_type):
     'exec_type',
     [Executor.StandaloneExecutorType.EXTERNAL, Executor.StandaloneExecutorType.SHARED],
 )
-def test_to_docker_compose_yaml(tmpdir, exec_type):
+@pytest.mark.parametrize(
+    'uses',
+    ['jinaai+docker://jina-ai/DummyHubExecutor'],
+)
+def test_to_docker_compose_yaml(tmpdir, exec_type, uses):
     compose_file = os.path.join(tmpdir, 'compose.yml')
     Executor.to_docker_compose_yaml(
         output_path=compose_file,
         port_expose=2020,
-        uses='jinahub+docker://DummyHubExecutor',
+        uses=uses,
         executor_type=exec_type,
     )