Merge pull request #115 from Paperspace/send-workspace-with-raw-post

Send workspace with raw post

Author: mkulaczkowski

paperspace/cli/jobs.py

@@ -84,7 +84,7 @@ def common_jobs_create_options(f):
         click.option("--useDockerfile", "useDockerfile", help="Flag: using Dockerfile"),
         click.option("--isPreemptible", "isPreemptible", help="Flag: isPreemptible"),
         click.option("--project", "project", help="Project name"),
-        click.option("--projectId", "projectHandle", help="Project ID"),
+        click.option("--projectId", "projectId", help="Project ID"),
         click.option("--startedByUserId", "startedByUserId", help="User ID"),
         click.option("--relDockerfilePath", "relDockerfilePath", help="Relative path to Dockerfile"),
         click.option("--registryUsername", "registryUsername", help="Docker registry username"),

paperspace/client.py

@@ -30,8 +30,8 @@ def get_path(self, url):
 
     def post(self, url, json=None, params=None, files=None, data=None):
         path = self.get_path(url)
-        logger.debug("POST request sent to: {} \n\theaders: {}\n\tjson: {}\n\tparams: {}\n\tdata: {}"
-                     .format(path, self.headers, json, params, data))
+        logger.debug("POST request sent to: {} \n\theaders: {}\n\tjson: {}\n\tparams: {}\n\tfiles: {}\n\tdata: {}"
+                     .format(path, self.headers, json, params, files, data))
         response = requests.post(path, json=json, params=params, headers=self.headers, files=files, data=data)
         logger.debug("Response status code: {}".format(response.status_code))
         logger.debug("Response content: {}".format(response.content))

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):
@@ -108,7 +108,7 @@ def _get_logs(self, job_id, line, limit):
             'jobId': job_id,
             'line': line,
             'limit': limit
-        };
+        }
         return self.api.get(self.base_url, params=params)
 
     def _log_logs_list(self, data, table, table_data, follow):
@@ -149,36 +149,39 @@ def __init__(self, workspace_handler=None, **kwargs):
 
     def execute(self, json_):
         url = "/jobs/createJob/"
-        files = None
+        data = None
+        self.set_project_if_not_provided(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 = self._get_files_dict(workspace_url)
+                data = self._get_multipart_data(json_)
             else:
                 json_["workspaceFileName"] = workspace_url
 
-        self.set_project(json_)
-
-        response = self.api.post(url, params=json_, files=files)
+        self.logger.log("Creating job...")
+        response = self.api.post(url, params=json_, data=data)
         self._log_message(response,
-                          "Job created",
+                          "Job created - ID: {id}",
                           "Unknown error while creating job")
 
-    @staticmethod
-    def _get_files_dict(workspace_url):
-        files = {"file": open(workspace_url, "rb")}
-        return files
+    def _get_multipart_data(self, json_):
+        archive_basename = self._workspace_handler.archive_basename
+        json_["workspaceFileName"] = archive_basename
+        job_data = self._get_files_dict(archive_basename)
+        monitor = MultipartEncoder(job_data).get_monitor()
+        self.api.headers["Content-Type"] = monitor.content_type
+        data = monitor
+        return data
+
+    def _get_files_dict(self, archive_basename):
+        job_data = {'file': (archive_basename, open(self._workspace_handler.archive_path, 'rb'), 'text/plain')}
+        return job_data
 
     @staticmethod
-    def set_project(json_):
-        project_id = json_.get("projectId", json_.get("projectHandle"))
-        if not project_id:
+    def set_project_if_not_provided(json_):
+        if not json_.get("projectId"):
             json_["project"] = "paperspace-python"
-        else:
-            json_["projectId"] = project_id
 
 
 class ArtifactsDestroyCommand(JobsCommandBase):

paperspace/workspace.py

@@ -15,8 +15,8 @@
 
 class MultipartEncoder(object):
     def __init__(self, fields):
-        s3_encoder = encoder.MultipartEncoder(fields=fields)
-        self.monitor = encoder.MultipartEncoderMonitor(s3_encoder, callback=self._create_callback(s3_encoder))
+        mp_encoder = encoder.MultipartEncoder(fields=fields)
+        self.monitor = encoder.MultipartEncoderMonitor(mp_encoder, callback=self._create_callback(mp_encoder))
 
     def get_monitor(self):
         return self.monitor
@@ -26,7 +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
 
@@ -105,13 +108,12 @@ def handle(self, input_data):
 
         # Should be removed as soon it won't be necessary by PS_API
         if workspace_path == 'none':
-            return workspace_path
-
+            return 'none'
         if workspace_archive:
             archive_path = os.path.abspath(workspace_archive)
         else:
             self.logger.log('Archiving your working directory for upload as your experiment workspace...'
-             '(See https://docs.paperspace.com/gradient/experiments/run-experiments for more information.)')
+                            '(See https://docs.paperspace.com/gradient/experiments/run-experiments for more information.)')
             archive_path = self._zip_workspace(workspace_path, ignore_files)
         self.archive_path = archive_path
         self.archive_basename = os.path.basename(archive_path)

