feat(azure): rebase based on Main not Common Code

Author: ikemilian-lewis

src/orb/infrastructure/scheduler/hostfactory/field_mappings.py

@@ -67,67 +67,6 @@ class HostFactoryFieldMappings:
             "launchTemplateId": "launch_template_id",
             "userDataScript": "user_data",
         },
-        # Azure-specific fields (only mapped when Azure provider is active)
-        "azure": {
-            # Override generic AWS-centric meanings for Azure templates
-            "vmType": "vm_size",
-            "vmTypes": "vm_sizes",
-            "keyName": "ssh_key_name",
-            "subnetId": "network_config.subnet_id",
-            "securityGroupIds": "network_config.network_security_group_id",
-            # Azure resource targeting
-            "resourceGroup": "resource_group",
-            "subscriptionId": "subscription_id",
-            # Azure VMSS / compute configuration
-            "vmSize": "vm_size",
-            "vmSizes": "vm_sizes",
-            "vmSizePreferences": "vm_size_preferences",
-            "vmssName": "vmss_name",
-            "orchestrationMode": "orchestration_mode",
-            "platformFaultDomainCount": "platform_fault_domain_count",
-            "singlePlacementGroup": "single_placement_group",
-            # Azure pricing / placement
-            "evictionPolicy": "eviction_policy",
-            "billingProfileMaxPrice": "billing_profile_max_price",
-            "spotPercentage": "spot_percentage",
-            "baseRegularPriorityCount": "base_regular_priority_count",
-            "vmssAllocationStrategy": "vmss_allocation_strategy",
-            "spotRestoreEnabled": "spot_restore_enabled",
-            "spotRestoreTimeout": "spot_restore_timeout",
-            "zoneBalance": "zone_balance",
-            "proximityPlacementGroupId": "proximity_placement_group_id",
-            "capacityReservationGroupId": "capacity_reservation_group_id",
-            # Azure storage / network / security
-            "osDisk": "os_disk",
-            "dataDisks": "data_disks",
-            "networkConfig": "network_config",
-            "securityType": "security_type",
-            "secureBootEnabled": "secure_boot_enabled",
-            "vtpmEnabled": "vtpm_enabled",
-            "encryptionAtHost": "encryption_at_host",
-            "diskEncryptionSetId": "disk_encryption_set_id",
-            # Azure identity / bootstrap
-            "adminUsername": "admin_username",
-            "sshKeyName": "ssh_key_name",
-            "sshPublicKeys": "ssh_public_keys",
-            "userAssignedIdentityIds": "user_assigned_identity_ids",
-            "systemAssignedIdentity": "system_assigned_identity",
-            "customData": "custom_data",
-            "extensionProfile": "extension_profile",
-            "upgradePolicyMode": "upgrade_policy_mode",
-            # Azure native spec / metadata
-            "providerApiSpec": "provider_api_spec",
-            "providerApiSpecFile": "provider_api_spec_file",
-            "nodeAttributes": "node_attributes",
-            # Azure CycleCloud
-            "clusterName": "cluster_name",
-            "nodeArray": "node_array",
-            "cyclecloudUrl": "cyclecloud_url",
-            "cyclecloudCredentialPath": "cyclecloud_credential_path",
-            "cyclecloudVerifySsl": "cyclecloud_verify_ssl",
-            "cyclecloudAuthMode": "cyclecloud_auth_mode",
-            "cyclecloudAadScope": "cyclecloud_aad_scope",
-        },
     }
 
     @classmethod

src/orb/infrastructure/scheduler/hostfactory/hostfactory_strategy.py

@@ -1027,36 +1027,6 @@ def _resolve_api_alias(self, raw_api: str) -> str:
 
     def _transform_machine_types_input(self, hf_data: dict) -> dict:
         """Transform HF vmType/vmTypes to internal machine_types."""
