azureml-registry-tools 0.1.0a6__py3-none-any.whl → 0.1.0a8__py3-none-any.whl

azureml/registry/data/validate_model_schema.py

@@ -0,0 +1,119 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+"""Validate model variant schema."""
+
+import argparse
+import yaml
+import sys
+import jsonschema
+from pathlib import Path
+from typing import List
+
+import azureml.assets as assets
+import azureml.assets.util as util
+from azureml.assets.util import logger
+
+
+def validate_model_schema(input_dirs: List[Path],
+                          schema_file: Path,
+                          asset_config_filename: str) -> bool:
+    """Validate model variant schema.
+
+    Args:
+        input_dirs (List[Path]): Directories containing assets.
+        schema_file (Path): File containing model variant schema.
+        asset_config_filename (str): Asset config filename to search for.
+
+    Returns:
+        bool: True on success.
+    """
+    # Load model schema from file
+    loaded_schema = {}
+    with open(schema_file, 'r') as file:
+        loaded_schema = yaml.safe_load(file)
+
+    # Create validator instance for collecting all errors
+    validator = jsonschema.Draft7Validator(loaded_schema)
+
+    asset_count = 0
+    model_count = 0
+    error_count = 0
+    for input_dir in input_dirs:
+        # Recursively find all files with the name matching asset_config_filename
+        for asset_config in util.find_assets(input_dir, asset_config_filename):
+            asset_count += 1
+            file_path = asset_config.spec_with_path
+            if asset_config.type == assets.AssetType.MODEL:
+                model_count += 1
+                # Validate the file against the schema
+                try:
+                    with open(file_path, "r") as f:
+                        spec_config = yaml.safe_load(f)
+
+                    # Collect all validation errors
+                    errors = list(validator.iter_errors(spec_config))
+
+                    if not errors:
+                        logger.print(f"{file_path} is valid.")
+                    else:
+                        logger.log_error(f"\n‼️{file_path} has {len(errors)} validation error(s):")
+
+                        for e in errors:
+                            # Get detailed error information for each error
+                            error_path = '.'.join(str(p) for p in e.path) if e.path else "root"
+                            line_info = ""
+
+                            # Get line number from jsonschema error if available
+                            if hasattr(e, 'lineno') and e.lineno is not None:
+                                line_info = f" at line {e.lineno}"
+                            else:
+                                # Try to find line number by looking at the path and instance
+                                try:
+                                    with open(file_path, "r") as f:
+                                        yaml_content = f.readlines()
+                                        yaml_lines = []
+                                        for idx, line in enumerate(yaml_content):
+                                            if error_path in line:
+                                                yaml_lines.append(f"line {idx+1}: {line.strip()}")
+                                        if yaml_lines:
+                                            line_info = "\nPossible location(s):\n  " + "\n  ".join(yaml_lines)
+                                except Exception:
+                                    pass
+
+                            schema_path = '.'.join(str(p) for p in e.schema_path)
+                            logger.print(f"⚠️ {file_path} is invalid at path '{error_path}'{line_info}:")
+                            logger.print(f"  Error: {e.message}")
+                            logger.print(f"  Instance: {e.instance}")
+                            logger.print(f"  Schema path: {schema_path}")
+                        error_count += 1
+                except Exception as e:
+                    logger.log_error(f"Error processing {file_path}: {str(e)}")
+                    error_count += 1
+
+    logger.print(f"Found {asset_count} total asset(s).")
+    logger.print(f"Found {error_count} model(s) with error(s) out of {model_count} total model(s)")
+    return error_count == 0
+
+
+if __name__ == "__main__":
+    # Handle command-line args
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input-dirs", required=True,
+                        help="Comma-separated list of directories containing assets")
+    parser.add_argument("-m", "--schema-file", default=Path(__file__).parent / "model.schema.json", type=Path, help="Model Schema file")
+    parser.add_argument("-a", "--asset-config-filename", default=assets.DEFAULT_ASSET_FILENAME,
+                        help="Asset config file name to search for")
+    args = parser.parse_args()
+
+    # Convert comma-separated values to lists
+    input_dirs = [Path(d) for d in args.input_dirs.split(",")]
+
+    # Validate against model schema
+    success = validate_model_schema(input_dirs=input_dirs,
+                                    schema_file=args.schema_file,
+                                    asset_config_filename=args.asset_config_filename)
+
+    if not success:
+        sys.exit(1)

