hiddenlayer-sdk 0.1.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

hiddenlayer/sdk/services/model_scan.py

@@ -1,19 +1,24 @@
+import json
 import os
 import random
+import tempfile
 import time
 import warnings
+import zipfile
 from datetime import datetime
 from pathlib import Path
 from typing import List, Optional, Union
 from uuid import uuid4
 
+from pydantic_core import ValidationError
+
 from hiddenlayer.sdk.constants import ScanStatus
-from hiddenlayer.sdk.enterprise.enterprise_model_scan_api import EnterpriseModelScanApi
-from hiddenlayer.sdk.models import ScanResults
-from hiddenlayer.sdk.rest.api import ModelScanApi, SensorApi
+from hiddenlayer.sdk.models import EmptyScanResults, Sarif, ScanResults
+from hiddenlayer.sdk.rest.api import ModelScanApi, ModelSupplyChainApi, SensorApi
 from hiddenlayer.sdk.rest.api_client import ApiClient
 from hiddenlayer.sdk.rest.models import MultipartUploadPart
 from hiddenlayer.sdk.rest.models.model import Model
+from hiddenlayer.sdk.rest.models.sarif210 import Sarif210
 from hiddenlayer.sdk.services.model import ModelAPI
 from hiddenlayer.sdk.utils import filter_path_objects, is_saas
 
