diff --git a/dask_cloudprovider/nebius/instances.py b/dask_cloudprovider/nebius/instances.py index e2a89d91..a8a6daf7 100644 --- a/dask_cloudprovider/nebius/instances.py +++ b/dask_cloudprovider/nebius/instances.py @@ -1,5 +1,6 @@ import dask +from dask_cloudprovider import __version__ from dask_cloudprovider.generic.vmcluster import ( VMCluster, VMInterface, @@ -66,7 +67,10 @@ def __init__( self.docker_image = docker_image self.server_platform = server_platform self.server_preset = server_preset - self.sdk = SDK(credentials=self.config.get("token")) + self.sdk = SDK( + credentials=self.config.get("token"), + user_agent_prefix=f"dask-cloudprovider/{__version__}", + ) self.disk_size = disk_size self.instance_id = None self.disk_id = None diff --git a/dask_cloudprovider/nebius/tests/test_nebius.py b/dask_cloudprovider/nebius/tests/test_nebius.py index 76798127..0d83c21e 100644 --- a/dask_cloudprovider/nebius/tests/test_nebius.py +++ b/dask_cloudprovider/nebius/tests/test_nebius.py @@ -1,10 +1,12 @@ +from unittest import mock + import pytest import dask nebius = pytest.importorskip("nebius") -from dask_cloudprovider.nebius.instances import NebiusCluster +from dask_cloudprovider.nebius.instances import NebiusCluster, NebiusInstance from dask.distributed import Client from distributed.core import Status @@ -52,6 +54,18 @@ async def test_init(): assert cluster.status == Status.created +def test_sdk_user_agent_prefix(): + with mock.patch( + "dask_cloudprovider.nebius.instances.__version__", "1.2.3" + ), mock.patch("dask_cloudprovider.nebius.instances.SDK") as sdk_cls: + NebiusInstance(cluster=mock.MagicMock(), config={"token": "token"}) + + sdk_cls.assert_called_once_with( + credentials="token", + user_agent_prefix="dask-cloudprovider/1.2.3", + ) + + @pytest.mark.asyncio @pytest.mark.external async def test_create_cluster(cluster): diff --git a/setup.py b/setup.py index 9ac9e953..903982af 100644 --- a/setup.py +++ b/setup.py @@ -17,7 +17,7 @@ "hetzner": ["hcloud>=1.10.0"], "ibm": ["ibm_code_engine_sdk>=3.1.0", "kubernetes>=25.3.0"], "openstack": ["openstacksdk>=3.3.0"], - "nebius": ["nebius>=0.2.0"], + "nebius": ["nebius>=0.2.55"], } extras_require["all"] = set(pkg for pkgs in extras_require.values() for pkg in pkgs)