Fix task.upload_artifact non standard dictionary will now revert to pickle. Issue #452

This commit is contained in:
allegroai 2021-10-02 21:42:18 +03:00
parent 7d1ebb360d
commit 35fe1c418c

View File

@ -341,6 +341,7 @@ class Artifacts(object):
except Exception: except Exception:
pass pass
store_as_pickle = False
artifact_type_data = tasks.ArtifactTypeData() artifact_type_data = tasks.ArtifactTypeData()
artifact_type_data.preview = '' artifact_type_data.preview = ''
override_filename_in_uri = None override_filename_in_uri = None
@ -380,21 +381,32 @@ class Artifacts(object):
elif isinstance(artifact_object, dict): elif isinstance(artifact_object, dict):
artifact_type = 'JSON' artifact_type = 'JSON'
artifact_type_data.content_type = 'application/json' artifact_type_data.content_type = 'application/json'
json_text = json.dumps(artifact_object, sort_keys=True, indent=4) # noinspection PyBroadException
override_filename_ext_in_uri = '.json' try:
override_filename_in_uri = name + override_filename_ext_in_uri json_text = json.dumps(artifact_object, sort_keys=True, indent=4)
fd, local_filename = mkstemp(prefix=quote(name, safe="") + '.', suffix=override_filename_ext_in_uri) except Exception:
os.write(fd, bytes(json_text.encode())) if not auto_pickle:
os.close(fd) raise
preview = preview or json_text LoggerRoot.get_base_logger().warning(
if len(preview) < self.max_preview_size_bytes: "JSON serialization of artifact \'{}\' failed, reverting to pickle".format(name))
artifact_type_data.preview = preview store_as_pickle = True
else: json_text = None
artifact_type_data.preview = '# full json too large to store, storing first {}kb\n{}'.format(
self.max_preview_size_bytes//1024, preview[:self.max_preview_size_bytes]
)
delete_after_upload = True if json_text is not None:
override_filename_ext_in_uri = '.json'
override_filename_in_uri = name + override_filename_ext_in_uri
fd, local_filename = mkstemp(prefix=quote(name, safe="") + '.', suffix=override_filename_ext_in_uri)
os.write(fd, bytes(json_text.encode()))
os.close(fd)
preview = preview or json_text
if len(preview) < self.max_preview_size_bytes:
artifact_type_data.preview = preview
else:
artifact_type_data.preview = '# full json too large to store, storing first {}kb\n{}'.format(
self.max_preview_size_bytes//1024, preview[:self.max_preview_size_bytes]
)
delete_after_upload = True
elif isinstance(artifact_object, pathlib_types): elif isinstance(artifact_object, pathlib_types):
# check if single file # check if single file
artifact_object = Path(artifact_object) artifact_object = Path(artifact_object)
@ -532,6 +544,13 @@ class Artifacts(object):
os.unlink(local_filename) os.unlink(local_filename)
raise raise
elif auto_pickle: elif auto_pickle:
# revert to pickling the object
store_as_pickle = True
else:
raise ValueError("Artifact type {} not supported".format(type(artifact_object)))
# revert to serializing the object with pickle
if store_as_pickle:
# if we are here it means we do not know what to do with the object, so we serialize it with pickle. # if we are here it means we do not know what to do with the object, so we serialize it with pickle.
artifact_type = 'pickle' artifact_type = 'pickle'
artifact_type_data.content_type = 'application/pickle' artifact_type_data.content_type = 'application/pickle'
@ -553,8 +572,6 @@ class Artifacts(object):
# cleanup and raise exception # cleanup and raise exception
os.unlink(local_filename) os.unlink(local_filename)
raise raise
else:
raise ValueError("Artifact type {} not supported".format(type(artifact_object)))
# verify preview not out of scope: # verify preview not out of scope:
if artifact_type_data.preview and len(artifact_type_data.preview) > (self.max_preview_size_bytes+1024): if artifact_type_data.preview and len(artifact_type_data.preview) > (self.max_preview_size_bytes+1024):