From 44ac40352706004802dd62e4b2f980be708cfe87 Mon Sep 17 00:00:00 2001
From: Daniel D <d.drizhuk@gmail.com>
Date: Mon, 3 Aug 2026 18:55:49 +0200
Subject: [PATCH] Add Dask Cloud Provider user agent to Nebius SDK

---
 dask_cloudprovider/nebius/instances.py         |  6 +++++-
 dask_cloudprovider/nebius/tests/test_nebius.py | 16 +++++++++++++++-
 setup.py                                       |  2 +-
 3 files changed, 21 insertions(+), 3 deletions(-)

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)