Fix error NoneType object is not subscriptable (Azure#38493)

azure-sdk · Nov 13, 2024 · e95bc99 · e95bc99
1 parent 7b18366
commit e95bc99
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py
@@ -171,7 +171,7 @@ def _from_rest_object(
         if not rest_obj:
             return None
 
-        workspace_object = Workspace._from_rest_object(rest_obj)
+        workspace_object = Workspace._from_rest_object(rest_obj, v2_service_context)
         if workspace_object is not None:
             return FeatureStore(
                 name=str(workspace_object.name),

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py
@@ -341,9 +341,13 @@ def _from_rest_object(
 
         if hasattr(rest_obj, "ml_flow_tracking_uri"):
             try:
-                from azureml.mlflow import get_mlflow_tracking_uri_v2
+                if v2_service_context:
+                    # v2_service_context is required (not None) in get_mlflow_tracking_uri_v2
+                    from azureml.mlflow import get_mlflow_tracking_uri_v2
 
-                mlflow_tracking_uri = get_mlflow_tracking_uri_v2(rest_obj, v2_service_context)
+                    mlflow_tracking_uri = get_mlflow_tracking_uri_v2(rest_obj, v2_service_context)
+                else:
+                    mlflow_tracking_uri = rest_obj.ml_flow_tracking_uri
             except ImportError:
                 mlflow_tracking_uri = rest_obj.ml_flow_tracking_uri
                 error_msg = (