Update tiledb-ml package to work with python 3.12

Author: ktsitsi

.github/workflows/ci.yml

@@ -9,10 +9,10 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9"]
-        ml-deps:
-          - "torch==2.1.2+cpu torchvision==0.16.2+cpu torchdata==0.7.1 tensorflow-cpu==2.15.1 "
-          - "torch==2.3.1+cpu torchvision==0.18.1+cpu torchdata==0.8.0 'tensorflow-cpu<2.16'"
+        python-version: ["3.12"]
+        # ml-deps:
+          # - "torch==2.1.2+cpu torchvision==0.16.2+cpu torchdata==0.7.1 tensorflow-cpu==2.15.1 "
+          # - "torch==2.3.1+cpu torchvision==0.18.1+cpu torchdata==0.8.0 'tensorflow-cpu<2.16'"
     env:
       run_coverage: ${{ github.ref == 'refs/heads/master' }}
 
@@ -28,14 +28,13 @@ jobs:
       uses: actions/cache@v4
       with:
         path: ~/.cache/pip
-        key: ${{ runner.os }}:ml-deps=[${{ matrix.ml-deps }}]
+        key: ${{ runner.os }} # :ml-deps=[${{ matrix.ml-deps }}]
 
     - name: Install dependencies
       run: |
         pip install --upgrade pip
-        pip install --extra-index-url https://download.pytorch.org/whl/cpu ${{ matrix.ml-deps }}
         pip install pytest-mock pytest-cov scikit-learn==1.0.2
-        pip install -e .[cloud]
+        pip install -e .[full]
 
     - name: Run pre-commit hooks
       run: |

examples/cloud/serverless_training/pytorch/README.md

