[azure][trace] Check trace Cosmos status and log warnings if not ready

src/promptflow-azure/promptflow/azure/_constants/_trace.py

@@ -6,3 +6,22 @@
 COSMOS_DB_SETUP_POLL_INTERVAL_SECOND = 30
 COSMOS_DB_SETUP_POLL_PRINT_INTERVAL_SECOND = 30
 COSMOS_DB_SETUP_RESOURCE_TYPE = "HOBO"
+
+
+class CosmosConfiguration:
+ NONE = "None"
+ READ_DISABLED = "ReadDisabled"
+ WRITE_DISABLED = "WriteDisabled"
+ DISABLED = "Disabled"
+ DIAGNOSTIC_DISABLED = "DiagnosticDisabled"
+ DATA_CLEANED = "DataCleaned"
+ ACCOUNT_DELETED = "AccountDeleted"
+
+
+class CosmosStatus:
+ NOT_EXISTS = "NotExists"
+ INITIALIZING = "Initializing"
+ INITIALIZED = "Initialized"
+ DELETING = "Deleting"
+ DELETED = "Deleted"
+ NOT_AVAILABLE = "NotAvailable"

src/promptflow-azure/promptflow/azure/_contracts/__init__.py

@@ -0,0 +1,3 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------

src/promptflow-azure/promptflow/azure/_contracts/_trace.py

@@ -0,0 +1,33 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from dataclasses import dataclass
+
+from promptflow.azure._constants._trace import CosmosConfiguration, CosmosStatus
+from promptflow.azure._restclient.flow.models import TraceCosmosMetaDto
+
+
+@dataclass
+class CosmosMetadata:
+ entity_id: str
+ configuration: str
+ status: str
+ database_name: str
+ resource_type: str
+
+ @staticmethod
+ def _from_rest_object(obj: TraceCosmosMetaDto) -> "CosmosMetadata":
+ return CosmosMetadata(
+ entity_id=obj.entity_id,
+ configuration=obj.trace_cosmos_configuration,
+ status=obj.trace_cosmos_status,
+ database_name=obj.database_name,
+ resource_type=obj.resource_type,
+ )
+
+ def is_disabled(self) -> bool:
+ return self.configuration == CosmosConfiguration.DISABLED
+
+ def is_ready(self) -> bool:
+ return not self.is_disabled() and self.status == CosmosStatus.INITIALIZED

src/promptflow-azure/promptflow/azure/_restclient/flow_service_caller.py

@@ -24,6 +24,11 @@
  COSMOS_DB_SETUP_POLL_PRINT_INTERVAL_SECOND,
  COSMOS_DB_SETUP_POLL_TIMEOUT_SECOND,
 )
+from promptflow.azure._constants._trace import (
+ COSMOS_DB_SETUP_POLL_INTERVAL_SECOND,
+ COSMOS_DB_SETUP_POLL_PRINT_INTERVAL_SECOND,
+ COSMOS_DB_SETUP_POLL_TIMEOUT_SECOND,
+)
 from promptflow.azure._restclient.flow import AzureMachineLearningDesignerServiceClient
 from promptflow.azure._utils.general import get_authorization, get_arm_token, get_aml_token
 from promptflow.exceptions import UserErrorException, PromptflowException, SystemErrorException
@@ -763,6 +768,22 @@ def init_workspace_cosmos(
  **kwargs,
  )
 
+ def get_workspace_cosmos_metadata(
+ self,
+ subscription_id: str,
+ resource_group_name: str,
+ workspace_name: str,
+ **kwargs,
+ ):
+ """Get Cosmos DB metadata."""
+ return self.caller.trace_sessions.get_trace_session_metadata_async(
+ subscription_id=subscription_id,
+ resource_group_name=resource_group_name,
+ workspace_name=workspace_name,
+ headers=self._get_headers(),
+ **kwargs,
+ )
+
  @_request_wrapper()
  def setup_workspace_cosmos(self, subscription_id, resource_group_name, workspace_name, body, **kwargs):
  """Setup Cosmos DB for workspace/project."""

