fix aihub

Author: AjitPadhi-Microsoft

infra/scripts/aihub_scripts/create_ai_hub.py

@@ -11,12 +11,6 @@
 )
 from azure.keyvault.secrets import SecretClient
 from azure.identity import DefaultAzureCredential
-import time
-import logging
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
 
 def get_secrets_from_kv(kv_name, secret_name):
     # Set the name of the Azure Key Vault
@@ -43,8 +37,6 @@ def get_secrets_from_kv(kv_name, secret_name):
 deployment_name = 'draftsinference-' + 'solutionname_to-be-replaced'
 solutionLocation = 'solutionlocation_to-be-replaced'
 
-logger.info(f"Starting AI Hub creation for {aihub_name} in {solutionLocation}")
-
 # Open AI Details
 open_ai_key = get_secrets_from_kv(key_vault_name, "AZURE-OPENAI-KEY")
 open_ai_res_name = (
@@ -70,133 +62,49 @@ def get_secrets_from_kv(kv_name, secret_name):
 # Credentials
 credential = DefaultAzureCredential()
 
-# Create an ML client without workspace_name for hub creation
+# Create an ML client
 ml_client = MLClient(
+    workspace_name=aihub_name,
     resource_group_name=resource_group_name,
     subscription_id=subscription_id,
     credential=credential,
 )
 
-try:
-    # Check if hub already exists
-    logger.info(f"Checking if AI Hub {aihub_name} already exists...")
-    try:
-        existing_hub = ml_client.workspaces.get(name=aihub_name)
-        logger.info(f"AI Hub {aihub_name} already exists, skipping creation")
-        created_hub = existing_hub
-    except Exception:
-        # Hub doesn't exist, create it
-        logger.info(f"Creating AI Hub {aihub_name}...")
-        
-        # construct a hub
-        my_hub = Hub(
-            name=aihub_name, 
-            location=solutionLocation, 
-            display_name=aihub_name,
-            description=f"AI Hub for {aihub_name}"
-        )
-
-        # Create hub with proper error handling
-        operation = ml_client.workspaces.begin_create(my_hub)
-        logger.info("Hub creation initiated, waiting for completion...")
-        
-        # Poll for completion with timeout
-        start_time = time.time()
-        timeout = 1800  # 30 minutes timeout
-        
-        while not operation.done():
-            elapsed = time.time() - start_time
-            if elapsed > timeout:
-                raise TimeoutError(f"Hub creation timed out after {elapsed/60:.1f} minutes")
-            
-            logger.info(f"Hub creation in progress... ({elapsed/60:.1f} minutes elapsed)")
-            time.sleep(30)  # Check every 30 seconds
-        
-        created_hub = operation.result()
-        logger.info(f"AI Hub {aihub_name} created successfully")
-
-    # Create ML client with the hub workspace
-    ml_client_hub = MLClient(
-        workspace_name=aihub_name,
-        resource_group_name=resource_group_name,
-        subscription_id=subscription_id,
-        credential=credential,
-    )
+# construct a hub
+my_hub = Hub(name=aihub_name, location=solutionLocation, display_name=aihub_name)
+
+created_hub = ml_client.workspaces.begin_create(my_hub).result()
+
+# construct the project
+my_project = Project(
+    name=project_name,
+    location=solutionLocation,
+    display_name=project_name,
+    hub_id=created_hub.id,
+)
+
+created_project = ml_client.workspaces.begin_create(workspace=my_project).result()
+
+open_ai_connection = AzureOpenAIConnection(
+    name="Azure_OpenAI",
+    api_key=open_ai_key,
+    api_version=openai_api_version,
+    azure_endpoint=f"https://{open_ai_res_name}.openai.azure.com/",
+    open_ai_resource_id=f"/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/providers/Microsoft.CognitiveServices/accounts/{open_ai_res_name}"
+)
+
+ml_client.connections.create_or_update(open_ai_connection)
+
+target = f"https://{ai_search_res_name}.search.windows.net/"
+
+# Create AI Search resource
+aisearch_connection = AzureAISearchConnection(
+    name="Azure_AISearch",
+    endpoint=target,
+    credentials=ApiKeyConfiguration(key=ai_search_key),
+)
+
+aisearch_connection.tags["ResourceId"] = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/providers/Microsoft.Search/searchServices/{ai_search_res_name}"
+aisearch_connection.tags["ApiVersion"] = "2024-05-01-preview"
 
-    # Check if project already exists
-    logger.info(f"Checking if AI Project {project_name} already exists...")
-    try:
-        existing_project = ml_client.workspaces.get(name=project_name)
-        logger.info(f"AI Project {project_name} already exists, skipping creation")
-        created_project = existing_project
-    except Exception:
-        # Project doesn't exist, create it
-        logger.info(f"Creating AI Project {project_name}...")
-        
-        # construct the project
-        my_project = Project(
-            name=project_name,
-            location=solutionLocation,
-            display_name=project_name,
-            hub_id=created_hub.id,
-            description=f"AI Project for {project_name}"
-        )
-
-        operation = ml_client.workspaces.begin_create(workspace=my_project)
-        logger.info("Project creation initiated, waiting for completion...")
-        
-        # Poll for completion
-        start_time = time.time()
-        while not operation.done():
-            elapsed = time.time() - start_time
-            if elapsed > timeout:
-                raise TimeoutError(f"Project creation timed out after {elapsed/60:.1f} minutes")
-            
-            logger.info(f"Project creation in progress... ({elapsed/60:.1f} minutes elapsed)")
-            time.sleep(30)
-        
-        created_project = operation.result()
-        logger.info(f"AI Project {project_name} created successfully")
-
-    # Create connections
-    logger.info("Creating AI connections...")
-    
-    # OpenAI connection
-    try:
-        logger.info("Creating Azure OpenAI connection...")
-        open_ai_connection = AzureOpenAIConnection(
-            name="Azure_OpenAI",
-            api_key=open_ai_key,
-            api_version=openai_api_version,
-            azure_endpoint=f"https://{open_ai_res_name}.openai.azure.com/",
-            open_ai_resource_id=f"/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/providers/Microsoft.CognitiveServices/accounts/{open_ai_res_name}"
-        )
-        ml_client_hub.connections.create_or_update(open_ai_connection)
-        logger.info("Azure OpenAI connection created successfully")
-    except Exception as e:
-        logger.warning(f"Failed to create OpenAI connection: {str(e)}")
-
-    # AI Search connection
-    try:
-        logger.info("Creating Azure AI Search connection...")
-        target = f"https://{ai_search_res_name}.search.windows.net/"
-        
-        aisearch_connection = AzureAISearchConnection(
-            name="Azure_AISearch",
-            endpoint=target,
-            credentials=ApiKeyConfiguration(key=ai_search_key),
-        )
-
-        aisearch_connection.tags["ResourceId"] = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/providers/Microsoft.Search/searchServices/{ai_search_res_name}"
-        aisearch_connection.tags["ApiVersion"] = "2024-05-01-preview"
-
-        ml_client_hub.connections.create_or_update(aisearch_connection)
-        logger.info("Azure AI Search connection created successfully")
-    except Exception as e:
-        logger.warning(f"Failed to create AI Search connection: {str(e)}")
-
-    logger.info("AI Hub setup completed successfully!")
-
-except Exception as e:
-    logger.error(f"Error during AI Hub creation: {str(e)}")
-    raise
+ml_client.connections.create_or_update(aisearch_connection)