nvidia-nat 1.2.0__py3-none-any.whl

nat/registry_handlers/rest/rest_handler.py

@@ -0,0 +1,237 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import base64
+import logging
+import os
+import shutil
+import subprocess
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+import httpx
+
+from nat.registry_handlers.registry_handler_base import AbstractRegistryHandler
+from nat.registry_handlers.schemas.headers import RequestHeaders
+from nat.registry_handlers.schemas.package import PackageNameVersionList
+from nat.registry_handlers.schemas.publish import Artifact
+from nat.registry_handlers.schemas.publish import PublishResponse
+from nat.registry_handlers.schemas.pull import PullRequestPackages
+from nat.registry_handlers.schemas.pull import PullResponse
+from nat.registry_handlers.schemas.remove import RemoveResponse
+from nat.registry_handlers.schemas.search import SearchQuery
+from nat.registry_handlers.schemas.search import SearchResponse
+from nat.registry_handlers.schemas.status import ActionEnum
+from nat.registry_handlers.schemas.status import StatusEnum
+
+logger = logging.getLogger(__name__)
+
+
+class RestRegistryHandler(AbstractRegistryHandler):
+    """A registry handler for interactions with a remote REST registry."""
+
+    def __init__(  # pylint: disable=R0917
+            self,
+            endpoint: str,
+            token: str,
+            timeout: int = 30,
+            publish_route: str = "",
+            pull_route: str = "",
+            search_route: str = "",
+            remove_route: str = ""):
+        super().__init__()
+        self._endpoint = endpoint.rstrip("/")
+        self._timeout = timeout
+        self._publish_route = publish_route.strip("/")
+        self._pull_route = pull_route.strip("/")
+        self._search_route = search_route.strip("/")
+        self._remove_route = remove_route.strip("/")
+        self._headers = RequestHeaders(Authorization=f"Bearer: {token}").model_dump(by_alias=True)
+
+    @asynccontextmanager
+    async def publish(self, artifact: Artifact) -> AsyncGenerator[PublishResponse]:
+        """Publishes a NAT artifact to a remote REST registry.
+
+        Args:
+            artifact (Artifact): An artifact that contain NAT plugin wheel and it's corrosponding discovery
+            metadata.
+
+        Yields:
+            Iterator[AsyncGenerator[PublishResponse]]: A response message that includes a completion status message.
+        """
+
+        try:
+
+            async with httpx.AsyncClient(headers=self._headers, timeout=self._timeout) as client:
+                response = await client.post(f"{self._endpoint}/{self._publish_route}",
+                                             content=artifact.artifact.model_dump_json())
+                response.raise_for_status()
+                response_dict = response.json()
+
+            validated_publish_response = PublishResponse(**response_dict)
+
+            yield validated_publish_response
+
+        except Exception as e:
+            msg = f"Error publishing package: {e}"
+            validated_publish_response = PublishResponse(status={
+                "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PUBLISH
+            })
+            logger.exception(validated_publish_response.status.message, exc_info=True)
+
+            yield validated_publish_response
+
+        finally:
+            logger.info("Execution complete.")
+
+    @asynccontextmanager
+    async def pull(self, packages: PullRequestPackages) -> AsyncGenerator[PullResponse]:
+        """Download and install NAT artifacts from a remote REST registry.
+
+        Args:
+            packages (PullRequestPackages): Parameters used to pull the NAT artifact.
+
+        Yields:
+            Iterator[AsyncGenerator[PullResponse]]: A response message that includes a the pulled packages and a
+                completion status message.
+        """
+
+        tmp_dir = "./.tmp/nat-pull"
+
+        try:
+            async with httpx.AsyncClient(headers=self._headers, timeout=self._timeout) as client:
+                packages_json = packages.model_dump_json()
+                response = await client.post(f"{self._endpoint}/{self._pull_route}", content=packages_json)
+                response.raise_for_status()
+
+            response_dict = response.json()
+            validated_pull_response = PullResponse(**response_dict)
+
+            if (validated_pull_response.status.status == StatusEnum.ERROR):
+                logger.error(validated_pull_response.status.message)
+                raise ValueError(f"Server error: {validated_pull_response.status.message}")
+
+            if (not os.path.exists(tmp_dir)):
+                os.mkdir(tmp_dir)
+
+            whl_paths = []
+
+            for package in validated_pull_response.packages:
+                whl_bytes = base64.b64decode(package.whl)
+                whl_path = os.path.join(tmp_dir, package.whl_name)
+
+                with open(whl_path, "wb") as f:
+                    f.write(whl_bytes)
+
+                whl_paths.append(whl_path)
+
+            cmd = ["uv", "pip", "install"]
+            cmd.extend(whl_paths)
+            result = subprocess.run(cmd, check=True)
+            result.check_returncode()
+
+            if (os.path.exists(tmp_dir)):
+                shutil.rmtree(tmp_dir)
+
+            yield validated_pull_response
+
+        except Exception as e:
+            msg = f"Error occured when installing packages: {e}"
+            logger.error(msg)
+            if (os.path.exists(tmp_dir)):
+                shutil.rmtree(tmp_dir)
+
+            validated_pull_response = PullResponse(status={
+                "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.PULL
+            })
+            logger.exception(validated_pull_response.status.message, exc_info=True)
+
+            yield validated_pull_response
+
+        finally:
+            logger.info("Execution complete.")
+
+    @asynccontextmanager
+    async def search(self, query: SearchQuery) -> AsyncGenerator[SearchResponse]:
+        """Searches a remote REST registry for relevant NAT components.
+
+        Args:
+            query (SearchQuery): Parameters of the search to be performed.
+
+        Yields:
+            Iterator[AsyncGenerator[SearchResponse]]: A response message that includes search
+                parameters and a completion status message.
+        """
+
+        try:
+            async with httpx.AsyncClient(headers=self._headers, timeout=self._timeout) as client:
+                query_json = query.model_dump_json()
+                response = await client.post(url=f"{self._endpoint}/{self._search_route}", content=query_json)
+                response.raise_for_status()
+
+            response_dict = response.json()
+            validated_search_response = SearchResponse(**response_dict)
+
+            yield validated_search_response
+
+        except Exception as e:
+            msg = f"Error searching for artifacts: {e}"
+            validated_search_response = SearchResponse(params=query,
+                                                       status={
+                                                           "status": StatusEnum.ERROR,
+                                                           "message": msg,
+                                                           "action": ActionEnum.SEARCH
+                                                       })
+            logger.exception(validated_search_response.status.message, exc_info=True)
+
+            yield validated_search_response
+
+        finally:
+            logger.info("Execution complete.")
+
+    @asynccontextmanager
+    async def remove(self, packages: PackageNameVersionList) -> AsyncGenerator[RemoveResponse]:
+        """Removes packages from a remote REST registry.
+
+        Args:
+            packages (PackageNameVersionList): The list of packages to remove.
+
+        Yields:
+            Iterator[AsyncGenerator[RemoveResponse]]: A response message that includes the packages and a
+                completion status message.
+        """
+
+        try:
+            async with httpx.AsyncClient(headers=self._headers, timeout=self._timeout) as client:
+                remove_json = packages.model_dump_json()
+                response = await client.post(url=f"{self._endpoint}/{self._remove_route}", content=remove_json)
+                response.raise_for_status()
+
+            response_dict = response.json()
+            validated_remove_response = RemoveResponse(**response_dict)
+
+            yield validated_remove_response
+
+        except Exception as e:
+            msg = f"Error removing artifacts: {e}"
+            validated_remove_response = RemoveResponse(status={
+                "status": StatusEnum.ERROR, "message": msg, "action": ActionEnum.REMOVE
+            })
+            logger.exception(validated_remove_response.status.message, exc_info=True)
+
+            yield validated_remove_response
+
+        finally:
+            logger.info("Execution complete.")

