Merge pull request #270 from Paperspace/model-upload-fix

model upload fix

Author: kossak

gradient/api_sdk/s3_uploader.py

@@ -79,6 +79,22 @@ def get_bucket_url(bucket_name, s3_fields):
         return url
 
 
+class S3PutFileUploader(S3FileUploader):
+    def _upload(self, url, data):
+        """Send data to S3 and raise exception if it was not a success
+
+        :param str url:
+        :param encoder.MultipartEncoderMonitor data:
+        """
+        file_path = data.encoder.fields['file'][0]
+        client = self._get_client(url)
+        client.headers = {"Content-Type": mimetypes.guess_type(file_path)[0] or ""}
+
+        response = client.put("", data=data)
+        if not response.ok:
+            raise sdk_exceptions.S3UploadFailedError(response)
+
+
 class S3ProjectFileUploader(object):
     def __init__(self, api_key, s3uploader=None, logger=None, ps_client_name=None):
         """
@@ -164,7 +180,7 @@ def __init__(self, api_key, multipart_encoder_cls=None, logger=None, ps_client_n
             api_key=api_key,
             ps_client_name=ps_client_name,
         )
-        self.s3uploader = s3uploader or S3FileUploader(
+        self.s3uploader = s3uploader or S3PutFileUploader(
             logger=self.logger,
             ps_client_name=ps_client_name,
             multipart_encoder_cls=self.multipart_encoder_cls

tests/functional/test_models.py

@@ -303,10 +303,12 @@ class TestModelUpload(object):
     UPDATE_TAGS_RESPONSE_JSON_200 = example_responses.UPDATE_TAGS_RESPONSE
 
     @mock.patch("gradient.api_sdk.clients.http_client.requests.get")
+    @mock.patch("gradient.api_sdk.clients.http_client.requests.put")
     @mock.patch("gradient.api_sdk.clients.http_client.requests.post")
     def test_should_send_post_request_when_models_update_command_was_used_with_basic_options(
-            self, post_patched, get_patched):
+            self, post_patched, put_patched, get_patched):
         post_patched.return_value = MockResponse(self.CREATE_MODEL_V2_REPONSE)
+        put_patched.return_value = MockResponse()
         get_patched.return_value = MockResponse(self.GET_PRESIGNED_URL_RESPONSE)
 
         runner = CliRunner()
@@ -324,10 +326,11 @@ def test_should_send_post_request_when_models_update_command_was_used_with_basic
                           files=None,
                           data=None,
                           params=self.BASE_PARAMS),
+            ])
+            put_patched.assert_has_calls([
                 mock.call(self.GET_PRESIGNED_URL_RESPONSE,
                           headers={"Content-Type": mock.ANY},
                           json=None,
-                          files=None,
                           params=None,
                           data=mock.ANY)
             ])
@@ -336,15 +339,17 @@ def test_should_send_post_request_when_models_update_command_was_used_with_basic
                                                 params=self.GET_PRESIGNED_URL_PARAMS,
                                                 json=None,
                                                 )
-            assert post_patched.call_args.kwargs["data"].encoder.fields["file"][0] == self.MODEL_FILE
+            assert put_patched.call_args.kwargs["data"].encoder.fields["file"][0] == self.MODEL_FILE
 
             assert EXPECTED_HEADERS["X-API-Key"] != "some_key"
 
     @mock.patch("gradient.api_sdk.clients.http_client.requests.get")
+    @mock.patch("gradient.api_sdk.clients.http_client.requests.put")
     @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, get_patched):
+            self, post_patched, put_patched, get_patched):
         post_patched.return_value = MockResponse(self.CREATE_MODEL_V2_REPONSE)
+        put_patched.return_value = MockResponse()
         get_patched.return_value = MockResponse(self.GET_PRESIGNED_URL_RESPONSE)
 
         runner = CliRunner()
@@ -362,19 +367,19 @@ def test_should_send_post_request_when_models_update_command_was_used_with_all_o
                           files=None,
                           data=None,
                           params=self.ALL_OPTIONS_PARAMS),
+            ])
+            put_patched.assert_has_calls([
                 mock.call(self.GET_PRESIGNED_URL_RESPONSE,
                           headers={"Content-Type": mock.ANY},
                           json=None,
-                          files=None,
                           params=None,
-                          data=mock.ANY)
-            ])
+                          data=mock.ANY)])
             get_patched.assert_called_once_with(self.GET_PRESIGNED_URL,
                                                 headers=EXPECTED_HEADERS,
                                                 params=self.GET_PRESIGNED_URL_PARAMS,
                                                 json=None,
                                                 )
-            assert post_patched.call_args.kwargs["data"].encoder.fields["file"][0] == self.MODEL_FILE
+            assert put_patched.call_args.kwargs["data"].encoder.fields["file"][0] == self.MODEL_FILE
 
             assert EXPECTED_HEADERS["X-API-Key"] != "some_key"
 
@@ -403,20 +408,20 @@ def test_should_replace_api_key_in_headers_when_api_key_parameter_was_used(
                           data=None,
                           params=self.ALL_OPTIONS_PARAMS
                           ),
+            ])
+            put_patched.assert_has_calls([
                 mock.call(self.GET_PRESIGNED_URL_RESPONSE,
                           headers={"Content-Type": mock.ANY},
-                          files=None,
                           json=None,
                           params=None,
                           data=mock.ANY)
             ])
-
             get_patched.assert_called_once_with(self.GET_PRESIGNED_URL,
                                                 headers=EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
                                                 params=self.GET_PRESIGNED_URL_PARAMS,
                                                 json=None,
                                                 )
-            assert post_patched.call_args.kwargs["data"].encoder.fields["file"][0] == self.MODEL_FILE
+            assert put_patched.call_args.kwargs["data"].encoder.fields["file"][0] == self.MODEL_FILE
 
             assert EXPECTED_HEADERS["X-API-Key"] != "some_key"
 
@@ -447,21 +452,18 @@ def test_should_read_options_from_yaml_file(
                     data=None,
                     params=self.ALL_OPTIONS_PARAMS
                 ),
-                mock.call(
-                    self.GET_PRESIGNED_URL_RESPONSE,
-                    headers={'Content-Type': mock.ANY},
-                    json=None,
-                    files=None,
-                    params=None,
-                    data=mock.ANY
-                )
             ])
+            put_patched.assert_called_once_with(self.GET_PRESIGNED_URL_RESPONSE,
+                                                headers={"Content-Type": ""},
+                                                json=None,
+                                                params=None,
+                                                data=mock.ANY)
             get_patched.assert_called_once_with(self.GET_PRESIGNED_URL,
                                                 headers=EXPECTED_HEADERS_WITH_CHANGED_API_KEY,
                                                 params=self.GET_PRESIGNED_URL_PARAMS,
                                                 json=None,
                                                 )
-            assert post_patched.call_args.kwargs["data"].encoder.fields["file"][0] == self.MODEL_FILE
+            assert put_patched.call_args.kwargs["data"].encoder.fields["file"][0] == self.MODEL_FILE
 
             assert EXPECTED_HEADERS["X-API-Key"] != "some_key"
 
@@ -511,15 +513,8 @@ def test_should_send_proper_data_and_tag_machine(
                           data=None,
                           params=self.BASE_PARAMS
                           ),
-                mock.call(
-                    self.GET_PRESIGNED_URL_RESPONSE,
-                    headers={"Content-Type": mock.ANY},
-                    files=None,
-                    json=None,
-                    params=None,
-                    data=mock.ANY,
-                ),
             ])
+
             get_patched.assert_has_calls(
                 [
                     mock.call(
@@ -532,6 +527,13 @@ def test_should_send_proper_data_and_tag_machine(
             )
             put_patched.assert_has_calls(
                 [
+                    mock.call(
+                        self.GET_PRESIGNED_URL_RESPONSE,
+                        headers={"Content-Type": mock.ANY},
+                        json=None,
+                        params=None,
+                        data=mock.ANY,
+                    ),
                     mock.call(
                         self.TAGS_URL,
                         headers=EXPECTED_HEADERS,