Client: execution graphs tracing using opentelemetry

client/quantum_serverless/core/constrants.py

@@ -9,3 +9,15 @@
 QS_EVENTS_REDIS_HOST: str = "QS_EVENTS_REDIS_HOST"
 QS_EVENTS_REDIS_PORT: str = "QS_EVENTS_REDIS_PORT"
 QS_EVENTS_REDIS_PASSWORD: str = "QS_EVENTS_REDIS_PASSWORD"
+
+
+# telemetry
+OT_PROGRAM_NAME = "OT_PROGRAM_NAME"
+OT_PROGRAM_NAME_DEFAULT = "unnamed_execution"
+OT_JAEGER_HOST = "OT_JAEGER_HOST"
+OT_JAEGER_HOST_KEY = "OT_JAEGER_HOST_KEY"
+OT_JAEGER_PORT_KEY = "OT_JAEGER_PORT_KEY"
+OT_TRACEPARENT_ID_KEY = "OT_TRACEPARENT_ID_KEY"
+OT_SPAN_DEFAULT_NAME = "entrypoint"
+OT_ATTRIBUTE_PREFIX = "qs"
+OT_LABEL_CALL_LOCATION = "qs.location"

client/quantum_serverless/core/decorators.py

@@ -28,23 +28,24 @@
     run_qiskit_remote
     get_refs_by_status
 """
+import functools
 import os
 from dataclasses import dataclass
-from typing import Optional, Dict, Any, Union, List
+from typing import Optional, Dict, Any, Union, List, Callable, Sequence
 
 import ray
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from qiskit import QuantumCircuit
+from ray.runtime_env import RuntimeEnv
 
 from quantum_serverless.core.constrants import (
-    META_TOPIC,
-    QS_EXECUTION_WORKLOAD_ID,
-    QS_EXECUTION_UID,
-)
-from quantum_serverless.core.events import (
-    RedisEventHandler,
-    ExecutionMessage,
-    EventHandler,
+    OT_ATTRIBUTE_PREFIX,
+    OT_JAEGER_HOST_KEY,
+    OT_JAEGER_PORT_KEY,
+    OT_TRACEPARENT_ID_KEY,
 )
 from quantum_serverless.core.state import StateHandler
+from quantum_serverless.core.tracing import get_tracer, _trace_env_vars
 from quantum_serverless.utils import JsonSerializable
 
 remote = ray.remote
@@ -76,10 +77,136 @@ def from_dict(cls, dictionary: dict):
         return Target(**dictionary)
 
 
+@dataclass
+class CircuitMeta:
+    """Circuit metainformation."""
+
+    n_qubits: int
+    depth: int
+
+    def to_seq(self) -> Sequence[int]:
+        """Converts meta to seq."""
+        return [self.n_qubits, self.depth]
+
+
+Numeric = Union[float, int]
+
+
+def fetch_execution_meta(*args, **kwargs) -> Dict[str, Sequence[Numeric]]:
+    """Extracts meta information from function arguments.
+
+    Meta information consist of metrics that describe circuits.
+    Metrics described in `CircuitMeta` class.
+
+    Args:
+        *args: arguments
+        **kwargs: named arguments
+
+    Returns:
+        return meta information dictionary
+    """
+
+    def fetch_meta(circuit: QuantumCircuit) -> CircuitMeta:
+        """Returns meta information from circuit."""
+        return CircuitMeta(n_qubits=circuit.num_qubits, depth=circuit.depth())
+
+    meta: Dict[str, Sequence[Numeric]] = {}
+
+    for idx, argument in enumerate(args):
+        if isinstance(argument, QuantumCircuit):
+            meta[f"{OT_ATTRIBUTE_PREFIX}.args.arg_{idx}"] = fetch_meta(
+                argument
+            ).to_seq()
+        elif isinstance(argument, list):
+            for sub_idx, sub_argument in enumerate(argument):
+                if isinstance(sub_argument, QuantumCircuit):
+                    meta[
+                        f"{OT_ATTRIBUTE_PREFIX}.args.arg_{idx}.{sub_idx}"
+                    ] = fetch_meta(sub_argument).to_seq()
+    for key, value in kwargs.items():
+        if isinstance(value, QuantumCircuit):
+            meta[f"{OT_ATTRIBUTE_PREFIX}.kwargs.{key}"] = fetch_meta(value).to_seq()
+        elif isinstance(value, list):
+            for sub_idx, sub_argument in enumerate(value):
+                if isinstance(sub_argument, QuantumCircuit):
+                    meta[f"{OT_ATTRIBUTE_PREFIX}.kwargs.{key}.{sub_idx}"] = fetch_meta(
+                        sub_argument
+                    ).to_seq()
+    return meta
+
+
+def _tracible_function(
+    name: str, target: Target, trace_id: Optional[str] = None
+) -> Callable:
+    """Wrap a function in an OTel span.
+
+    Args:
+        name: name of function
+        target: target for function
+        trace_id: trace id
+
+    Returns:
+        traced function
+    """
+
+    def decorator(func: Callable):
+        @functools.wraps(func)
+        def wraps(*args, **kwargs):
+            tracer = get_tracer(
+                func.__module__,
+                agent_host=os.environ.get(OT_JAEGER_HOST_KEY, None),
+                agent_port=int(os.environ.get(OT_JAEGER_PORT_KEY, 6831)),
+            )
+            ctx = TraceContextTextMapPropagator().extract(
+                {
+                    TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: trace_id  # pylint:disable=protected-access
+                }
+            )
+
+            circuits_meta = fetch_execution_meta(*args, **kwargs)
+
+            with tracer.start_as_current_span(name, context=ctx) as rollspan:
+                # TODO: add serverless package version # pylint: disable=fixme
+                rollspan.set_attribute(
+                    f"{OT_ATTRIBUTE_PREFIX}.meta.function_name", name
+                )
+                rollspan.set_attribute(
+                    f"{OT_ATTRIBUTE_PREFIX}.meta.stack_layer", "quantum_serverless"
+                )
+
+                rollspan.set_attribute(
+                    f"{OT_ATTRIBUTE_PREFIX}.resources.cpu", target.cpu
+                )
+                rollspan.set_attribute(
+                    f"{OT_ATTRIBUTE_PREFIX}.resources.memory", target.mem
+                )
+                rollspan.set_attribute(
+                    f"{OT_ATTRIBUTE_PREFIX}.resources.gpu", target.gpu
+                )
+
+                resources = target.resources or {}
+                for resource_name, resource_value in resources.items():
+                    rollspan.set_attribute(
+                        f"{OT_ATTRIBUTE_PREFIX}.resources.{resource_name}",
+                        resource_value,
+                    )
+
+                if target.pip is not None:
+                    rollspan.set_attribute("requirements", target.pip)
+
+                for meta_key, meta_value in circuits_meta.items():
+                    rollspan.set_attribute(meta_key, meta_value)
+
+                return func(*args, **kwargs)
+
+        return wraps
+
+    return decorator
+
+
 def run_qiskit_remote(
     target: Optional[Union[Dict[str, Any], Target]] = None,
     state: Optional[StateHandler] = None,
-    events: Optional[EventHandler] = None,
 ):
     """Wraps local function as remote executable function.
     New function will return reference object when called.
