Merge branch 'main' of https://github.com/microsoft/promptflow into orchestrator

Author: lalala123123

.github/actions/step_sdk_setup/action.yml

@@ -26,7 +26,7 @@ runs:
       pip list
       python ./setup.py bdist_wheel
       $package = Get-ChildItem ./dist | ? { $_.Name.Contains('.whl')}
-      pip install $($package.FullName + "[azure,executable,azureml-serving]")
+      pip install $($package.FullName + "[azure,executable,azureml-serving,executor-service]")
       echo "########### pip freeze (After) ###########"
       pip freeze
     working-directory: ${{ inputs.scriptPath }}

.github/workflows/promptflow-executor-e2e-test.yml

@@ -100,7 +100,7 @@ jobs:
       run: |
         Set-PSDebug -Trace 1
         pip install -r ${{ github.workspace }}/src/promptflow/dev_requirements.txt
-        gci ./promptflow -Recurse | % {if ($_.Name.Contains('.whl')) {python -m pip install "$($_.FullName)"}}
+        gci ./promptflow -Recurse | % {if ($_.Name.Contains('.whl')) {python -m pip install "$($_.FullName)[executor-service]"}}
         gci ./promptflow-tools -Recurse | % {if ($_.Name.Contains('.whl')) {python -m pip install $_.FullName}}
         pip freeze
     - name: Azure Login

.github/workflows/promptflow-executor-unit-test.yml

@@ -105,7 +105,7 @@ jobs:
       run: |
         Set-PSDebug -Trace 1
         pip install -r ${{ github.workspace }}/src/promptflow/dev_requirements.txt
-        gci ./promptflow -Recurse | % {if ($_.Name.Contains('.whl')) {python -m pip install "$($_.FullName)"}}
+        gci ./promptflow -Recurse | % {if ($_.Name.Contains('.whl')) {python -m pip install "$($_.FullName)[executor-service]"}}
         gci ./promptflow-tools -Recurse | % {if ($_.Name.Contains('.whl')) {python -m pip install $_.FullName}}
         pip freeze
     - name: Azure Login

src/promptflow/CHANGELOG.md

@@ -5,6 +5,7 @@
 
 ### Features Added
 
+- [SDK/CLI][azure] Support specify compute instance as session compute in run.yaml
 
 ### Bugs Fixed
 

src/promptflow/dev_requirements.txt

@@ -29,3 +29,6 @@ httpx
 # test dummy flow run in notebook, give a minimal version for vulnerability issue
 ipykernel>=6.27.1
 papermill>=2.5.0
+
+# test executor server in local
+uvicorn>=0.27.0

src/promptflow/promptflow/_cli/_pf/_experiment.py

@@ -13,9 +13,8 @@
 )
 from promptflow._cli._utils import activate_action, exception_handler
 from promptflow._sdk._constants import get_list_view_type
-from promptflow._sdk._load_functions import load_common
 from promptflow._sdk._pf_client import PFClient
-from promptflow._sdk.entities._experiment import Experiment, ExperimentTemplate
+from promptflow._sdk.entities._experiment import Experiment
 from promptflow._utils.logger_utils import get_cli_sdk_logger
 
 logger = get_cli_sdk_logger()
@@ -176,9 +175,11 @@ def dispatch_experiment_commands(args: argparse.Namespace):
 
 @exception_handler("Create experiment")
 def create_experiment(args: argparse.Namespace):
+    from promptflow._sdk._load_functions import _load_experiment_template
+
     template_path = args.template
     logger.debug("Loading experiment template from %s", template_path)
-    template = load_common(ExperimentTemplate, source=template_path)
+    template = _load_experiment_template(source=template_path)
     logger.debug("Creating experiment from template %s", template.dir_name)
     experiment = Experiment.from_template(template, name=args.name)
     logger.debug("Creating experiment %s", experiment.name)

src/promptflow/promptflow/_constants.py

@@ -50,10 +50,8 @@ class AvailableIDE:
 
 _ENV_PF_INSTALLER = "PF_INSTALLER"
 
-# trace related
-TRACE_SESSION_ID_OP_CTX_NAME = "pf_trace_session_id"
-
 
