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:
pass
store_as_pickle = False
artifact_type_data = tasks.ArtifactTypeData()
artifact_type_data.preview = ''
override_filename_in_uri = None
@ -380,7 +381,18 @@ class Artifacts(object):
elif isinstance(artifact_object, dict):
artifact_type = 'JSON'
artifact_type_data.content_type = 'application/json'
# noinspection PyBroadException
try:
json_text = json.dumps(artifact_object, sort_keys=True, indent=4)
except Exception:
if not auto_pickle:
raise
LoggerRoot.get_base_logger().warning(
"JSON serialization of artifact \'{}\' failed, reverting to pickle".format(name))
store_as_pickle = True
json_text = None
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)
@ -532,6 +544,13 @@ class Artifacts(object):
os.unlink(local_filename)
raise
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.
artifact_type = 'pickle'
artifact_type_data.content_type = 'application/pickle'
@ -553,8 +572,6 @@ class Artifacts(object):
# cleanup and raise exception
os.unlink(local_filename)
raise
else:
raise ValueError("Artifact type {} not supported".format(type(artifact_object)))
# verify preview not out of scope:
if artifact_type_data.preview and len(artifact_type_data.preview) > (self.max_preview_size_bytes+1024):