@@ -31,24 +36,22 @@ EXCLUDE_FILE_TYPES = [
 
 class ModelScanAPI:
     def __init__(self, api_client: ApiClient) -> None:
-        self.is_saas = is_saas(api_client.configuration.host)
         self._api_client = api_client
 
-        if self.is_saas:
-            self._model_scan_api = ModelScanApi(api_client=api_client)
-            self._model_api = ModelAPI(api_client=api_client)
-            self._sensor_api = SensorApi(
-                api_client=api_client
-            )  # lower level api of ModelAPI
-        else:
-            self._model_scan_api = EnterpriseModelScanApi(api_client=api_client)
+        self._model_supply_chain_api = ModelSupplyChainApi(api_client=api_client)
+        self._model_api = ModelAPI(api_client=api_client)
+        self._sensor_api = SensorApi(
+            api_client=api_client
+        )  # lower level api of ModelAPI
+
+        self._model_scan_api = ModelScanApi(api_client=api_client)
 
     def scan_file(
         self,
         *,
         model_name: str,
         model_path: Union[str, os.PathLike],
-        threads: int = 1,
+        model_version: Optional[int] = None,
         chunk_size: int = 16,
         wait_for_results: bool = True,
     ) -> ScanResults:
@@ -57,67 +60,45 @@ class ModelScanAPI:
 
         :param model_name: Name of the model to be shown on the HiddenLayer UI
         :param model_path: Local path to the model file.
-        :param threads: Number of threads used to upload the file, defaults to 1.
+        :param model_version: Version of the model to be shown on the HiddenLayer UI.
         :param chunk_size: Number of chunks of the file to upload at once, defaults to 4.
         :param wait_for_results: True whether to wait for the scan to finish, defaults to True.
 
         :returns: Scan Results
         """
 
-        warnings.warn(
-            "Use of the threads parameter is deprecated and will be removed in version 0.2.0.",
-            category=DeprecationWarning,
-            stacklevel=2,
-        )
-
         file_path = Path(model_path)
 
-        # Can combine the 2 paths when SaaS API and Enterprise APIs are in sync
-        if self.is_saas:
-            filesize = file_path.stat().st_size
-            sensor = self._model_api.create(model_name=model_name)
-            upload = self._sensor_api.begin_multipart_upload(filesize, sensor.sensor_id)
-
-            with open(file_path, "rb") as f:
-                for i in range(0, len(upload.parts), chunk_size):
-                    group: List[MultipartUploadPart] = upload.parts[i : i + chunk_size]
-                    for part in group:
-                        read_amount = part.end_offset - part.start_offset
-                        f.seek(int(part.start_offset))
-                        part_data = f.read(int(read_amount))
-
-                        # The SaaS multipart upload returns a upload url for each part
-                        # So there is no specified route
-                        self._api_client.call_api(
-                            "PUT",
-                            part.upload_url,
-                            body=part_data,
-                            header_params={"Content-Type": "application/octet-binary"},
-                        )
-
-            self._sensor_api.complete_multipart_upload(
-                sensor.sensor_id, upload.upload_id
-            )
-
-            self._model_scan_api.scan_model(sensor.sensor_id)
-        else:
-            with open(file_path, "rb") as f:
-                data = f.read()
-
-            sensor = Model(
-                sensor_id=str(uuid4()),
-                created_at=datetime.now(),
-                tenant_id="0000",
-                plaintext_name=model_name,
-                active=True,
-                version=1,
-            )
-
-            self._model_scan_api: EnterpriseModelScanApi
-            self._model_scan_api.scan_model(sensor.sensor_id, data)
-            model_name = sensor.sensor_id
-
-        scan_results = self.get_scan_results(model_name=model_name)
+        filesize = file_path.stat().st_size
+        sensor = self._model_api.create_or_get(
+            model_name=model_name, model_version=model_version
+        )
+        upload = self._sensor_api.begin_multipart_upload(sensor.sensor_id, filesize)
+
+        with open(file_path, "rb") as f:
+            for i in range(0, len(upload.parts), chunk_size):
+                group: List[MultipartUploadPart] = upload.parts[i : i + chunk_size]
+                for part in group:
+                    read_amount = part.end_offset - part.start_offset
+                    f.seek(int(part.start_offset))
+                    part_data = f.read(int(read_amount))
+
+                    # The SaaS multipart upload returns a upload url for each part
+                    # So there is no specified route
+                    self._api_client.call_api(
+                        "PUT",
+                        part.upload_url,
+                        body=part_data,
+                        header_params={"Content-Type": "application/octet-binary"},
+                    )
+
+        self._sensor_api.complete_multipart_upload(sensor.sensor_id, upload.upload_id)
+
+        self._model_scan_api.scan_model(sensor.sensor_id)
+
+        scan_results = self.get_scan_results(
+            model_name=model_name, model_version=model_version
+        )
 
         base_delay = 0.1  # seconds
         retries = 0
@@ -129,13 +110,13 @@ class ModelScanAPI:
                     0, 1
                 )  # exponential back off retry
                 time.sleep(delay)
-                scan_results = self.get_scan_results(model_name=model_name)
+                scan_results = self.get_scan_results(
+                    model_name=model_name, model_version=model_version
+                )
                 print(f"{file_path.name} scan status: {scan_results.status}")
 
-        scan_results = ScanResults.from_scanresultsv2(scan_results_v2=scan_results)
         scan_results.file_name = file_path.name
         scan_results.file_path = str(file_path)
-        scan_results.sensor_id = sensor.sensor_id
 
         return scan_results
 
@@ -145,8 +126,8 @@ class ModelScanAPI:
         model_name: str,
         bucket: str,
         key: str,
+        model_version: Optional[int] = None,
         s3_client: Optional[object] = None,
-        threads: int = 1,
         chunk_size: int = 4,
         wait_for_results: bool = True,
     ) -> ScanResults:
@@ -156,9 +137,9 @@ class ModelScanAPI:
         :param model_name: Name of the model to be shown on the HiddenLayer UI.
         :param bucket: Name of the s3 bucket where the model file is stored.
         :param key: Path to the model file on s3.
+        :param model_version: Version of the model to be shown on the HiddenLayer UI.
         :param wait_for_results: True whether to wait for the scan to finish, defaults to True.
         :param s3_client: boto3 s3 client.
-        :param threads: Number of threads used to upload the file, defaults to 1.
         :param chunk_size: Number of chunks of the file to upload at once, defaults to 4.
         :param wait_for_results: True whether to wait for the scan to finish, defaults to True.
 
@@ -174,7 +155,7 @@ class ModelScanAPI:
                 )
         """
         try:
-            import boto3
+            import boto3  # type: ignore
         except ImportError:
             raise ImportError("Python package boto3 is not installed.")
 
@@ -184,14 +165,14 @@ class ModelScanAPI:
         file_name = key.split("/")[-1]
 
         try:
-            s3_client.download_file(bucket, key, f"/tmp/{file_name}")
+            s3_client.download_file(bucket, key, f"/tmp/{file_name}")  # type: ignore
         except Exception as e:
             raise RuntimeError(f"Couldn't download model s3://{bucket}/{key}: {e}")
 
         return self.scan_file(
             model_path=f"/tmp/{file_name}",
             model_name=model_name,
-            threads=threads,
+            model_version=model_version,
             chunk_size=chunk_size,
             wait_for_results=wait_for_results,
         )
@@ -203,9 +184,9 @@ class ModelScanAPI:
         account_url: str,
         container: str,
         blob: str,
+        model_version: Optional[int] = None,
         blob_service_client: Optional[object] = None,
         credential: Optional[object] = None,
-        threads: int = 1,
         chunk_size: int = 4,
         wait_for_results: bool = True,
     ) -> ScanResults:
@@ -216,10 +197,10 @@ class ModelScanAPI:
         :param account_url: Azure Blob url of where the file is stored.
         :param container: Azure Blob container containing the model file.
         :param blob: Path to the model file inside the Azure blob container.
+        :param model_version: Version of the model to be shown on the HiddenLayer UI.
         :param blob_service_client: BlobServiceClient object. Defaults to creating one using DefaultCredential().
         :param credential: Credential to be passed to the BlobServiceClient object, can be a credential object, SAS key, etc.
             Defaults to `DefaultCredential`
-        :param threads: Number of threads used to upload the file, defaults to 1.
         :param chunk_size: Number of chunks of the file to upload at once, defaults to 4.
         :param wait_for_results: True whether to wait for the scan to finish, defaults to True.
 
@@ -250,11 +231,11 @@ class ModelScanAPI:
             credential = DefaultAzureCredential()
 
         if not blob_service_client:
-            blob_service_client = BlobServiceClient(account_url, credential=credential)
+            blob_service_client = BlobServiceClient(account_url, credential=credential)  # type: ignore
 
         file_name = blob.split("/")[-1]
 
-        blob_client = blob_service_client.get_blob_client(
+        blob_client = blob_service_client.get_blob_client(  # type: ignore
             container=container, blob=blob
         )
 
@@ -271,7 +252,7 @@ class ModelScanAPI:
         return self.scan_file(
             model_path=f"/tmp/{file_name}",
             model_name=model_name,
-            threads=threads,
+            model_version=model_version,
             chunk_size=chunk_size,
             wait_for_results=wait_for_results,
         )
@@ -289,15 +270,15 @@ class ModelScanAPI:
         force_download: bool = False,
         hf_token: Optional[Union[str, bool]] = None,
         # HL parameters
-        threads: int = 1,
         chunk_size: int = 4,
         wait_for_results: bool = True,
-    ) -> List[ScanResults]:
+    ) -> ScanResults:
         """
         Scans a model on HuggingFace.
 
         Note: Requires the `huggingface_hub` pip package to be installed.
 
+        :param repo_id: The HuggingFace repository id.
         :param revision: An optional Git revision id which can be a branch name, a tag, or a commit hash.
         :param local_dir: If provided, the downloaded files will be placed under this directory.
         :param allow_file_patterns: If provided, only files matching at least one pattern are scanned.
@@ -306,7 +287,6 @@ class ModelScanAPI:
         :param hf_token: A token to be used for the download.
             If True, the token is read from the HuggingFace config folder.
             If a string, it’s used as the authentication token.
-        :param threads: Number of threads used to upload the file, defaults to 1.
         :param chunk_size: Number of chunks of the file to upload at once, defaults to 4.
         :param wait_for_results: True whether to wait for the scan to finish, defaults to True.
 
@@ -337,52 +317,117 @@ class ModelScanAPI:
         )
 
         return self.scan_folder(
+            model_name=repo_id,
             path=local_dir,
             allow_file_patterns=allow_file_patterns,
             ignore_file_patterns=ignore_file_patterns,
-            threads=threads,
             chunk_size=chunk_size,
             wait_for_results=wait_for_results,
         )
 
-    def get_scan_results(self, *, model_name: str) -> ScanResults:
+    def get_scan_results(
+        self,
+        *,
+        model_name: str,
+        model_version: Optional[int] = None,
+    ) -> ScanResults:
         """
         Get results from a model scan.
 
         :param model_name: Name of the model.
