Add deployment setup to clarifai model upload CLI command

clarifai/runners/models/model_builder.py

@@ -7,6 +7,7 @@
 import sys
 import tarfile
 import time
+import webbrowser
 from string import Template
 from unittest.mock import MagicMock
 
@@ -16,6 +17,7 @@
 from google.protobuf import json_format
 
 from clarifai.client.base import BaseClient
+from clarifai.client.user import User
 from clarifai.runners.models.model_class import ModelClass
 from clarifai.runners.utils.const import (
     AMD_PYTHON_BASE_IMAGE,
@@ -1108,4 +1110,189 @@ def upload_model(folder, stage, skip_dockerfile):
         )
 
     input("Press Enter to continue...")
-    builder.upload_model_version()
+    model_version = builder.upload_model_version()
+
+    # Ask user if they want to deploy the model
+    deploy_model = input("Do you want to deploy the model? (y/n): ")
+    if deploy_model.lower() != 'y':
+        logger.info("Model uploaded successfully. Skipping deployment setup.")
+        return
+
+    # Setup deployment for the uploaded model
+    setup_deployment_for_model(builder)
+
+def setup_deployment_for_model(builder):
+    """
+    Set up deployment for a model after upload.
+
+    :param builder: The ModelBuilder instance that has uploaded the model.
+    """
+
+    model = builder.config.get('model')
+    user_id = model.get('user_id')
+    app_id = model.get('app_id')
+    model_id = model.get('id')
+
+    # Set up the API client with the user's credentials
+    user = User(
+        user_id=user_id,
+        pat=builder.client.pat,
+        base_url=builder.client.base
+    )
+
+    # Step 1: Check for available compute clusters and let user choose or create a new one
+    logger.info("Checking for available compute clusters...")
+    compute_clusters = list(user.list_compute_clusters())
+
+    compute_cluster = None
+    if compute_clusters:
+        logger.info("Available compute clusters:")
+        for i, cc in enumerate(compute_clusters):
+            logger.info(f"{i+1}. {cc.id} ({cc.description if hasattr(cc, 'description') else 'No description'})")
+
+        choice = input(f"Choose a compute cluster (1-{len(compute_clusters)}) or 'n' to create a new one: ")
+        if choice.lower() == 'n':
+            create_new_cc = True
+        else:
+            try:
+                idx = int(choice) - 1
+                if 0 <= idx < len(compute_clusters):
+                    compute_cluster = compute_clusters[idx]
+                    create_new_cc = False
+                else:
+                    logger.info("Invalid choice. Creating a new compute cluster.")
+                    create_new_cc = True
+            except ValueError:
+                logger.info("Invalid choice. Creating a new compute cluster.")
+                create_new_cc = True
+    else:
+        logger.info("No compute clusters found.")
+        create_new_cc = True
+
+    if create_new_cc:
+        # Provide URL to create a new compute cluster
+        url_helper = ClarifaiUrlHelper()
+        compute_cluster_url = f"{url_helper.ui}/settings/compute/new"
+        logger.info(f"Please create a new compute cluster by visiting: {compute_cluster_url}")
+
+        # Ask if they want to open the URL in browser
+        open_browser = input("Do you want to open the compute cluster creation page in your browser? (y/n): ")
+        if open_browser.lower() == 'y':
+            try:
+                webbrowser.open(compute_cluster_url)
+            except Exception as e:
+                logger.error(f"Failed to open browser: {e}")
+
+        input("After creating the compute cluster, press Enter to continue...")
+
+        # Re-fetch the compute clusters list after user has created one
+        logger.info("Re-checking for available compute clusters...")
+        compute_clusters = list(user.list_compute_clusters())
+
+        if not compute_clusters:
+            logger.info("No compute clusters found. Please make sure you have created a compute cluster and try again.")
+            return
+
+        # Show the updated list and let user choose
+        logger.info("Available compute clusters:")
+        for i, cc in enumerate(compute_clusters):
+            logger.info(f"{i+1}. {cc.id} ({cc.description if hasattr(cc, 'description') else 'No description'})")
+
+        choice = input(f"Choose a compute cluster (1-{len(compute_clusters)}): ")
+        try:
+            idx = int(choice) - 1
+            if 0 <= idx < len(compute_clusters):
+                compute_cluster = compute_clusters[idx]
+            else:
+                logger.info("Invalid choice. Aborting deployment setup.")
+                return
+        except ValueError:
+            logger.info("Invalid choice. Aborting deployment setup.")
+            return
+
+    # Step 2: Check for available nodepools and let user choose or create a new one
+    logger.info(f"Checking for available nodepools in compute cluster '{compute_cluster.id}'...")
+    nodepools = list(compute_cluster.list_nodepools())
+
+    nodepool = None
+    if nodepools:
+        logger.info("Available nodepools:")
+        for i, np in enumerate(nodepools):
+            logger.info(f"{i+1}. {np.id} ({np.description if hasattr(np, 'description') else 'No description'})")
+
+        choice = input(f"Choose a nodepool (1-{len(nodepools)}) or 'n' to create a new one: ")
+        if choice.lower() == 'n':
+            create_new_np = True
+        else:
+            try:
+                idx = int(choice) - 1
+                if 0 <= idx < len(nodepools):
+                    nodepool = nodepools[idx]
+                    create_new_np = False
+                else:
+                    logger.info("Invalid choice. Creating a new nodepool.")
+                    create_new_np = True
+            except ValueError:
+                logger.info("Invalid choice. Creating a new nodepool.")
+                create_new_np = True
+    else:
+        logger.info("No nodepools found in this compute cluster.")
+        create_new_np = True
+
+    if create_new_np:
+        # Provide URL to create a new nodepool
+        url_helper = ClarifaiUrlHelper()
+        nodepool_url = f"{url_helper.ui}/settings/compute/{compute_cluster.id}/nodepools/new"
+        logger.info(f"Please create a new nodepool by visiting: {nodepool_url}")
+
+        # Ask if they want to open the URL in browser
+        open_browser = input("Do you want to open the nodepool creation page in your browser? (y/n): ")
+        if open_browser.lower() == 'y':
+            try:
+                webbrowser.open(nodepool_url)
+            except Exception as e:
+                logger.error(f"Failed to open browser: {e}")
+
+        input("After creating the nodepool, press Enter to continue...")
+
+        # Re-fetch the nodepools list after user has created one
+        logger.info(f"Re-checking for available nodepools in compute cluster '{compute_cluster.id}'...")
+        nodepools = list(compute_cluster.list_nodepools())
+
+        if not nodepools:
+            logger.info("No nodepools found. Please make sure you have created a nodepool in the selected compute cluster and try again.")
+            return
+
+        # Show the updated list and let user choose
+        logger.info("Available nodepools:")
+        for i, np in enumerate(nodepools):
+            logger.info(f"{i+1}. {np.id} ({np.description if hasattr(np, 'description') else 'No description'})")
+
+        choice = input(f"Choose a nodepool (1-{len(nodepools)}): ")
+        try:
+            idx = int(choice) - 1
+            if 0 <= idx < len(nodepools):
+                nodepool = nodepools[idx]
+            else:
+                logger.info("Invalid choice. Aborting deployment setup.")
+                return
+        except ValueError:
+            logger.info("Invalid choice. Aborting deployment setup.")
+            return
+
+    # Step 3: Help create a new deployment by providing URL
+    # Provide URL to create a new deployment
+    url_helper = ClarifaiUrlHelper()
+    deployment_url = f"{url_helper.ui}/settings/compute/deployments/new?computeClusterId={compute_cluster.id}&nodePoolId={nodepool.id}"
+    logger.info(f"Please create a new deployment by visiting: {deployment_url}")
+
+    # Ask if they want to open the URL in browser
+    open_browser = input("Do you want to open the deployment creation page in your browser? (y/n): ")
+    if open_browser.lower() == 'y':
+        try:
+            webbrowser.open(deployment_url)
+        except Exception as e:
+            logger.error(f"Failed to open browser: {e}")
+
+    logger.info("After creating the deployment, your model will be ready for inference!")
+    logger.info(f"You can always return to view your deployments at: {deployment_url}")