@@ -6,7 +6,7 @@ Firstly, we ingest all training images and labels in TileDB arrays and register
 and finally, we serverless-ly get some predictions using the trained model. In case you want to run the example, you will need a TileDB-Cloud account as described
 [here](https://docs.tiledb.com/cloud/tutorials/start-here). After signing up, you should export your username and password 
 as environmental variables (**TILEDB_USER_NAME**, **TILEDB_PASSWD**), in order to run ingestion, model training and prediction UDFs. Moreover,
-please add your TileDB namespace and your **S3** bucket in each script.
+please add your TileDB teamspace and your **S3** bucket in each script.
 
 # Steps
 

examples/cloud/serverless_training/pytorch/data_ingestion.py

@@ -8,7 +8,7 @@
 TILEDB_USER_NAME = os.environ.get("TILEDB_USER_NAME")
 TILEDB_PASSWD = os.environ.get("TILEDB_PASSWD")
 
-# Your TileDB namespace
+# Your TileDB workspace/teamspace
 TILEDB_WORKSPACE = "your_tiledb_WORKSPACE"
 TILEDB_TEAMSPACE = "your_tiledb_TEAMSPACE"
 

examples/cloud/serverless_training/pytorch/model_training.py

@@ -1,6 +1,7 @@
 import os
 
 import tiledb.client
+from tiledb.ml.readers.types import ArrayParams
 
 # Your TileDB username and password, exported as environmental variables
 TILEDB_USER_NAME = os.environ.get("TILEDB_USER_NAME")
@@ -13,9 +14,9 @@
 # Your S3 bucket
 S3_BUCKET = "your_s3_bucket"
 
-IMAGES_URI = f"tiledb://{TILEDB_WORKSPACE}/{TILEDB_TEAMSPACE}/s3://{S3_BUCKET}/mnist_images"
-LABELS_URI = f"tiledb://{TILEDB_WORKSPACE}/{TILEDB_TEAMSPACE}/s3://{S3_BUCKET}/mnist_labels"
-MODEL_URI = f"tiledb://{TILEDB_WORKSPACE}/{TILEDB_TEAMSPACE}/s3://{S3_BUCKET}/mnist_model"
+IMAGES_URI = f"tiledb://{TILEDB_WORKSPACE}/{TILEDB_TEAMSPACE}/mnist_images"
+LABELS_URI = f"tiledb://{TILEDB_WORKSPACE}/{TILEDB_TEAMSPACE}/mnist_labels"
+MODEL_URI = f"tiledb://{TILEDB_WORKSPACE}/{TILEDB_TEAMSPACE}/mnist_model"
 
 # The size of each slice from a image and label TileDB arrays.
 IO_BATCH_SIZE = 20000
@@ -50,8 +51,20 @@ def forward(self, x: torch.Tensor) -> Any:
             logits = self.linear_relu_stack(x)
             return logits
 
+    def do_random_noise(img, mag=0.1):
+        noise = np.random.uniform(-1, 1,img.shape)*mag
+        img = img + noise
+        img = np.clip(img,0,1)
+        return img
+
     with tiledb.open(IMAGES_URI) as x, tiledb.open(LABELS_URI) as y:
-        train_loader = PyTorchTileDBDataLoader(x, y, batch_size=IO_BATCH_SIZE)
+        train_loader = PyTorchTileDBDataLoader(
+            ArrayParams(x, fn=do_random_noise),
+            ArrayParams(y),
+            batch_size=IO_BATCH_SIZE,
+            num_workers=0,
+            shuffle_buffer_size=256,
+        )
 
         net = Net(shape=(28, 28))
         criterion = nn.CrossEntropyLoss()
@@ -96,7 +109,7 @@ def forward(self, x: torch.Tensor) -> Any:
 
         model = PyTorchTileDBModel(
             uri="mnist_model",
-            namespace=TILEDB_NAMESPACE,
+            teamspace=TILEDB_TEAMSPACE,
             model=net,
             optimizer=optimizer,
         )

setup.py

@@ -1,7 +1,7 @@
 import setuptools
 
-tensorflow = ["tensorflow>=2.6, <2.16.0"]
-pytorch = ["torch>=1.11", "torchdata"]
+tensorflow = ["tensorflow>=2.6"]
+pytorch = ["torch>=1.11", "torchdata>=0.6.0,<0.8.0"]
 sklearn = ["scikit-learn>=1.0"]
 cloud = ["tiledb-client"]
 full = sorted({"torchvision", *tensorflow, *pytorch, *sklearn, *cloud})

tests/models/test_cloud_utils.py

@@ -1,61 +1,10 @@
 import pytest
 
 from tiledb.ml.models._cloud_utils import (
-    get_cloud_uri,
-    get_s3_prefix,
     update_file_properties,
 )
 
-
 class TestCloudUtils:
-    def test_get_s3_prefix(self, mocker):
-
-        assert get_s3_prefix(None) is None
-
-        profile = mocker.patch(
-            "tiledb.client.user_profile",
-            return_value=mocker.Mock(username="foo", default_s3_path="bar"),
-        )
-        profile.default_s3_path = "bar"
-        assert get_s3_prefix("foo") == "bar/ml_models"
-
-        profile = mocker.patch(
-            "tiledb.client.user_profile",
-            return_value=mocker.Mock(username="foo", default_s3_path=None),
-        )
-        profile.default_s3_path = "bar"
-        assert get_s3_prefix("foo") is None
-
-        org = mocker.patch(
-            "tiledb.client.organization",
-            return_value=mocker.Mock(default_s3_path="orgbar"),
-        )
-        org.default_s3_path = "orgbar"
-        assert get_s3_prefix("nofoo") == "orgbar/ml_models"
-
-        org = mocker.patch(
-            "tiledb.client.organization",
-            return_value=mocker.Mock(
-                default_s3_path=None,
-            ),
-        )
-        org.default_s3_path = "orgbar"
-        assert get_s3_prefix("nofoo") is None
-
-    def test_get_cloud_uri(self, mocker):
-        mocker.patch(
-            "tiledb.ml.models._cloud_utils.get_s3_prefix", return_value="s3://"
-        )
-
-        assert "tiledb://test_namespace/s3://tiledb_array" == get_cloud_uri(
-            uri="tiledb_array", namespace="test_namespace"
-        )
-
-        mocker.patch("tiledb.ml.models._cloud_utils.get_s3_prefix", return_value=None)
-        with pytest.raises(ValueError) as ex:
-            get_cloud_uri(uri="tiledb_array", namespace="test_namespace")
-
-        assert "You must set the default s3 prefix path for ML models" in str(ex.value)
 
     def test_update_file_properties(self, mocker):
         mock_tiledb_cloud_update_file_properties = mocker.patch(

tests/models/test_pytorch_models.py

@@ -103,6 +103,8 @@ class TestPyTorchModel:
         ],
     )
     def test_save(self, tmpdir, net, optimizer):
+        if optimizer.__name__ == "Muon":
+            pytest.skip("Muon optimizer is 2D only")
         model = net()
         saved_optimizer = optimizer(model.parameters(), lr=0.001)
         tiledb_array = os.path.join(tmpdir, "model_array")