+# trace related
 class SpanFieldName:
     NAME = "name"
     CONTEXT = "context"
@@ -75,7 +73,7 @@ class SpanContextFieldName:
 
 
 class SpanStatusFieldName:
-    CODE = "code"
+    STATUS_CODE = "status_code"
 
 
 class SpanAttributeFieldName:
@@ -89,6 +87,11 @@ class SpanAttributeFieldName:
     FLOW_ID = "flow_id"
     RUN = "run"
     EXPERIMENT = "experiment"
+    LINE_RUN_ID = "line_run_id"
+    REFERENCED_LINE_RUN_ID = "referenced.line_run_id"
+    COMPLETION_TOKEN_COUNT = "__computed__.cumulative_token_count.completion"
+    PROMPT_TOKEN_COUNT = "__computed__.cumulative_token_count.prompt"
+    TOTAL_TOKEN_COUNT = "__computed__.cumulative_token_count.total"
 
 
 class SpanResourceAttributesFieldName:

src/promptflow/promptflow/_core/operation_context.py

@@ -2,7 +2,6 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 import copy
-
 from contextvars import ContextVar
 from typing import Dict, Mapping
 
@@ -30,6 +29,14 @@ def _add_otel_attributes(self, key, value):
         attributes[key] = value
         self[OperationContext._OTEL_ATTRIBUTES] = attributes
 
+    def _remove_otel_attributes(self, keys: list):
+        if isinstance(keys, str):
+            keys = [keys]
+        attributes = self.get(OperationContext._OTEL_ATTRIBUTES, {})
+        for key in keys:
+            attributes.pop(key, None)
+        self[OperationContext._OTEL_ATTRIBUTES] = attributes
+
     def _get_otel_attributes(self):
         return self.get(OperationContext._OTEL_ATTRIBUTES, {})
 

src/promptflow/promptflow/_core/tracer.py

@@ -10,9 +10,10 @@
 from collections.abc import Iterator
 from contextvars import ContextVar
 from datetime import datetime
+from threading import Lock
 from typing import Callable, Dict, List, Optional
 
-from opentelemetry import trace
+import opentelemetry.trace as otel_trace
 from opentelemetry.trace.status import StatusCode
 
 from promptflow._core.generator_proxy import GeneratorProxy, generate_from_proxy
@@ -24,7 +25,8 @@
 
 from .thread_local_singleton import ThreadLocalSingleton
 
-open_telemetry_tracer = trace.get_tracer("promptflow")
+
+open_telemetry_tracer = otel_trace.get_tracer("promptflow")
 
 
 class Tracer(ThreadLocalSingleton):
@@ -153,6 +155,46 @@ def _format_error(error: Exception) -> dict:
         }
 
 
+class TokenCollector():
+    _lock = Lock()
+
+    def __init__(self):
+        self._span_id_to_tokens = {}
+
+    def collect_openai_tokens(self, span, output):
+        span_id = span.get_span_context().span_id
+        if not inspect.isgenerator(output) and hasattr(output, "usage") and output.usage is not None:
+            tokens = {
+                f"__computed__.cumulative_token_count.{k.split('_')[0]}": v for k, v in output.usage.dict().items()
+            }
+            if tokens:
+                with self._lock:
+                    self._span_id_to_tokens[span_id] = tokens
+
+    def collect_openai_tokens_for_parent_span(self, span):
+        tokens = self.try_get_openai_tokens(span.get_span_context().span_id)
+        if tokens:
+            if not hasattr(span, "parent") or span.parent is None:
+                return
+            parent_span_id = span.parent.span_id
+            with self._lock:
+                if parent_span_id in self._span_id_to_tokens:
+                    merged_tokens = {
+                        key: self._span_id_to_tokens[parent_span_id].get(key, 0) + tokens.get(key, 0)
+                        for key in set(self._span_id_to_tokens[parent_span_id]) | set(tokens)
+                    }
+                    self._span_id_to_tokens[parent_span_id] = merged_tokens
+                else:
+                    self._span_id_to_tokens[parent_span_id] = tokens
+
+    def try_get_openai_tokens(self, span_id):
+        with self._lock:
+            return self._span_id_to_tokens.get(span_id, None)
+
+
+token_collector = TokenCollector()
+
+
 def _create_trace_from_function_call(
     f, *, args=None, kwargs=None, args_to_ignore: Optional[List[str]] = None, trace_type=TraceType.FUNCTION
 ):