tests/runners/test_model_deployment_setup.py

@@ -0,0 +1,176 @@
+import unittest
+from unittest.mock import patch, MagicMock
+
+
+class TestModelDeploymentSetup(unittest.TestCase):
+    """Test case for the model deployment setup functionality."""
+
+    @patch('clarifai.client.user.User')
+    @patch('clarifai.client.compute_cluster.ComputeCluster')
+    @patch('clarifai.client.nodepool.Nodepool')
+    @patch('builtins.input')
+    @patch('clarifai.utils.logging.logger')
+    def test_deployment_setup_existing_resources(self, mock_logger, mock_input, 
+                                               mock_nodepool, mock_compute_cluster, 
+                                               mock_user):
+        """Test the deployment setup with existing compute cluster and nodepool."""
+        # Import the function here to avoid circular imports in tests
+        from clarifai.runners.models.model_builder import setup_deployment_for_model
+
+        # Setup mocks
+        builder = MagicMock()
+        builder.user_id = "test-user"
+        builder.app_id = "test-app"
+        builder.model_id = "test-model"
+        builder._client._session.pat = "test-pat"
+        builder._client._session.base = "https://api.clarifai.com"
+
+        # Mock User
+        mock_user_instance = mock_user.return_value
+        mock_compute_cluster_obj = MagicMock()
+        mock_compute_cluster_obj.id = "test-cc"
+        mock_compute_cluster_obj.description = "Test compute cluster"
+        mock_user_instance.list_compute_clusters.return_value = [mock_compute_cluster_obj]
+
+        # Mock ComputeCluster
+        mock_compute_cluster_instance = mock_compute_cluster.return_value
+        mock_nodepool_obj = MagicMock()
+        mock_nodepool_obj.id = "test-np"
+        mock_nodepool_obj.description = "Test nodepool"
+        mock_compute_cluster_instance.list_nodepools.return_value = [mock_nodepool_obj]
+
+        # Mock Nodepool
+        mock_nodepool_instance = mock_nodepool.return_value
+
+        # Mock input responses
+        mock_input.side_effect = ["1", "1", "test-deployment", "n"]
+
+        # Call the function
+        setup_deployment_for_model(builder)
+
+        # Verify user was initialized correctly
+        mock_user.assert_called_once_with(
+            user_id="test-user",
+            pat="test-pat",
+            base_url="https://api.clarifai.com"
+        )
+
+        # Verify compute cluster selection
+        mock_user_instance.list_compute_clusters.assert_called_once()
+        mock_compute_cluster.assert_called_once_with(
+            compute_cluster_id=mock_compute_cluster_obj.id,
+            user_id="test-user",
+            pat="test-pat",
+            base_url="https://api.clarifai.com"
+        )
+
+        # Verify nodepool selection
+        mock_compute_cluster_instance.list_nodepools.assert_called_once()
+
+        # Verify deployment creation
+        mock_nodepool.assert_called_once_with(
+            nodepool_id=mock_nodepool_obj.id,
+            user_id="test-user",
+            pat="test-pat",
+            base_url="https://api.clarifai.com"
+        )
+
+        # Verify deployment config
+        expected_deployment_config = {
+            "deployment": {
+                "id": "test-deployment",
+                "description": "Deployment for test-model",
+                "worker": {
+                    "model": {
+                        "id": "test-model",
+                        "user_id": "test-user",
+                        "app_id": "test-app",
+                    }
+                },
+                "nodepools": [
+                    {
+                        "id": "test-np",
+                        "compute_cluster": {
+                            "id": "test-cc",
+                            "user_id": "test-user"
+                        }
+                    }
+                ]
+            }
+        }
+        mock_nodepool_instance.create_deployment.assert_called_once_with(expected_deployment_config)
+
+        # Verify logging messages
+        mock_logger.info.assert_any_call("Checking for available compute clusters...")
+        mock_logger.info.assert_any_call("Available compute clusters:")
+        mock_logger.info.assert_any_call("Checking for available nodepools in compute cluster 'test-cc'...")
+        mock_logger.info.assert_any_call("Available nodepools:")
+        mock_logger.info.assert_any_call("Creating deployment 'test-deployment'...")
+        mock_logger.info.assert_any_call("Deployment 'test-deployment' created successfully.")
+
+    @patch('clarifai.client.user.User')
+    @patch('builtins.input')
+    @patch('clarifai.utils.logging.logger')
+    def test_deployment_setup_create_new_compute_cluster(self, mock_logger, mock_input, mock_user):
+        """Test the deployment setup with creating a new compute cluster."""
+        # Import the function here to avoid circular imports in tests
+        from clarifai.runners.models.model_builder import setup_deployment_for_model
+
+        # Setup mocks
+        builder = MagicMock()
+        builder.user_id = "test-user"
+        builder.app_id = "test-app"
+        builder.model_id = "test-model"
+        builder._client._session.pat = "test-pat"
+        builder._client._session.base = "https://api.clarifai.com"
+
+        # Mock User
+        mock_user_instance = mock_user.return_value
+        # No existing compute clusters
+        mock_user_instance.list_compute_clusters.return_value = []
+
+        # Mock creating compute cluster
+        mock_cc = MagicMock()
+        mock_cc.id = "new-cc"
+        mock_user_instance.create_compute_cluster.return_value = mock_cc
+
+        # Abort after creating compute cluster to simplify test
+        mock_input.side_effect = ["new-cc"]
+
+        # Abort after creating the compute cluster
+        mock_cc.list_nodepools = MagicMock(side_effect=Exception("Test exception"))
+
+        # This should abort after creating the compute cluster due to the exception
+        with self.assertRaises(Exception):
+            setup_deployment_for_model(builder)
+
+        # Verify user was initialized correctly
+        mock_user.assert_called_once_with(
+            user_id="test-user",
+            pat="test-pat",
+            base_url="https://api.clarifai.com"
+        )
+
+        # Verify compute cluster creation
+        mock_user_instance.list_compute_clusters.assert_called_once()
+        mock_user_instance.create_compute_cluster.assert_called_once()
+
+        # Verify expected config for compute cluster creation
+        expected_cc_config = {
+            "compute_cluster": {
+                "id": "new-cc",
+                "cluster_type": "users",
+                "description": "Compute cluster for test-model",
+            }
+        }
+        # Check that the config matches (ignoring the specific object instance)
+        actual_config = mock_user_instance.create_compute_cluster.call_args[0][0]
+        self.assertEqual(actual_config["compute_cluster"]["id"], expected_cc_config["compute_cluster"]["id"])
+        self.assertEqual(actual_config["compute_cluster"]["cluster_type"], expected_cc_config["compute_cluster"]["cluster_type"])
+        self.assertEqual(actual_config["compute_cluster"]["description"], expected_cc_config["compute_cluster"]["description"])
+
+        # Verify logging messages
+        mock_logger.info.assert_any_call("Checking for available compute clusters...")
+        mock_logger.info.assert_any_call("No compute clusters found.")
+        mock_logger.info.assert_any_call("Creating new compute cluster 'new-cc'...")
+        mock_logger.info.assert_any_call("Compute cluster 'new-cc' created successfully.")