Repository: jina-ai/serve
Pull request: fix: use start_as_current_span to maintain tracing operation chain
**Goals:**
<!---https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword--->

- resolves #5386 
- the python SDK for whatever reason doesn't consider the span started from the origin request context as the direct child span. This causes inconsistency when the trace is viewed.


Before:
![before-start-as-current-span](https://user-images.githubusercontent.com/9018501/201644800-d0b7ccb3-6971-4436-ac4c-c4ec2e4a8249.png)


After:

![after-start-as-span](https://user-images.githubusercontent.com/9018501/201644837-d5b30a66-1f9e-4c61-923c-8154f6cb0c79.png)



Problem statement:
bug: worker runtime request handler span appears out of order
**Describe the bug**
<!-- A clear and concise description of what the bug is. -->
The worker runtime request handler starts the span to record the executor request using `trace.start_as_span` which doesn't sequence the span correctly if the executor requests method creates another span using the `tracing_context`.

**Describe how you solve it**
<!-- copy past your code/pull request link -->
To correctly sequence the span order, the worker runtime request handler must start it's span using `trace.start_as_current_span`
---

<!-- 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:

Commit message: style: fix overload and cli autocomplete

Patch to review:
diff --git a/jina/serve/executors/__init__.py b/jina/serve/executors/__init__.py
index b0019612d2f1a..d3eaf50257789 100644
--- a/jina/serve/executors/__init__.py
+++ b/jina/serve/executors/__init__.py
@@ -1,9 +1,9 @@
 import contextlib
+import copy
 import inspect
 import multiprocessing
 import os
 import threading
-import copy
 import warnings
 from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
@@ -14,7 +14,10 @@
 from jina.importer import ImportExtensions
 from jina.jaml import JAML, JAMLCompatible, env_var_regex, internal_var_regex
 from jina.logging.logger import JinaLogger
-from jina.serve.executors.decorators import avoid_concurrent_lock_cls, _init_requests_by_class
+from jina.serve.executors.decorators import (
+    _init_requests_by_class,
+    avoid_concurrent_lock_cls,
+)
 from jina.serve.executors.metas import get_executor_taboo
 from jina.serve.helper import store_init_kwargs, wrap_func
 from jina.serve.instrumentation import MetricsTimer
@@ -57,7 +60,7 @@ def register_class(cls):
             arg_spec = inspect.getfullargspec(cls.__init__)
 
             if not arg_spec.varkw and not __args_executor_init__.issubset(
-                    arg_spec.args
+                arg_spec.args
             ):
                 raise TypeError(
                     f'{cls.__init__} does not follow the full signature of `Executor.__init__`, '
@@ -117,12 +120,12 @@ def __init__(awesomeness=5):
     """
 
     def __init__(
-            self,
-            metas: Optional[Dict] = None,
-            requests: Optional[Dict] = None,
-            runtime_args: Optional[Dict] = None,
-            workspace: Optional[str] = None,
-            **kwargs,
+        self,
+        metas: Optional[Dict] = None,
+        requests: Optional[Dict] = None,
+        runtime_args: Optional[Dict] = None,
+        workspace: Optional[str] = None,
+        **kwargs,
     ):
         """`metas` and `requests` are always auto-filled with values from YAML config.
 
@@ -160,12 +163,12 @@ def _add_runtime_args(self, _runtime_args: Optional[Dict]):
 
     def _init_monitoring(self):
         if (
-                hasattr(self.runtime_args, 'metrics_registry')
-                and self.runtime_args.metrics_registry
+            hasattr(self.runtime_args, 'metrics_registry')
+            and self.runtime_args.metrics_registry
         ):
             with ImportExtensions(
-                    required=True,
-                    help_text='You need to install the `prometheus_client` to use the montitoring functionality of jina',
+                required=True,
+                help_text='You need to install the `prometheus_client` to use the montitoring functionality of jina',
             ):
                 from prometheus_client import Summary
 
@@ -296,7 +299,7 @@ def _add_metas(self, _metas: Optional[Dict]):
                 if not hasattr(target, k):
                     if isinstance(v, str):
                         if not (
-                                env_var_regex.findall(v) or internal_var_regex.findall(v)
+                            env_var_regex.findall(v) or internal_var_regex.findall(v)
                         ):
                             setattr(target, k, v)
                         else:
@@ -348,10 +351,10 @@ async def __acall__(self, req_endpoint: str, **kwargs):
             return await self.__acall_endpoint__(__default_endpoint__, **kwargs)
 
     async def __acall_endpoint__(
-            self, req_endpoint, tracing_context: Optional['Context'], **kwargs
+        self, req_endpoint, tracing_context: Optional['Context'], **kwargs
     ):
         async def exec_func(
-                summary, histogram, histogram_metric_labels, tracing_context
+            summary, histogram, histogram_metric_labels, tracing_context
         ):
             with MetricsTimer(summary, histogram, histogram_metric_labels):
                 if iscoroutinefunction(func):
@@ -379,7 +382,9 @@ async def exec_func(
         }
 
         if self.tracer:
-            with self.tracer.start_span(req_endpoint, context=tracing_context) as _:
+            with self.tracer.start_as_current_span(
+                req_endpoint, context=tracing_context
+            ) as _:
                 from opentelemetry.propagate import extract
                 from opentelemetry.trace.propagation.tracecontext import (
                     TraceContextTextMapPropagator,
@@ -409,10 +414,10 @@ def workspace(self) -> Optional[str]:
         :return: returns the workspace of the current shard of this Executor.
         """
         workspace = (
-                getattr(self.runtime_args, 'workspace', None)
-                or getattr(self.metas, 'workspace')
-                or self._init_workspace
-                or __cache_path__
+            getattr(self.runtime_args, 'workspace', None)
+            or getattr(self.metas, 'workspace')
+            or self._init_workspace
+            or __cache_path__
         )
         if workspace:
             complete_workspace = os.path.join(workspace, self.metas.name)
@@ -435,13 +440,13 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     @classmethod
     def from_hub(
-            cls: Type[T],
-            uri: str,
-            context: Optional[Dict[str, Any]] = None,
-            uses_with: Optional[Dict] = None,
-            uses_metas: Optional[Dict] = None,
-            uses_requests: Optional[Dict] = None,
-            **kwargs,
+        cls: Type[T],
+        uri: str,
+        context: Optional[Dict[str, Any]] = None,
+        uses_with: Optional[Dict] = None,
+        uses_metas: Optional[Dict] = None,
+        uses_requests: Optional[Dict] = None,
+        **kwargs,
     ) -> T:
         """Construct an Executor from Hub.
 
@@ -493,12 +498,12 @@ def from_hub(
 
     @classmethod
     def serve(
-            cls,
-            uses_with: Optional[Dict] = None,
-            uses_metas: Optional[Dict] = None,
-            uses_requests: Optional[Dict] = None,
-            stop_event: Optional[Union[threading.Event, multiprocessing.Event]] = None,
-            **kwargs,
+        cls,
+        uses_with: Optional[Dict] = None,
+        uses_metas: Optional[Dict] = None,
+        uses_requests: Optional[Dict] = None,
+        stop_event: Optional[Union[threading.Event, multiprocessing.Event]] = None,
+        **kwargs,
     ):
         """Serve this Executor in a temporary Flow. Useful in testing an Executor in remote settings.
 
@@ -531,16 +536,16 @@ class StandaloneExecutorType(BetterEnum):
 
     @staticmethod
     def to_kubernetes_yaml(
-            uses: str,
-            output_base_path: str,
-            k8s_namespace: Optional[str] = None,
-            executor_type: Optional[
-                StandaloneExecutorType
-            ] = StandaloneExecutorType.EXTERNAL,
-            uses_with: Optional[Dict] = None,
-            uses_metas: Optional[Dict] = None,
-            uses_requests: Optional[Dict] = None,
-            **kwargs,
+        uses: str,
+        output_base_path: str,
+        k8s_namespace: Optional[str] = None,
+        executor_type: Optional[
+            StandaloneExecutorType
+        ] = StandaloneExecutorType.EXTERNAL,
+        uses_with: Optional[Dict] = None,
+        uses_metas: Optional[Dict] = None,
+        uses_requests: Optional[Dict] = None,
+        **kwargs,
     ):
         """
         Converts the Executor into a set of yaml deployments to deploy in Kubernetes.
@@ -568,23 +573,23 @@ def to_kubernetes_yaml(
             output_base_path=output_base_path,
             k8s_namespace=k8s_namespace,
             include_gateway=executor_type
-                            == BaseExecutor.StandaloneExecutorType.EXTERNAL,
+            == BaseExecutor.StandaloneExecutorType.EXTERNAL,
         )
 
     to_k8s_yaml = to_kubernetes_yaml
 
     @staticmethod
     def to_docker_compose_yaml(
-            uses: str,
-            output_path: Optional[str] = None,
-            network_name: Optional[str] = None,
-            executor_type: Optional[
-                StandaloneExecutorType
-            ] = StandaloneExecutorType.EXTERNAL,
-            uses_with: Optional[Dict] = None,
-            uses_metas: Optional[Dict] = None,
-            uses_requests: Optional[Dict] = None,
-            **kwargs,
+        uses: str,
+        output_path: Optional[str] = None,
+        network_name: Optional[str] = None,
+        executor_type: Optional[
+            StandaloneExecutorType
+        ] = StandaloneExecutorType.EXTERNAL,
+        uses_with: Optional[Dict] = None,
+        uses_metas: Optional[Dict] = None,
+        uses_requests: Optional[Dict] = None,
+        **kwargs,
     ):
         """
         Converts the Executor into a yaml file to run with `docker-compose up`
@@ -609,11 +614,11 @@ def to_docker_compose_yaml(
             output_path=output_path,
             network_name=network_name,
             include_gateway=executor_type
-                            == BaseExecutor.StandaloneExecutorType.EXTERNAL,
+            == BaseExecutor.StandaloneExecutorType.EXTERNAL,
         )
 
     def monitor(
-            self, name: Optional[str] = None, documentation: Optional[str] = None
+        self, name: Optional[str] = None, documentation: Optional[str] = None
     ) -> Optional[MetricsTimer]:
         """
         Get a given prometheus metric, if it does not exist yet, it will create it and store it in a buffer.
diff --git a/jina_cli/autocomplete.py b/jina_cli/autocomplete.py
index 3cfc083f768d9..cc1e7140445f6 100644
--- a/jina_cli/autocomplete.py
+++ b/jina_cli/autocomplete.py
@@ -210,10 +210,9 @@
         'hub': ['--help', 'new', 'push', 'pull', 'status', 'list'],
         'cloud login': ['--help'],
         'cloud logout': ['--help'],
-        'cloud deploy': ['--help', '--name', '--workspace', '--env-file'],
-        'cloud list': ['--help', '--status'],
-        'cloud logs': ['--help', '--executor'],
-        'cloud status': ['--help'],
+        'cloud deploy': ['--help'],
+        'cloud list': ['--help', '--phase', '--name'],
+        'cloud status': ['--help', '--verbose'],
         'cloud remove': ['--help'],
         'cloud new': ['--help'],
         'cloud survey': ['--help'],
@@ -225,7 +224,6 @@
             'logout',
             'deploy',
             'list',
-            'logs',
             'status',
             'remove',
             'new',
diff --git a/tests/integration/instrumentation/flow.yml b/tests/integration/instrumentation/flow.yml
new file mode 100644
index 0000000000000..fa837fe55ece5
--- /dev/null
+++ b/tests/integration/instrumentation/flow.yml
@@ -0,0 +1,8 @@
+jtype: Flow
+version: '1'
+with:
+  tracing: true
+  traces_exporter_host: localhost
+  traces_exporter_port: 4317
+executors:
+  - uses: ExecutorTestWithTracing
diff --git a/tests/integration/instrumentation/test_flow_instrumentation.py b/tests/integration/instrumentation/test_flow_instrumentation.py
index f24616e9f7936..4f8b2dc1f7ceb 100644
--- a/tests/integration/instrumentation/test_flow_instrumentation.py
+++ b/tests/integration/instrumentation/test_flow_instrumentation.py
@@ -192,10 +192,10 @@ def test_flow_metrics(
     )
     assert len(receiving_request_seconds_metrics) > 0
     assert receiving_request_seconds_exported_jobs == {
-            'gateway/rep-0',
-            'executor0/head',
-            'executor0/shard-0/rep-0',
-            'executor0/shard-1/rep-0',
+        'gateway/rep-0',
+        'executor0/head',
+        'executor0/shard-0/rep-0',
+        'executor0/shard-1/rep-0',
     }
 
     (
@@ -224,10 +224,10 @@ def test_flow_metrics(
     )
     assert len(sent_response_bytes_metrics) > 0
     assert sent_response_bytes_exported_jobs == {
-            'gateway/rep-0',
-            'executor0/shard-0/rep-0',
-            'executor0/shard-1/rep-0',
-            'executor0/head'
+        'gateway/rep-0',
+        'executor0/shard-0/rep-0',
+        'executor0/shard-1/rep-0',
+        'executor0/head',
     }
 
     (
@@ -237,7 +237,10 @@ def test_flow_metrics(
         prometheus_client, 'jina_number_of_pending_requests'
     )
     assert len(number_of_pending_requests_metrics) > 0
-    assert number_of_pending_requests_exported_jobs == {'gateway/rep-0', 'executor0/head'}
+    assert number_of_pending_requests_exported_jobs == {
+        'gateway/rep-0',
+        'executor0/head',
+    }
 
     (
         failed_requests_metrics,
@@ -245,10 +248,10 @@ def test_flow_metrics(
     ) = get_metrics_and_exported_jobs_by_name(prometheus_client, 'jina_failed_requests')
     assert len(failed_requests_metrics) > 0
     assert failed_requests_exported_jobs == {
-            'gateway/rep-0',
-            'executor0/head',
-            'executor0/shard-0/rep-0',
-            'executor0/shard-1/rep-0',
+        'gateway/rep-0',
+        'executor0/head',
+        'executor0/shard-0/rep-0',
+        'executor0/shard-1/rep-0',
     }
 
     (
@@ -258,7 +261,12 @@ def test_flow_metrics(
         prometheus_client, 'jina_successful_requests'
     )
     assert len(successful_requests_metrics) > 0
-    assert successful_requests_exported_jobs == {'gateway/rep-0', 'executor0/shard-0/rep-0', 'executor0/shard-1/rep-0', 'executor0/head'}
+    assert successful_requests_exported_jobs == {
+        'gateway/rep-0',
+        'executor0/shard-0/rep-0',
+        'executor0/shard-1/rep-0',
+        'executor0/head',
+    }
 
     (
         received_request_bytes_metrics,
@@ -267,7 +275,12 @@ def test_flow_metrics(
         prometheus_client, 'jina_received_request_bytes'
     )
     assert len(received_request_bytes_metrics) > 0
-    assert received_request_bytes_exported_jobs == {'gateway/rep-0', 'executor0/shard-0/rep-0', 'executor0/shard-1/rep-0', 'executor0/head'}
+    assert received_request_bytes_exported_jobs == {
+        'gateway/rep-0',
+        'executor0/shard-0/rep-0',
+        'executor0/shard-1/rep-0',
+        'executor0/head',
+    }
 
     (
         process_requests_seconds_metrics,
@@ -276,7 +289,10 @@ def test_flow_metrics(
         prometheus_client, 'jina_process_request_seconds'
     )
     assert len(process_requests_seconds_metrics) > 0
-    assert process_requests_seconds_exported_jobs == {'executor0/shard-0/rep-0', 'executor0/shard-1/rep-0'}
+    assert process_requests_seconds_exported_jobs == {
+        'executor0/shard-0/rep-0',
+        'executor0/shard-1/rep-0',
+    }
 
     # filter by attributes/labels
     (
@@ -288,7 +304,10 @@ def test_flow_metrics(
         {'executor_endpoint': '/search'},
     )
     assert len(process_requests_seconds_search_endpoint) > 0
-    assert process_requests_seconds_search_endpoint_exported_jobs == {'executor0/shard-0/rep-0', 'executor0/shard-1/rep-0'}
+    assert process_requests_seconds_search_endpoint_exported_jobs == {
+        'executor0/shard-0/rep-0',
+        'executor0/shard-1/rep-0',
+    }
 
     (
         process_requests_seconds_executor,
@@ -299,7 +318,10 @@ def test_flow_metrics(
         {'executor': 'ExecutorFailureWithTracing'},
     )
     assert len(process_requests_seconds_executor) > 0
-    assert process_requests_seconds_executor_exported_jobs == {'executor0/shard-0/rep-0', 'executor0/shard-1/rep-0'}
+    assert process_requests_seconds_executor_exported_jobs == {
+        'executor0/shard-0/rep-0',
+        'executor0/shard-1/rep-0',
+    }
 
     (
         process_requests_seconds_runtime,
diff --git a/tests/integration/instrumentation/test_worker_instrumentation.py b/tests/integration/instrumentation/test_worker_instrumentation.py
new file mode 100644
index 0000000000000..33680673936dc
--- /dev/null
+++ b/tests/integration/instrumentation/test_worker_instrumentation.py
@@ -0,0 +1,44 @@
+import time
+
+from jina import Flow
+from tests.integration.instrumentation import ExecutorTestWithTracing, get_traces
+
+
+def test_span_order(jaeger_port, otlp_collector, otlp_receiver_port):
+    # cannot use the pythonic f = Flow(...) because an actual gateway process is not started
+    f = Flow.load_config('flow.yml')
+
+    with f:
+        from jina import DocumentArray
+
+        f.post(f'/search', DocumentArray.empty(), continue_on_error=True)
+        # give some time for the tracing and metrics exporters to finish exporting.
+        # the client is slow to export the data
+        time.sleep(8)
+
+    traces = get_traces(jaeger_port, 'executor0/rep-0')
+    with open('tmp.json', 'w') as f:
+        import json
+
+        json.dump(traces, f)
+    process_single_data_span_ids = set()
+    search_request_parent_span_ids = set()
+    for trace in traces:
+        for span in trace['spans']:
+            if (
+                span['operationName']
+                == '/jina.JinaSingleDataRequestRPC/process_single_data'
+            ):
+                process_single_data_span_ids.add(span['spanID'])
+
+            if span['operationName'] == '/search':
+                references = span.get('references', [])
+                for ref in references:
+                    search_request_parent_span_ids.add(ref.get('spanID', ''))
+
+    assert any(
+        [
+            search_span in process_single_data_span_ids
+            for search_span in search_request_parent_span_ids
+        ]
+    )


Code review comments:
diff_hunk: @@ -379,7 +382,9 @@ async def exec_func(
         }
 
         if self.tracer:
-            with self.tracer.start_span(req_endpoint, context=tracing_context) as _:
+            with self.tracer.start_as_current_span(
+                req_endpoint, context=tracing_context
+            ) as _:
path: jina/serve/executors/__init__.py
text: @user1:
```suggestion
            ):
```

diff_hunk: @@ -0,0 +1,44 @@
+import time
+
+from jina import Flow
+from tests.integration.instrumentation import ExecutorTestWithTracing, get_traces
+
+
+def test_span_order(jaeger_port, otlp_collector, otlp_receiver_port):
+    # cannot use the pythonic f = Flow(...) because an actual gateway process is not started
path: tests/integration/instrumentation/test_worker_instrumentation.py
text: @user1:
I do not understand this

@author:
This might be a bug or something that is part of the orchestrate layer but if I start as flow using:
```python
f = Flow(tracing=True, traces_exporter_host='localhost', traces_exporter_port=4317)

with f:
  f.post(...)
```

This doesn't produce the proper chain (v.s. flow.yml) of requests to the executor. 

@author:
I tested this again on my local machine and I cannot reproduce this issue. Maybe something was broken in the test that caused me to use teh `flow.yml`.

@user1:
can u get back to the Pythonic way then?

@author:
Yes, the code has already been replaced with the pythonic code.