-        provider_type = self.field_mapper.provider_type
-
-        if provider_type == "azure":
-            primary_vm_type = hf_data.get("vmType")
-            if "vmTypes" in hf_data:
-                raw_vm_types = hf_data["vmTypes"]
-                candidate_sizes: list[str] = []
-                if isinstance(raw_vm_types, dict):
-                    candidate_sizes = [str(vm_size) for vm_size in raw_vm_types.keys()]
-
-                if candidate_sizes:
-                    if primary_vm_type:
-                        candidate_sizes = [
-                            str(primary_vm_type),
-                            *[
-                                vm_size
-                                for vm_size in candidate_sizes
-                                if vm_size != str(primary_vm_type)
-                            ],
-                        ]
-
-                    primary_vm_size, *fallback_vm_sizes = candidate_sizes
-                    result: dict[str, Any] = {"vm_size": primary_vm_size}
-                    if fallback_vm_sizes:
-                        result["vm_sizes"] = fallback_vm_sizes
-                    return result
-            if primary_vm_type:
-                return {"vm_size": primary_vm_type}
-            return {}
-
         if "vmType" in hf_data:
             return {"machine_types": {hf_data["vmType"]: 1}}
         elif "vmTypes" in hf_data:

src/orb/providers/azure/configuration/template_extension.py

@@ -22,6 +22,7 @@
     AzureVmSizePreference,
     AzureVMSSOrchestrationMode,
 )
+from orb.providers.azure.services.spot_placement_planner import PlacementSplitStrategy
 
 
 class AzureTemplateExtensionConfig(BaseModel):
@@ -72,6 +73,21 @@ class AzureTemplateExtensionConfig(BaseModel):
     spot_restore_timeout: Optional[str] = Field(
         None, description="ISO 8601 Spot Try-Restore timeout"
     )
+    spot_placement_score_enabled: Optional[bool] = Field(
+        None, description="Enable Azure Spot Placement Score planning before launch"
+    )
+    placement_split_strategy: Optional[PlacementSplitStrategy] = Field(
+        None, description="How Spot Placement Score launches split capacity"
+    )
+    placement_primary_share_percent: Optional[int] = Field(
+        None, description="Capacity percentage assigned to the top placement candidate"
+    )
+    placement_regions: Optional[list[str]] = Field(
+        None, description="Azure regions considered for Spot Placement Score planning"
+    )
+    placement_zones: Optional[list[str]] = Field(
+        None, description="Azure zones considered for Spot Placement Score planning"
+    )
     zones: Optional[list[str]] = Field(None, description="Availability zones")
     zone_balance: Optional[bool] = Field(None, description="Enable zone balancing")
     proximity_placement_group_id: Optional[AzureProximityPlacementGroupId] = Field(
@@ -146,9 +162,9 @@ class AzureTemplateExtensionConfig(BaseModel):
     )
 
     # VM configuration
