Merge branch 'master' of github.com:Paperspace/gradient-cli into PS-12720-Make_workspace_optional_for_v2

Author: BartoszCki

gradient/api_sdk/clients/base_client.py

@@ -1,25 +1,124 @@
+import copy
+
+from gradient.api_sdk.repositories.common import BaseRepository
 from .. import logger as sdk_logger
+from ..repositories.tags import ListTagRepository, UpdateTagRepository
+from ...exceptions import ReceivingDataFailedError
 
 
 class BaseClient(object):
     def __init__(
-            self, api_key,
+            self,
+            api_key,
+            ps_client_name=None,
             logger=sdk_logger.MuteLogger()
     ):
         """
         Base class. All client classes inherit from it.
 
-        An API key can be created at paperspace.com after you sign in to your account. After obtaining it, you can set
-        it in the CLI using the command::
-
-            gradient apiKey XXXXXXXXXXXXXXXXXXX
-
-        or you can provide your API key in any command, for example::
-
-            gradient experiments run ... --apiKey XXXXXXXXXXXXXXXXXXX
-
         :param str api_key: your API key
+        :param str ps_client_name:
         :param sdk_logger.Logger logger:
         """
         self.api_key = api_key
+        self.ps_client_name = ps_client_name
         self.logger = logger
+
+    KNOWN_TAGS_ENTITIES = [
+        "project", "job", "notebook", "experiment", "deployment", "mlModel", "machine",
+    ]
+    entity = ""
+
+    def _validate_entities(self, entity):
+        """
+        Method to validate if passed entity is correct
+        :param entity:
+        :return:
+        """
+        if entity not in self.KNOWN_TAGS_ENTITIES:
+            raise ReceivingDataFailedError("Not known entity type provided")
+
+    @staticmethod
+    def merge_tags(entity_id, entity_tags, new_tags):
+        result_tags = []
+        if entity_tags:
+            entity_tags = entity_tags[0].get(entity_id, [])
+
+            result_tags = entity_tags + new_tags
+        else:
+            result_tags += new_tags
+        return sorted(list(set(result_tags)))
+
+    @staticmethod
+    def diff_tags(entity_id, entity_tags, tags_to_remove):
+        result_tags = []
+        if entity_tags:
+            entity_tags = entity_tags[0].get(entity_id, [])
+            entity_tags = set(entity_tags) - set(tags_to_remove)
+            result_tags = sorted(list(entity_tags))
+
+        return result_tags
+
+    def add_tags(self, entity_id, entity, tags):
+        """
+        Add tags to entity.
+        :param entity_id:
+        :param entity:
+        :param tags:
+        :return:
+        """
+        self._validate_entities(entity)
+
+        list_tag_repository = self.build_repository(ListTagRepository)
+        entity_tags = list_tag_repository.list(entity=entity, entity_ids=[entity_id])
+
+        if entity_tags:
+            tags = self.merge_tags(entity_id, entity_tags, tags)
+
+        update_tag_repository = self.build_repository(UpdateTagRepository)
+        update_tag_repository.update(entity=entity, entity_id=entity_id, tags=tags)
+
+    def remove_tags(self, entity_id, entity, tags):
+        """
+        Remove tags from entity.
+        :param str entity_id:
+        :param str entity:
+        :param list[str] tags: list of tags to remove from entity
+        :return:
+        """
+        self._validate_entities(entity)
+
+        list_tag_repository = self.build_repository(ListTagRepository)
+        entity_tags = list_tag_repository.list(entity=entity, entity_ids=[entity_id])
+
+        if entity_tags:
+            entity_tags = self.diff_tags(entity_id, entity_tags, tags)
+
+            update_tag_repository = self.build_repository(UpdateTagRepository)
+            update_tag_repository.update(entity=entity, entity_id=entity_id, tags=entity_tags)
+
+    def list_tags(self, entity_ids, entity):
+        """
+        List tags for entity
+        :param list[str] entity_ids:
+        :param str entity:
+        :return:
+        """
+        self._validate_entities(entity)
+
+        list_tag_repository = self.build_repository(ListTagRepository)
+        entity_tags = list_tag_repository.list(entity=entity, entity_ids=entity_ids)
+        return entity_tags
+
+    def build_repository(self, repository_class, *args, **kwargs):
+        """
+        :param type[BaseRepository] repository_class:
+        :rtype: BaseRepository
+        """
+
+        if self.ps_client_name is not None and kwargs.get("ps_client_name") is None:
+            kwargs = copy.deepcopy(kwargs)
+            kwargs["ps_client_name"] = self.ps_client_name
+
+        repository = repository_class(*args, api_key=self.api_key, logger=self.logger, **kwargs)
+        return repository

gradient/api_sdk/clients/clusters.py

@@ -0,0 +1,19 @@
+from gradient.api_sdk import repositories
+
+from gradient.api_sdk.clients.base_client import BaseClient
+
+
+class ClustersClient(BaseClient):
+    def list(self, limit=11, offset=0, **kwargs):
+        """
+        Get a list of clusters for your team
+
+        :param int limit: how many element to return on request
+        :param int offset: from what position we should return clusters
+
+        :return: clusters
+        :rtype: list
+        """
+        repository = repositories.ListClusters(api_key=self.api_key, logger=self.logger)
+        clusters = repository.list(limit=limit, offset=offset)
+        return clusters