@@ -174,6 +176,8 @@ def test_file_properties(self, tmpdir, net):
         ],
     )
     def test_tensorboard_callback(self, tmpdir, net, optimizer):
+        if optimizer.__name__ == "Muon":
+            pytest.skip("Muon optimizer is 2D only")
         model = net()
         saved_optimizer = optimizer(model.parameters(), lr=0.001)
         tiledb_array = os.path.join(tmpdir, "model_array")
@@ -219,21 +223,21 @@ def test_get_cloud_uri_call_for_models_on_tiledb_cloud(self, tmpdir, mocker):
             "tiledb.ml.models._base.get_cloud_uri", return_value=uri
         )
 
-        _ = PyTorchTileDBModel(uri=uri, namespace="test_namespace", model=model)
+        _ = PyTorchTileDBModel(uri=uri, teamspace="test_teamspace", model=model)
 
-        mock_get_cloud_uri.assert_called_once_with(uri, "test_namespace")
+        mock_get_cloud_uri.assert_called_once_with(uri, "test_teamspace")
 
-    def test_get_s3_prefix_call_for_models_on_tiledb_cloud(self, tmpdir, mocker):
-        model = Net()
-        uri = os.path.join(tmpdir, "model_array")
+    # def test_get_s3_prefix_call_for_models_on_tiledb_cloud(self, tmpdir, mocker):
+    #     model = Net()
+    #     uri = os.path.join(tmpdir, "model_array")
 
-        mock_get_s3_prefix = mocker.patch(
-            "tiledb.ml.models._cloud_utils.get_s3_prefix", return_value="s3 prefix"
-        )
+    #     mock_get_s3_prefix = mocker.patch(
+    #         "tiledb.ml.models._cloud_utils.get_s3_prefix", return_value="s3 prefix"
+    #     )
 
-        _ = PyTorchTileDBModel(uri=uri, namespace="test_namespace", model=model)
+    #     _ = PyTorchTileDBModel(uri=uri, teamspace="test_teamspace", model=model)
 
-        mock_get_s3_prefix.assert_called_once_with("test_namespace")
+    #     mock_get_s3_prefix.assert_called_once_with("test_teamspace")
 
     def test_update_file_properties_call(self, tmpdir, mocker):
         model = Net()
@@ -242,7 +246,7 @@ def test_update_file_properties_call(self, tmpdir, mocker):
         mocker.patch("tiledb.ml.models._base.get_cloud_uri", return_value=uri)
 
         tiledb_obj = PyTorchTileDBModel(
-            uri=uri, namespace="test_namespace", model=model
+            uri=uri, teamspace="test_teamspace", model=model
         )
 
         mock_update_file_properties = mocker.patch(

tests/models/test_sklearn_models.py

@@ -82,21 +82,9 @@ def test_get_cloud_uri_call_for_models_on_tiledb_cloud(self, tmpdir, mocker):
             "tiledb.ml.models._base.get_cloud_uri", return_value=uri
         )
 
-        _ = SklearnTileDBModel(uri=uri, namespace="test_namespace", model=model)
+        _ = SklearnTileDBModel(uri=uri, teamspace="test_teamspace", model=model)
 
-        mock_get_cloud_uri.assert_called_once_with(uri, "test_namespace")
-
-    def test_get_s3_prefix_call_for_models_on_tiledb_cloud(self, tmpdir, mocker):
-        model = sklearn.linear_model.LinearRegression()
-        uri = os.path.join(tmpdir, "model_array")
-
-        mock_get_s3_prefix = mocker.patch(
-            "tiledb.ml.models._cloud_utils.get_s3_prefix", return_value="s3 prefix"
-        )
-
-        _ = SklearnTileDBModel(uri=uri, namespace="test_namespace", model=model)
-
-        mock_get_s3_prefix.assert_called_once_with("test_namespace")
+        mock_get_cloud_uri.assert_called_once_with(uri, "test_teamspace")
 
     def test_update_file_properties_call(self, tmpdir, mocker):
         model = sklearn.linear_model.LinearRegression()
@@ -105,7 +93,7 @@ def test_update_file_properties_call(self, tmpdir, mocker):
         mocker.patch("tiledb.ml.models._base.get_cloud_uri", return_value=uri)
 
         tiledb_obj = SklearnTileDBModel(
-            uri=uri, namespace="test_namespace", model=model
+            uri=uri, teamspace="test_teamspace", model=model
         )
 
         mock_update_file_properties = mocker.patch(