feat: support jinaai executor scheme

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',
+    ['jinahub+sandbox://DummyHubExecutor', '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', ['jinahub+sandbox://Hello', '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,14 @@ 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,
+        'jinahub+docker://HubBeforeExecutor',
+        '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])
@@ -396,10 +403,29 @@ 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',
+    [
+        'jinahub+docker://HubExecutor',
+        'docker://docker_image:latest',
+        'jinaai+docker://jina-ai/HubExecutor',
+    ],
+)
+@pytest.mark.parametrize(
+    'uses_before',
+    [
+        None,
+        'jinahub+docker://HubBeforeExecutor',
+        'jinaai+docker://jina-ai/HubBeforeExecutor',
+    ],
+)
+@pytest.mark.parametrize(
+    'uses_after',
+    [
+        None,
+        'jinahub+docker://HubAfterExecutor',
+        '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'])

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', ['jinahub://DummyExecutor', 'jinaai://jina-ai/DummyExecutor']
+)
+def test_get_image_name(mocker, monkeypatch, uses):
     mock = mocker.Mock()
 
     def _mock_fetch(
@@ -71,8 +74,6 @@ def _mock_fetch(
 
     monkeypatch.setattr(HubIO, 'fetch_meta', _mock_fetch)
 
-    uses = 'jinahub://DummyExecutor'
-
     image_name = get_image_name(uses)
 
     assert image_name == 'jinahub/DummyExecutor'

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

@@ -28,7 +28,14 @@ def namespace_equal(
 
 
 @pytest.mark.parametrize('shards', [1, 5])
-@pytest.mark.parametrize('uses_before', [None, 'jinahub+docker://HubBeforeExecutor'])
+@pytest.mark.parametrize(
+    'uses_before',
+    [
+        None,
+        'jinahub+docker://HubBeforeExecutor',
+        '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])
@@ -373,10 +380,29 @@ 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',
+    [
+        'jinahub+docker://HubExecutor',
+        'docker://docker_image:latest',
+        'jinaai+docker://jina-ai/HubExecutor',
+    ],
+)
+@pytest.mark.parametrize(
+    'uses_before',
+    [
+        None,
+        'jinahub+docker://HubBeforeExecutor',
+        'jinaai+docker://jina-ai/HubBeforeExecutor',
+    ],
+)
+@pytest.mark.parametrize(
+    'uses_after',
+    [
+        None,
+        'jinahub+docker://HubAfterExecutor',
+        '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'])

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',
+    ['jinahub+sandbox://DummyHubExecutor', '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,16 @@ 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',
+    [
+        'jinahub+sandbox://DummyHubExecutor',
+        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', ['jinahub+docker://DummyExecutor', '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', ['jinahub://DummyHubExecutor', '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',
+    ['jinahub+docker://DummyHubExecutor', '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,
     )
 

tests/unit/test_cli.py

@@ -39,11 +39,14 @@ def test_cli_help():
     subprocess.check_call(['jina', 'help', 'deployment'])
 
 
-def test_cli_hub():
+@pytest.mark.parametrize(
+    'uses', ['jinahub://DummyHubExecutor', 'jinaai://jina-ai/DummyHubExecutor']
+)
+def test_cli_hub(uses):
     subprocess.check_call(['jina', 'hub', '--help'])
     for cmd in ['new', 'status', 'pull', 'push']:
         subprocess.check_call(['jina', 'hub', cmd, '--help'])
-    subprocess.check_call(['jina', 'hub', 'pull', 'jinahub://DummyHubExecutor'])
+    subprocess.check_call(['jina', 'hub', 'pull', uses])
 
 
 def test_cli_warn_unknown_args():