Add limit for logging images (#1046)

src/neptune/new/attributes/series/file_series.py

@@ -35,6 +35,7 @@
 )
 from neptune.new.internal.utils import base64_encode
 from neptune.new.internal.utils.iteration import get_batches
+from neptune.new.internal.utils.limits import image_size_exceeds_limit_for_logging
 from neptune.new.types import File
 from neptune.new.types.atoms.file import FileType
 from neptune.new.types.series.file_series import FileSeries as FileSeriesVal
@@ -74,15 +75,20 @@ def _get_base64_image_content(file: File) -> str:
             if not os.path.exists(file.path):
                 raise FileNotFound(file.path)
             with open(file.path, "rb") as image_file:
-                file = File.from_stream(image_file)
+                file_content = File.from_stream(image_file).content
+        else:
+            file_content = file.content
 
-        ext = imghdr.what("", h=file.content)
+        ext = imghdr.what("", h=file_content)
         if not ext:
             raise OperationNotSupported(
-                "FileSeries supports only image files for now. " "Other file types will be implemented in future."
+                "FileSeries supports only image files for now. Other file types will be implemented in future."
             )
 
-        return base64_encode(file.content)
+        if image_size_exceeds_limit_for_logging(len(file_content)):
+            file_content = b""
+
+        return base64_encode(file_content)
 
     def download(self, destination: Optional[str]):
         target_dir = self._get_destination(destination)

src/neptune/new/internal/utils/limits.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 #
 import logging
+import warnings
 
 _logger = logging.getLogger(__name__)
 
@@ -22,6 +23,7 @@
 _IMAGE_SIZE_LIMIT_MB = 15
 _IN_MEMORY_SIZE_LIMIT_MB = 32
 _STREAM_SIZE_LIMIT_MB = 32
+_LOGGED_IMAGE_SIZE_LIMIT_MB = 15
 
 BYTES_IN_MB = 1024 * 1024
 
@@ -51,6 +53,18 @@ def image_size_exceeds_limit(content_size):
     return False
 
 
+def image_size_exceeds_limit_for_logging(content_size):
+    if content_size > _LOGGED_IMAGE_SIZE_LIMIT_MB * BYTES_IN_MB:
+        warnings.warn(
+            f"You are attempting to log an image that is {content_size / BYTES_IN_MB:.2f}MB large. "
+            f"Neptune supports logging images smaller than {_LOGGED_IMAGE_SIZE_LIMIT_MB}MB. "
+            "Resize or increase compression of this image.",
+            category=UserWarning,
+        )
+        return True
+    return False
+
+
 def file_size_exceeds_limit(content_size):
     if content_size > _IN_MEMORY_SIZE_LIMIT_MB * BYTES_IN_MB:
         _logger.warning(

tests/neptune/new/attributes/series/test_file_series.py

@@ -13,9 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+from unittest import mock
 
 # pylint: disable=protected-access
 import numpy
+import pytest
 from mock import (
     MagicMock,
     call,
@@ -186,3 +188,22 @@ def test_assign_raise_not_image(self):
                 attr.assign([file])
             with self.assertRaises(OperationNotSupported):
                 attr.assign([saved_file])
+
+    @mock.patch("neptune.new.internal.utils.limits._LOGGED_IMAGE_SIZE_LIMIT_MB", (10**-3))
+    def test_image_limit(self):
+        """Test if we prohibit logging images greater than mocked 1KB limit size"""
+        # given
+        path = self._random_path()
+        op_processor = MagicMock()
+        exp = self._create_run(processor=op_processor)
+        attr = FileSeries(exp, path)
+
+        file = File.as_image(numpy.random.rand(100, 100) * 255)
+        with create_file(file.content, binary_mode=True) as tmp_filename:
+            saved_file = File(tmp_filename)
+
+            # when
+            with pytest.warns(expected_warning=UserWarning, match=".* Neptune supports logging images smaller than .*"):
+                attr.assign([file])
+            with pytest.warns(expected_warning=UserWarning, match=".* Neptune supports logging images smaller than .*"):
+                attr.assign([saved_file])