-    vm_size: str = Field(
-        "Standard_D4s_v5",
-        description="Default Azure VM size",
+    vm_size: Optional[str] = Field(
+        default=None,
+        description="Explicit Azure VM size default",
     )
     vm_sizes: Optional[list[str]] = Field(
         default=None,
@@ -184,17 +200,28 @@ class AzureTemplateExtensionConfig(BaseModel):
     def to_template_defaults(self) -> dict[str, Any]:
         """Convert extension to a dict of default values for template creation."""
         defaults: dict[str, Any] = {
-            "vm_size": self.vm_size,
             "priority": self.priority,
             "admin_username": self.admin_username,
             "node_attributes": self.node_attributes,
         }
+        if self.vm_size:
+            defaults["vm_size"] = self.vm_size
         if self.vm_sizes:
             defaults["vm_sizes"] = self.vm_sizes
         if self.vm_size_preferences:
             defaults["vm_size_preferences"] = self.vm_size_preferences
         if self.vmss_allocation_strategy:
             defaults["vmss_allocation_strategy"] = self.vmss_allocation_strategy
+        if self.spot_placement_score_enabled is not None:
+            defaults["spot_placement_score_enabled"] = self.spot_placement_score_enabled
+        if self.placement_split_strategy:
+            defaults["placement_split_strategy"] = self.placement_split_strategy
+        if self.placement_primary_share_percent is not None:
+            defaults["placement_primary_share_percent"] = self.placement_primary_share_percent
+        if self.placement_regions:
+            defaults["placement_regions"] = self.placement_regions
+        if self.placement_zones:
+            defaults["placement_zones"] = self.placement_zones
         if self.os_disk_size_gb is not None:
             defaults["os_disk"] = {
                 "storage_account_type": self.os_disk_type or "Premium_LRS",

src/orb/providers/azure/defaults_loader.py

@@ -3,18 +3,32 @@
 from __future__ import annotations
 
 from orb.domain.base.ports.provider_defaults_loader_port import ProviderDefaultsLoaderPort
+from orb.providers.azure.capabilities import get_supported_api_capabilities
 from orb.providers.azure.registration import get_azure_extension_defaults
 
 
+_HANDLER_CLASSES = {
+    "VMSS": "VMSSHandler",
+    "VMSSUniform": "VMSSHandler",
+    "SingleVM": "SingleVMHandler",
+    "CycleCloud": "CycleCloudHandler",
+}
+
+
 class AzureDefaultsLoader:
     """Loads Azure provider defaults from the provider-owned template extension."""
 
     def load_defaults(self) -> dict:
         """Return Azure provider defaults contributed by the Azure provider."""
+        handlers = {
+            api: {"handler_class": _HANDLER_CLASSES[api], **capabilities}
+            for api, capabilities in get_supported_api_capabilities().items()
+        }
         return {
             "provider": {
                 "provider_defaults": {
                     "azure": {
+                        "handlers": handlers,
                         "template_defaults": get_azure_extension_defaults(),
                     }
                 }

src/orb/providers/azure/domain/template/azure_template_aggregate.py

@@ -22,10 +22,9 @@
 
 from typing import Any, Optional
 
-from pydantic import AliasChoices, ConfigDict, Field, model_validator
+from pydantic import AliasChoices, ConfigDict, Field, field_validator, model_validator
 
 from orb.domain.template.template_aggregate import Template
-from orb.domain.base.value_objects import AllocationStrategy
 from orb.providers.azure.domain.template.value_objects import (
     AzureAllocationStrategy,
     AzureCapacityReservationGroupId,
@@ -45,6 +44,7 @@
     AzureVmSizePreference,
     AzureVMSSOrchestrationMode,
 )
+from orb.providers.azure.services.spot_placement_planner import PlacementSplitStrategy
 
 
 
@@ -129,6 +129,11 @@ class AzureTemplate(Template):
         description="Primary Azure VM size, e.g. 'Standard_D4s_v5'.",
         validation_alias=AliasChoices("vm_size", "vmSize"),
     )
+    vm_sizes: list[str] = Field(
+        default_factory=list,
+        description="Additional unranked Azure VM size candidates.",
+        validation_alias=AliasChoices("vm_sizes", "vmSizes"),
+    )
     vm_size_preferences: list[AzureVmSizePreference] = Field(
         default_factory=list,
         description=(
@@ -201,6 +206,37 @@ class AzureTemplate(Template):
         description="ISO 8601 duration for spot restore timeout, e.g. 'PT1H'.",
         validation_alias=AliasChoices("spot_restore_timeout", "spotRestoreTimeout"),
     )
+    spot_placement_score_enabled: bool = Field(
+        default=False,
+        description="Enable Azure Spot Placement Score planning before launch.",
+        validation_alias=AliasChoices(
+            "spot_placement_score_enabled", "spotPlacementScoreEnabled"
+        ),
+    )
+    placement_split_strategy: PlacementSplitStrategy = Field(
+        default=PlacementSplitStrategy.HYBRID,
+        description="How spot placement score launches split capacity across candidates.",
+        validation_alias=AliasChoices("placement_split_strategy", "placementSplitStrategy"),
+    )
+    placement_primary_share_percent: int = Field(
+        default=80,
+        ge=0,
+        le=100,
+        description="Capacity percentage assigned to the top spot placement candidate.",
+        validation_alias=AliasChoices(
+            "placement_primary_share_percent", "placementPrimarySharePercent"
+        ),
+    )
+    placement_regions: list[str] = Field(
+        default_factory=list,
+        description="Azure regions considered for spot placement score planning.",
+        validation_alias=AliasChoices("placement_regions", "placementRegions"),
+    )
+    placement_zones: list[str] = Field(
+        default_factory=list,
+        description="Azure zones considered for spot placement score planning.",
+        validation_alias=AliasChoices("placement_zones", "placementZones"),
+    )
 
     # ------------------------------------------------------------------
     # Availability
@@ -404,6 +440,10 @@ class AzureTemplate(Template):
         description="AAD scope used to resolve a bearer token for CycleCloud.",
         validation_alias=AliasChoices("cyclecloud_aad_scope", "cyclecloudAadScope"),
     )
+    version: Optional[str] = Field(
+        default=None,
+        description="Template catalog or DTO version metadata.",
+    )
 
     # ------------------------------------------------------------------
     # Initialisation
@@ -418,6 +458,14 @@ def __init__(self, **data: Any) -> None:
     # Validators
     # ------------------------------------------------------------------
 
+    @field_validator("vm_sizes", mode="before")
+    @classmethod
+    def normalize_vm_sizes(cls, value: Any) -> Any:
+        """Accept HostFactory ``vmTypes`` dictionaries as Azure VM size lists."""
+        if isinstance(value, dict):
+            return [str(vm_size) for vm_size in value.keys()]
+        return value
+
     @model_validator(mode="before")
     @classmethod
     def apply_implied_defaults(cls, data: Any) -> Any:
@@ -438,6 +486,17 @@ def apply_implied_defaults(cls, data: Any) -> Any:
             return data
 
         data = dict(data)
+        raw_provider_config = data.pop("provider_config", None)
+        if raw_provider_config is not None:
+            if hasattr(raw_provider_config, "model_dump"):
+                provider_config = raw_provider_config.model_dump(mode="json")
+            elif isinstance(raw_provider_config, dict):
+                provider_config = dict(raw_provider_config)
+            else:
+                provider_config = {}
+            for key, value in provider_config.items():
+                if value is not None and data.get(key) in (None, ""):
+                    data[key] = value
 
         location = data.get("location")
         region = data.get("region")
@@ -458,12 +517,16 @@ def apply_implied_defaults(cls, data: Any) -> Any:
                 data["max_instances"] = data["max_number"]
             data.pop("max_number", None)
 
-        if data.get("allocation_strategy") in (None, ""):
-            data["allocation_strategy"] = (
-                AllocationStrategy.PRICE_CAPACITY_OPTIMIZED.value
-                if data.get("price_type") == "spot"
-                else AllocationStrategy.LOWEST_PRICE.value
+        if data.get("allocation_strategy") not in (None, ""):
+            raise ValueError(
+                "Azure templates do not support allocation_strategy; use "
+                "vmss_allocation_strategy for Azure VMSS instance mix or "
+                "spot_placement_score_enabled for Spot Placement Score planning"
             )
+        # The shared Template base still owns an AWS-shaped allocation_strategy
+        # field and mutates it when absent. Azure does not use that contract, so
+        # set an inert value here after rejecting explicit non-empty input.
+        data["allocation_strategy"] = ""
 
         if data.get("spot_percentage") is not None and data.get("priority") in (
             None,
@@ -495,11 +558,11 @@ def validate_azure_template(self) -> "AzureTemplate":
                     "provider_api 'VMSSUniform' requires orchestration_mode 'Uniform'"
                 )
 
-        if self.allocation_strategy == AllocationStrategy.SPOT_PLACEMENT_SCORE.value:
+        if self.spot_placement_score_enabled:
             candidate_sizes = self.candidate_vm_sizes
             if len(candidate_sizes) < 2:
                 raise ValueError(
-                    "spotPlacementScore allocation strategy requires at least two candidate vm sizes"
+                    "spot_placement_score_enabled requires at least two candidate vm sizes"
                 )
 
         if self.vm_sizes and self.vm_size_preferences:

src/orb/providers/azure/domain/template/value_objects.py

@@ -14,7 +14,7 @@
     model_serializer,
 )
 
-from orb.domain.base.value_objects import AllocationStrategy, PriceType, ValueObject
+from orb.domain.base.value_objects import PriceType, ValueObject
 
 
 # ---------------------------------------------------------------------------
@@ -88,18 +88,6 @@ class AzureAllocationStrategy(str, Enum):
     CAPACITY_OPTIMIZED = "CapacityOptimized"
     PRIORITIZED = "Prioritized"
 
-    @classmethod
-    def from_core(cls, strategy: AllocationStrategy) -> "AzureAllocationStrategy":
-        """Map a core AllocationStrategy to Azure-specific value."""
-        mapping = {
-            AllocationStrategy.LOWEST_PRICE: cls.LOWEST_PRICE,
-            AllocationStrategy.CAPACITY_OPTIMIZED: cls.CAPACITY_OPTIMIZED,
-            AllocationStrategy.CAPACITY_OPTIMIZED_PRIORITIZED: cls.PRIORITIZED,
-            AllocationStrategy.PRICE_CAPACITY_OPTIMIZED: cls.CAPACITY_OPTIMIZED,
-            AllocationStrategy.DIVERSIFIED: cls.LOWEST_PRICE,
-        }
-        return mapping.get(strategy, cls.LOWEST_PRICE)
-
     def to_arm_value(self) -> str:
         """Return the value expected by the ARM / REST API."""
         return self.value