nat/registry_handlers/schemas/headers.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import logging
+
+from pydantic import BaseModel
+from pydantic import Field
+
+logger = logging.getLogger(__name__)
+
+
+class RequestHeaders(BaseModel):
+    """Represents a data model for REST registry handler request headers.
+
+    Args:
+        accept (str): Specifies the media types the client can accept. Defaults to 'application/json'
+        content_type (str): Describes the format of the request body data. Defaults to 'application/json'
+        authorization (str): Contains authentication credentials for accessing a protected resource.
+    """
+
+    accept: str = Field("application/json", alias="Accept")
+    content_type: str = Field("application/json", alias="Content-Type")
+    authorization: str = Field(..., alias="Authorization")
+
+
+class ResponseHeaders(BaseModel):
+    """Placehoder data model for REST registry handler resopnse headers.
+    """
+
+    pass

nat/registry_handlers/schemas/package.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import logging
+
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class WheelData(BaseModel):
+    """Data model containing a built python wheel and its corresponding metadata.
+
+    Args:
+        package_root (str): The path to the package root directory containing the pyproject.toml file.
+        package_name (str): The name of the python package.
+        toml_project (dict): A dictionary containing data about the python project.
+        toml_dependencies (set): The list of dependencies provided in the pyproject.toml file.
+        toml_nat_packages (set): The NAT plugins listed in the pyproject.toml.
+        union_dependencies (set): The union of toml_dependencies and toml_nat_packages.
+        whl_path (str): The path to the package wheel file.
+        whl_base64 (str): Base64 encoded string of the wheel file.
+        whl_version (str): The version representing the wheel file.
+    """
+
+    package_root: str
+    package_name: str
+    toml_project: dict
+    toml_dependencies: set
+    toml_nat_packages: set
+    union_dependencies: set
+    whl_path: str
+    whl_base64: str
+    whl_version: str
+
+
+class PackageNameVersion(BaseModel):
+    """Represents a data model containing a package name and version.
+
+    Args:
+        name (str): Package name, excluding the version.
+        version (str | None): The package version, excluding the name. Defaults to None.
+    """
+
+    name: str
+    version: str | None = None
+
+
+class PackageNameVersionList(BaseModel):
+    """Represents a data model containing a list of `PackageNameVersion` packages.
+
+    Args:
+        packages (list[PackageNameVersion]): A list of `PackageNameVersion` models. Defaults to `[]`
+    """
+
+    packages: list[PackageNameVersion] = []

