azureml-registry-tools 0.1.0a2__py3-none-any.whl → 0.1.0a4__py3-none-any.whl

azureml/registry/_rest_client/arm_client.py

@@ -1,16 +1,14 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
-from azure.identity import DefaultAzureCredential
-from .base_rest_client import BaseRestClient
-import time
+from .base_rest_client import BaseAzureRestClient
 from json.decoder import JSONDecodeError
 
 DEFAULT_API_VERSION = "2025-04-01"  # Default API version for Azure Resource Manager
 
 
-class ArmClient(BaseRestClient):
-    """Simple Azure Resource Manager (ARM) client leveraging BaseRestClient for GET, PATCH, PUT, and DELETE operations.
+class ArmClient(BaseAzureRestClient):
+    """Simple Azure Resource Manager (ARM) client leveraging BaseAzureRestClient for GET, PATCH, PUT, and DELETE operations.
 
     Handles authentication via Bearer token (pass as api_key) and supports standard ARM resource operations.
     """
@@ -25,27 +23,8 @@ class ArmClient(BaseRestClient):
             backoff_factor (int): Backoff factor for retries.
         """
         base_url = "https://management.azure.com"
-        self._credential = None
-        self._token_expires_on = None
-        if api_key is None:
-            # Use DefaultAzureCredential for authentication if no API key is provided
-            self._credential = DefaultAzureCredential()
-            token = self._credential.get_token("https://management.azure.com/.default")
-            api_key = token.token
-            self._token_expires_on = token.expires_on
         super().__init__(base_url, api_key=api_key, max_retries=max_retries, backoff_factor=backoff_factor)
 
-    def _refresh_api_key_if_needed(self) -> None:
-        """Refresh the API key if using DefaultAzureCredential and the token is close to expiration."""
-        if self._credential is not None:
-            now = int(time.time())
-            # Refresh if less than 10 minutes (600 seconds) left
-            if not self._token_expires_on or self._token_expires_on - now < 600:
-                token = self._credential.get_token("https://management.azure.com/.default")
-                self.api_key = token.token
-                self.session.headers.update({'Authorization': f'Bearer {self.api_key}'})
-                self._token_expires_on = token.expires_on
-
     def get_resource(self, resource_id, api_version=DEFAULT_API_VERSION, **kwargs) -> object:
         """
         Get an ARM resource by its resource ID.

azureml/registry/_rest_client/base_rest_client.py

@@ -4,6 +4,11 @@
 import requests
 import time
 import logging
+from azure.identity import DefaultAzureCredential
+from diskcache import Cache
+import tempfile
+
+ENABLE_CACHE = True  # Set to False to disable diskcache usage for API key
 
 
 class BaseRestClient:
@@ -136,3 +141,61 @@ class BaseRestClient:
             requests.Response: The HTTP response object.
         """
         return self._request_with_retry('DELETE', url, **kwargs)