+        :param model_version: Version of the model. When the model version is not specified, the scan results for the latest version will be returned.
 
         :returns: Scan results.
         """
 
-        if self.is_saas:
-            model = self._model_api.get(model_name=model_name)
-            sensor_id = model.sensor_id
-        else:
-            sensor_id = model_name
+        response = self._sensor_api.sensor_sor_api_v3_model_cards_query_get(
+            model_name_eq=model_name, limit=1
+        )
+        model_id = response.results[0].model_id
 
-        scan_results_v2 = self._model_scan_api.scan_status(sensor_id)
+        scans = self._model_supply_chain_api.model_scan_api_v3_scan_query(
+            model_ids=[model_id], latest_per_model_version_only=True
+        )
+        if scans.total == 0:
+            return EmptyScanResults()
+
+        if scans.items is None:
+            return EmptyScanResults()
+
+        scan = scans.items[0]
+        if model_version:
+            scan = next(
+                (
+                    s
+                    for s in scans.items
+                    if s.inventory.model_version == str(model_version)
+                ),
+                None,
+            )
+        if not scan:
+            return EmptyScanResults()
+
+        scan_report = (
+            self._model_supply_chain_api.model_scan_api_v3_scan_model_version_id_get(
+                scan.scan_id
+            )
+        )
 
-        return ScanResults.from_scanresultsv2(
-            scan_results_v2=scan_results_v2, sensor_id=sensor_id
+        return ScanResults.from_scanreportv3(
+            scan_report_v3=scan_report, model_id=model_id
         )
 
+    def get_sarif_results(
+        self,
+        *,
+        model_name: str,
+        model_version: Optional[int] = None,
+    ) -> Optional[Sarif]:
+        """
+        Get sarif results from a model scan.
+
+        :param model_name: Name of the model.
+        :param model_version: Version of the model. When the model version is not specified, the scan results for the latest version will be returned.
+
+        :returns: Scan results.
+        """
+        scan = self.get_scan_results(model_name=model_name, model_version=model_version)
+        if scan.scan_id == "":
+            return None
+
+        # Unfortunately, the generated code for the API doesn't directly support modifying the Accept header
+        # in order to enable us to get the Sarif results
+        # Here we will reach in to the request serialization process. The 2nd element in the tuple is the headers
+        # where we will modify the Accept header to application/sarif+json
+        request = self._model_supply_chain_api._model_scan_api_v3_scan_model_version_id_get_serialize(
+            scan.scan_id, None, None, None, None, 0
+        )
+        request[2]["Accept"] = "application/sarif+json"
+        response = self._api_client.call_api(*request)
+        response.read()
+        return self._api_client.response_deserialize(
+            response_data=response, response_types_map={"200": Sarif}
+        ).data  # type: ignore
+
     def scan_folder(
         self,
         *,
+        model_name: str,
         path: Union[str, os.PathLike],
+        model_version: Optional[int] = None,
         allow_file_patterns: Optional[List[str]] = None,
         ignore_file_patterns: Optional[List[str]] = None,
-        threads: int = 1,
         chunk_size: int = 4,
         wait_for_results: bool = True,
-    ) -> List[ScanResults]:
+    ) -> ScanResults:
         """
         Submits all files in a directory and its sub directories to be scanned.
 
+        :param model_name: Name of the model to be shown on the HiddenLayer UI.
         :param path: Path to the folder on disk to be scanned.
+        :param model_version: Version of the model to be shown on the HiddenLayer UI.
         :param allow_file_patterns: If provided, only files matching at least one pattern are scanned.
         :param ignore_file_patterns: If provided, files matching any of the patterns are not scanned.
-        :param threads: Number of threads used to upload the file, defaults to 1.
         :param chunk_size: Number of chunks of the file to upload at once, defaults to 4.
         :param wait_for_results: True whether to wait for the scan to finish, defaults to True.
 