tests/functional/test_jobs.py

@@ -366,3 +366,86 @@ 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",
+        "workspaceFileName": 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

tests/functional/test_run.py

@@ -17,27 +17,32 @@ class TestRunCommand(object):
 
     @mock.patch("paperspace.client.requests.post")
     @mock.patch("paperspace.workspace.WorkspaceHandler._zip_workspace")
+    @mock.patch("paperspace.workspace.MultipartEncoder.get_monitor")
     @mock.patch("paperspace.commands.jobs.CreateJobCommand._get_files_dict")
-    def test_run_simple_file_with_args(self, get_files_patched, workspace_zip_patched, post_patched):
+    def test_run_simple_file_with_args(self, get_files_patched, get_moniror_patched, workspace_zip_patched, post_patched):
         get_files_patched.return_value = mock.MagicMock()
         workspace_zip_patched.return_value = '/foo/bar'
         post_patched.return_value = MockResponse(status_code=200)
 
+        mock_monitor = mock.MagicMock()
+        mock_monitor.content_type = "mock/multipart"
+
+        get_moniror_patched.return_value = mock_monitor
+
         runner = CliRunner()
         result = runner.invoke(cli.cli, [self.command_name] + self.common_commands + ["/myscript.py", "a", "b"])
 
         expected_headers = self.headers.copy()
         expected_headers.update({
-            'Content-Type': "multipart/form-data"
+            'Content-Type': "mock/multipart"
         })
         post_patched.assert_called_with(self.url,
                                         params={'name': u'test', 'projectId': u'projectId',
                                                 'workspaceFileName': 'bar',
                                                 'command': 'python{} myscript.py a b'.format(str(sys.version_info[0])),
-                                                'projectHandle': u'projectId',
                                                 'container': u'paperspace/tensorflow-python'},
-                                        data=None,
-                                        files=mock.ANY,
+                                        data=mock.ANY,
+                                        files=None,
                                         headers=expected_headers,
                                         json=None)
 
@@ -55,7 +60,6 @@ def test_run_python_command_with_args_and_no_workspace(self, post_patched):
                                                 'workspaceFileName': 'none',
                                                 'workspace': 'none',
                                                 'command': 'python{} -c print(foo)'.format(str(sys.version_info[0])),
-                                                'projectHandle': u'projectId',
                                                 'container': u'paperspace/tensorflow-python'},
                                         data=None,
                                         files=None,
@@ -79,7 +83,6 @@ def test_run_shell_command_with_args_with_s3_workspace(self, workspace_zip_patch
                                                 'workspaceFileName': 's3://bucket/object',
                                                 'workspaceUrl': 's3://bucket/object',
                                                 'command': 'echo foo',
-                                                'projectHandle': u'projectId',
                                                 'container': u'paperspace/tensorflow-python'},
                                         data=None,
                                         files=None,