azureml/registry/data/validate_model_variant_schema.py

@@ -0,0 +1,85 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+"""Validate model variant schema."""
+
+import argparse
+import yaml
+import sys
+import jsonschema
+from jsonschema import validate
+from pathlib import Path
+from typing import List
+
+import azureml.assets as assets
+import azureml.assets.util as util
+from azureml.assets.util import logger
+
+
+def validate_model_variant_schema(input_dirs: List[Path],
+                                  model_variant_schema_file: Path,
+                                  asset_config_filename: str) -> bool:
+    """Validate model variant schema.
+
+    Args:
+        input_dirs (List[Path]): Directories containing assets.
+        model_variant_schema_file (Path): File containing model variant schema.
+        asset_config_filename (str): Asset config filename to search for.
+
+    Returns:
+        bool: True on success.
+    """
+    # Load variantInfo schema from file
+    model_variant_info_schema = {}
+    with open(model_variant_schema_file, 'r') as file:
+        model_variant_info_schema = yaml.safe_load(file)
+
+    asset_count = 0
+    model_count = 0
+    error_count = 0
+    for input_dir in input_dirs:
+        for asset_config in util.find_assets(input_dir, asset_config_filename):
+            asset_count += 1
+            if asset_config.type == assets.AssetType.MODEL:
+                model_count += 1
+                # Extract model variant info from spec
+                variant_info = None
+                with open(asset_config.spec_with_path, "r") as f:
+                    spec_config = yaml.safe_load(f)
+                    variant_info = spec_config.get("variantInfo")
+
+                if variant_info is not None:
+                    logger.print(f"Found variantInfo in spec {asset_config.spec_with_path}. "
+                                 f"Validating variantInfo against schema: {variant_info}")
+                    # Validate data
+                    try:
+                        validate(instance=variant_info, schema=model_variant_info_schema)
+                        logger.print("variantInfo is valid.")
+                    except jsonschema.exceptions.ValidationError as e:
+                        logger.log_error(f"variantInfo is invalid for {asset_config.spec_with_path}: {e.message}")
+                        error_count += 1
+
+    logger.print(f"Found {asset_count} total asset(s).")
+    logger.print(f"Found {error_count} model(s) with error(s) out of {model_count} total model(s)")
+    return error_count == 0
+
+
+if __name__ == "__main__":
+    # Handle command-line args
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input-dirs", required=True, help="Comma-separated list of directories containing assets")
+    parser.add_argument("-m", "--model-variant-schema-file", default=Path(__file__).parent / "model-variant.schema.json", type=Path, help="Model Variant Schema file")
+    parser.add_argument("-a", "--asset-config-filename", default=assets.DEFAULT_ASSET_FILENAME, help="Asset config file name to search for")
+    args = parser.parse_args()
+
+    # Convert comma-separated values to lists
+    input_dirs = [Path(d) for d in args.input_dirs.split(",")]
+
+    # Validate variantInfo against model variant schema
+    success = validate_model_variant_schema(input_dirs=input_dirs,
+                                            model_variant_schema_file=args.model_variant_schema_file,
+                                            asset_config_filename=args.asset_config_filename)
+
+    if not success:
+        sys.exit(1)

azureml/registry/mgmt/asset_management.py

