clarifai 11.3.0rc2__py3-none-any.whl → 11.4.0__py3-none-any.whl

clarifai/client/compute_cluster.py

@@ -14,184 +14,218 @@ from clarifai.utils.logging import logger
 
 
 class ComputeCluster(Lister, BaseClient):
-  """ComputeCluster is a class that provides access to Clarifai API endpoints related to Compute Cluster information."""
-
-  def __init__(self,
-               compute_cluster_id: str = None,
-               user_id: str = None,
-               base_url: str = "https://api.clarifai.com",
-               pat: str = None,
-               token: str = None,
-               root_certificates_path: str = None,
-               **kwargs):
-    """Initializes an ComputeCluster object.
-
-    Args:
-        compute_cluster_id (str): The ComputeCluster ID for the ComputeCluster to interact with.
-        user_id (str): The user ID of the user.
-        base_url (str): Base API url. Default "https://api.clarifai.com"
-        pat (str): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
-        token (str): A session token for authentication. Accepts either a session token or a pat. Can be set as env var CLARIFAI_SESSION_TOKEN
-        root_certificates_path (str): Path to the SSL root certificates file, used to establish secure gRPC connections.
-        **kwargs: Additional keyword arguments to be passed to the compute cluster.
-    """
-    self.kwargs = {**kwargs, 'id': compute_cluster_id, 'user_id': user_id}
-    self.compute_cluster_info = resources_pb2.ComputeCluster(**self.kwargs)
-    self.logger = logger
-    BaseClient.__init__(
+    """ComputeCluster is a class that provides access to Clarifai API endpoints related to Compute Cluster information."""
+
+    def __init__(
+        self,
+        compute_cluster_id: str = None,
+        user_id: str = None,
+        base_url: str = "https://api.clarifai.com",
+        pat: str = None,
+        token: str = None,
+        root_certificates_path: str = None,
+        **kwargs,
+    ):
+        """Initializes an ComputeCluster object.
+
+        Args:
+            compute_cluster_id (str): The ComputeCluster ID for the ComputeCluster to interact with.
+            user_id (str): The user ID of the user.
+            base_url (str): Base API url. Default "https://api.clarifai.com"
+            pat (str): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
+            token (str): A session token for authentication. Accepts either a session token or a pat. Can be set as env var CLARIFAI_SESSION_TOKEN
+            root_certificates_path (str): Path to the SSL root certificates file, used to establish secure gRPC connections.
+            **kwargs: Additional keyword arguments to be passed to the compute cluster.
+        """
+        self.kwargs = {**kwargs, 'id': compute_cluster_id, 'user_id': user_id}
+        self.compute_cluster_info = resources_pb2.ComputeCluster(**self.kwargs)
+        self.logger = logger
+        BaseClient.__init__(
+            self,
+            user_id=self.user_id,
+            base=base_url,
+            pat=pat,
+            token=token,
+            root_certificates_path=root_certificates_path,
+        )
+        Lister.__init__(self)
+
+    def list_nodepools(
+        self, page_no: int = None, per_page: int = None
+    ) -> Generator[Nodepool, None, None]:
+        """Lists all the available nodepools of the compute cluster.
+
+        Args:
+            compute_cluster_id (str): The compute cluster ID to list the nodepools.
+            page_no (int): The page number to list.
+            per_page (int): The number of items per page.
+
+        Yields:
+            Nodepool: Nodepool objects for the nodepools in the compute cluster.
+
+        Example:
+            >>> from clarifai.client.compute_cluster import ComputeCluster
+            >>> compute_cluster = ComputeCluster(compute_cluster_id="compute_cluster_id", user_id="user_id")
+            >>> all_nodepools = list(compute_cluster.list_nodepools())
+
+        Note:
+            Defaults to 16 per page if page_no is specified and per_page is not specified.
+            If both page_no and per_page are None, then lists all the resources.
+        """
+        request_data = dict(user_app_id=self.user_app_id, compute_cluster_id=self.id)
+        all_nodepools_info = self.list_pages_generator(
+            self.STUB.ListNodepools,
+            service_pb2.ListNodepoolsRequest,
+            request_data,
+            per_page=per_page,
+            page_no=page_no,
+        )
+
+        for nodepool_info in all_nodepools_info:
+            yield Nodepool.from_auth_helper(auth=self.auth_helper, **nodepool_info)
+
+    def _process_nodepool_config(self, nodepool_config: str) -> Dict[str, Any]:
+        assert "nodepool" in nodepool_config, "nodepool info not found in the config file"
+        nodepool = nodepool_config['nodepool']
+        assert "instance_types" in nodepool, "instance_types not found in the config file"
+        assert "node_capacity_type" in nodepool, "node_capacity_type not found in the config file"
+        assert "min_instances" in nodepool, "min_instances not found in the config file"
+        assert "max_instances" in nodepool, "max_instances not found in the config file"
+        nodepool['compute_cluster'] = resources_pb2.ComputeCluster(
+            id=self.id, user_id=self.user_id
+        )
+        nodepool['node_capacity_type'] = resources_pb2.NodeCapacityType(
+            capacity_types=[
+                capacity_type for capacity_type in nodepool['node_capacity_type']['capacity_types']
+            ]
+        )
+        instance_types = []
+        for instance_type in nodepool['instance_types']:
+            if 'compute_info' in instance_type:
+                instance_type['compute_info'] = resources_pb2.ComputeInfo(
+                    **instance_type['compute_info']
+                )
+            instance_types.append(resources_pb2.InstanceType(**instance_type))
+        nodepool['instance_types'] = instance_types
+        if "visibility" in nodepool:
+            nodepool["visibility"] = resources_pb2.Visibility(**nodepool["visibility"])
+        return nodepool
+
+    def create_nodepool(
         self,
-        user_id=self.user_id,
-        base=base_url,
-        pat=pat,
-        token=token,
-        root_certificates_path=root_certificates_path)
-    Lister.__init__(self)
-
-  def list_nodepools(self, page_no: int = None,
-                     per_page: int = None) -> Generator[Nodepool, None, None]:
-    """Lists all the available nodepools of the compute cluster.
-
-    Args:
-        compute_cluster_id (str): The compute cluster ID to list the nodepools.
-        page_no (int): The page number to list.
-        per_page (int): The number of items per page.
-
-    Yields:
-        Nodepool: Nodepool objects for the nodepools in the compute cluster.
-
-    Example:
-        >>> from clarifai.client.compute_cluster import ComputeCluster
-        >>> compute_cluster = ComputeCluster(compute_cluster_id="compute_cluster_id", user_id="user_id")
-        >>> all_nodepools = list(compute_cluster.list_nodepools())
-
-    Note:
-        Defaults to 16 per page if page_no is specified and per_page is not specified.
-        If both page_no and per_page are None, then lists all the resources.
-    """
-    request_data = dict(user_app_id=self.user_app_id, compute_cluster_id=self.id)
-    all_nodepools_info = self.list_pages_generator(
-        self.STUB.ListNodepools,
-        service_pb2.ListNodepoolsRequest,
-        request_data,
-        per_page=per_page,
-        page_no=page_no)
-
-    for nodepool_info in all_nodepools_info:
-      yield Nodepool.from_auth_helper(auth=self.auth_helper, **nodepool_info)
-
-  def _process_nodepool_config(self, config_filepath: str) -> Dict[str, Any]:
-    with open(config_filepath, "r") as file:
-      nodepool_config = yaml.safe_load(file)
-
-    assert "nodepool" in nodepool_config, "nodepool info not found in the config file"
-    nodepool = nodepool_config['nodepool']
-    assert "instance_types" in nodepool, "region not found in the config file"
-    assert "node_capacity_type" in nodepool, "managed_by not found in the config file"
-    assert "max_instances" in nodepool, "cluster_type not found in the config file"
-    nodepool['compute_cluster'] = resources_pb2.ComputeCluster(id=self.id, user_id=self.user_id)
-    nodepool['node_capacity_type'] = resources_pb2.NodeCapacityType(capacity_types=[
-        capacity_type for capacity_type in nodepool['node_capacity_type']['capacity_types']
-    ])
-    instance_types = []
-    for instance_type in nodepool['instance_types']:
-      if 'compute_info' in instance_type:
-        instance_type['compute_info'] = resources_pb2.ComputeInfo(**instance_type['compute_info'])
-      instance_types.append(resources_pb2.InstanceType(**instance_type))
-    nodepool['instance_types'] = instance_types
-    if "visibility" in nodepool:
-      nodepool["visibility"] = resources_pb2.Visibility(**nodepool["visibility"])
-    return nodepool
-
-  def create_nodepool(self, config_filepath: str, nodepool_id: str = None) -> Nodepool:
-    """Creates a nodepool for the compute cluster.
-
-    Args:
-        config_filepath (str): The path to the nodepool config file.
-        nodepool_id (str): New nodepool ID for the nodepool to create.
-
-    Returns:
-        Nodepool: A Nodepool object for the specified nodepool ID.
-
-    Example:
-        >>> from clarifai.client.compute_cluster import ComputeCluster
-        >>> compute_cluster = ComputeCluster(compute_cluster_id="compute_cluster_id", user_id="user_id")
-        >>> nodepool = compute_cluster.create_nodepool(config_filepath="config.yml")
-    """
-    if not os.path.exists(config_filepath):
-      raise UserError(f"Nodepool config file not found at {config_filepath}")
-
-    nodepool_config = self._process_nodepool_config(config_filepath)
-
-    if 'id' in nodepool_config:
-      if nodepool_id is None:
-        nodepool_id = nodepool_config['id']
-      nodepool_config.pop('id')
-
-    request = service_pb2.PostNodepoolsRequest(
-        user_app_id=self.user_app_id,
-        compute_cluster_id=self.id,
-        nodepools=[resources_pb2.Nodepool(id=nodepool_id, **nodepool_config)])
-    response = self._grpc_request(self.STUB.PostNodepools, request)
-    if response.status.code != status_code_pb2.SUCCESS:
-      raise Exception(response.status)
-    self.logger.info("\nNodepool created\n%s", response.status)
-
-    return Nodepool.from_auth_helper(self.auth_helper, nodepool_id=nodepool_id)
-
-  def nodepool(self, nodepool_id: str) -> Nodepool:
-    """Returns a Nodepool object for the existing nodepool ID.
-
-    Args:
-        nodepool_id (str): The nodepool ID for the nodepool to interact with.
-
-    Returns:
-        Nodepool: A Nodepool object for the existing nodepool ID.
-
-    Example:
-        >>> from clarifai.client.compute_cluster import ComputeCluster
-        >>> compute_cluster = ComputeCluster(compute_cluster_id="compute_cluster_id", user_id="user_id")
-        >>> nodepool = compute_cluster.nodepool(nodepool_id="nodepool_id")
-    """
-    request = service_pb2.GetNodepoolRequest(
-        user_app_id=self.user_app_id, compute_cluster_id=self.id, nodepool_id=nodepool_id)
-    response = self._grpc_request(self.STUB.GetNodepool, request)
-
-    if response.status.code != status_code_pb2.SUCCESS:
-      raise Exception(response.status)
-    dict_response = MessageToDict(response, preserving_proto_field_name=True)
-    kwargs = self.process_response_keys(dict_response[list(dict_response.keys())[1]],
-                                        list(dict_response.keys())[1])
-
-    return Nodepool.from_auth_helper(auth=self.auth_helper, **kwargs)
-
-  def delete_nodepools(self, nodepool_ids: List[str]) -> None:
-    """Deletes list of nodepools for the compute cluster.
-
-    Args:
-        nodepool_ids (List[str]): The nodepool IDs of the compute cluster to delete.
-
-    Example:
-        >>> from clarifai.client.compute_cluster import ComputeCluster
-        >>> compute_cluster = ComputeCluster(compute_cluster_id="compute_cluster_id", user_id="user_id")
-        >>> compute_cluster.delete_nodepools(nodepool_ids=["nodepool_id1", "nodepool_id2"])
-    """
-    assert isinstance(nodepool_ids, list), "nodepool_ids param should be a list"
-
-    request = service_pb2.DeleteNodepoolsRequest(
-        user_app_id=self.user_app_id, compute_cluster_id=self.id, ids=nodepool_ids)
-    response = self._grpc_request(self.STUB.DeleteNodepools, request)
-
-    if response.status.code != status_code_pb2.SUCCESS:
-      raise Exception(response.status)
-    self.logger.info("\nNodepools Deleted\n%s", response.status)
-
-  def __getattr__(self, name):
-    return getattr(self.compute_cluster_info, name)
-
-  def __str__(self):
-    init_params = [param for param in self.kwargs.keys()]
-    attribute_strings = [
-        f"{param}={getattr(self.compute_cluster_info, param)}" for param in init_params
-        if hasattr(self.compute_cluster_info, param)
-    ]
-    return f"Clarifai Compute Cluster Details: \n{', '.join(attribute_strings)}\n"
+        config_filepath: str = None,
+        nodepool_id: str = None,
+        nodepool_config: Dict[str, Any] = None,
+    ) -> Nodepool:
+        """Creates a nodepool for the compute cluster.
+
+        Args:
+            config_filepath (str): The path to the nodepool config file.
+            nodepool_id (str): New nodepool ID for the nodepool to create.
+            nodepool_config (Dict[str, Any]) = nodepool_config or {}
+
+        Returns:
+            Nodepool: A Nodepool object for the specified nodepool ID.
+
+        Example:
+            >>> from clarifai.client.compute_cluster import ComputeCluster
+            >>> compute_cluster = ComputeCluster(compute_cluster_id="compute_cluster_id", user_id="user_id")
+            >>> nodepool = compute_cluster.create_nodepool(config_filepath="config.yml")
+        """
+
+        if config_filepath is not None:
+            assert nodepool_config is None, (
+                "nodepool_config has to be None if config_filepath is provided"
+            )
+
+            if not os.path.exists(config_filepath):
+                raise UserError(f"Nodepool config file not found at {config_filepath}")
+            with open(config_filepath, "r") as file:
+                nodepool_config = yaml.safe_load(file)
+        elif nodepool_config is not None:
+            assert isinstance(nodepool_config, dict), (
+                "nodepool_config should be a dictionary if provided."
+            )
+        else:
+            raise AssertionError("Either config_filepath or nodepool_config must be provided.")
+
+        nodepool_config = self._process_nodepool_config(nodepool_config)
+
+        if 'id' in nodepool_config:
+            if nodepool_id is None:
+                nodepool_id = nodepool_config['id']
+            nodepool_config.pop('id')
+
+        request = service_pb2.PostNodepoolsRequest(
+            user_app_id=self.user_app_id,
+            compute_cluster_id=self.id,
+            nodepools=[resources_pb2.Nodepool(id=nodepool_id, **nodepool_config)],
+        )
+        response = self._grpc_request(self.STUB.PostNodepools, request)
+        if response.status.code != status_code_pb2.SUCCESS:
+            raise Exception(response.status)
+        self.logger.info("\nNodepool created\n%s", response.status)
+
+        dict_response = MessageToDict(response.nodepools[0], preserving_proto_field_name=True)
+        kwargs = self.process_response_keys(dict_response, 'nodepool')
+        return Nodepool.from_auth_helper(auth=self.auth_helper, **kwargs)
+
+    def nodepool(self, nodepool_id: str) -> Nodepool:
+        """Returns a Nodepool object for the existing nodepool ID.
+
+        Args:
+            nodepool_id (str): The nodepool ID for the nodepool to interact with.
+
+        Returns:
+            Nodepool: A Nodepool object for the existing nodepool ID.
+
+        Example:
+            >>> from clarifai.client.compute_cluster import ComputeCluster
+            >>> compute_cluster = ComputeCluster(compute_cluster_id="compute_cluster_id", user_id="user_id")
+            >>> nodepool = compute_cluster.nodepool(nodepool_id="nodepool_id")
+        """
+        request = service_pb2.GetNodepoolRequest(
+            user_app_id=self.user_app_id, compute_cluster_id=self.id, nodepool_id=nodepool_id
+        )
+        response = self._grpc_request(self.STUB.GetNodepool, request)
+
+        if response.status.code != status_code_pb2.SUCCESS:
+            raise Exception(response.status)
+        dict_response = MessageToDict(response, preserving_proto_field_name=True)
+
+        kwargs = self.process_response_keys(dict_response['nodepool'], 'nodepool')
+        return Nodepool.from_auth_helper(auth=self.auth_helper, **kwargs)
+
+    def delete_nodepools(self, nodepool_ids: List[str]) -> None:
+        """Deletes list of nodepools for the compute cluster.
+
+        Args:
+            nodepool_ids (List[str]): The nodepool IDs of the compute cluster to delete.
+
+        Example:
+            >>> from clarifai.client.compute_cluster import ComputeCluster
+            >>> compute_cluster = ComputeCluster(compute_cluster_id="compute_cluster_id", user_id="user_id")
+            >>> compute_cluster.delete_nodepools(nodepool_ids=["nodepool_id1", "nodepool_id2"])
+        """
+        assert isinstance(nodepool_ids, list), "nodepool_ids param should be a list"
+
+        request = service_pb2.DeleteNodepoolsRequest(
+            user_app_id=self.user_app_id, compute_cluster_id=self.id, ids=nodepool_ids
+        )
+        response = self._grpc_request(self.STUB.DeleteNodepools, request)
+
+        if response.status.code != status_code_pb2.SUCCESS:
+            raise Exception(response.status)
+        self.logger.info("\nNodepools Deleted\n%s", response.status)
+
+    def __getattr__(self, name):
+        return getattr(self.compute_cluster_info, name)
+
+    def __str__(self):
+        init_params = [param for param in self.kwargs.keys()]
+        attribute_strings = [
+            f"{param}={getattr(self.compute_cluster_info, param)}"
+            for param in init_params
+            if hasattr(self.compute_cluster_info, param)
+        ]
+        return f"Clarifai Compute Cluster Details: \n{', '.join(attribute_strings)}\n"