+
+
+class BaseAzureRestClient(BaseRestClient):
+    """Base Azure REST client that handles Azure authentication and token refresh, with diskcache for API key."""
+
+    if ENABLE_CACHE:
+        _cache_dir = tempfile.gettempdir() + '/.azureml_registry_token_cache'
+        _cache = Cache(_cache_dir)
+    else:
+        _cache = None
+
+    def __init__(self, base_url: str, api_key: str = None, max_retries: int = 5, backoff_factor: int = 1):
+        """
+        Initialize the BaseAzureRestClient.
+
+        Args:
+            base_url (str): The base URL for the REST API.
+            api_key (str, optional): Bearer token for authentication. If None, uses DefaultAzureCredential.
+            max_retries (int): Maximum number of retries for failed requests.
+            backoff_factor (int): Backoff factor for retry delays.
+        """
+        self._credential = None
+        self._token_expires_on = None
+        cache_key = f"azureml_api_key_{base_url}"
+        cache_expiry_key = f"azureml_api_key_expiry_{base_url}"
+        cached_token = None
+        cached_expiry = None
+        now = int(time.time())
+        if ENABLE_CACHE and self._cache is not None:
+            cached_token = self._cache.get(cache_key)
+            cached_expiry = self._cache.get(cache_expiry_key)
+        if api_key is None and cached_token and cached_expiry and cached_expiry > now:
+            api_key = cached_token
+            self._token_expires_on = cached_expiry
+        super().__init__(base_url, api_key=api_key, max_retries=max_retries, backoff_factor=backoff_factor)
+        # Only after super().__init__ is self.session available
+        if api_key is None:
+            self._credential = DefaultAzureCredential()
+            self._refresh_api_key_if_needed()
+            # Ensure self.api_key is set for future use
+            api_key = self.api_key
+
+    def _refresh_api_key_if_needed(self) -> None:
+        if self._credential is None:
+            return
+        now = int(time.time())
+        if not getattr(self, '_token_expires_on', None) or self._token_expires_on - now < 600:
+            token = self._credential.get_token("https://management.azure.com/.default")
+            self.api_key = token.token
+            self.session.headers.update({'Authorization': f'Bearer {self.api_key}'})
+            self._token_expires_on = token.expires_on
+            # Store in diskcache if enabled
+            if ENABLE_CACHE and self._cache is not None:
+                base_url = getattr(self, 'base_url', 'default')
+                cache_key = f"azureml_api_key_{base_url}"
+                cache_expiry_key = f"azureml_api_key_expiry_{base_url}"
+                self._cache.set(cache_key, self.api_key, expire=(self._token_expires_on - now))
+                self._cache.set(cache_expiry_key, self._token_expires_on, expire=(self._token_expires_on - now))

azureml/registry/_rest_client/registry_management_client.py

@@ -1,18 +1,16 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
-from azure.identity import DefaultAzureCredential
-from .base_rest_client import BaseRestClient
-import time
+from .base_rest_client import BaseAzureRestClient
 
 
-class RegistryManagementClient(BaseRestClient):
+class RegistryManagementClient(BaseAzureRestClient):
     """Python client for RegistrySyndicationManifestController (excluding S2S APIs).
 
     Handles authentication, token refresh, and provides methods for manifest and registry management.
     """
 
-    def __init__(self, registry_name: str, primary_region: str = "eastus2euap", api_key: str = None, max_retries: int = 5, backoff_factor: int = 1) -> None:
+    def __init__(self, registry_name: str, primary_region: str = None, api_key: str = None, max_retries: int = 5, backoff_factor: int = 1) -> None:
         """
         Initialize the RegistryManagementClient.
 
@@ -23,29 +21,26 @@ class RegistryManagementClient(BaseRestClient):
             max_retries (int): Maximum number of retries for failed requests.
             backoff_factor (int): Backoff factor for retry delays.
         """
+        if primary_region is None:
+            # Resolve the primary region if not provided
+            primary_region = self.resolve_registry_primary_region(registry_name)
         base_url = f"https://{primary_region}.api.azureml.ms"
-        self._credential = None
-        self._token_expires_on = None
-        if api_key is None:
-            # Use DefaultAzureCredential for authentication if no API key is provided
-            self._credential = DefaultAzureCredential()
-            token = self._credential.get_token("https://management.azure.com/.default")
-            api_key = token.token
-            self._token_expires_on = token.expires_on
         super().__init__(base_url, api_key=api_key, max_retries=max_retries, backoff_factor=backoff_factor)
         self.registry_name = registry_name
 
-    def _refresh_api_key_if_needed(self) -> None:
-        """Refresh the API key if using DefaultAzureCredential and the token is close to expiration."""
-        # Only refresh if using DefaultAzureCredential
-        if self._credential is not None:
-            now = int(time.time())
-            # Refresh if less than 10 minutes (600 seconds) left
-            if not self._token_expires_on or self._token_expires_on - now < 600:
-                token = self._credential.get_token("https://management.azure.com/.default")
-                self.api_key = token.token
-                self.session.headers.update({'Authorization': f'Bearer {self.api_key}'})
-                self._token_expires_on = token.expires_on
+    @staticmethod
+    def resolve_registry_primary_region(registry_name: str) -> str:
+        """
+        Resolve the primary region for the given registry name.
+
+        Args:
+            registry_name (str): The name of the AzureML registry.
+
+        Returns:
+            str: The primary region for the registry.
+        """
+        discovery = RegistryManagementClient(registry_name, primary_region="eastus").discovery()
+        return discovery.get('primaryRegion', 'eastus')
 
     def create_or_update_manifest(self, manifest_dto: dict) -> dict:
         """
@@ -110,5 +105,15 @@ class RegistryManagementClient(BaseRestClient):
         """
         self._refresh_api_key_if_needed()
         url = f"{self.base_url}/registrymanagement/v1.0/registries/{self.registry_name}/discovery"
