feat(models): add model create command (#347)

Roger Lam · web-flow · commit cd9e4c5350f9 · 2021-05-11T21:49:30.000-04:00
This adds the ability to create a model with a dataset ref.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -2,6 +2,10 @@ version: 2.1
 
 orbs:
   release-tools: paperspace/release-tools@0.1.1
+  docker-tools: paperspace/docker-tools@0.0.4
+
+_docker_image: &docker_image paperspace/gradient-sdk
+_workspace_root: &workspace_root .
 
 workflows:
   master:
@@ -21,13 +25,39 @@ workflows:
           filters:
             branches:
               only: master
+      - docker-tools/build_and_push:
+          name: build_and_push_master
+          context: docker-deploy
+          docker_username: ${DOCKER_USERNAME}
+          docker_password: ${DOCKER_PASSWORD}
+          workspace_root: *workspace_root
+          docker_image: *docker_image
+          docker_tag: 0.0.0-latest
+          requires:
+            - test
+          filters:
+            branches:
+              only: master
 
   pr:
     jobs:
       - test:
           filters:
             branches:
               ignore: master
+      - docker-tools/build_and_push:
+          name: build_and_push
+          context: docker-deploy
+          docker_username: ${DOCKER_USERNAME}
+          docker_password: ${DOCKER_PASSWORD}
+          workspace_root: *workspace_root
+          docker_image: *docker_image
+          requires:
+            - test
+          filters:
+            branches:
+              ignore: master
+
 
   release:
     jobs:
@@ -37,6 +67,17 @@ workflows:
               only: /.*/
             branches:
               ignore: /.*/
+      - docker-tools/tag:
+          name: tag
+          context: docker-deploy
+          docker_username: ${DOCKER_USERNAME}
+          docker_password: ${DOCKER_PASSWORD}
+          docker_image: *docker_image
+          filters:
+            tags:
+              only: /.*/
+            branches:
+              ignore: /.*/
 
 executors:
   python-tox:
@@ -99,3 +140,5 @@ jobs:
           name: upload to pypi
           command: |
             python -m twine upload dist/*
+
+  build-docker:
diff --git a/Dockerfile b/Dockerfile
@@ -0,0 +1,14 @@
+FROM python:3.8
+
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir gradient
+
+ENV PAPERSPACE_API_KEY your_api_key_value
+ENV PAPERSPACE_WEB_URL https://console.paperspace.com
+ENV PAPERSPACE_CONFIG_HOST https://api.paperspace.io
+ENV PAPERSPACE_CONFIG_LOG_HOST https://logs.paperspace.io
+ENV PAPERSPACE_CONFIG_EXPERIMENTS_HOST https://services.paperspace.io/experiments/v1/
+ENV PAPERSPACE_CONFIG_EXPERIMENTS_HOST_V2 https://services.paperspace.io/experiments/v2/
+ENV PAPERSPACE_CONFIG_SERVICE_HOST https://services.paperspace.io
+
+ENTRYPOINT ["gradient"]
diff --git a/gradient/api_sdk/clients/model_client.py b/gradient/api_sdk/clients/model_client.py
@@ -61,6 +61,38 @@ def upload(self, path, name, model_type, model_summary=None, notes=None, tags=No
 
         return model_id
 
+    def create(self, name, model_type, dataset_ref, model_summary=None, notes=None, tags=None, project_id=None):
+        """Create model
+
+        :param str name: Model name
+        :param str model_type: Model Type
+        :param str dataset_ref: Dataset ref to associate a model with
+        :param dict|None model_summary: Dictionary describing model parameters like loss, accuracy, etc.
+        :param str|None notes: Optional model description
+        :param list[str] tags: List of tags
+        :param str|None project_id: ID of a project
+
+        :return: ID of new model
+        :rtype: str
+        """
+
+        model = models.Model(
+            name=name,
+            model_type=model_type,
+            dataset_ref=dataset_ref,
+            summary=json.dumps(model_summary) if model_summary else None,
+            notes=notes,
+            project_id=project_id,
+        )
+
+        repository = self.build_repository(repositories.CreateModel)
+        model_id = repository.create(model)
+
+        if tags:
+            self.add_tags(entity_id=model_id, tags=tags)
+
+        return model_id
+
     def get(self, model_id):
         """Get model instance
 
diff --git a/gradient/api_sdk/models/model.py b/gradient/api_sdk/models/model.py
@@ -18,6 +18,7 @@ class Model(object):
     :param str deployment_state:
     :param str summary:
     :param str detail:
+    :param str dataset_ref:
     """
     id = attr.ib(type=str, default=None)
     name = attr.ib(type=str, default=None)
@@ -32,6 +33,7 @@ class Model(object):
     summary = attr.ib(type=dict, default=None)
     detail = attr.ib(type=dict, default=None)
     notes = attr.ib(type=str, default=None)
+    dataset_ref = attr.ib(type=str, default=None)
 
 
 @attr.s
diff --git a/gradient/api_sdk/repositories/__init__.py b/gradient/api_sdk/repositories/__init__.py
@@ -15,7 +15,7 @@
 from .machine_types import ListMachineTypes
 from .machines import CheckMachineAvailability, CreateMachine, CreateResource, StartMachine, StopMachine, \
     RestartMachine, GetMachine, UpdateMachine, GetMachineUtilization
-from .models import DeleteModel, ListModels, UploadModel, GetModel, ListModelFiles
+from .models import DeleteModel, ListModels, UploadModel, GetModel, ListModelFiles, CreateModel
 from .notebooks import CreateNotebook, DeleteNotebook, GetNotebook, ListNotebooks, GetNotebookMetrics, ListNotebookMetrics, \
     StreamNotebookMetrics, StopNotebook, StartNotebook, ForkNotebook, ListNotebookArtifacts, ListNotebookLogs
 from .projects import CreateProject, ListProjects, DeleteProject, GetProject
diff --git a/gradient/api_sdk/repositories/models.py b/gradient/api_sdk/repositories/models.py
@@ -107,6 +107,24 @@ def _delete_model(self, model_id):
         repository.delete(model_id)
 
 
+class CreateModel(GetBaseModelsApiUrlMixin, CreateResource):
+    SERIALIZER_CLS = serializers.Model
+    HANDLE_FIELD = "id"
+
+    def get_request_url(self, **kwargs):
+        return "/mlModels/createModelV2"
+
+    def _get_request_params(self, kwargs):
+        return kwargs
+
+    def _get_request_json(self, instance_dict):
+        return None
+
+    def create(self, instance, data=None, path=None):
+        model_id = super(CreateModel, self).create(instance, data=data, path=path)
+        return model_id
+
+
 class GetModel(GetBaseModelsApiUrlMixin, GetResource):
     SERIALIZER_CLS = serializers.Model
 
diff --git a/gradient/api_sdk/repositories/workflows.py b/gradient/api_sdk/repositories/workflows.py
@@ -125,7 +125,6 @@ def _send_create_request(self, **kwargs):
 
         response = client.post(url, json=json_)
         gradient_response = http_client.GradientResponse.interpret_response(response)
-
         json_formatted_str = json.dumps(gradient_response.data, indent=4)
         return gradient_response
 
diff --git a/gradient/api_sdk/serializers/model.py b/gradient/api_sdk/serializers/model.py
@@ -19,6 +19,7 @@ class Model(BaseSchema):
     summary = marshmallow.fields.Dict()
     detail = marshmallow.fields.Dict()
     notes = marshmallow.fields.Str()
+    dataset_ref = marshmallow.fields.Str(dump_to="datasetRef", load_from="datasetRef")
 
 
 class ModelFileSchema(BaseSchema):
diff --git a/gradient/cli/models.py b/gradient/cli/models.py
@@ -60,6 +60,69 @@ def delete_model(api_key, model_id, options_file):
     command.execute(model_id=model_id)
 
 
+@models_group.command("create", help="Create a model from an url or dataset id")
+@click.option(
+    "--name",
+    "name",
+    required=True,
+    help="Model name",
+    cls=common.GradientOption,
+)
+@click.option(
+    "--modelType",
+    "model_type",
+    required=True,
+    type=ChoiceType(constants.MODEL_TYPES_MAP, case_sensitive=False),
+    help="Model type",
+    cls=common.GradientOption,
+)
+@click.option(
+    "--datasetRef",
+    "dataset_ref",
+    required=True,
+    help="Dataset ref to associate a model with",
+    cls=common.GradientOption,
+)
+@click.option(
+    "--projectId",
+    "project_id",
+    help="ID of a project",
+    cls=common.GradientOption,
+)
+@click.option(
+    "--modelSummary",
+    "model_summary",
+    type=json_string,
+    help="Model summary",
+    cls=common.GradientOption,
+)
+@click.option(
+    "--notes",
+    "notes",
+    help="Additional notes",
+    cls=common.GradientOption,
+)
+@click.option(
+    "--tag",
+    "tags",
+    multiple=True,
+    help="One or many tags that you want to add to experiment",
+    cls=common.GradientOption
+)
+@click.option(
+    "--tags",
+    "tags_comma",
+    help="Separated by comma tags that you want add to experiment",
+    cls=common.GradientOption
+)
+@common.api_key_option
+@common.options_file
+def create_model(api_key, options_file, **model):
+    model["tags"] = validate_comma_split_option(model.pop("tags_comma"), model.pop("tags"))
+    command = models_commands.CreateModel(api_key=api_key)
+    command.execute(**model)
+
+
 @models_group.command("upload", help="Upload a model file or directory")
 @click.argument(
     "PATH",
diff --git a/gradient/commands/models.py b/gradient/commands/models.py
@@ -49,6 +49,16 @@ def execute(self, model_id, *args, **kwargs):
         self.logger.log("Model deleted")
 
 
+class CreateModel(GetModelsClientMixin, BaseCommand):
+    SPINNER_MESSAGE = "Creating model"
+
+    def execute(self, **kwargs):
+        with halo.Halo(text=self.SPINNER_MESSAGE, spinner="dots"):
+            model_id = self.client.create(**kwargs)
+
+        self.logger.log("Model created with ID: {}".format(model_id))
+
+
 class UploadModel(GetModelsClientMixin, BaseCommand):
     SPINNER_MESSAGE = "Uploading model"
 
diff --git a/tests/functional/test_models.py b/tests/functional/test_models.py
@@ -233,6 +233,78 @@ def test_should_print_proper_message_when_wrong_api_key_was_used(self, post_patc
 
         assert result.output == "Failed to delete resource: Invalid API token\n"
 
+class TestModelCreate(object):
+    URL = "https://api.paperspace.io/mlModels/createModelV2"
+    BASE_COMMAND = [
+        "models", "create",
+        "--name", "some_name",
+        "--modelType", "custom",
+        "--datasetRef", "dsr8k5qzn401lb5:latest",
+    ]
+    BASE_PARAMS = {
+        "name": "some_name",
+        "modelType": "Custom",
+        "datasetRef": "dsr8k5qzn401lb5:latest"
+    }
+    COMMAND_WITH_ALL_OPTIONS = [
+        "models", "create",
+        "--name", "some_name",
+        "--modelType", "tensorflow",
+        "--modelSummary", """{"key": "value"}""",
+        "--datasetRef", "dsr8k5qzn401lb5:latest",
+        "--notes", "some notes",
+        "--projectId", "some_project_id",
+    ]
+    ALL_OPTIONS_PARAMS = {
+        "name": "some_name",
+        "modelType": "Tensorflow",
+        "summary": """{"key": "value"}""",
+        "notes": "some notes",
+        "projectId": "some_project_id",
+        "datasetRef": "dsr8k5qzn401lb5:latest",
+    }
+
+
+    EXPECTED_STDOUT = "Model created with ID: some_model_id\n"
+
+    CREATE_MODEL_V2_REPONSE = example_responses.MODEL_CREATE_RESPONSE_JSON_V2
+
+    @mock.patch("gradient.api_sdk.clients.http_client.requests.post")
+    def test_should_send_post_request_when_models_create_command_was_used_with_basic_options(
+            self, post_patched):
+        post_patched.return_value = MockResponse(self.CREATE_MODEL_V2_REPONSE)
+
+        runner = CliRunner()
+        result = runner.invoke(cli.cli, self.BASE_COMMAND)
+
+        assert result.output == self.EXPECTED_STDOUT, result.exc_info
+        post_patched.assert_has_calls([
+            mock.call(self.URL,
+                        headers=EXPECTED_HEADERS,
+                        json=None,
+                        files=None,
+                        data=None,
+                        params=self.BASE_PARAMS),
+        ])
+
+    @mock.patch("gradient.api_sdk.clients.http_client.requests.post")
+    def test_should_send_post_request_when_models_update_command_was_used_with_all_options(
+            self, post_patched):
+        post_patched.return_value = MockResponse(self.CREATE_MODEL_V2_REPONSE)
+
+        runner = CliRunner()
+        result = runner.invoke(cli.cli, self.COMMAND_WITH_ALL_OPTIONS)
+
+        assert result.output == self.EXPECTED_STDOUT, result.exc_info
+        post_patched.assert_has_calls([
+            mock.call(self.URL,
+                        headers=EXPECTED_HEADERS,
+                        json=None,
+                        files=None,
+                        data=None,
+                        params=self.ALL_OPTIONS_PARAMS),
+        ])
+
 
 class TestModelUpload(object):
     URL = "https://api.paperspace.io/mlModels/createModelV2"
