feat: migration from python-sdk to pysdk provided by nebius and added cpu-d3 support

.github/actions/nebius_runner_create/action.yaml

@@ -25,6 +25,9 @@ inputs:
   vm_name:
     required: true
     description: "vm name"
+  vm_platform_id:
+    required: true
+    description: "vm platform id"
   vm_preset:
     required: true
     description: "vm cpu"
@@ -94,8 +97,7 @@ runs:
     - name: install dependencies
       shell: bash
       run: |
-        pip install https://github.com/librarian/python-sdk/releases/download/v0.1.1/nebiusai-0.1.1-py3-none-any.whl
-        pip install PyGithub==2.2.0
+        pip install PyGithub==2.5.0 nebius
     - name: create vm
       id: create-vm
       shell: bash
@@ -106,6 +108,7 @@ runs:
           --service-account-key "${SERVICE_ACCOUNT_KEY}" \
           --name "${VM_NAME}" \
           --parent-id "${VM_PARENT_ID}" \
+          --platform-id "${VM_PLATFORM_ID}" \
           --preset "${VM_PRESET}" --disk-size "${VM_DISK_SIZE}" \
           --disk-type "${VM_DISK_TYPE}" \
           --subnet-id "${VM_SUBNET}" \
@@ -120,6 +123,7 @@ runs:
         REPO: ${{ inputs.repo }}
         VM_NAME: ${{ inputs.vm_name }}
         VM_PARENT_ID: ${{ inputs.vm_parent_id }}
+        VM_PLATFORM_ID: ${{ inputs.vm_platform_id }}
         VM_PRESET: ${{ inputs.vm_preset }}
         VM_DISK_SIZE: ${{ inputs.vm_disk_size }}
         VM_DISK_TYPE: ${{ inputs.vm_disk_type }}

.github/actions/nebius_runner_remove/action.yaml

@@ -34,8 +34,7 @@ runs:
     - name: install dependencies
       shell: bash
       run: |
-        pip install https://github.com/librarian/python-sdk/releases/download/v0.1.1/nebiusai-0.1.1-py3-none-any.whl
-        pip install PyGithub==2.2.0
+        pip install PyGithub==2.5.0 nebius
     - name: remove vm
       id: remove-vm
       shell: bash

.github/actions/nebius_threads_calculator/action.yaml

@@ -33,6 +33,30 @@ runs:
         test_threads=0
         if [[ $tests_size == *"large"* ]]; then
           case "$vm_preset" in
+            "256vcpu-1024gb")
+                test_threads=64
+                build_threads=256
+            ;;
+            "224vcpu-896gb")
+                test_threads=56
+                build_threads=224
+            ;;
+            "192vcpu-768gb")
+                test_threads=48
+                build_threads=192
+            ;;
+            "160vcpu-640gb")
+                test_threads=40
+                build_threads=160
+            ;;
+            "128vcpu-512gb")
+                test_threads=32
+                build_threads=128
+            ;;
+            "96vcpu-384gb")
+                test_threads=24
+                build_threads=96
+            ;;
             "80vcpu-320gb")
                 test_threads=16
                 build_threads=80
@@ -72,6 +96,30 @@ runs:
           esac
         else
           case "$vm_preset" in
+          "256vcpu-1024gb")
+                test_threads=64
+                build_threads=256
+            ;;
+            "224vcpu-896gb")
+                test_threads=56
+                build_threads=224
+            ;;
+            "192vcpu-768gb")
+                test_threads=48
+                build_threads=192
+            ;;
+            "160vcpu-640gb")
+                test_threads=40
+                build_threads=160
+            ;;
+            "128vcpu-512gb")
+                test_threads=32
+                build_threads=128
+            ;;
+            "96vcpu-384gb")
+                test_threads=28
+                build_threads=96
+            ;;
             "80vcpu-320gb")
                 test_threads=24
                 build_threads=80

.github/actions/prepare/action.yaml

@@ -11,11 +11,10 @@ runs:
           tzdata python3-dev python3-pip antlr3 libidn11-dev file distcc strace qemu-kvm \
           qemu-utils dpkg-dev atop pigz pbzip2 xz-utils pixz gdb unzip
         sudo apt-get remove -y unattended-upgrades
-        sudo pip install https://github.com/librarian/python-sdk/releases/download/v0.1.1/nebiusai-0.1.1-py3-none-any.whl
         # shellcheck disable=SC2102
         sudo pip3 install pytest pytest-timeout pytest-xdist setproctitle grpcio grpcio-tools \
              PyHamcrest tornado xmltodict pyarrow boto3 psutil PyGithub==2.5.0 pyinstaller \
-             cryptography==41.0.7 protobuf==4.25.6 packaging six pyyaml rapidgzip pyOpenSSL typing-extensions==4.10.0
+             cryptography protobuf packaging six pyyaml rapidgzip pyOpenSSL typing-extensions nebius
     - name: add user github to kvm group if exists
       shell: bash
       run: |

.github/scripts/github-runner.sh

@@ -70,14 +70,14 @@ pyarrow
 boto3
 psutil
 PyGithub==2.5.0
-cryptography==41.0.7
-protobuf==4.25.6
-pyOpenSSL==24.2.1
+cryptography
+protobuf
+pyOpenSSL
 packaging
 rapidgzip
-typing-extensions==4.10.0
+typing-extensions
+nebius
 EOF
