content: #  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
  #
  #  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:
  #
  #       https://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.
  """AzureML definitions."""
  
  from typing import Optional
  
  from pydantic import Field, model_validator
  
  from zenml.config.base_settings import BaseSettings
  from zenml.logger import get_logger
  from zenml.utils.enum_utils import StrEnum
  
  logger = get_logger(__name__)
  
  
  class AzureMLComputeTypes(StrEnum):
      """Enum for different types of compute on AzureML."""
  
      SERVERLESS = "serverless"
      COMPUTE_INSTANCE = "compute-instance"
      COMPUTE_CLUSTER = "compute-cluster"
  
  
  class AzureMLComputeSettings(BaseSettings):
      """Settings for the AzureML compute.
  
      These settings adjust the compute resources that will be used by the
      pipeline execution.
  
      There are three possible use cases for this implementation:
  
          1. Serverless compute (default behavior):
              - The `mode` is set to `serverless` (default behavior).
              - All the other parameters become irrelevant and will throw a
              warning if set.
  
          2. Compute instance:
              - The `mode` is set to `compute-instance`.
              - In this case, users have to provide a `compute-name`.
                  - If a compute instance exists with this name, this instance
                  will be used and all the other parameters become irrelevant
                  and will throw a warning if set.
                  - If a compute instance does not already exist, ZenML will
                  create it. You can use the parameters `compute_size` and
                  `idle_type_before_shutdown_minutes` for this operation.
  
          3. Compute cluster:
              - The `mode` is set to `compute-cluster`.
              - In this case, users have to provide a `compute-name`.
                  - If a compute cluster exists with this name, this instance
                  will be used and all the other parameters become irrelevant
                  and will throw a warning if set.
                  - If a compute cluster does not already exist, ZenML will
                  create it. You can set the additional parameters for this
                  operation.
      """
  
      # Mode for compute
      mode: AzureMLComputeTypes = AzureMLComputeTypes.SERVERLESS
  
      # Common Configuration for Compute Instances and Clusters
      compute_name: Optional[str] = None
      size: Optional[str] = None
  
      # Additional configuration for a Compute Instance
      idle_time_before_shutdown_minutes: Optional[int] = None
  
      # Additional configuration for a Compute Cluster
      idle_time_before_scaledown_down: Optional[int] = None
      location: Optional[str] = None
      min_instances: Optional[int] = None
      max_instances: Optional[int] = None
      tier: Optional[str] = None
  
      # Container resource configuration
      shm_size: Optional[str] = Field(
          default=None,
          description="Size of the shared memory block for the container. Must be "
          "specified in format supported by Docker (e.g., '2g', '200g', '1024m'). "
          "Useful for PyTorch dataloaders with multiple workers, large model caching, "
          "and RAM-based operations. Defaults to AzureML's default (typically 64MB)",
      )
  
      @model_validator(mode="after")
      def azureml_settings_validator(self) -> "AzureMLComputeSettings":
          """Checks whether the right configuration is set based on mode.
  
          Returns:
              the instance itself.
  
          Raises:
              AssertionError: if a name is not provided when working with
                  instances and clusters.
          """
          viable_configuration_fields = {
              AzureMLComputeTypes.SERVERLESS: {"mode", "shm_size"},
              AzureMLComputeTypes.COMPUTE_INSTANCE: {
                  "mode",
                  "compute_name",
                  "size",
                  "idle_time_before_shutdown_minutes",
                  "shm_size",
              },
              AzureMLComputeTypes.COMPUTE_CLUSTER: {
                  "mode",
                  "compute_name",
                  "size",
                  "idle_time_before_scaledown_down",
                  "location",
                  "min_instances",
                  "max_instances",
                  "tier",
                  "shm_size",
              },
          }
          viable_fields = viable_configuration_fields[self.mode]
  
          for field in self.model_fields_set:
              if (
                  field not in viable_fields
                  and field in AzureMLComputeSettings.model_fields
              ):
                  logger.warning(
                      f"In the {self.__class__.__name__} settings, the mode of "
                      f"operation is set to {self.mode}. In this mode, you can "
                      f"not configure the parameter '{field}'. This "
                      "configuration will be ignored."
                  )
  
          if (
              self.mode == AzureMLComputeTypes.COMPUTE_INSTANCE
              or self.mode == AzureMLComputeTypes.COMPUTE_CLUSTER
          ):
              assert self.compute_name is not None, (
                  "When you are working with compute instances and clusters, "
                  "please define a name for the compute target."
              )
  
          return self
  