src/promptflow-azure/promptflow/azure/_utils/_tracing.py

@@ -6,12 +6,11 @@
 from azure.core.exceptions import ResourceNotFoundError
 from azure.identity import AzureCliCredential
 
-from promptflow._constants import AzureWorkspaceKind, CosmosDBContainerName
+from promptflow._constants import AzureWorkspaceKind
 from promptflow._sdk._utilities.general_utils import extract_workspace_triad_from_trace_provider
 from promptflow._utils.logger_utils import get_cli_sdk_logger
 from promptflow.azure import PFClient
 from promptflow.azure._constants._trace import COSMOS_DB_SETUP_RESOURCE_TYPE
-from promptflow.azure._restclient.flow_service_caller import FlowRequestException
 from promptflow.exceptions import ErrorTarget, UserErrorException
 
 _logger = get_cli_sdk_logger()
@@ -61,13 +60,16 @@ def validate_trace_destination(value: str) -> None:
  _logger.debug("Resource type is valid.")
 
  # the workspace Cosmos DB is initialized
- # try to retrieve the token from PFS; if failed, call PFS init API and start polling
+ # if not, call PFS setup API and start polling
  _logger.debug("Validating workspace Cosmos DB is initialized...")
  pf_client = PFClient(ml_client=ml_client)
- try:
- pf_client._traces._get_cosmos_db_token(container_name=CosmosDBContainerName.SPAN)
- _logger.debug("The workspace Cosmos DB is already initialized.")
- except FlowRequestException:
+ cosmos_metadata = pf_client._traces._get_cosmos_metadata()
+ # raise error if the Cosmos DB is disabled
+ if cosmos_metadata.is_disabled():
+ error_message = "The workspace Cosmos DB is disabled, please enable it first."
+ _logger.error(error_message)
+ raise _create_trace_destination_value_user_error(error_message)
+ if not cosmos_metadata.is_ready():
  # print here to let users aware this operation as it's kind of time consuming
  init_cosmos_msg = (
  "The workspace Cosmos DB is not initialized yet, "
@@ -76,6 +78,8 @@ def validate_trace_destination(value: str) -> None:
  print(init_cosmos_msg)
  _logger.debug(init_cosmos_msg)
  pf_client._traces._setup_cosmos_db(resource_type=COSMOS_DB_SETUP_RESOURCE_TYPE)
+ else:
+ _logger.debug("The workspace Cosmos DB is available.")
  _logger.debug("The workspace Cosmos DB is initialized.")
 
  _logger.debug("pf.config.trace.destination is valid.")

src/promptflow-azure/promptflow/azure/operations/_trace_operations.py

@@ -7,6 +7,7 @@
 from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationScope, _ScopeDependentOperations
 
 from promptflow._sdk._telemetry import ActivityType, WorkspaceTelemetryMixin, monitor_operation
+from promptflow.azure._contracts._trace import CosmosMetadata
 from promptflow.azure._restclient.flow.models import TraceDbSetupRequest
 from promptflow.azure._restclient.flow_service_caller import FlowServiceCaller
 
@@ -32,6 +33,7 @@ def __init__(
 
  @monitor_operation(activity_name="pfazure.traces._init_cosmos_db", activity_type=ActivityType.INTERNALCALL)
  def _init_cosmos_db(self) -> Optional[Dict]:
+ # this API is deprecated and will be removed in the future
  resp = self._service_caller.init_workspace_cosmos(
  subscription_id=self._operation_scope.subscription_id,
  resource_group_name=self._operation_scope.resource_group_name,
@@ -63,3 +65,11 @@ def _setup_cosmos_db(self, resource_type: str) -> None:
  body=body,
  )
  return
+
+ def _get_cosmos_metadata(self) -> CosmosMetadata:
+ rest_obj = self._service_caller.get_workspace_cosmos_metadata(
+ subscription_id=self._operation_scope.subscription_id,
+ resource_group_name=self._operation_scope.resource_group_name,
+ workspace_name=self._operation_scope.workspace_name,
+ )
+ return CosmosMetadata._from_rest_object(rest_obj)