From 8df1c3d684fc3006b27215f9c50f4061bbad0e87 Mon Sep 17 00:00:00 2001
From: Zhengfei Wang <zhengfei.wang@microsoft.com>
Date: Thu, 11 Apr 2024 23:32:49 +0800
Subject: [PATCH] raise exception when pfs inovke error

---
 src/promptflow-devkit/promptflow/_sdk/_errors.py  | 6 ++++++
 src/promptflow-devkit/promptflow/_sdk/_tracing.py | 5 +++++
 2 files changed, 11 insertions(+)

diff --git a/src/promptflow-devkit/promptflow/_sdk/_errors.py b/src/promptflow-devkit/promptflow/_sdk/_errors.py
index 3ce415be56a..4a5f564fee4 100644
--- a/src/promptflow-devkit/promptflow/_sdk/_errors.py
+++ b/src/promptflow-devkit/promptflow/_sdk/_errors.py
@@ -271,3 +271,9 @@ def __init__(
             'please install it by running "pip install promptflow-azure" with your version.'
         )
         super().__init__(message=msg, no_personal_data_message=msg, **kwargs)
+
+
+class PromptFlowServiceInvokeError(SDKError):
+    """Exception raised if prompt flow service invoke failed."""
+
+    pass
diff --git a/src/promptflow-devkit/promptflow/_sdk/_tracing.py b/src/promptflow-devkit/promptflow/_sdk/_tracing.py
index 0661cea282c..91e0b8c9c1d 100644
--- a/src/promptflow-devkit/promptflow/_sdk/_tracing.py
+++ b/src/promptflow-devkit/promptflow/_sdk/_tracing.py
@@ -30,6 +30,7 @@
     AzureMLWorkspaceTriad,
     ContextAttributeKey,
 )
+from promptflow._sdk._errors import PromptFlowServiceInvokeError
 from promptflow._sdk._service.utils.utils import (
     add_executable_script_to_env_path,
     get_port_from_config,
@@ -258,6 +259,10 @@ def start_trace_with_devkit(
 
     # invoke prompt flow service
     pfs_port = _invoke_pf_svc()
+    if not is_pfs_service_healthy(pfs_port):
+        error_msg = "prompt flow service is not healthy, please check the logs for more details."
+        logger.error(error_msg)
+        raise PromptFlowServiceInvokeError(error_msg)
 
     _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