Repository: microsoft/promptflow
File: src/promptflow-devkit/promptflow/_sdk/_tracing.py (Python)

Code before:
    run: typing.Optional[str] = None,
) -> None:
    # honor and set attributes if user has specified
    if isinstance(attrs, dict):
        _inject_attrs_to_op_ctx(attrs)

    # experiment related attributes, pass from environment
    env_tracing_ctx = os.environ.get(PF_TRACE_CONTEXT, None)
    logger.debug("Read tracing context from environment: %s", env_tracing_ctx)
    env_attrs = dict(json.loads(env_tracing_ctx)).get(PF_TRACE_CONTEXT_ATTR) if env_tracing_ctx else dict()
    exp = env_attrs.get(ContextAttributeKey.EXPERIMENT, None)
    ref_line_run_id = env_attrs.get(ContextAttributeKey.REFERENCED_LINE_RUN_ID, None)
    op_ctx = OperationContext.get_instance()
    # remove `referenced.line_run_id` from context to avoid stale value set by previous node
    if ref_line_run_id is None:
        op_ctx._remove_otel_attributes(SpanAttributeFieldName.REFERENCED_LINE_RUN_ID)
    else:
        op_ctx._add_otel_attributes(SpanAttributeFieldName.REFERENCED_LINE_RUN_ID, ref_line_run_id)

    # local to cloud feature
    ws_triad = _get_ws_triad_from_pf_config()
    if ws_triad is not None and not _is_azure_ext_installed():
        warning_msg = (
            "`trace.provider` is configured in prompt flow config, while Azure extension is not installed; "
            "in this case, traces will not be exported to cloud. If you expect to export traces to cloud, "
            "please install the Azure extension with command `pip install promptflow-azure`, "
            "and manually restart prompt flow service with command `pf service restart` to make it take effect."
        )
        logger.warning(warning_msg)

    # invoke prompt flow service
    pfs_port = _invoke_pf_svc()

    _inject_res_attrs_to_environ(pfs_port=pfs_port, collection=collection, exp=exp, ws_triad=ws_triad)
    # instrument openai and setup exporter to pfs here for flex mode
    inject_openai_api()
    setup_exporter_to_pfs()
    # print tracing url(s)
    _print_tracing_url_from_local(pfs_port=pfs_port, collection=collection, exp=exp, run=run)
    _print_tracing_url_from_azure_portal(ws_triad=ws_triad, collection=collection, exp=exp, run=run)


def setup_exporter_to_pfs() -> None:
    # get resource attributes from environment
    # For local trace, collection is the only identifier for name and id
    # For cloud trace, we use collection here as name and collection_id for id
    collection = os.getenv(TraceEnvironmentVariableName.COLLECTION, None)
    # Only used for runtime
    collection_id = os.getenv(TraceEnvironmentVariableName.COLLECTION_ID, None)
    exp = os.getenv(TraceEnvironmentVariableName.EXPERIMENT, None)
    # local to cloud scenario: workspace triad in resource.attributes

Review comment (suggestion):
Last time I heard from @brynn-code  `pip install promptflow[azure]` is the recommended way to install.

Code after:
    run: typing.Optional[str] = None,
) -> None:
    # honor and set attributes if user has specified
    if isinstance(attrs, dict):
        _inject_attrs_to_op_ctx(attrs)

    # experiment related attributes, pass from environment
    env_tracing_ctx = os.environ.get(PF_TRACE_CONTEXT, None)
    logger.debug("Read tracing context from environment: %s", env_tracing_ctx)
    env_attrs = dict(json.loads(env_tracing_ctx)).get(PF_TRACE_CONTEXT_ATTR) if env_tracing_ctx else dict()
    exp = env_attrs.get(ContextAttributeKey.EXPERIMENT, None)
    ref_line_run_id = env_attrs.get(ContextAttributeKey.REFERENCED_LINE_RUN_ID, None)
    op_ctx = OperationContext.get_instance()
    # remove `referenced.line_run_id` from context to avoid stale value set by previous node
    if ref_line_run_id is None:
        op_ctx._remove_otel_attributes(SpanAttributeFieldName.REFERENCED_LINE_RUN_ID)
    else:
        op_ctx._add_otel_attributes(SpanAttributeFieldName.REFERENCED_LINE_RUN_ID, ref_line_run_id)

    # local to cloud feature
    ws_triad = _get_ws_triad_from_pf_config()
    if ws_triad is not None and not _is_azure_ext_installed():
        warning_msg = (
            "Azure extension is not installed, though export to cloud is configured, "
            "traces cannot be exported to cloud. To fix this, please run `pip install promptflow-azure` "
            "and restart prompt flow service."
        )
        logger.warning(warning_msg)

    # invoke prompt flow service
    pfs_port = _invoke_pf_svc()

    _inject_res_attrs_to_environ(pfs_port=pfs_port, collection=collection, exp=exp, ws_triad=ws_triad)
    # instrument openai and setup exporter to pfs here for flex mode
    inject_openai_api()
    setup_exporter_to_pfs()
    # print tracing url(s)
    _print_tracing_url_from_local(pfs_port=pfs_port, collection=collection, exp=exp, run=run)
    _print_tracing_url_from_azure_portal(ws_triad=ws_triad, collection=collection, exp=exp, run=run)


def setup_exporter_to_pfs() -> None:
    # get resource attributes from environment
    # For local trace, collection is the only identifier for name and id
    # For cloud trace, we use collection here as name and collection_id for id
    collection = os.getenv(TraceEnvironmentVariableName.COLLECTION, None)
    # Only used for runtime
    collection_id = os.getenv(TraceEnvironmentVariableName.COLLECTION_ID, None)
    exp = os.getenv(TraceEnvironmentVariableName.EXPERIMENT, None)
    # local to cloud scenario: workspace triad in resource.attributes
    workspace_triad = None