nat/registry_handlers/schemas/publish.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import logging
+
+from pydantic import BaseModel
+
+from nat.data_models.component import ComponentEnum
+from nat.data_models.discovery_metadata import DiscoveryMetadata
+from nat.registry_handlers.schemas.status import StatusMessage
+
+logger = logging.getLogger(__name__)
+
+
+class BuiltArtifact(BaseModel):
+    """A NAT artifact including base64 encoded string of wheel package and corrosponding discovery metadata.
+
+    Args:
+        whl (str): A base64 encoded string of a NAT package wheel (.whl).
+
+        metadata (dict[ComponentEnum, list[DiscoveryMetadata]]): Provides rich discover metadata for developers to
+        quickly find useful components.
+    """
+
+    whl: str
+    metadata: dict[ComponentEnum, list[DiscoveryMetadata]]
+
+
+class Artifact(BaseModel):
+    """A NAT artifact including base64 encoded string of wheel package and corrosponding discovery metadata.
+
+    Args:
+        artifact (BuiltArtifact): A NAT artifact including base64 encoded string of wheel package and
+        corrosponding discovery metadata.
+
+        whl_path (str): A local path to the built wheel package.
+    """
+
+    artifact: BuiltArtifact | None = None
+    whl_path: str
+
+
+class PublishResponse(BaseModel):
+    """The expected response from a publish request denoting status information.
+
+    Args:
+        status (StatusMessage): Provides metadata describing the success or errors that occurred when
+            making a publish request.
+    """
+
+    status: StatusMessage
+
+
+# Compatibility aliases with previous releases
+BuiltAIQArtifact = BuiltArtifact
+AIQArtifact = Artifact

nat/registry_handlers/schemas/pull.py

@@ -0,0 +1,82 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import logging
+
+from pydantic import BaseModel
+
+from nat.registry_handlers.schemas.package import PackageNameVersion
+from nat.registry_handlers.schemas.status import StatusMessage
+
+logger = logging.getLogger(__name__)
+
+
+class PulledPackage(BaseModel):
+    """Represents a data model of a pulled package containing the package wheel and its name.
+
+    Args:
+        whl (str): Base64 encoded string of the NAT python package wheel (.whl).
+        whl_name (str): A string representing the wheel filename.
+    """
+
+    whl: str
+    whl_name: str
+
+
+class PullResponse(BaseModel):
+    """
+    Represents a data model of the expected respones from a NAT pull request, including detailed status
+    information.
+
+    Args:
+        packages (list[PulledPackage]): A list of pulled packages included in the pull request.
+        status (StatusMessage): Provides metadata describing the success or errors that occurred when making to pull in
+        a package.
+    """
+
+    packages: list[PulledPackage] = []
+    status: StatusMessage
+
+
+class PullPackageWhl(BaseModel):
+    """Local path to wheel (.whl) file.
+
+    Args:
+        whl_path (str): The local path the wheel (.whl) file.
+    """
+
+    whl_path: str
+
+
+class PullRequestPackage(BaseModel):
+    """Represents all data for a single package needed to download an install its components.
+
+    Args:
+        package (typing.Union[PackageNameVersion, PullPackageWhl]): Attributes of a single package necessary
+            to download and install its components.
+    """
+
+    package: PackageNameVersion | PullPackageWhl
+
+
+class PullRequestPackages(BaseModel):
+    """Represents a list of all packages th download and install in the local NAT environment.
+
+    Args:
+        packages (list[typing.Union[PackageNameVersion, PullPackageWhl]]): A list of packages that can be
+            downloaded and installed in the local NAT environment.
+    """
+
+    packages: list[PackageNameVersion | PullPackageWhl]

