diff --git a/examples/artifacts_toy.py b/examples/artifacts_toy.py index da6f8e8e..b9558f32 100644 --- a/examples/artifacts_toy.py +++ b/examples/artifacts_toy.py @@ -31,7 +31,8 @@ task.upload_artifact('Numpy Eye', np.eye(100, 100)) # add and upload Image (stored as .png file) im = Image.open('samples/dancing.jpg') task.upload_artifact('pillow_image', im) - +# add and upload a folder, artifact_object should be the folder path +task.upload_artifact('local folder', artifact_object='samples/') # do something sleep(1.) diff --git a/trains/binding/artifacts.py b/trains/binding/artifacts.py index 318dafec..92eb38f8 100644 --- a/trains/binding/artifacts.py +++ b/trains/binding/artifacts.py @@ -114,17 +114,36 @@ class Artifact(object): self._metadata = dict(artifact_api_object.display_data) if artifact_api_object.display_data else {} self._preview = artifact_api_object.type_data.preview if artifact_api_object.type_data else None - def get_local_copy(self): + def get_local_copy(self, extract_archive=True): """ + :param bool extract_archive: If True and artifact is of type 'archive' (compressed folder) + The returned path will be a temporary folder containing the archive content :return: a local path to a downloaded copy of the artifact """ from trains.storage.helper import StorageHelper - return StorageHelper.get_local_copy(self.url) + local_path = StorageHelper.get_local_copy(self.url) + if local_path and extract_archive and self.type == 'archive': + try: + temp_folder = mkdtemp(prefix='artifact_', suffix='.archive_'+self.name) + ZipFile(local_path).extractall(path=temp_folder) + except Exception: + try: + Path(temp_folder).rmdir() + except Exception: + pass + return local_path + try: + Path(local_path).unlink() + except Exception: + pass + return temp_folder + + return local_path def __repr__(self): return str({'name': self.name, 'size': self.size, 'type': self.type, 'mode': self.mode, 'url': self.url, 'hash': self.hash, 'timestamp': self.timestamp, - 'metadata': self.metadata, 'preview': self.preview,}) + 'metadata': self.metadata, 'preview': self.preview, }) class Artifacts(object): @@ -294,7 +313,7 @@ class Artifacts(object): os.close(fd) artifact_object = zip_file - artifact_type = 'zip' + artifact_type = 'archive' artifact_type_data.content_type = mimetypes.guess_type(artifact_object)[0] local_filename = artifact_object delete_after_upload = True diff --git a/trains/storage/helper.py b/trains/storage/helper.py index 9cb5b848..5881f071 100644 --- a/trains/storage/helper.py +++ b/trains/storage/helper.py @@ -569,6 +569,8 @@ class StorageHelper(object): # we won't think we have the entire file temp_local_path = '{}_{}{}'.format(local_path, time(), self._temp_download_suffix) obj = self._get_object(remote_path) + if not obj: + return None # object size in bytes total_size_mb = -1 diff --git a/trains/task.py b/trains/task.py index fe27529f..c91a9b7e 100644 --- a/trains/task.py +++ b/trains/task.py @@ -686,10 +686,10 @@ class Task(_Task): :param str name: Artifact name. Notice! it will override previous artifact if name already exists :param object artifact_object: Artifact object to upload. Currently supports: - string / pathlib2.Path are treated as path to artifact file to upload - If wildcard or a folder is passed, zip file containing the local files will be created and uploaded. - - dict will be stored as .json, - - pandas.DataFrame will be stored as .csv.gz (compressed CSV file), - - numpy.ndarray will be stored as .npz, + If wildcard or a folder is passed, zip file containing the local files will be created and uploaded + - dict will be stored as .json file and uploaded + - pandas.DataFrame will be stored as .csv.gz (compressed CSV file) and uploaded + - numpy.ndarray will be stored as .npz and uploaded - PIL.Image will be stored to .png file and uploaded :param dict metadata: Simple key/value dictionary to store on the artifact :param bool delete_after_upload: If True local artifact will be deleted