send workspace with multipart encoder

Author: kossak

paperspace/commands/jobs.py

@@ -6,7 +6,7 @@
 from paperspace.commands import common
 from paperspace.exceptions import BadResponseError
 from paperspace.utils import get_terminal_lines
-from paperspace.workspace import WorkspaceHandler
+from paperspace.workspace import WorkspaceHandler, MultipartEncoder
 
 
 class JobsCommandBase(common.CommandBase):
@@ -130,32 +130,40 @@ def __init__(self, workspace_handler=None, **kwargs):
 
     def execute(self, json_):
         url = "/jobs/createJob/"
-        files = None
+        data = None
+        self.set_project(json_)
+
         workspace_url = self._workspace_handler.handle(json_)
         if workspace_url:
             if self._workspace_handler.archive_path:
-                archive_basename = self._workspace_handler.archive_basename
-                json_["workspaceFileName"] = archive_basename
-                self.api.headers["Content-Type"] = "multipart/form-data"
-                files = {"file": open(workspace_url, "rb")}
+                data = self._get_multipart_data(json_)
             else:
-                json_["workspaceFileName"] = workspace_url
-
-        self.set_project(json_)
+                data["workspaceFileName"] = workspace_url
 
         self.logger.log("Creating job...")
-        response = self.api.post(url, params=json_, files=files)
+        response = self.api.post(url, params=json_, data=data)
         self._log_message(response,
-                          "Job created",
+                          "Job created - ID: {id}",
                           "Unknown error while creating job")
 
+    def _get_multipart_data(self, json_):
+        archive_basename = self._workspace_handler.archive_basename
+        json_["workspaceFileName"] = archive_basename
+        job_data = {'file': (archive_basename, open(self._workspace_handler.archive_path, 'rb'), 'text/plain')}
+        monitor = MultipartEncoder(job_data).get_monitor()
+        self.api.headers["Content-Type"] = monitor.content_type
+        data = monitor
+        return data
+
     @staticmethod
     def set_project(json_):
-        project_id = json_.get("projectId", json_.get("projectHandle"))
-        if not project_id:
-            json_["project"] = "paperspace-python"
+        if json_.get("projectId"):
+            return
+
+        if json_.get("projectHandle"):
+            json_["projectId"] = json_.pop("projectHandle")
         else:
-            json_["projectId"] = project_id
+            json_["project"] = "paperspace-python"
 
 
 class ArtifactsDestroyCommand(JobsCommandBase):

paperspace/workspace.py

@@ -26,8 +26,10 @@ def _create_callback(encoder_obj):
         bar = progressbar.ProgressBar(max_value=encoder_obj.len)
 
         def callback(monitor):
-            bar.update(monitor.bytes_read)
-
+            if monitor.bytes_read == bar.max_value:
+                bar.finish()
+            else:
+                bar.update(monitor.bytes_read)
         return callback
 
 

tests/functional/test_jobs.py

@@ -342,3 +342,85 @@ def test_should_send_valid_get_request_with_valid_param_for_a_list_of_artifacts_
                                        params={"jobId": job_id,
                                                param: True})
         assert result.exit_code == 0
+
+
+class TestJobsCreate(object):
+    URL = "https://api.paperspace.io"
+    EXPECTED_HEADERS = default_headers.copy()
+    EXPECTED_HEADERS_WITH_CHANGED_API_KEY = default_headers.copy()
+    EXPECTED_HEADERS_WITH_CHANGED_API_KEY["X-API-Key"] = "some_key"
+    BASIC_OPTIONS_COMMAND = [
+        "jobs", "create",
+        "--name", "exp1",
+        "--projectId", "testHandle",
+        "--container", "testContainer",
+        "--machineType", "testType",
+        "--command", "testCommand",
+        "--workspaceUrl", "some-workspace",
+    ]
+    FULL_OPTIONS_COMMAND = [
+        "jobs", "create",
+        "--name", "exp1",
+        "--ports", 4567,
+        "--workspaceUrl", "wsp.url",
+        "--workingDirectory", "/work/dir/",
+        "--artifactDirectory", "/artifact/dir/",
+        "--clusterId", 42,
+        "--experimentEnv", '{"key":"val"}',
+        "--projectId", "testHandle",
+        "--container", "testContainer",
+        "--machineType", "testType",
+        "--command", "testCommand",
+        "--containerUser", "conUser",
+        "--registryUsername", "userName",
+        "--registryPassword", "passwd",
+        "--apiKey", "some_key",
+    ]
+    BASIC_OPTIONS_REQUEST = {
+        "name": u"exp1",
+        "projectId": u"testHandle",
+        "container": u"testContainer",
+        "machineType": u"testType",
+        "command": u"testCommand",
+        "workspaceUrl": u"some-workspace",
+    }
+    FULL_OPTIONS_REQUEST = {
+        "name": u"exp1",
+        "ports": 4567,
+        "workspaceUrl": u"wsp.url",
+        "workingDirectory": u"/work/dir/",
+        "artifactDirectory": u"/artifact/dir/",
+        "clusterId": 42,
+        "experimentEnv": {u"key": u"val"},
+        "projectHandle": u"testHandle",
+        "container": u"testContainer",
+        "machineType": u"testType",
+        "command": u"testCommand",
+        "containerUser": u"conUser",
+        "registryUsername": u"userName",
+        "registryPassword": u"passwd",
+    }
+    RESPONSE_JSON_200 = {"id": "sadkfhlskdjh", "message": "success"}
+    RESPONSE_CONTENT_200 = b'{"handle":"sadkfhlskdjh","message":"success"}\n'
+    EXPECTED_STDOUT = u'Creating job...\nJob created - ID: sadkfhlskdjh\n'
+
+    RESPONSE_JSON_404_PROJECT_NOT_FOUND = {"details": {"handle": "wrong_handle"}, "error": "Project not found"}
+    RESPONSE_CONTENT_404_PROJECT_NOT_FOUND = b'{"details":{"handle":"wrong_handle"},"error":"Project not found"}\n'
+    EXPECTED_STDOUT_PROJECT_NOT_FOUND = "Project not found\nhandle: wrong_handle\n"
+
+    @mock.patch("paperspace.client.requests.post")
+    def test_should_send_proper_data_and_print_message_when_create_job_was_run_with_basic_options(self, post_patched):
+        post_patched.return_value = MockResponse(self.RESPONSE_JSON_200, 200, self.RESPONSE_CONTENT_200)
+
+        runner = CliRunner()
+        result = runner.invoke(cli.cli, self.BASIC_OPTIONS_COMMAND)
+
+        post_patched.assert_called_once_with(self.URL + '/jobs/createJob/',
+                                             headers=self.EXPECTED_HEADERS,
+                                             json=None,
+                                             params=self.BASIC_OPTIONS_REQUEST,
+                                             files=None,
+                                             data=None)
+
+        assert result.output == self.EXPECTED_STDOUT
+        assert result.exit_code == 0