-        response = self.get(url)
-        return response.json()
+        try:
+            response = self.get(url)
+            return response.json()
+        except Exception as ex:
+            # Special handling for HTTP errors with status_code attribute
+            if hasattr(ex, 'response') and ex.response is not None:
+                status_code = ex.response.status_code
+                if status_code == 403:
+                    raise RuntimeError(f"Received 403 Forbidden. This may indicate that the registry '{self.registry_name}' does not exist or you do not have access.")
+                else:
+                    raise RuntimeError(f"Failed to get discovery information: {status_code} {ex.response.text}")
+            raise RuntimeError(f"Error occurred while trying to get discovery information: {ex}")

azureml/registry/tools/create_or_update_assets.py

@@ -5,6 +5,7 @@
 """Create or update assets."""
 
 import azure
+import copy
 import json
 import os
 import re
@@ -14,7 +15,8 @@ from git import InvalidGitRepositoryError, Repo
 from pathlib import Path
 from typing import List, Union
 
-from azure.ai.ml import MLClient, load_data, load_model
+from azure.ai.ml import MLClient, load_component, load_data, load_model
+from azure.ai.ml.entities._component.component import Component
 from azure.ai.ml.entities._assets._artifacts.data import Data
 from azure.ai.ml.entities._assets._artifacts.model import Model
 from azure.identity import DefaultAzureCredential
