Account Operations

def list_supported_images(account_list_supported_images_options=None, custom_headers=None, raw=False, **operation_config):
    """
    List all virtual machine images supported by Azure Batch service.
    
    This operation returns only images and node agent SKUs that are available
    for use on compute nodes in batch pools.
    
    Args:
        account_list_supported_images_options: Additional options for listing images
        custom_headers: Custom headers to include in request
        raw: Return raw response if True
        
    Returns:
        ItemPaged[ImageInformation]: Paginated list of supported VM images
    """

def list_pool_node_counts(account_list_pool_node_counts_options=None, custom_headers=None, raw=False, **operation_config):
    """
    Get the number of nodes in each state, grouped by pool.
    
    Args:
        account_list_pool_node_counts_options: Additional options for listing node counts
        custom_headers: Custom headers to include in request  
        raw: Return raw response if True
        
    Returns:
        ItemPaged[PoolNodeCounts]: Paginated list of pool node count information
    """

# List all supported images
supported_images = client.account.list_supported_images()
for image in supported_images:
    print(f"Image: {image.image_reference.publisher}/{image.image_reference.offer}")
    print(f"  SKU: {image.image_reference.sku}")
    print(f"  Node Agent SKU: {image.node_agent_sku_id}")
    print(f"  OS Type: {image.os_type}")
    print(f"  Verification Type: {image.verification_type}")
    
    # Check capabilities
    if image.capabilities:
        print(f"  Capabilities: {', '.join(image.capabilities)}")
    
    print("---")

# Filter images for specific requirements
linux_images = []
windows_images = []

for image in client.account.list_supported_images():
    if image.os_type.lower() == "linux":
        linux_images.append(image)
    elif image.os_type.lower() == "windows":
        windows_images.append(image)

print(f"Found {len(linux_images)} Linux images and {len(windows_images)} Windows images")

# Find Ubuntu images
ubuntu_images = []
for image in client.account.list_supported_images():
    if ("ubuntu" in image.image_reference.offer.lower() or 
        "ubuntu" in image.image_reference.sku.lower()):
        ubuntu_images.append(image)

print("Available Ubuntu images:")
for image in ubuntu_images:
    print(f"  {image.image_reference.publisher}/{image.image_reference.offer}")
    print(f"    SKU: {image.image_reference.sku}")
    print(f"    Node Agent: {image.node_agent_sku_id}")

# Find the latest Ubuntu LTS image
ubuntu_lts_images = [img for img in ubuntu_images if "lts" in img.image_reference.sku.lower()]
if ubuntu_lts_images:
    latest_ubuntu = ubuntu_lts_images[0]  # Typically the first is the latest
    print(f"\nRecommended Ubuntu LTS image:")
    print(f"  Publisher: {latest_ubuntu.image_reference.publisher}")
    print(f"  Offer: {latest_ubuntu.image_reference.offer}")
    print(f"  SKU: {latest_ubuntu.image_reference.sku}")
    print(f"  Node Agent SKU ID: {latest_ubuntu.node_agent_sku_id}")

# Get node count information for all pools
pool_node_counts = client.account.list_pool_node_counts()

total_dedicated = 0
total_low_priority = 0

for pool_counts in pool_node_counts:
    print(f"Pool: {pool_counts.pool_id}")
    print(f"  Dedicated nodes:")
    print(f"    Creating: {pool_counts.dedicated.creating}")
    print(f"    Idle: {pool_counts.dedicated.idle}")
    print(f"    Running: {pool_counts.dedicated.running}")
    print(f"    Rebooting: {pool_counts.dedicated.rebooting}")
    print(f"    Reimaging: {pool_counts.dedicated.reimaging}")
    print(f"    Total: {pool_counts.dedicated.total}")
    
    print(f"  Low-priority nodes:")
    print(f"    Creating: {pool_counts.low_priority.creating}")
    print(f"    Idle: {pool_counts.low_priority.idle}")
    print(f"    Running: {pool_counts.low_priority.running}")
    print(f"    Preempted: {pool_counts.low_priority.preempted}")
    print(f"    Total: {pool_counts.low_priority.total}")
    
    total_dedicated += pool_counts.dedicated.total
    total_low_priority += pool_counts.low_priority.total
    print("---")

print(f"Account totals:")
print(f"  Total dedicated nodes: {total_dedicated}")
print(f"  Total low-priority nodes: {total_low_priority}")
print(f"  Total nodes: {total_dedicated + total_low_priority}")

def find_compatible_images(os_type, required_capabilities=None):
    """Find VM images compatible with specific requirements."""
    compatible_images = []
    required_capabilities = required_capabilities or []
    
    for image in client.account.list_supported_images():
        if image.os_type.lower() != os_type.lower():
            continue
            
        # Check if image has required capabilities
        image_capabilities = image.capabilities or []
        if all(cap in image_capabilities for cap in required_capabilities):
            compatible_images.append(image)
    
    return compatible_images

# Find Linux images with GPU support
gpu_linux_images = find_compatible_images("linux", ["cuda", "docker"])
print(f"Found {len(gpu_linux_images)} Linux images with GPU/Docker support")

# Find Windows images
windows_images = find_compatible_images("windows")
print(f"Found {len(windows_images)} Windows images")

def get_recommended_image(os_type, prefer_latest=True):
    """Get recommended image for a specific OS type."""
    images = find_compatible_images(os_type)
    
    if not images:
        return None
    
    # Sort by publisher, offer, sku for consistency
    images.sort(key=lambda x: (x.image_reference.publisher, 
                              x.image_reference.offer, 
                              x.image_reference.sku))
    
    if prefer_latest:
        return images[-1]  # Assume last is latest
    return images[0]

# Get recommended images
recommended_linux = get_recommended_image("linux")
recommended_windows = get_recommended_image("windows")

if recommended_linux:
    print(f"Recommended Linux image: {recommended_linux.image_reference.offer} {recommended_linux.image_reference.sku}")
    
if recommended_windows:
    print(f"Recommended Windows image: {recommended_windows.image_reference.offer} {recommended_windows.image_reference.sku}")

class ImageInformation:
    """Information about a supported VM image."""
    def __init__(self):
        self.node_agent_sku_id: str
        self.image_reference: ImageReference
        self.os_type: str  # linux, windows
        self.capabilities: List[str]
        self.batch_support_end_of_life: datetime.datetime
        self.verification_type: str  # verified, unverified

class ImageReference:
    """Reference to a specific VM image."""
    def __init__(self):
        self.publisher: str
        self.offer: str
        self.sku: str
        self.version: str
        self.virtual_machine_image_id: str  # For custom images

class PoolNodeCounts:
    """Node count information for a pool."""
    def __init__(self):
        self.pool_id: str
        self.dedicated: NodeCounts
        self.low_priority: NodeCounts

class NodeCounts:
    """Count of nodes in different states."""
    def __init__(self):
        self.creating: int
        self.idle: int
        self.offline: int
        self.preempted: int
        self.rebooting: int
        self.reimaging: int
        self.running: int
        self.starting: int
        self.start_task_failed: int
        self.leaving_pool: int
        self.unknown: int
        self.unusable: int
        self.waiting_for_start_task: int
        self.total: int

class AccountListSupportedImagesOptions:
    """Options for listing supported images."""
    def __init__(self):
        self.filter: str
        self.max_results: int
        self.timeout: int

class AccountListPoolNodeCountsOptions:
    """Options for listing pool node counts."""
    def __init__(self):
        self.filter: str
        self.max_results: int
        self.timeout: int