@@ -390,6 +435,8 @@ class ModelScanAPI:
         """
 
         model_path = Path(path)
+        filename = tempfile.NamedTemporaryFile().name + ".zip"
+
         ignore_file_patterns = (
             EXCLUDE_FILE_TYPES + ignore_file_patterns
             if ignore_file_patterns
@@ -402,13 +449,14 @@ class ModelScanAPI:
             ignore_patterns=ignore_file_patterns,
         )
 
-        return [
-            self.scan_file(
-                model_name=str(file),
-                model_path=file,
-                threads=threads,
-                chunk_size=chunk_size,
-                wait_for_results=wait_for_results,
-            )
-            for file in files
-        ]
+        with zipfile.ZipFile(filename, "a") as zipf:
+            for file in files:
+                zipf.write(file, os.path.relpath(file, model_path))
+
+        return self.scan_file(
+            model_name=model_name,
+            model_version=model_version,
+            model_path=filename,
+            chunk_size=chunk_size,
+            wait_for_results=wait_for_results,
+        )

hiddenlayer/sdk/version.py

@@ -1 +1 @@
-VERSION = "0.1.2"
+VERSION = "1.1.0"

{hiddenlayer_sdk-0.1.2.dist-info → hiddenlayer_sdk-1.1.0.dist-info}/METADATA

@@ -1,10 +1,10 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: hiddenlayer-sdk
-Version: 0.1.2
+Version: 1.1.0
 Summary: Official HiddenLayer Python SDK
 Author-email: HiddenLayer Integrations Team <integrations@hiddenlayer.com>
 Maintainer-email: HiddenLayer Integrations Team <integrations@hiddenlayer.com>
-License: Apache License
+License:                                  Apache License
                                    Version 2.0, January 2004
                                 http://www.apache.org/licenses/
         
@@ -222,20 +222,20 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: requests
-Requires-Dist: pydantic >=1.10.9
-Requires-Dist: python-dateutil >=2.0.0
-Requires-Dist: numpy >=1.0.0
-Provides-Extra: aws
-Requires-Dist: boto3 >=1.0.0 ; extra == 'aws'
-Provides-Extra: azure
-Requires-Dist: azure-storage-blob ; extra == 'azure'
-Requires-Dist: azure-identity ; extra == 'azure'
+Requires-Dist: pydantic>=1.10.9
+Requires-Dist: python-dateutil>=2.0.0
+Requires-Dist: numpy>=1.0.0
 Provides-Extra: dev
-Requires-Dist: ruff >=0.2.2 ; extra == 'dev'
-Requires-Dist: pytest ; extra == 'dev'
-Requires-Dist: pandas ; extra == 'dev'
+Requires-Dist: ruff>=0.2.2; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
+Provides-Extra: aws
+Requires-Dist: boto3>=1.0.0; extra == "aws"
 Provides-Extra: hf
-Requires-Dist: huggingface-hub ; extra == 'hf'
+Requires-Dist: huggingface_hub; extra == "hf"
+Provides-Extra: azure
+Requires-Dist: azure-storage-blob; extra == "azure"
+Requires-Dist: azure-identity; extra == "azure"
 
 # HiddenLayer SDK Python (Beta)
 
@@ -277,25 +277,44 @@ hl_client = HiddenlayerServiceClient(
 )
 ```
 
-If you are using the Enterprise version of the production, you can instantiate the `HiddenlayerServiceClient` as follows:
+### Scanning Models
 
 ```python
-from hiddenlayer import HiddenlayerServiceClient
-
-hl_client = HiddenlayerServiceClient(
-  host="https://your.hiddenlayer.enterprise.url",
+hl_client.model_scanner.scan_file(
+  model_name="name_of_the_model",
+  model_path="path/to/model/file.pkl"
 )
 ```
 
-### Scanning Models
+If you would like to specify the version associated with the model, you can pass that in as well. Note that the model version must not already exist.
 
 ```python
 hl_client.model_scanner.scan_file(
   model_name="name_of_the_model",
+  model_version=1,
   model_path="path/to/model/file.pkl"
 )
 ```
 
+### Scanning Folders
+
+```python
+hl_client.model_scanner.scan_folder(
+  model_name="name_of_the_model",
+  path="path/to/model/"
+)
+```
+
+If you would like to specify the version associated with the model, you can pass that in as well. Note that the model version must not already exist.
+
+```python
+hl_client.model_scanner.scan_folder(
+  model_name="name_of_the_model",
+  model_version=1,
+  path="path/to/model/"
+)
+```
+
 ### Using AIDR for Predictive Models
 
 > Note: This is only supported using the SaaS version of the platform.