gradient/api_sdk/clients/deployment_client.py

@@ -24,6 +24,7 @@ class DeploymentsClient(BaseClient):
         )
     """
     HOST_URL = config.config.CONFIG_HOST
+    entity = "deployment"
 
     def create(
             self,
@@ -46,7 +47,7 @@ def create(
             cluster_id=None,
             auth_username=None,
             auth_password=None,
-            use_vpc=False,
+            tags=None,
     ):
         """
         Method to create a Deployment instance.
@@ -89,7 +90,7 @@ def create(
         :param str cluster_id: cluster ID
         :param str auth_username: Username
         :param str auth_password: Password
-        :param bool use_vpc:
+        :param list[str] tags: List of tags
 
         :returns: Created deployment id
         :rtype: str
@@ -117,8 +118,10 @@ def create(
             auth_password=auth_password,
         )
 
-        repository = repositories.CreateDeployment(api_key=self.api_key, logger=self.logger)
-        deployment_id = repository.create(deployment, use_vpc=use_vpc)
+        repository = self.build_repository(repositories.CreateDeployment)
+        deployment_id = repository.create(deployment)
+        if tags:
+            self.add_tags(entity_id=deployment_id, entity=self.entity, tags=tags)
         return deployment_id
 
     def get(self, deployment_id):
@@ -129,11 +132,11 @@ def get(self, deployment_id):
         :return: Deployment instance
         :rtype: models.Deployment
         """
-        repository = repositories.GetDeployment(self.api_key, logger=self.logger)
+        repository = self.build_repository(repositories.GetDeployment)
         deployment = repository.get(deployment_id=deployment_id)
         return deployment
 
-    def start(self, deployment_id, use_vpc=False):
+    def start(self, deployment_id):
         """
         Start deployment
 
@@ -142,13 +145,12 @@ def start(self, deployment_id, use_vpc=False):
             gradient deployments start --id <your-deployment-id>
 
         :param str deployment_id: Deployment ID
-        :param bool use_vpc:
         """
 
-        repository = repositories.StartDeployment(api_key=self.api_key, logger=self.logger)
-        repository.start(deployment_id, use_vpc=use_vpc)
+        repository = self.build_repository(repositories.StartDeployment)
+        repository.start(deployment_id)
 
-    def stop(self, deployment_id, use_vpc=False):
+    def stop(self, deployment_id):
         """
         Stop deployment
 
@@ -157,29 +159,31 @@ def stop(self, deployment_id, use_vpc=False):
             gradient deployments stop --id <your-deployment-id>
 
         :param deployment_id: Deployment ID
-        :param bool use_vpc:
         """
 
-        repository = repositories.StopDeployment(api_key=self.api_key, logger=self.logger)
-        repository.stop(deployment_id, use_vpc=use_vpc)
+        repository = self.build_repository(repositories.StopDeployment)
+        repository.stop(deployment_id)
 
-    def list(self, state=None, project_id=None, model_id=None, use_vpc=False):
+    def list(self, state=None, project_id=None, model_id=None, tags=None):
         """
         List deployments with optional filtering
 
-        :param str state:
-        :param str project_id:
-        :param str model_id:
-        :param bool use_vpc:
+        :param str state: state to filter deployments
+        :param str project_id: project ID to filter deployments
+        :param str model_id: model ID to filter deployments
+        :param list[str]|tuple[str] tags: tags to filter deployments with OR
+
+        :returns: List of Deployment model instances
+        :rtype: list[models.Deployment]
         """
 
-        repository = repositories.ListDeployments(api_key=self.api_key, logger=self.logger)
-        deployments = repository.list(state=state, project_id=project_id, model_id=model_id, use_vpc=use_vpc)
+        repository = self.build_repository(repositories.ListDeployments)
+        deployments = repository.list(state=state, project_id=project_id, model_id=model_id, tags=tags)
         return deployments
 
-    def delete(self, deployment_id, use_vpc=False):
-        repository = repositories.DeleteDeployment(api_key=self.api_key, logger=self.logger)
-        repository.delete(deployment_id, use_vpc=use_vpc)
+    def delete(self, deployment_id):
+        repository = self.build_repository(repositories.DeleteDeployment)
+        repository.delete(deployment_id)
 
     def update(
             self,
@@ -204,7 +208,6 @@ def update(
             cluster_id=None,
             auth_username=None,
             auth_password=None,
-            use_vpc=False,
     ):
         deployment = models.Deployment(
             deployment_type=deployment_type,
@@ -229,5 +232,5 @@ def update(
             auth_password=auth_password,
         )
 
-        repository = repositories.UpdateDeployment(api_key=self.api_key, logger=self.logger)
-        repository.update(deployment_id, deployment, use_vpc=use_vpc)
+        repository = self.build_repository(repositories.UpdateDeployment)
+        repository.update(deployment_id, deployment)