@@ -110,40 +237,33 @@ def run_qiskit_remote(
     else:
         remote_target = Target.from_dict(target)
 
-    runtime_env: Dict[str, Any] = {"env_vars": remote_target.env_vars}
-    if remote_target.pip is not None:
-        runtime_env["pip"] = remote_target.pip
-
     def decorator(function):
         def wrapper(*args, **kwargs):
             # inject state as an argument when passed in decorator
             if state is not None:
                 args = tuple([state] + list(args))
 
-            # inject execution meta emitter
-            if events is not None:
-                emitter = events
-            else:
-                emitter = RedisEventHandler.from_env_vars()
-            if emitter is not None:
-                emitter.publish(
-                    META_TOPIC,
-                    message=ExecutionMessage(
-                        workload_id=os.environ.get(QS_EXECUTION_WORKLOAD_ID),
-                        uid=os.environ.get(QS_EXECUTION_UID),
-                        layer="qs",
-                        function_meta={"name": function.__name__},
-                        resources=remote_target.to_dict(),
-                    ).to_dict(),
-                )
+            # tracing
+            traced_env_vars = _trace_env_vars(
+                remote_target.env_vars or {}, location="on decoration"
+            )
+            traced_function = _tracible_function(
+                name=function.__name__,
+                target=remote_target,
+                trace_id=traced_env_vars.get(OT_TRACEPARENT_ID_KEY),
+            )(function)
+
+            # runtime env
+            runtime_env = RuntimeEnv(pip=remote_target.pip, env_vars=traced_env_vars)
 
+            # remote function
             result = ray.remote(
                 num_cpus=remote_target.cpu,
                 num_gpus=remote_target.gpu,
                 resources=remote_target.resources,
                 memory=remote_target.mem,
                 runtime_env=runtime_env,
-            )(function).remote(*args, **kwargs)
+            )(traced_function).remote(*args, **kwargs)
 
             return result
 

client/quantum_serverless/core/provider.py

@@ -34,6 +34,9 @@
 import ray
 from ray.dashboard.modules.job.sdk import JobSubmissionClient
 
+from quantum_serverless.core.constrants import OT_PROGRAM_NAME
+
+from quantum_serverless.core.tracing import _trace_env_vars
 from quantum_serverless.core.job import Job
 from quantum_serverless.core.program import Program
 from quantum_serverless.exception import QuantumServerlessException
@@ -79,7 +82,9 @@ def job_client(self) -> Optional[JobSubmissionClient]:
         return None
 
     def context(self, **kwargs):
-        """Return context allocated for this compute_resource."""
+        """Returns context allocated for this compute_resource."""
+        _trace_env_vars({}, location="on context allocation")
+
         init_args = {
             **kwargs,
             **{
@@ -238,13 +243,16 @@ def run_program(self, program: Program) -> Job:
             arguments = " ".join(arg_list)
         entrypoint = f"python {program.entrypoint} {arguments}"
 
+        # set program name so OT can use it as parent span name
+        env_vars = {**(program.env_vars or {}), **{OT_PROGRAM_NAME: program.name}}
+
         job_id = job_client.submit_job(
             entrypoint=entrypoint,
             submission_id=f"qs_{uuid4()}",
             runtime_env={
                 "working_dir": program.working_dir,
                 "pip": program.dependencies,
-                "env_vars": program.env_vars,
+                "env_vars": env_vars,
             },
         )
         return Job(job_id=job_id, job_client=job_client)