Merge branch 'development' of https://github.com/mlrun/functions into convert-to-uv

Author: Eyal-Danieli

modules/src/vllm_module/vllm_module.ipynb

@@ -17,7 +17,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import mlrun\n",
+    "import mlrun\n"
    ]
   },
   {
@@ -170,34 +170,31 @@
     "body = {\n",
     "    \"model\": vllm_module.model,\n",
     "    \"messages\": [{\"role\": \"user\", \"content\": \"what are the 3 countries with the most gpu as far as you know\"}],\n",
-    "    \"max_tokens\": vllm_module.max_tokens,   # start smaller for testing\n",
+    "    \"max_tokens\": vllm_module.max_tokens,     # start smaller for testing\n",
     "}\n",
     "\n",
     "resp = app.invoke(path=\"/v1/chat/completions\", body=body)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 22,
    "id": "a459d5f8-dad0-4735-94c2-3801d4f94bb5",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Raw response keys: dict_keys(['id', 'object', 'created', 'model', 'choices', 'service_tier', 'system_fingerprint', 'usage', 'prompt_logprobs', 'prompt_token_ids', 'kv_transfer_params'])\n",
       "\n",
       "assistant:\n",
       "\n",
-      "As of July 2023, the three countries with the most scientists in articles that explain or discuss GPU contributions to AI are the United States, China, and India.\n",
-      "The number of scientists is not the best measure of the number of GPUs. According to Practical Deep Learning forמות, China has 6,307 GPU-equipped tens of thousands of compute servers, the number of GPUs in the top 100 supercomputers is 6,492 (19th largest: 10,363), and the per capita number of GPUs invested by research institutions is high. From the total value perspective, the annual procurement increase of GPUs in China is estimated to be more than $40 billion. Similarly, the United States and India have significantly higher prices than China purely due to price controls.\n",
-      "In summary, there is limited data to support the claim that GPU prices vary significantly between the three countries. However, China has a significant number of GPUs in use, and its computational resources are some of the largest in the world.\n"
+      "As of the most commonly cited estimates, the three countries with the largest GPU capacity for AI workloads are the United States, China, and India.\n"
      ]
     }
    ],
    "source": [
-    "data = resp.json()\n",
+    "data = resp\n",
     "assistant_text = data[\"choices\"][0][\"message\"][\"content\"]\n",
     "\n",
     "print(\"\\nassistant:\\n\")\n",
@@ -207,7 +204,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1de85b32-9c91-4609-9a63-0b38ed4fde65",
+   "id": "957b5d21-7ade-4131-9100-878652c477fc",
    "metadata": {},
    "outputs": [],
    "source": []

modules/src/vllm_module/vllm_module.py

@@ -52,21 +52,12 @@ def __init__(
             mem: str = "10G",
             port: int = 8000,
             dtype: str = "auto",
-            tensor_parallel_size: Optional[int] = None,
             uvicorn_log_level: str = "info",
             max_tokens: int = 500,
     ):
         if gpus < 1:
             raise ValueError("gpus must be >= 1")
 
-        if tensor_parallel_size is not None:
-            if tensor_parallel_size < 1:
-                raise ValueError("tensor_parallel_size must be >= 1")
-            if tensor_parallel_size > gpus:
-                raise ValueError(
-                    f"tensor_parallel_size ({tensor_parallel_size}) cannot be greater than gpus ({gpus})"
-                )
-
 
 
         if node_selector is None:
@@ -87,7 +78,6 @@ def __init__(
         self.node_selector = node_selector
         self.port = port
         self.dtype = dtype
-        self.tensor_parallel_size = tensor_parallel_size
         self.uvicorn_log_level = uvicorn_log_level
         self.max_tokens = max_tokens
 
@@ -117,8 +107,7 @@ def __init__(
             args += ["--uvicorn-log-level", self.uvicorn_log_level]
 
         if self.gpus > 1:
-            tps = self.tensor_parallel_size or self.gpus
-            args += ["--tensor-parallel-size", str(tps)]
+            args += ["--tensor-parallel-size", str(gpus)]
 
             # For more than one GPU you should create a share volume for the multiple GPUs
             self.vllm_app.spec.volumes = [{"name": "dshm", "emptyDir": {"medium": "Memory"}}]

steps/README.md

@@ -4,5 +4,7 @@
 ## Catalog
 
 <!-- AUTOGEN:START (do not edit below) -->
-_No items found_
+| Name | Description | Class Name | Categories |
+| --- | --- | --- | --- |
+| [verify_schema](https://github.com/mlrun/functions/tree/development/steps/src/verify_schema) | Verifies the event is aligned with the provided schema | VerifySchema | data-preparation, model-serving, utilities |
 <!-- AUTOGEN:END -->

steps/src/.gitkeep

@@ -0,0 +1,20 @@
+apiVersion: v1
+categories:
+  - data-preparation
+  - model-serving
+  - utilities
+description: Verifies the event is aligned with the provided schema
+example: verify_schema.ipynb
+generationDate: 2025-12-29:11-59    
+hidden: false                                       
+labels:
+  author: Iguazio
+mlrunVersion: 1.10.0
+name: verify_schema                        
+className: VerifySchema
+defaultHandler:     
+spec:
+  filename: verify_schema.py              
+  image: mlrun/mlrun                                
+  requirements:
+version: 1.0.0                                      

steps/src/verify_schema/item.yaml

@@ -0,0 +1,66 @@
+# Copyright 2025 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from verify_schema import VerifySchema
+
+class TestVerifySchema:
+    def test_verify_schema(self):
+        schema = ["id", "name", "active"]
+        verifier = VerifySchema(schema=schema, allow_unexpected_keys=False)
+
+        # Test with valid event
+        event = {
+            "id": 1,
+            "name": "Test Event",
+            "active": True
+        }
+        result = verifier.do(event)
+        assert result == event
+
+        # Test with missing key
+        event_missing_key = {
+            "id": 1,
+            "name": "Test Event"
+        }
+        try:
+            verifier.do(event_missing_key)
+        except KeyError as e:
+            assert "missing keys {'active'} in event" in str(e)
+
+        # Test with unexpected key
+        event_unexpected_key = {
+            "id": 1,
+            "name": "Test Event",
+            "active": True,
+            "extra": "unexpected"
+        }
+        try:
+            verifier.do(event_unexpected_key)
+        except KeyError as e:
+            assert "unexpected keys {'extra'} in event" in str(e)
+
+    def test_verify_schema_allow_unexpected(self):
+        schema = ["id", "name", "active"]
+        verifier = VerifySchema(schema=schema, allow_unexpected_keys=True)
+
+        # Test with valid event and unexpected key
+        event = {
+            "id": 1,
+            "name": "Test Event",
+            "active": True,
+            "extra": "unexpected"
+        }
+        result = verifier.do(event)
+        assert result == event

steps/src/verify_schema/test_verify_schema.py

@@ -0,0 +1,33 @@
+{
+ "cells": [
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": "",
+   "id": "556b36b9b89d0515"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

steps/src/verify_schema/verify_schema.ipynb

@@ -0,0 +1,40 @@
+# Copyright 2025 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+class VerifySchema:
+    """
+    This step validates that an event dictionary contains exactly the keys defined in the schema,
+    raising a KeyError if any are missing or unexpected.
+    """
+
+    def __init__(self, schema: list, allow_unexpected_keys: bool = False):
+        self.schema = schema
+        self.allow_unexpected_keys = allow_unexpected_keys
+
+    def do(self, event: dict):
+        # Check if all keys in the expected schema are present in the event
+        missing = set(self.schema) - set(event)
+        if missing:
+            raise KeyError(f"Schema verification failed: missing keys {missing} in event: {event}")
+
+        if self.allow_unexpected_keys:
+            return event
+
+        # Check if there are any unexpected keys in the event
+        unexpected = set(event) - set(self.schema)
+        if unexpected:
+            raise KeyError(f"Schema verification failed: unexpected keys {unexpected} in event: {event}")
+
+        return event