Upload in-memory files using copy stored on disk

CHANGELOG.md

@@ -1,3 +1,8 @@
+## [UNRELEASED]  neptune-client 0.16.11
+
+### Changes
+- Upload in-memory files using copy stored on disk ([#1052](https://github.com/neptune-ai/neptune-client/pull/1052))
+
 ## neptune-client 0.16.10
 
 ### Features

src/neptune/new/attributes/atoms/file.py

@@ -13,20 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import os
 from typing import Optional
 
 from neptune.new.attributes.atoms.atom import Atom
-from neptune.new.internal.operation import (
-    UploadFile,
-    UploadFileContent,
-)
-from neptune.new.internal.utils import (
-    base64_encode,
-    verify_type,
-)
+from neptune.new.internal.operation import UploadFile
+from neptune.new.internal.utils import verify_type
 from neptune.new.types.atoms.file import File as FileVal
-from neptune.new.types.atoms.file import FileType
 
 # pylint: disable=protected-access
 
@@ -35,16 +27,11 @@ class File(Atom):
     def assign(self, value: FileVal, wait: bool = False) -> None:
         verify_type("value", value, FileVal)
 
-        if value.file_type is FileType.LOCAL_FILE:
-            operation = UploadFile(self._path, ext=value.extension, file_path=os.path.abspath(value.path))
-        elif value.file_type is FileType.IN_MEMORY:
-            operation = UploadFileContent(
-                self._path,
-                ext=value.extension,
-                file_content=base64_encode(value.content),
-            )
-        else:
-            raise ValueError(f"Unexpected FileType: {value.file_type}")
+        operation = UploadFile.of_file(
+            value=value,
+            attribute_path=self._path,
+            upload_path=self._container._op_processor._operation_storage.upload_path,
+        )
 
         with self._container.lock():
             self._enqueue_operation(operation, wait)

src/neptune/new/internal/backends/hosted_neptune_backend.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import itertools
 import logging
 import os
 import re
@@ -472,16 +473,22 @@ def execute_operations(
         )
 
         errors.extend(artifact_operations_errors)
-        preprocessed_operations.other_operations.extend(assign_artifact_operations)
 
         errors.extend(
             self._execute_operations(
                 container_id,
                 container_type,
-                operations=preprocessed_operations.other_operations,
+                operations=itertools.chain(assign_artifact_operations, preprocessed_operations.other_operations),
             )
         )
 
+        for op in itertools.chain(
+            preprocessed_operations.upload_operations,
+            assign_artifact_operations,
+            preprocessed_operations.other_operations,
+        ):
+            op.clean()
+
         return (
             operations_preprocessor.processed_ops_count + dropped_count,
             errors,
@@ -607,7 +614,7 @@ def _execute_operations(
         self,
         container_id: UniqueId,
         container_type: ContainerType,
-        operations: List[Operation],
+        operations: Iterable[Operation],
     ) -> List[MetadataInconsistency]:
         kwargs = {
             "experimentId": container_id,

src/neptune/new/internal/operation.py

@@ -14,8 +14,10 @@
 # limitations under the License.
 #
 import abc
+import os
 from dataclasses import dataclass
 from datetime import datetime
+from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Generic,
@@ -30,6 +32,10 @@
 from neptune.common.exceptions import InternalClientError
 from neptune.new.exceptions import MalformedOperation
 from neptune.new.internal.container_type import ContainerType
+from neptune.new.types.atoms.file import (
+    File,
+    FileType,
+)
 
 if TYPE_CHECKING:
     from neptune.new.attributes.attribute import Attribute
@@ -53,6 +59,9 @@ class Operation:
     def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
         pass
 
+    def clean(self):
+        pass
+
     # pylint: disable=unused-argument
     def to_dict(self) -> dict:
         return {"type": self.__class__.__name__, "path": self.path}
@@ -180,6 +189,33 @@ class UploadFile(Operation):
 
     ext: str
     file_path: str
+    clean_after_upload: bool = False
+
+    @classmethod
+    def of_file(cls, value: File, attribute_path: List[str], upload_path: Path):
+        if value.file_type is FileType.LOCAL_FILE:
+            operation = UploadFile(
+                path=attribute_path,
+                ext=value.extension,
+                file_path=os.path.abspath(value.path),
+            )
+        elif value.file_type is FileType.IN_MEMORY:
+            tmp_file_path = cls.get_upload_path(attribute_path, value.extension, upload_path)
+            # pylint: disable=protected-access
+            value._save(tmp_file_path)
+            operation = UploadFile(
+                path=attribute_path,
+                ext=value.extension,
+                file_path=os.path.abspath(tmp_file_path),
+                clean_after_upload=True,
+            )
+        else:
+            raise ValueError(f"Unexpected FileType: {value.file_type}")
+        return operation
+
+    def clean(self):
+        if self.clean_after_upload:
+            os.remove(self.file_path)
 
     def accept(self, visitor: "OperationVisitor[Ret]") -> Ret:
         return visitor.visit_upload_file(self)
@@ -188,11 +224,20 @@ def to_dict(self) -> dict:
         ret = super().to_dict()
         ret["ext"] = self.ext
         ret["file_path"] = self.file_path
+        ret["clean_after_upload"] = self.clean_after_upload
         return ret
 
     @staticmethod
     def from_dict(data: dict) -> "UploadFile":
-        return UploadFile(data["path"], data["ext"], data["file_path"])
+        return UploadFile(data["path"], data["ext"], data["file_path"], data.get("clean_after_upload", False))
+
+    @staticmethod
+    def get_upload_path(attribute_path: List[str], extension: str, upload_path: Path):
+        now = datetime.now()
+        tmp_file_name = (
+            f"{'_'.join(attribute_path)}-{now.timestamp()}-{now.strftime('%Y-%m-%d_%H.%M.%S.%f')}.{extension}"
+        )
+        return upload_path / tmp_file_name
 
 
 @dataclass

src/neptune/new/internal/operation_processors/async_operation_processor.py

@@ -95,7 +95,7 @@ def __init__(
     def _init_data_path(container_id: UniqueId, container_type: ContainerType):
         now = datetime.now()
         container_dir = f"{NEPTUNE_DATA_DIRECTORY}/{ASYNC_DIRECTORY}/{create_dir_name(container_type, container_id)}"
-        data_path = f"{container_dir}/exec-{now.timestamp()}-{now}"
+        data_path = f"{container_dir}/exec-{now.timestamp()}-{now.strftime('%Y-%m-%d_%H.%M.%S.%f')}"
         data_path = data_path.replace(" ", "_").replace(":", ".")
         return data_path
 

src/neptune/new/internal/operation_processors/operation_storage.py

@@ -17,6 +17,7 @@
     "OperationStorage",
 ]
 
+import os
 from pathlib import Path
 
 from neptune.new.constants import NEPTUNE_DATA_DIRECTORY
@@ -29,6 +30,10 @@ class OperationStorage:
     def __init__(self, data_path: str):
         self._data_path = Path(data_path)
 
+        # initialize directories
+        os.makedirs(self.data_path, exist_ok=True)
+        os.makedirs(self.upload_path, exist_ok=True)
+
     @property
     def data_path(self) -> Path:
         return self._data_path

src/neptune/new/internal/operation_processors/sync_operation_processor.py

@@ -42,8 +42,7 @@ def __init__(self, container_id: UniqueId, container_type: ContainerType, backen
     def _init_data_path(container_id: UniqueId, container_type: ContainerType):
         now = datetime.now()
         container_dir = f"{NEPTUNE_DATA_DIRECTORY}/{SYNC_DIRECTORY}/{create_dir_name(container_type, container_id)}"
-        data_path = f"{container_dir}/exec-{now.timestamp()}-{now}"
-        data_path = data_path.replace(" ", "_").replace(":", ".")
+        data_path = f"{container_dir}/exec-{now.timestamp()}-{now.strftime('%Y-%m-%d_%H.%M.%S.%f')}"
         return data_path
 
     def enqueue_operation(self, op: Operation, wait: bool) -> None:

tests/neptune/new/attributes/atoms/test_file.py

@@ -21,9 +21,12 @@
     BytesIO,
     StringIO,
 )
+from pathlib import Path
+from unittest.mock import PropertyMock
 
 from mock import MagicMock
 
+from e2e_tests.utils import tmp_context
 from neptune.common.utils import IS_WINDOWS
 from neptune.new.attributes.atoms.file import (
     File,
@@ -35,43 +38,64 @@
 )
 from neptune.new.internal.operation import (
     UploadFile,
-    UploadFileContent,
     UploadFileSet,
 )
-from neptune.new.internal.utils import base64_encode
+from neptune.new.types.atoms.file import FileType
 from tests.neptune.new.attributes.test_attribute_base import TestAttributeBase
 
 
 class TestFile(TestAttributeBase):
     @unittest.skipIf(IS_WINDOWS, "Windows behaves strangely")
     def test_assign(self):
+        def get_tmp_uploaded_file(tmp_upload_dir):
+            """Get tmp file to uploaded from `upload_path`
+            - here's assumption that we upload only one file per one path in test"""
+            uploaded_files = os.listdir(tmp_upload_dir)
+            assert len(uploaded_files) == 1
+            return f"{tmp_upload_dir}/{uploaded_files[0]}"
+
         a_text = "Some text stream"
         a_binary = b"Some binary stream"
         value_and_operation_factory = [
             (
                 FileVal("other/../other/file.txt"),
-                lambda _: UploadFile(_, "txt", os.getcwd() + "/other/file.txt"),
+                lambda attribute_path, _: UploadFile(
+                    attribute_path, ext="txt", file_path=os.getcwd() + "/other/file.txt"
+                ),
             ),
             (
                 FileVal.from_stream(StringIO(a_text)),
-                lambda _: UploadFileContent(_, "txt", base64_encode(a_text.encode("utf-8"))),
+                lambda attribute_path, tmp_uploaded_file: UploadFile(
+                    attribute_path, ext="txt", file_path=tmp_uploaded_file, clean_after_upload=True
+                ),
             ),
             (
                 FileVal.from_stream(BytesIO(a_binary)),
-                lambda _: UploadFileContent(_, "bin", base64_encode(a_binary)),
+                lambda attribute_path, tmp_uploaded_file: UploadFile(
+                    attribute_path, ext="bin", file_path=tmp_uploaded_file, clean_after_upload=True
+                ),
             ),
         ]
 
         for value, operation_factory in value_and_operation_factory:
-            processor = MagicMock()
-            exp, path, wait = (
-                self._create_run(processor),
-                self._random_path(),
-                self._random_wait(),
-            )
-            var = File(exp, path)
-            var.assign(value, wait=wait)
-            processor.enqueue_operation.assert_called_once_with(operation_factory(path), wait)
+            with tmp_context() as tmp_upload_dir:
+                processor = MagicMock()
+                processor._operation_storage = PropertyMock(upload_path=Path(tmp_upload_dir))
+                exp, path, wait = (
+                    self._create_run(processor),
+                    self._random_path(),
+                    self._random_wait(),
+                )
+                var = File(exp, path)
+                var.assign(value, wait=wait)
+
+                if value.file_type is not FileType.LOCAL_FILE:
+                    tmp_uploaded_file = get_tmp_uploaded_file(tmp_upload_dir)
+                    self.assertTrue(os.path.exists(tmp_uploaded_file))
+                else:
+                    tmp_uploaded_file = None
+
+                processor.enqueue_operation.assert_called_once_with(operation_factory(path, tmp_uploaded_file), wait)
 
     def test_assign_type_error(self):
         values = [55, None, []]

tests/neptune/new/test_imports.py

@@ -66,7 +66,6 @@
     File,
     FileVal,
     UploadFile,
-    UploadFileContent,
 )
 from neptune.new.attributes.atoms.float import (
     AssignFloat,