@@ -121,38 +123,52 @@ def find_assets(input_dir: Path,
 
 
 def merge_yamls(existing_asset_file_name: str,
-                updated_asset_file_name: str) -> dict:
-    """Overwrite values from updated asset to existing asset with the exception of storage path.
+                updated_asset: AssetSpec) -> dict:
+    """Preserve storage value from existing asset to updated asset.
 
     Args:
         existing_asset_file_name (str): Name of file containing details of asset currently in registry.
-        updated_asset_file_name (str): Name of file containing updated asset.
-
+        updated_asset (AssetSpec): AssetSpec containing updated asset info.
     Returns:
         dict: Merged asset data.
     """
-    with open(existing_asset_file_name, "r") as existing_asset_file, open(updated_asset_file_name, "r") as updated_asset_file:
-        existing_asset_yaml = yaml.safe_load(existing_asset_file)
-        updated_asset_yaml = yaml.safe_load(updated_asset_file)
-    merged_asset = {**existing_asset_yaml, **updated_asset_yaml}
-    merged_asset["path"] = existing_asset_yaml["path"]
+    with open(existing_asset_file_name, "r") as existing_asset_file, open(updated_asset.file_name_with_path, "r") as updated_asset_file:
+        existing_asset_dict = yaml.safe_load(existing_asset_file)
+        updated_asset_dict = yaml.safe_load(updated_asset_file)
+
+    merged_asset = copy.deepcopy(updated_asset_dict)
+
+    if updated_asset.type == AssetType.MODEL or updated_asset.type == AssetType.DATA:
+        storage_key = "path"
+        # "path" is always required for models, use the existing value
+        if existing_asset_dict.get(storage_key) != updated_asset_dict.get(storage_key):
+            print(f'"{storage_key}" is immutable, using existing value from asset in registry: {existing_asset_dict[storage_key]}')
+            merged_asset[storage_key] = existing_asset_dict[storage_key]
+    elif updated_asset.type == AssetType.COMPONENT:
+        storage_key = "code"
+        if existing_asset_dict.get(storage_key) is None and updated_asset_dict.get(storage_key) is not None:
+            print(f'"{storage_key}" is immutable and cannot be added. Please create a new version to include "{storage_key}"')
+        elif existing_asset_dict.get(storage_key) != updated_asset_dict.get(storage_key):
+            print(f'"{storage_key}" is immutable, using existing value from asset in registry: {existing_asset_dict[storage_key]}')
+            merged_asset[storage_key] = existing_asset_dict[storage_key]
 
     return merged_asset
 
 
-def merge_assets(existing_asset: Union[Data, Model],
+def merge_assets(existing_asset: Union[Component, Data, Model],
                  asset: AssetSpec):
-    """Overwrite values from updated to existing asset while preserving existing asset path.
+    """Prepare asset for update by preserving storage value from existing asset.
 
     Args:
-        existing_asset (Union[Data, Model]): Asset currently in registry.
+        existing_asset (Union[Component, Data, Model]): Asset currently in registry.
         asset (AssetSpec): AssetSpec containing updated asset info.
     """
     with tempfile.NamedTemporaryFile(suffix=".yaml", delete=False) as existing_asset_temp_file:
         existing_asset_temp_file_name = existing_asset_temp_file.name
 
     existing_asset.dump(existing_asset_temp_file_name)
-    merged_result = merge_yamls(existing_asset_temp_file_name, asset.file_name_with_path)
+
+    merged_result = merge_yamls(existing_asset_temp_file_name, asset)
 
     with tempfile.NamedTemporaryFile(suffix=".yaml", delete=False) as merged_asset_temp_file:
         merged_asset_temp_file_name = merged_asset_temp_file.name
@@ -164,6 +180,8 @@ def merge_assets(existing_asset: Union[Data, Model],
         merged_asset = load_model(merged_asset_temp_file_name)
     elif asset.type == AssetType.DATA:
         merged_asset = load_data(merged_asset_temp_file_name)
+    elif asset.type == AssetType.COMPONENT:
+        merged_asset = load_component(merged_asset_temp_file_name)
 
     asset.set_asset_obj_to_create_or_update(merged_asset)
 
@@ -289,7 +307,7 @@ def create_or_update_assets(input_dir: Path,
                     raise
 
             if existing_asset:
-                if asset.type == AssetType.MODEL or asset.type == AssetType.DATA:
+                if asset.type in [AssetType.COMPONENT, AssetType.DATA, AssetType.MODEL]:
                     try:
                         merge_assets(existing_asset, asset)
                     except Exception as e:

{azureml_registry_tools-0.1.0a2.dist-info → azureml_registry_tools-0.1.0a4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: azureml-registry-tools
-Version: 0.1.0a2
+Version: 0.1.0a4
 Summary: AzureML Registry tools and CLI
 Author: Microsoft Corp
 License: https://aka.ms/azureml-sdk-license
@@ -8,6 +8,7 @@ Requires-Python: >=3.9,<3.14
 License-File: LICENSE.txt
 Requires-Dist: azure-identity<2.0
 Requires-Dist: ruamel-yaml<0.19,>=0.17.21
+Requires-Dist: diskcache~=5.6
 Requires-Dist: azure-ai-ml<2.0
 Requires-Dist: azureml-assets<2.0
 Dynamic: author

{azureml_registry_tools-0.1.0a2.dist-info → azureml_registry_tools-0.1.0a4.dist-info}/RECORD

@@ -4,20 +4,20 @@ azureml/registry/_cli/__init__.py,sha256=dJ020ynrregpUaNGu8xVLLmX3cC9DAjaBMP1qnl
 azureml/registry/_cli/registry_syndication_cli.py,sha256=87hxSu4iwFOdwfukSGqc0fjJEpSpQhDRlthiCGCDwQ8,10003
 azureml/registry/_cli/repo2registry_cli.py,sha256=4CDv4tYWLTjVgdXIcoKA9iu_gOv_Z_xn05wSANkdmdg,5378
 azureml/registry/_rest_client/__init__.py,sha256=Eh0vB8AVlwkZlHLO4DCGHyyfRYYgeTeVjiUAX_WujG0,219
-azureml/registry/_rest_client/arm_client.py,sha256=dgm9nDnmGI1NIZaopXCIXmViM8sSAipZETfuOWnQRa0,5638
-azureml/registry/_rest_client/base_rest_client.py,sha256=cXFYvYgSYSveEeY511VIlDpmXLOY9Jat_E-Zs1ehMyA,5164
-azureml/registry/_rest_client/registry_management_client.py,sha256=JGBohCVpL8XR4rNuScOKt7tNQnjw203pqlTSLKVjYgM,4833
+azureml/registry/_rest_client/arm_client.py,sha256=KVSj0V1bN-vCUV3fBbIdZDpWTLKYT6qdWIZtDjx2la4,4516
+azureml/registry/_rest_client/base_rest_client.py,sha256=JcmhckfR-ZdUtEh0EeQIKfe63Zjowa1fLzmVeWpeLno,8173
+azureml/registry/_rest_client/registry_management_client.py,sha256=P2aijKHdHvsSZ507-06Rq_n9pNJ_xci1oQKDGBJjm-I,4986
 azureml/registry/mgmt/__init__.py,sha256=LMhqcEC8ItmmpKZljElGXH-6olHlT3SLl0dJU01OvuM,226
 azureml/registry/mgmt/create_manifest.py,sha256=N9wRmjAKO09A3utN_lCUsM_Ufpj7PL0SJz-XHPHWuyM,9528
 azureml/registry/mgmt/syndication_manifest.py,sha256=ma1_vUEHHtkRcgiFMQanlXZajvKodynylxfblTHuleE,6195
 azureml/registry/tools/__init__.py,sha256=IAuWWpGfZm__pAkBIxmpJz84QskpkxBr0yDk1TUSnkE,223
 azureml/registry/tools/config.py,sha256=tjPaoBsWtPXBL8Ww1hcJtsr2SuIjPKt79dR8iovcebg,3639
-azureml/registry/tools/create_or_update_assets.py,sha256=NexSY22bpN2eHPqZZ21eDIMerwq3SzVk_4fBV5_GTBY,14787
+azureml/registry/tools/create_or_update_assets.py,sha256=Q-_BV7KWn1huQn5JriKT_8xJNoQQ_HK5wCftrq9DepA,15988
 azureml/registry/tools/registry_utils.py,sha256=zgYlCiOONtQJ4yZ9wg8tKVoE8dh6rrjB8hYBGhpV9-0,1403
 azureml/registry/tools/repo2registry_config.py,sha256=eXp_tU8Jyi30g8xGf7wbpLgKEPpieohBANKxMSLzq7s,4873
-azureml_registry_tools-0.1.0a2.dist-info/licenses/LICENSE.txt,sha256=n20rxwp7_NGrrShv9Qvcs90sjI1l3Pkt3m-5OPCWzgs,845
-azureml_registry_tools-0.1.0a2.dist-info/METADATA,sha256=SJpiD2SsYXqJi_DUSLk-mquFZgbnb54Sd_mrP6H-GBI,491
-azureml_registry_tools-0.1.0a2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-azureml_registry_tools-0.1.0a2.dist-info/entry_points.txt,sha256=iRUkAeQidMnO6RQzpLqMUBTcyYtNzAfSin9WnSdVGLw,147
-azureml_registry_tools-0.1.0a2.dist-info/top_level.txt,sha256=ZOeEa0TAXo6i5wOjwBoqfIGEuxOcKuscGgNSpizqREY,8
-azureml_registry_tools-0.1.0a2.dist-info/RECORD,,
+azureml_registry_tools-0.1.0a4.dist-info/licenses/LICENSE.txt,sha256=n20rxwp7_NGrrShv9Qvcs90sjI1l3Pkt3m-5OPCWzgs,845
+azureml_registry_tools-0.1.0a4.dist-info/METADATA,sha256=XCJDpX9oa2LTU8pUOVQOGjizcLtrlRYa1f0fECIa9tM,521
+azureml_registry_tools-0.1.0a4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+azureml_registry_tools-0.1.0a4.dist-info/entry_points.txt,sha256=iRUkAeQidMnO6RQzpLqMUBTcyYtNzAfSin9WnSdVGLw,147
+azureml_registry_tools-0.1.0a4.dist-info/top_level.txt,sha256=ZOeEa0TAXo6i5wOjwBoqfIGEuxOcKuscGgNSpizqREY,8
+azureml_registry_tools-0.1.0a4.dist-info/RECORD,,