@@ -205,6 +247,14 @@ def get_node_name_from_context():
     return None
 
 
+def enrich_span_with_context(span):
+    try:
+        attrs_from_context = OperationContext.get_instance()._get_otel_attributes()
+        span.set_attributes(attrs_from_context)
+    except Exception as e:
+        logging.warning(f"Failed to enrich span with context: {e}")
+
+
 def enrich_span_with_trace(span, trace):
     try:
         span.set_attributes(
@@ -215,8 +265,7 @@ def enrich_span_with_trace(span, trace):
                 "node_name": get_node_name_from_context(),
             }
         )
-        attrs_from_context = OperationContext.get_instance()._get_otel_attributes()
-        span.set_attributes(attrs_from_context)
+        enrich_span_with_context(span)
     except Exception as e:
         logging.warning(f"Failed to enrich span with trace: {e}")
 
@@ -235,6 +284,9 @@ def enrich_span_with_output(span, output):
     try:
         serialized_output = serialize_attribute(output)
         span.set_attribute("output", serialized_output)
+        tokens = token_collector.try_get_openai_tokens(span.get_span_context().span_id)
+        if tokens:
+            span.set_attributes(tokens)
     except Exception as e:
         logging.warning(f"Failed to enrich span with output: {e}")
 
@@ -306,12 +358,16 @@ async def wrapped(*args, **kwargs):
                 Tracer.push(trace)
                 enrich_span_with_input(span, trace.inputs)
                 output = await func(*args, **kwargs)
+                if trace_type == TraceType.LLM:
+                    token_collector.collect_openai_tokens(span, output)
                 enrich_span_with_output(span, output)
                 span.set_status(StatusCode.OK)
-                return Tracer.pop(output)
+                output = Tracer.pop(output)
             except Exception as e:
                 Tracer.pop(None, e)
                 raise
+        token_collector.collect_openai_tokens_for_parent_span(span)
+        return output
 
     wrapped.__original_function = func
 
@@ -351,12 +407,16 @@ def wrapped(*args, **kwargs):
                 Tracer.push(trace)
                 enrich_span_with_input(span, trace.inputs)
                 output = func(*args, **kwargs)
+                if trace_type == TraceType.LLM:
+                    token_collector.collect_openai_tokens(span, output)
                 enrich_span_with_output(span, output)
                 span.set_status(StatusCode.OK)
-                return Tracer.pop(output)
+                output = Tracer.pop(output)
             except Exception as e:
                 Tracer.pop(None, e)
                 raise
+        token_collector.collect_openai_tokens_for_parent_span(span)
+        return output
 
     wrapped.__original_function = func
 

src/promptflow/promptflow/_sdk/_constants.py

@@ -70,6 +70,7 @@ def _prepare_home_dir() -> Path:
 SERVICE_CONFIG_FILE = "pf.yaml"
 PF_SERVICE_PORT_FILE = "pfs.port"
 PF_SERVICE_LOG_FILE = "pfs.log"
+PF_TRACE_CONTEXT = "PF_TRACE_CONTEXT"
 
 LOCAL_MGMT_DB_PATH = (HOME_PROMPT_FLOW_DIR / "pf.sqlite").resolve()
 LOCAL_MGMT_DB_SESSION_ACQUIRE_LOCK_PATH = (HOME_PROMPT_FLOW_DIR / "pf.sqlite.lock").resolve()
@@ -135,6 +136,7 @@ def _prepare_home_dir() -> Path:
 TRACE_MGMT_DB_PATH = (HOME_PROMPT_FLOW_DIR / "trace.sqlite").resolve()
 TRACE_MGMT_DB_SESSION_ACQUIRE_LOCK_PATH = (HOME_PROMPT_FLOW_DIR / "trace.sqlite.lock").resolve()
 SPAN_TABLENAME = "span"
+PFS_MODEL_DATETIME_FORMAT = "iso8601"
 
 
 class CustomStrongTypeConnectionConfigs:
@@ -425,3 +427,10 @@ class ExperimentNodeRunStatus(object):
     COMPLETED = "Completed"
     FAILED = "Failed"
     CANCELED = "Canceled"
+
+
+class ExperimentContextKey:
+    EXPERIMENT = "experiment"
+    # Note: referenced id not used for lineage, only for evaluation
+    REFERENCED_LINE_RUN_ID = "referenced.line_run_id"
+    REFERENCED_RUN_ID = "referenced.run_id"

src/promptflow/promptflow/_sdk/_load_functions.py

@@ -11,6 +11,7 @@
 from .._utils.yaml_utils import load_yaml
 from .entities import Run
 from .entities._connection import CustomConnection, _Connection
+from .entities._experiment import ExperimentTemplate
 from .entities._flow import Flow
 
 logger = get_cli_sdk_logger()
@@ -105,6 +106,15 @@ def load_connection(
     source: Union[str, PathLike, IO[AnyStr]],
     **kwargs,
 ):
+    """Load connection from YAML file or .env file.
+
+    :param source: The local yaml source of a connection or .env file. Must be a path to a local file.
+        If the source is a path, it will be open and read.
+        An exception is raised if the file does not exist.
+    :type source: Union[PathLike, str]
+    :return: A Connection object
+    :rtype: Connection
+    """
     if Path(source).name.endswith(".env"):
         return _load_env_to_connection(source, **kwargs)
     return load_common(_Connection, source, **kwargs)
@@ -132,3 +142,24 @@ def _load_env_to_connection(
         return CustomConnection(name=name, secrets=data)
     except Exception as e:
         raise Exception(f"Load entity error: {e}") from e
+
+
+def _load_experiment_template(
+    source: Union[str, PathLike, IO[AnyStr]],
+    **kwargs,
+):
+    """Load experiment template from YAML file.
+
+    :param source: The local yaml source of an experiment template. Must be a path to a local file.
+        If the source is a path, it will be open and read.
+        An exception is raised if the file does not exist.
+    :type source: Union[PathLike, str]
+    :return: An ExperimentTemplate object
+    :rtype: ExperimentTemplate
+    """
+    source_path = Path(source)
+    if source_path.is_dir():
+        source = source_path / "flow.exp.yaml"
+    if not source_path.exists():
+        raise FileNotFoundError(f"Experiment template file {source.resolve().absolute().as_posix()} not found.")
+    return load_common(ExperimentTemplate, source=source)

src/promptflow/promptflow/_sdk/_orm/trace.py

@@ -2,9 +2,10 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
+import copy
 import typing
 
-from sqlalchemy import TEXT, Column, Index
+from sqlalchemy import TEXT, Column, Index, text
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import declarative_base
 
@@ -65,3 +66,40 @@ def list(
             if parent_span_id is not None:
                 stmt = stmt.filter(Span.parent_span_id == parent_span_id)
             return [span for span in stmt.all()]
+
+
+class LineRun:
+    """Line run is an abstraction of spans, which is not persisted in the database."""
+
+    @staticmethod
+    def list(
+        session_id: typing.Optional[str] = None,
+    ) -> typing.List[typing.List[Span]]:
+        with trace_mgmt_db_session() as session:
+            stmt = session.query(Span)
+            if session_id is not None:
+                stmt = stmt.filter(Span.session_id == session_id)
+                # other filters, e.g., experiment, run, path, etc.
+                stmt = stmt.filter(
+                    text("json_extract(json_extract(span.content, '$.attributes'), '$.framework') = 'promptflow'")
+                )
+            else:
+                # TODO: fully support query
+                raise NotImplementedError
+            stmt = stmt.order_by(Span.trace_id)
+            line_runs = []
+            current_spans: typing.List[Span] = []
+            span: Span
+            for span in stmt.all():
+                if len(current_spans) == 0:
+                    current_spans.append(span)
+                    continue
+                current_trace_id = current_spans[0].trace_id
+                if span.trace_id == current_trace_id:
+                    current_spans.append(span)
+                    continue
+                line_runs.append(copy.deepcopy(current_spans))
+                current_spans = [span]
+            if len(current_spans) > 0:
+                line_runs.append(copy.deepcopy(current_spans))
+            return line_runs