nat/registry_handlers/schemas/remove.py

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import logging
+
+from pydantic import BaseModel
+
+from nat.registry_handlers.schemas.package import PackageNameVersion
+from nat.registry_handlers.schemas.status import StatusMessage
+
+logger = logging.getLogger(__name__)
+
+
+class RemoveResponse(BaseModel):
+    """Represents a data model for the expected response from a remove request, including packages and status metadata.
+
+    Args:
+        packages (list[PackageNameVersion]): A list of packages that are to be removed from a remote registry.
+        status (StatusMessage): Provides metadata describing the success or errors that occurred when making a remove
+        request.
+    """
+
+    packages: list[PackageNameVersion] = []
+    status: StatusMessage

nat/registry_handlers/schemas/search.py

@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import logging
+from enum import Enum
+
+from pydantic import BaseModel
+
+from nat.data_models.component import ComponentEnum
+from nat.registry_handlers.schemas.status import StatusMessage
+
+logger = logging.getLogger(__name__)
+
+
+class SearchFields(str, Enum):
+    ALL = "all"
+    PACKAGE = "package"
+    VERSION = "version"
+    COMPONENT_NAME = "component_name"
+    DESCRIPTION = "description"
+    DEVELOPER_NOTES = "developer_notes"
+
+
+class VisualizeFields(str, Enum):
+    PACKAGE = "package"
+    VERSION = "version"
+    COMPONENT_TYPE = "component_type"
+    COMPONENT_NAME = "component_name"
+    DESCRIPTION = "description"
+
+
+class SearchQuery(BaseModel):
+    """Represents the search criteria that will be used to discover useful NAT components.
+
+    Args:
+        query (str): A query string used to find useful NAT components.
+        fields (list[SearchFields]): The list of fields used when applying the query string.
+        component_types (list[ComponentEnum]): NAT components types to filter search results.
+        top_k (int): Specifies the number of search results to provide.
+    """
+
+    query: str = "*"
+    fields: list[SearchFields] = [SearchFields.ALL]
+    component_types: list[ComponentEnum]
+    top_k: int = 10
+
+
+class SearchResponseItem(BaseModel):
+    """Represents an individual item in the search response, including elements of it's discovery metadata.
+
+    Args:
+        package (str): The name of the NAT package that includes the component.
+        version (str): The version of the NAT package that includes the component.
+        component_type (ComponentEnum): Type of NAT component this item represents.
+        description (str): A description of this NAT component.
+        developer_notes (str): Additional details that would help a developer use this component.
+    """
+
+    package: str
+    version: str
+    component_type: ComponentEnum
+    component_name: str
+    description: str
+    developer_notes: str
+
+
+class SearchResponse(BaseModel):
+    """Represents a data model of the expected search response.
+
+    Args:
+        results (list[SearchResponseItem]): A list of results that matched the search criteria.
+        params (SearchQuery): The search criterial that produced these search results.
+        status (StatusMessage): Provides metadata describing the success or errors that occurred when making the search
+        request.
+    """
+
+    results: list[SearchResponseItem] = []
+    params: SearchQuery
+    status: StatusMessage

nat/registry_handlers/schemas/status.py

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+import logging
+from enum import Enum
+
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class ActionEnum(str, Enum):
+    PUBLISH = "publish"
+    PULL = "pull"
+    REMOVE = "remove"
+    SEARCH = "search"
+
+
+class StatusEnum(str, Enum):
+    SUCCESS = "success"
+    ERROR = "error"
+
+
+class StatusMessage(BaseModel):
+    """Represents a data model to record success or error when performing registry interactions.
+
+    Args:
+        status (StatusEnum): Represents the outcome (success or error) of the registry interaction.
+        action: (ActionEnum): Represents the type of registry action that was taken.
+        message: (str): Provides a more detailed status message for the registry interaction.
+    """
+
+    status: StatusEnum
+    action: ActionEnum
+    message: str