Copy refactorings (#1365)

* s3-specific copy

* client_kwargs

* lint

* no custom put/get

* optimize copy shape

* no conditional write

* changelog

* better progress bar for check-equality

* ci

* ci

* ci

* coverage

* recheck_cached=open

* return attachments

* return attachments in methods

* changelog

* fixes for attachments

* downgrade pytest-cov

* pytest-cov

* merge

* uv.lock

* expose unallowed dtypes as constant

* add mag arg to add_mag_as_ref

* reset deps

Author: normanrz

webknossos/Changelog.md

@@ -13,6 +13,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 [Commits](https://github.com/scalableminds/webknossos-libs/compare/v2.5.0...HEAD)
 
 ### Breaking Changes
+- The `endpoint_url` in `UPath` objects is now stored directly in the `storage_options` dict, instead of being stored in the `storage_options["client_kwargs"]` dict. [#1365](https://github.com/scalableminds/webknossos-libs/pull/1365)
 - Removed Docker image `scalableminds/webknossos-cli` build. It will not be released further. [#1376](https://github.com/scalableminds/webknossos-libs/pull/1376)
 - `Dataset.add_layer_as_ref(remote_layer.path)` does not work anymore. Please use `Dataset.add_layer_as_ref(remote_layer)` instead. [#1371](https://github.com/scalableminds/webknossos-libs/pull/1371])
 - Due to the refactoring, the imports of various classes have changed. Please update your code accordingly. [#1371](https://github.com/scalableminds/webknossos-libs/pull/1371])
@@ -28,6 +29,8 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 - Dataset.add_layer_as_ref() now accepts RemoteLayer objects as well as Layer objects. [#1371](https://github.com/scalableminds/webknossos-libs/pull/1371])
 
 ### Changed
+- Only use fs-based copy for mags on local file systems and S3. Other protocols, e.g. memory and HTTP, have caveats that break fs-based copying. [#1365](https://github.com/scalableminds/webknossos-libs/pull/1365)
+- The `add_*` methods in `Attachments` now return the created attachment objects, similar to `add_layer` and `add_mag`. [#1365](https://github.com/scalableminds/webknossos-libs/pull/1365)
 - The returned dataset ID from dataset exploration is now used. [#1378](https://github.com/scalableminds/webknossos-libs/pull/1378) 
 - Refactored the architecture, by introducing RemoteLayers, RemoteSegmentationLayers and their abstract base classes. [#1371](https://github.com/scalableminds/webknossos-libs/pull/1371])
 - Updated the api version of the webknossos-api to 12. [#1371](https://github.com/scalableminds/webknossos-libs/pull/1371])

webknossos/tests/constants.py

@@ -22,7 +22,7 @@
     "s3://testoutput",
     key=MINIO_ROOT_USER,
     secret=MINIO_ROOT_PASSWORD,
-    client_kwargs={"endpoint_url": f"http://localhost:{MINIO_PORT}"},
+    endpoint_url=f"http://localhost:{MINIO_PORT}",
 )
 
 

webknossos/tests/dataset/test_attachments.py

@@ -46,7 +46,7 @@ def test_attachments(tmp_upath: UPath) -> None:
     seg_layer.attachments.add_agglomerate(
         UPath(
             "s3://bucket/agglomerate.zarr",
-            client_kwargs={"endpoint_url": "https://s3.eu-central-1.amazonaws.com"},
+            endpoint_url="https://s3.eu-central-1.amazonaws.com",
         ),
         name="identity",
         data_format=AttachmentDataFormat.Zarr3,
@@ -242,7 +242,7 @@ def test_remote_layer(tmp_upath: UPath) -> None:
 
     mesh_path = UPath(
         "s3://bucket/meshfile.zarr",
-        client_kwargs={"endpoint_url": "https://s3.eu-central-1.amazonaws.com"},
+        endpoint_url="https://s3.eu-central-1.amazonaws.com",
     )
 
     seg_layer.attachments.add_mesh(

webknossos/tests/dataset/test_dataset.py

@@ -2010,6 +2010,40 @@ def test_add_mag_as_ref(data_format: DataFormat, output_path: UPath) -> None:
     assert (ds_path / "color" / "4").exists()
 
 
+@pytest.mark.parametrize("data_format,output_path", DATA_FORMATS_AND_OUTPUT_PATHS)
+def test_add_mag_as_ref_with_mag(data_format: DataFormat, output_path: UPath) -> None:
+    ds_path = prepare_dataset_path(data_format, output_path, "original")
+    new_path = prepare_dataset_path(data_format, output_path, "with_ref")
+
+    original_ds = Dataset(ds_path, voxel_size=(1, 1, 1))
+    original_layer = original_ds.add_layer(
+        "color",
+        COLOR_CATEGORY,
+        dtype_per_channel="uint8",
+        bounding_box=BoundingBox((0, 0, 0), (10, 20, 30)),
+    )
+    original_layer.add_mag(1).write(
+        data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8)
+    )
+
+    ds = Dataset(new_path, voxel_size=(1, 1, 1))
+    layer = ds.add_layer(
+        "color",
+        COLOR_CATEGORY,
+        dtype_per_channel="uint8",
+        bounding_box=BoundingBox((6, 6, 6), (10, 20, 30)),
+    )
+    layer.add_mag_as_ref(original_layer.get_mag(1), mag="2")
+
+    assert list(layer.mags.values())[0].mag == Mag("2")
+    assert list(layer.mags.values())[0]._properties.path == dump_path(
+        ds_path / "color" / "1", new_path
+    )
+
+    assure_exported_properties(ds)
+    assure_exported_properties(original_ds)
+
+
 @pytest.mark.parametrize("data_format", [DataFormat.Zarr, DataFormat.Zarr3])
 def test_remote_add_symlink_layer(data_format: DataFormat) -> None:
     src_dataset_path = copy_simple_dataset(data_format, REMOTE_TESTOUTPUT_DIR)

webknossos/tests/test_utils.py

@@ -110,27 +110,27 @@ def test_dump_path(tmp_path: Path) -> None:
     dataset_path = tmp_path / "test_dataset"
     path = UPath(
         "s3://bucket/test.txt",
-        client_kwargs={"endpoint_url": "https://s3.amazonaws.com"},
+        endpoint_url="https://s3.amazonaws.com",
     )
     assert dump_path(path, dataset_path) == "s3://s3.amazonaws.com/bucket/test.txt"
 
     # s3 relative
     dataset_path = UPath(
         "s3://bucket/test_dataset",
-        client_kwargs={"endpoint_url": "https://s3.amazonaws.com"},
+        endpoint_url="https://s3.amazonaws.com",
     )
     path = dataset_path / "test.txt"
     assert dump_path(path, dataset_path) == "./test.txt"
 
     # s3 dataset path is a prefix
     dataset_path = UPath(
         "s3://bucket/test_dataset",
-        client_kwargs={"endpoint_url": "https://s3.amazonaws.com"},
+        endpoint_url="https://s3.amazonaws.com",
     )
     path = (
         UPath(
             "s3://bucket/test_dataset_longer",
-            client_kwargs={"endpoint_url": "https://s3.amazonaws.com"},
+            endpoint_url="https://s3.amazonaws.com",
         )
         / "test.txt"
     )
@@ -142,7 +142,7 @@ def test_dump_path(tmp_path: Path) -> None:
     # s3 with ..
     dataset_path = UPath(
         "s3://bucket/test_dataset",
-        client_kwargs={"endpoint_url": "https://s3.amazonaws.com"},
+        endpoint_url="https://s3.amazonaws.com",
     )
     path = dataset_path / ".." / "test_dataset2" / "test.txt"
     assert (
@@ -163,7 +163,7 @@ def test_dump_path(tmp_path: Path) -> None:
     dataset_path = (
         UPath(
             "s3://bucket/test_dataset",
-            client_kwargs={"endpoint_url": "https://s3.amazonaws.com"},
+            endpoint_url="https://s3.amazonaws.com",
         )
         / ".."
         / "test_dataset2"
@@ -173,7 +173,7 @@ def test_dump_path(tmp_path: Path) -> None:
 
     path = UPath(
         "s3://bucket/test_dataset/test.txt",
-        client_kwargs={"endpoint_url": "https://s3.amazonaws.com"},
+        endpoint_url="https://s3.amazonaws.com",
     )
     assert (
         dump_path(path, dataset_path)
@@ -183,7 +183,7 @@ def test_dump_path(tmp_path: Path) -> None:
     path = (
         UPath(
             "s3://bucket/",
-            client_kwargs={"endpoint_url": "https://s3.amazonaws.com"},
+            endpoint_url="https://s3.amazonaws.com",
         )
         / "test_dataset2"
         / "test.txt"

webknossos/webknossos/cli/_utils.py

@@ -193,7 +193,7 @@ def parse_path(value: str) -> UPath:
     if value.startswith("s3://") and "S3_ENDPOINT_URL" in environ:
         return UPath(
             value,
-            client_kwargs={"endpoint_url": environ["S3_ENDPOINT_URL"]},
+            endpoint_url=environ["S3_ENDPOINT_URL"],
         )
 
     return UPath(value)

webknossos/webknossos/cli/check_equality.py

@@ -103,7 +103,7 @@ def main(
             f"The datasets {source} and {target} are equal \
     (with regard to the layers: {layer_names})"
         )
-    except AssertionError as err:
+    except RuntimeError as err:
         print(f"The datasets are not equal: {err}")
         exit(1)
 
@@ -118,25 +118,32 @@ def compare_layers(
     layer_name = source_layer.name
     logger.info("Checking layer_name: %s", layer_name)
 
-    assert source_layer.bounding_box == target_layer.bounding_box, (
-        f"The bounding boxes of '{layer_name}' layer of source and target \
+    if source_layer.bounding_box != target_layer.bounding_box:
+        raise RuntimeError(
+            f"The bounding boxes of '{layer_name}' layer of source and target \
 are not equal: {source_layer.bounding_box} != {target_layer.bounding_box}"
-    )
+        )
 
     source_mags = set(source_layer.mags.keys())
     target_mags = set(target_layer.mags.keys())
 
-    assert source_mags == target_mags, (
-        f"The mags of '{layer_name}' layer of source and target are not equal: \
+    if source_mags != target_mags:
+        raise RuntimeError(
+            f"The mags of '{layer_name}' layer of source and target are not equal: \
 {source_mags} != {target_mags}"
-    )
+        )
 
     for mag in source_mags:
         source_mag = source_layer.mags[mag]
         target_mag = target_layer.mags[mag]
 
         logger.info("Start verification of %s in mag %s", layer_name, mag)
         with get_executor_for_args(args=executor_args) as executor:
-            assert source_mag.content_is_equal(target_mag, executor=executor), (
-                f"The contents of {source_mag} and {target_mag} differ."
-            )
+            if not source_mag.content_is_equal(
+                target_mag,
+                executor=executor,
+                progress_desc=f"Comparing {layer_name}/{mag}",
+            ):
+                raise RuntimeError(
+                    f"The contents of {source_mag} and {target_mag} differ."
+                )

webknossos/webknossos/dataset/dataset.py

@@ -112,6 +112,25 @@
 
 logger = logging.getLogger(__name__)
 
+_ALLOWED_COLOR_LAYER_DTYPES = (
+    "uint8",
+    "uint16",
+    "uint32",
+    "int8",
+    "int16",
+    "int32",
+    "float32",
+)
+_ALLOWED_SEGMENTATION_LAYER_DTYPES = (
+    "uint8",
+    "uint16",
+    "uint32",
+    "uint64",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+)
 
 SAFE_LARGE_XY: int = 10_000_000_000  # 10 billion
 
@@ -1082,36 +1101,17 @@ def add_layer(
 
         # assert that the dtype_per_channel is supported by webknossos
         if category == COLOR_CATEGORY:
-            color_dtypes = (
-                "uint8",
-                "uint16",
-                "uint32",
-                "int8",
-                "int16",
-                "int32",
-                "float32",
-            )
-            if dtype_per_channel.name not in color_dtypes:
+            if dtype_per_channel.name not in _ALLOWED_COLOR_LAYER_DTYPES:
                 raise ValueError(
                     f"Cannot add color layer with dtype {dtype_per_channel.name}. "
-                    f"Supported dtypes are: {', '.join(color_dtypes)}."
+                    f"Supported dtypes are: {', '.join(_ALLOWED_COLOR_LAYER_DTYPES)}."
                     "For an overview of supported dtypes, see https://docs.webknossos.org/webknossos/data/upload_ui.html",
                 )
         else:
-            segmentation_dtypes = (
-                "uint8",
-                "uint16",
-                "uint32",
-                "uint64",
-                "int8",
-                "int16",
-                "int32",
-                "int64",
-            )
-            if dtype_per_channel.name not in segmentation_dtypes:
+            if dtype_per_channel.name not in _ALLOWED_SEGMENTATION_LAYER_DTYPES:
                 raise ValueError(
                     f"Cannot add segmentation layer with dtype {dtype_per_channel.name}. "
-                    f"Supported dtypes are: {', '.join(segmentation_dtypes)}."
+                    f"Supported dtypes are: {', '.join(_ALLOWED_SEGMENTATION_LAYER_DTYPES)}."
                     "For an overview of supported dtypes, see https://docs.webknossos.org/webknossos/data/upload_ui.html",
                 )