def test_upload_multiple_files_in_file_set_attribute(self, upload_raw_data_mock):
# given
exp_uuid = str(uuid.uuid4())
swagger_mock = self._get_swagger_mock()
upload_raw_data_mock.return_value = b"null"
swagger_mock.api.getUploadConfig.return_value.response.return_value.result.chunkSize = (
10
)
# when
with NamedTemporaryFile("w") as temp_file_1:
with NamedTemporaryFile("w") as temp_file_2:
upload_file_set_attribute(
swagger_client=swagger_mock,
container_id=exp_uuid,
attribute="some/attribute",
file_globs=[temp_file_1.name, temp_file_2.name],
reset=True,
multipart_config=self.multipart_config,
)
# then
upload_raw_data_mock.assert_called_once_with(
http_client=swagger_mock.swagger_spec.http_client,
url="https://ui.neptune.ai/uploadFileSetTar",
data=mock.ANY,
headers={"Content-Type": "application/octet-stream"},
query_params={
"experimentId": str(exp_uuid),
"attribute": "some/attribute",
"reset": "True",
},
)
def test_upload_single_file_in_file_set_attribute(self, upload_loop_chunk_mock):
# given
exp_uuid = uuid.uuid4()
swagger_mock = self._get_swagger_mock()
upload_loop_chunk_mock.return_value = b"null"
chunk_size = 5 * 1024 * 1024
swagger_mock.api.getUploadConfig.return_value.response.return_value.result.chunkSize = (
chunk_size
)
# when
with NamedTemporaryFile("w") as temp_file:
with open(temp_file.name, "wb") as handler:
handler.write(self.get_random_bytes(2 * chunk_size))
upload_file_set_attribute(
swagger_client=swagger_mock,
container_id=str(exp_uuid),
attribute="some/attribute",
file_globs=[temp_file.name],
reset=True,
multipart_config=self.multipart_config,
)
# then
upload_loop_chunk_mock.assert_has_calls(
[
call(
mock.ANY,
mock.ANY,
http_client=swagger_mock.swagger_spec.http_client,
query_params={
"experimentId": str(exp_uuid),
"attribute": "some/attribute",
"reset": "True",
"path": os.path.basename(temp_file.name),
},
url="https://ui.neptune.ai/uploadFileSetChunk",
),
call(
mock.ANY,
mock.ANY,
http_client=swagger_mock.swagger_spec.http_client,
query_params={
"experimentId": str(exp_uuid),
"attribute": "some/attribute",
"reset": "False",
"path": os.path.basename(temp_file.name),
},
url="https://ui.neptune.ai/uploadFileSetChunk",
),
]
)
def _execute_upload_operation(self, experiment: Experiment,
upload_operation: alpha_operation.Operation):
experiment_id = experiment.internal_id
try:
if isinstance(upload_operation, alpha_operation.UploadFile):
alpha_hosted_file_operations.upload_file_attribute(
swagger_client=self.leaderboard_swagger_client,
container_id=experiment_id,
attribute=alpha_path_utils.path_to_str(
upload_operation.path),
source=upload_operation.file_path,
ext=upload_operation.ext,
multipart_config=self._client_config.multipart_config,
)
elif isinstance(upload_operation,
alpha_operation.UploadFileContent):
alpha_hosted_file_operations.upload_file_attribute(
swagger_client=self.leaderboard_swagger_client,
container_id=experiment_id,
attribute=alpha_path_utils.path_to_str(
upload_operation.path),
source=base64_decode(upload_operation.file_content),
ext=upload_operation.ext,
multipart_config=self._client_config.multipart_config,
)
elif isinstance(upload_operation, alpha_operation.UploadFileSet):
alpha_hosted_file_operations.upload_file_set_attribute(
swagger_client=self.leaderboard_swagger_client,
container_id=experiment_id,
attribute=alpha_path_utils.path_to_str(
upload_operation.path),
file_globs=upload_operation.file_globs,
reset=upload_operation.reset,
multipart_config=self._client_config.multipart_config,
)
else:
raise NeptuneException(
"Upload operation in neither File or FileSet")
except alpha_exceptions.NeptuneException as e:
raise NeptuneException(e) from e
return None
def test_upload_single_small_file_in_file_set_attribute(self, upload_raw_data):
# given
exp_uuid = uuid.uuid4()
swagger_mock = self._get_swagger_mock()
upload_raw_data.return_value = json.dumps(
{
"errors": [],
}
)
data = b"testdata"
# when
with create_file(content=data, binary_mode=True) as filename:
upload_file_set_attribute(
swagger_client=swagger_mock,
container_id=str(exp_uuid),
attribute="some/attribute",
file_globs=[filename],
reset=True,
multipart_config=self.multipart_config,
)
# then
swagger_mock.api.fileSetFileMultipartUploadStart.assert_not_called()
swagger_mock.api.fileSetFileMultipartUploadFinish.assert_not_called()
swagger_mock.api.fileSetFileMultipartUploadPart.assert_not_called()
swagger_mock.api.fileSetFileUpload.assert_not_called()
swagger_mock.api.fileAtomMultipartUploadStart.assert_not_called()
swagger_mock.api.fileAtomMultipartUploadFinish.assert_not_called()
swagger_mock.api.fileAtomMultipartUploadPart.assert_not_called()
swagger_mock.api.fileAtomUpload.assert_not_called()
upload_raw_data.assert_called_once_with(
data=data,
http_client=swagger_mock.swagger_spec.http_client,
url="https://ui.neptune.ai/attributes/storage/fileset/upload",
query_params={
"subPath": os.path.basename(filename),
"experimentIdentifier": str(exp_uuid),
"attribute": "some/attribute",
},
)
def test_missing_files_or_directory(self, upload_raw_data_mock):
# given
exp_uuid = str(uuid.uuid4())
swagger_mock = self._get_swagger_mock()
upload_raw_data_mock.return_value = b"null"
# when
with NamedTemporaryFile("w") as temp_file_1:
with NamedTemporaryFile("w") as temp_file_2:
with TemporaryDirectory() as temp_dir:
upload_file_set_attribute(
swagger_client=swagger_mock,
container_id=exp_uuid,
attribute="some/attribute",
file_globs=[
temp_file_1.name,
temp_file_2.name,
os.path.abspath("missing_file"),
temp_dir,
],
reset=True,
multipart_config=self.multipart_config,
)
# then
upload_raw_data_mock.assert_called_once_with(
http_client=swagger_mock.swagger_spec.http_client,
url="https://ui.neptune.ai/uploadFileSetTar",
data=mock.ANY,
headers={"Content-Type": "application/octet-stream"},
query_params={
"experimentId": str(exp_uuid),
"attribute": "some/attribute",
"reset": "True",
},
)
def _execute_upload_operations(
self,
container_id: str,
container_type: ContainerType,
upload_operations: List[Operation],
) -> List[NeptuneException]:
errors = list()
if self._client_config.has_feature(OptionalFeatures.MULTIPART_UPLOAD):
multipart_config = self._client_config.multipart_config
# collect delete operations and execute them
attributes_to_reset = [
DeleteAttribute(op.path)
for op in upload_operations
if isinstance(op, UploadFileSet) and op.reset
]
if attributes_to_reset:
errors.extend(
self._execute_operations(
container_id, container_type, operations=attributes_to_reset
)
)
else:
multipart_config = None
for op in upload_operations:
if isinstance(op, UploadFile):
upload_errors = upload_file_attribute(
swagger_client=self.leaderboard_client,
container_id=container_id,
attribute=path_to_str(op.path),
source=op.file_path,
ext=op.ext,
multipart_config=multipart_config,
)
if upload_errors:
errors.extend(upload_errors)
elif isinstance(op, UploadFileContent):
upload_errors = upload_file_attribute(
swagger_client=self.leaderboard_client,
container_id=container_id,
attribute=path_to_str(op.path),
source=base64_decode(op.file_content),
ext=op.ext,
multipart_config=multipart_config,
)
if upload_errors:
errors.extend(upload_errors)
elif isinstance(op, UploadFileSet):
upload_errors = upload_file_set_attribute(
swagger_client=self.leaderboard_client,
container_id=container_id,
attribute=path_to_str(op.path),
file_globs=op.file_globs,
reset=op.reset,
multipart_config=multipart_config,
)
if upload_errors:
errors.extend(upload_errors)
else:
raise InternalClientError("Upload operation in neither File or FileSet")
return errors
def test_upload_single_big_file_in_file_set_attribute(self, upload_raw_data):
# given
exp_uuid = uuid.uuid4()
swagger_mock = self._get_swagger_mock()
upload_id = "placeholder"
set_expected_result(
swagger_mock.api.fileSetFileMultipartUploadStart,
{
"uploadId": upload_id,
"errors": [],
},
)
upload_raw_data.return_value = json.dumps(
{
"errors": [],
}
)
data = self.get_random_bytes(8 * 2 ** 20) # 8 MB
chunk_size = self.multipart_config.min_chunk_size
# when
with create_file(content=data, binary_mode=True) as filename:
upload_file_set_attribute(
swagger_client=swagger_mock,
container_id=str(exp_uuid),
attribute="some/attribute",
file_globs=[filename],
reset=True,
multipart_config=self.multipart_config,
)
# then
swagger_mock.api.fileSetFileMultipartUploadPart.assert_not_called()
swagger_mock.api.fileSetFileUpload.assert_not_called()
swagger_mock.api.fileAtomMultipartUploadStart.assert_not_called()
swagger_mock.api.fileAtomMultipartUploadFinish.assert_not_called()
swagger_mock.api.fileAtomMultipartUploadPart.assert_not_called()
swagger_mock.api.fileAtomUpload.assert_not_called()
swagger_mock.api.fileSetFileMultipartUploadStart.assert_called_once_with(
attribute="some/attribute",
experimentIdentifier=str(exp_uuid),
totalLength=len(data),
subPath=os.path.basename(filename),
)
swagger_mock.api.fileSetFileMultipartUploadFinish.assert_called_once_with(
attribute="some/attribute",
experimentIdentifier=str(exp_uuid),
subPath=os.path.basename(filename),
uploadId=upload_id,
)
upload_raw_data.assert_has_calls(
[
call(
data=data[:chunk_size],
http_client=swagger_mock.swagger_spec.http_client,
url="https://ui.neptune.ai/attributes/storage/fileset/upload/part",
headers={"X-Range": f"bytes=0-{chunk_size - 1}/{len(data)}"},
query_params={
"uploadPartIdx": 0,
"uploadId": upload_id,
"subPath": os.path.basename(filename),
"experimentIdentifier": str(exp_uuid),
"attribute": "some/attribute",
},
),
call(
data=data[chunk_size:],
http_client=swagger_mock.swagger_spec.http_client,
url="https://ui.neptune.ai/attributes/storage/fileset/upload/part",
headers={
"X-Range": f"bytes={chunk_size}-{len(data) - 1}/{len(data)}"
},
query_params={
"uploadPartIdx": 1,
"uploadId": upload_id,
"subPath": os.path.basename(filename),
"experimentIdentifier": str(exp_uuid),
"attribute": "some/attribute",
},
),
]
)