-sudo pip3 install https://github.com/librarian/python-sdk/releases/download/v0.1.1/nebiusai-0.1.1-py3-none-any.whl
 sudo pip3 install -r /tmp/requirements.txt
 curl -L "https://github.com/ccache/ccache/releases/download/v${CCACHE_VERSION}/ccache-${CCACHE_VERSION}-linux-${OS_ARCH}.tar.xz" | sudo tar -xJ -C /usr/local/bin/ --strip-components=1 --no-same-owner "ccache-${CCACHE_VERSION}-linux-${OS_ARCH}/ccache"
 sudo apt-get remove -y unattended-upgrades

.github/scripts/nebius-manage-images.py

@@ -1,22 +1,42 @@
 import os
-import json
-import grpc
+import asyncio
 import math
 import logging
 import argparse
 from github import Github
-from datetime import datetime, timezone
-from nebius.compute.v1.image_pb2 import Image, ImageStatus
-from nebius.compute.v1.image_service_pb2 import ListImagesRequest
+from nebius.sdk import SDK
+from nebius.aio.service_error import RequestError
+from nebius.api.nebius.compute.v1 import (
+    ListImagesRequest,
+    ImageServiceClient,
+)
 
-from nebius.compute.v1.image_service_pb2_grpc import ImageServiceStub
-import nebius.compute.v1.image_service_pb2 as image_service_pb2
-from nebiusai import SDK, RetryInterceptor, backoff_linear_with_jitter
+SENSITIVE_DATA_VALUES = {}
+if os.environ.get("GITHUB_TOKEN"):
+    SENSITIVE_DATA_VALUES["github_token"] = os.environ.get("GITHUB_TOKEN")
 
-logging.basicConfig(
-    level=logging.INFO, format="%(asctime)s: %(levelname)s: %(message)s"
-)
-logger = logging.getLogger(__name__)
+
+class MaskingFormatter(logging.Formatter):
+    @staticmethod
+    def mask_sensitive_data(msg):
+        # Iterate over the patterns and replace sensitive data with '***'
+        for pattern_name, pattern in SENSITIVE_DATA_VALUES.items():
+            msg = msg.replace(pattern, f"[{pattern_name}=***]")
+        return msg
+
+    def format(self, record):
+        original = logging.Formatter.format(self, record)
+        return self.mask_sensitive_data(original)
+
+
+formatter = MaskingFormatter("%(asctime)s: %(levelname)s: %(message)s")
+console_handler = logging.StreamHandler()
+console_handler.setLevel(logging.INFO)
+console_handler.setFormatter(formatter)
+
+logger = logging.getLogger()
+logger.addHandler(console_handler)
+logger.setLevel(logging.INFO)
 
 # they won't appear in the list of images
 # but they will be protected from deletion anyway
@@ -25,10 +45,6 @@
 ]
 
 
-def status_to_string(image: Image) -> str:
-    return ImageStatus.State.Name(image.status.state)
-
-
 def convert_size(size_bytes):
     if size_bytes == 0:
         return "0 B"
@@ -39,7 +55,7 @@ def convert_size(size_bytes):
     return f"{s} {size_name[i]}"
 
 
-def main(
+async def main(
     sdk: SDK,
     github_token: str,
     github_repository: str,
@@ -68,19 +84,21 @@ def main(
     else:
         logger.info("Would set %s (new) = %s", image_variable_name, new_image_id)
 
+    service = ImageServiceClient(sdk)
     request = ListImagesRequest(
         parent_id=parent_id,
         filter=f"family IN ('{image_family_name}') AND status = 'READY'",
     )
+    try:
+        response = await service.list(request)
+    except RequestError as e:
+        logger.error("Failed to list images: %s", e)
+        return
 
-    client = sdk.client(image_service_pb2, ImageServiceStub)
-    response = client.List(request)
     candidate_images = []
     for image in response.items:
-        status = status_to_string(image)
-        created_at = datetime.fromtimestamp(
-            image.metadata.created_at.seconds, tz=timezone.utc
-        )
+        status = image.status.state.name
+        created_at = image.metadata.created_at
         storage_size = convert_size(image.status.storage_size_bytes)
         min_disk_size = convert_size(image.status.min_disk_size_bytes)
         prefix = " (PROTECTED)"
@@ -158,25 +176,15 @@ def main(
     args = parser.parse_args()
     logger.info(args)
 
-    interceptor = RetryInterceptor(
-        max_retry_count=30,
-        retriable_codes=[grpc.StatusCode.UNAVAILABLE],
-        back_off_func=backoff_linear_with_jitter(5, 0),
-    )
-
-    with open(args.service_account_key, "r") as fp:
-        sdk = SDK(
-            service_account_key=json.load(fp),
-            endpoint=args.api_endpoint,
-            interceptor=interceptor,
+    sdk = SDK(credentials_file_name=args.service_account_key)
+    asyncio.run(
+        main(
+            sdk=sdk,
+            github_token=args.github_token,
+            github_repository=args.github_repo,
+            new_image_id=args.new_image_id,
+            image_variable_name=args.image_variable_name,
+            update_image_id=args.update_image_id,
+            parent_id=args.parent_id,
         )
-
-    main(
-        sdk=sdk,
-        github_token=args.github_token,
-        github_repository=args.github_repo,
-        new_image_id=args.new_image_id,
-        image_variable_name=args.image_variable_name,
-        update_image_id=args.update_image_id,
-        parent_id=args.parent_id,
     )