@@ -0,0 +1,260 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""Asset management commands for registry-mgmt CLI."""
+
+import sys
+import shutil
+import tempfile
+import yaml
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import List
+
+from azure.ai.ml import MLClient
+from azure.identity import DefaultAzureCredential
+
+from azureml.registry.data.validate_model_schema import validate_model_schema
+from azureml.registry.data.validate_model_variant_schema import validate_model_variant_schema
+from azureml.registry.mgmt.registry_config import RegistryConfig
+
+# Windows compatibility patch - must be applied before importing azureml.assets
+from subprocess import run
+
+
+def patched_run_command(cmd: List[str]):
+    """Run command with shell=True for Windows compatibility."""
+    result = run(cmd, capture_output=True, encoding=sys.stdout.encoding, errors="ignore", shell=True)
+    return result
+
+
+# Apply patch before importing azureml.assets
+import azureml.assets.publish_utils as publish_utils  # noqa: E402
+publish_utils.run_command = patched_run_command
+
+import azureml.assets as assets  # noqa: E402
+import azureml.assets.util as util  # noqa: E402
+from azureml.assets.config import AssetConfig, AssetType, AzureBlobstoreAssetPath  # noqa: E402
+from azureml.assets.publish_utils import create_asset  # noqa: E402
+from azureml.assets.validate_assets import validate_assets  # noqa: E402
+
+
+def validate_model(asset_path: Path) -> bool:
+    """Validate model.
+
+    Args:
+        asset_path (Path): Path to the asset folder to validate
+
+    Returns:
+        bool: True if validation passes, False otherwise
+    """
+    errors = 0
+
+    print("⚙️ [VALIDATION #1]: Validate assets...")
+    if not validate_assets(asset_path, assets.DEFAULT_ASSET_FILENAME):
+        print("❌ [FAILED] Validation #1: validate_assets\n\n")
+        errors += 1
+    else:
+        print("✅ [PASSED] Validation #1: validate_assets passed\n")
+
+    # Model variant schema validation
+    model_variant_schema_file = Path(__file__).parent.parent / "data" / "model-variant.schema.json"
+
+    print("⚙️ [VALIDATION #2]: Validating model variant schema...")
+    if not validate_model_variant_schema(input_dirs=[asset_path], model_variant_schema_file=model_variant_schema_file,
+                                         asset_config_filename=assets.DEFAULT_ASSET_FILENAME):
+        print("❌ [FAILED] Validation #2: validate_model_variant_schema\n")
+        errors += 1
+    else:
+        print("✅ [PASSED] Validation #2: validate_model_variant_schema passed\n")
+
+    # Model schema validation
+    model_schema_file = Path(__file__).parent.parent / "data" / "model.schema.json"
+
+    print("⚙️ [VALIDATION #3]: Validating model schema...")
+    if not validate_model_schema(input_dirs=[asset_path], schema_file=model_schema_file,
+                                 asset_config_filename=assets.DEFAULT_ASSET_FILENAME):
+        print("❌ [FAILED] Validation #3: validate_model_schema\n")
+        errors += 1
+    else:
+        print("✅ [PASSED] Validation #3: validate_model_schema passed\n")
+
+    if errors != 0:
+        return False
+
+    print("🎉 [VALIDATION COMPLETE] All validations passed!\n")
+    return True
+
+
+def set_storage_and_sas(asset: AssetConfig, storage_config: dict):
+    """Use storage configuration and generate/set SAS token.
+
+    Args:
+        asset (AssetConfig): Asset configuration object to modify
+        storage_config (dict): Storage configuration dictionary
+    """
+    if not storage_config:
+        # No storage overrides provided, skip storage configuration
+        return
+
+    print("Overriding storage configuration with provided values...")
+    extra_config = asset.extra_config_as_object()
+    extra_config._path = AzureBlobstoreAssetPath(
+        storage_name=storage_config["storage_name"],
+        container_name=storage_config["container_name"],
+        container_path=storage_config["container_path"]
+    )
+    _ = extra_config.path.get_uri(token_expiration=timedelta(hours=1))
+
+
+def build_mutable_asset(base_asset: AssetConfig, mutable_asset_dir: str) -> AssetConfig:
+    """Build a mutable copy of the asset in a temporary directory.
+
+    Args:
+        base_asset (AssetConfig): Base asset configuration to copy
+        mutable_asset_dir (str): Directory path for the mutable asset copy
+
+    Returns:
+        AssetConfig: Mutable asset configuration object
+    """
+    common_dir, _ = util.find_common_directory(base_asset.release_paths)
+
+    # Convert string paths to Path objects and ensure they're absolute
+    common_dir = Path(common_dir).resolve()
+    mutable_asset_dir = Path(mutable_asset_dir).resolve()
+    base_asset_file = base_asset.file_name_with_path.resolve()
+    base_spec_file = base_asset.spec_with_path.resolve()
+
+    shutil.copytree(common_dir, mutable_asset_dir, dirs_exist_ok=True)
+
+    # Reference asset files in mutable directory
+    asset_config_file = mutable_asset_dir / base_asset_file.relative_to(common_dir)
+    spec_config_file = mutable_asset_dir / base_spec_file.relative_to(common_dir)
+
+    # Autoincrement version for mutable asset
+    with open(spec_config_file, "r") as f:
+        spec_config = yaml.safe_load(f)
+        spec_config["version"] = datetime.now().strftime("%Y%m%d%H%M%S")
+
+    with open(spec_config_file, "w") as f:
+        yaml.dump(spec_config, f)
+
+    mutable_asset = AssetConfig(asset_config_file)
+
+    return mutable_asset
+
+
+def create_or_update_asset(readonly_asset: AssetConfig, config: RegistryConfig):
+    """Create or update an asset in the AzureML registry.
+
+    Args:
+        readonly_asset (AssetConfig): Asset configuration to create or update
+        config (RegistryConfig): Registry configuration settings
+    """
+    print("[CREATING/UPDATING ASSET]")
+    print(f"Using registry configuration from: {config.config_path}")
+    # Create ML client
+    ml_client = MLClient(
+        subscription_id=config.subscription_id,
+        resource_group_name=config.resource_group,
+        registry_name=config.registry_name,
+        credential=DefaultAzureCredential(),
+    )
+
+    with tempfile.TemporaryDirectory() as mutable_asset_dir:
+        mutable_asset = build_mutable_asset(base_asset=readonly_asset, mutable_asset_dir=mutable_asset_dir)
+        # autoincrement version
+        try:
+            set_storage_and_sas(mutable_asset, config.storage_config)
+            success = create_asset(mutable_asset, config.registry_name, ml_client)
+        except Exception as e:
+            print(f"Failed to create/update asset: {e}")
+            raise
+
+        if not success:
+            print(f"Failed to create/update asset: create_asset 'success' returned {success}")
+            raise
+
+        print("\n[VALIDATE YOUR ASSET IN THE UI HERE]")
+        print(f" - Model Catalog link: https://ai.azure.com/explore/models/{mutable_asset.name}/version/{mutable_asset.version}/registry/{config.registry_name}?tid={config.tenant_id}")
+        print(f" - Azure Portal link: https://ml.azure.com/registries/{config.registry_name}/models/{mutable_asset.name}/version/{mutable_asset.version}?tid={config.tenant_id}")
+
+
+def asset_validate(asset_path: Path, dry_run: bool = False) -> bool:
+    """Validate an asset at the specified path.
+
+    Args:
+        asset_path (Path): Path to the asset folder to validate
+        dry_run (bool): If True, perform a dry run without side effects
+
+    Returns:
+        bool: True if validation passes, False otherwise
+    """
+    if dry_run:
+        print(f"[DRY RUN] Would validate asset at: {asset_path}")
+        return True
+
+    asset_path = asset_path.resolve()
+    print(f"[VALIDATION] Begin validating for asset at: {asset_path}...")
+
+    # Check if asset path exists
+    if not asset_path.exists():
+        print(f"❌ [ERROR]: Asset path {asset_path} does not exist")
+        return False
+
+    # Check for exactly one asset
+    asset_count = len(util.find_assets([asset_path], assets.DEFAULT_ASSET_FILENAME))
+    if asset_count != 1:
+        print(f"❌ [ERROR]: Expected exactly one asset in {asset_path}, found {asset_count}")
+        return False
+
+    # Load asset configuration
+    readonly_asset = assets.AssetConfig(asset_path / assets.DEFAULT_ASSET_FILENAME)
+
+    # Check asset type
+    if readonly_asset.type != AssetType.MODEL:
+        print(f"❌ [ERROR]: Asset type {readonly_asset.type} is not supported for validation. "
+              f"Only models are currently supported.")
+        return False
+
+    # Perform validation
+    return validate_model(readonly_asset.file_path)
+
+
+def asset_deploy(asset_path: Path, config_path: Path, dry_run: bool = False) -> bool:
+    """Deploy an asset using configuration file.
+
+    Args:
+        asset_path (Path): Path to the asset folder to deploy
+        config_path (Path): Path to configuration file
+        dry_run (bool): If True, perform a dry run without deploying
+
+    Returns:
+        bool: True if deployment succeeds, False otherwise
+    """
+    try:
+        config = RegistryConfig(config_path)
+    except Exception as e:
+        print(f"❌ [ERROR]: Configuration validation failed: {e}")
+        return False
+
+    if dry_run:
+        print(f"[DRY RUN] Would deploy asset at {asset_path} to registry {config.registry_name}")
+        return True
+
+    asset_path = asset_path.resolve()
+
+    # Validate asset before deployment
+    if not asset_validate(asset_path, dry_run=False):
+        print("❌ [ERROR]: Asset validation failed. Asset deployment aborted.")
+        return False
+
+    # Load asset configuration
+    readonly_asset = assets.AssetConfig(asset_path / assets.DEFAULT_ASSET_FILENAME)
+
+    try:
+        create_or_update_asset(readonly_asset, config)
+        return True
+    except Exception as e:
+        print(f"❌ [ERROR]: Failed to deploy asset: {e}")
+        return False

azureml/registry/mgmt/create_asset_template.py

@@ -0,0 +1,87 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""Asset template creation commands for registry-mgmt CLI."""
+
+from pathlib import Path
+from .create_model_spec import generate_model_spec_content
+
+
+def asset_template(folder_path: Path, dry_run: bool = False) -> bool:
+    """Create asset template files in the specified folder.
+
+    Args:
+        folder_path (Path): Path to the folder where template files will be created
+        dry_run (bool): If True, perform a dry run without creating files
+
+    Returns:
+        bool: True if template creation succeeds, False otherwise
+    """
+    if dry_run:
+        print(f"[DRY RUN] Would create asset template files in: {folder_path}")
+        return True
+
+    folder_path = folder_path.resolve()
+    print(f"Creating asset template files in {folder_path} ...")
+
+    # Check if folder path exists, create if it doesn't
+    if not folder_path.exists():
+        try:
+            print(f"Folder path does not exist, creating directory {folder_path} ...")
+            folder_path.mkdir(parents=True, exist_ok=True)
+            print(f"Created directory: {folder_path}")
+        except Exception as e:
+            print(f"[ERROR] Failed to create directory {folder_path}: {e}")
+            return False
+
+    # Get the data directory path (relative to this module)
+    data_dir = Path(__file__).parent.parent / "data"
+
+    # List of template files to create
+    template_files = ["asset.yaml", "spec.yaml", "model.yaml", "notes.md", "evaluation.md", "description.md"]
+
+    # Create each template file
+    try:
+        for output_name in template_files:
+            output_path = folder_path / output_name
+
+            # Special handling for spec.yaml (generate from schema)
+            if output_name == "spec.yaml":
+                schema_path = data_dir / "model.schema.json"
+                if not schema_path.exists():
+                    print(f"[ERROR] Schema file not found: {schema_path}")
+                    return False
+
+                try:
+                    spec_content = generate_model_spec_content(schema_path)
+                    with open(output_path, "w", encoding="utf-8") as output_file:
+                        output_file.write(spec_content)
+                    print(f"Created {output_path}")
+                except Exception as e:
+                    print(f"[ERROR] Failed to generate spec.yaml: {e}")
+                    return False
+            else:
+                # Handle other template files from data/ folder
+                template_name = f"{output_name}.template"
+                template_path = data_dir / template_name
+
+                # Check if template file exists
+                if not template_path.exists():
+                    print(f"[ERROR] Template file not found: {template_path}")
+                    return False
+
+                # Read template content and write to output file
+                with open(template_path, "r", encoding="utf-8") as template_file:
+                    content = template_file.read()
+
+                with open(output_path, "w", encoding="utf-8") as output_file:
+                    output_file.write(content)
+
+                print(f"Created {output_path}")
+
+        print(f"Created {len(template_files)} template files in {folder_path}")
+        return True
+
+    except Exception as e:
+        print(f"[ERROR] Failed to create template files: {e}")
+        return False