zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/integrations/aws/service_connectors/aws_service_connector.py

@@ -24,6 +24,7 @@ IAM role)
 - STS Session tokens
 
 """
+
 import base64
 import datetime
 import json
@@ -36,7 +37,7 @@ from aws_profile_manager import Common  # type: ignore[import-untyped]
 from botocore.client import BaseClient
 from botocore.exceptions import BotoCoreError, ClientError
 from botocore.signers import RequestSigner
-from pydantic import Field, SecretStr
+from pydantic import Field
 
 from zenml.constants import (
     DOCKER_REGISTRY_RESOURCE_TYPE,
@@ -64,22 +65,24 @@ from zenml.service_connectors.service_connector import (
     ServiceConnector,
 )
 from zenml.utils.enum_utils import StrEnum
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 logger = get_logger(__name__)
 
-EKS_KUBE_API_TOKEN_EXPIRATION = 60
+EKS_KUBE_API_TOKEN_EXPIRATION = 15  # 15 minutes
 DEFAULT_IAM_ROLE_TOKEN_EXPIRATION = 3600  # 1 hour
 DEFAULT_STS_TOKEN_EXPIRATION = 43200  # 12 hours
+BOTO3_SESSION_EXPIRATION_BUFFER = 15  # 15 minutes
 
 
 class AWSSecretKey(AuthenticationConfig):
     """AWS secret key credentials."""
 
-    aws_access_key_id: SecretStr = Field(
+    aws_access_key_id: PlainSerializedSecretStr = Field(
         title="AWS Access Key ID",
         description="An AWS access key ID associated with an AWS account or IAM user.",
     )
-    aws_secret_access_key: SecretStr = Field(
+    aws_secret_access_key: PlainSerializedSecretStr = Field(
         title="AWS Secret Access Key",
     )
 
@@ -87,7 +90,7 @@ class AWSSecretKey(AuthenticationConfig):
 class STSToken(AWSSecretKey):
     """AWS STS token."""
 
-    aws_session_token: SecretStr = Field(
+    aws_session_token: PlainSerializedSecretStr = Field(
         title="AWS Session Token",
     )
 
@@ -104,28 +107,6 @@ class AWSBaseConfig(AuthenticationConfig):
     )
 
 
-class AWSImplicitConfig(AWSBaseConfig):
-    """AWS implicit configuration."""
-
-    profile_name: Optional[str] = Field(
-        default=None,
-        title="AWS Profile Name",
-    )
-
-
-class AWSSecretKeyConfig(AWSBaseConfig, AWSSecretKey):
-    """AWS secret key authentication configuration."""
-
-
-class STSTokenConfig(AWSBaseConfig, STSToken):
-    """AWS STS token authentication configuration."""
-
-    expires_at: Optional[datetime.datetime] = Field(
-        default=None,
-        title="AWS STS Token Expiration",
-    )
-
-
 class AWSSessionPolicy(AuthenticationConfig):
     """AWS session IAM policy configuration."""
 
@@ -142,6 +123,27 @@ class AWSSessionPolicy(AuthenticationConfig):
     )
 
 
+class AWSImplicitConfig(AWSBaseConfig, AWSSessionPolicy):
+    """AWS implicit configuration."""
+
+    profile_name: Optional[str] = Field(
+        default=None,
+        title="AWS Profile Name",
+    )
+    role_arn: Optional[str] = Field(
+        default=None,
+        title="Optional AWS IAM Role ARN to assume",
+    )
+
+
+class AWSSecretKeyConfig(AWSBaseConfig, AWSSecretKey):
+    """AWS secret key authentication configuration."""
+
+
+class STSTokenConfig(AWSBaseConfig, STSToken):
+    """AWS STS token authentication configuration."""
+
+
 class IAMRoleAuthenticationConfig(AWSSecretKeyConfig, AWSSessionPolicy):
     """AWS IAM authentication config."""
 
@@ -241,13 +243,39 @@ the ZenML server is running (e.g. an EKS cluster). The IAM role permissions may
 need to be adjusted to allows listing and accessing/describing the AWS resources
 that the connector is configured to access.
 
+An IAM role may optionally be specified to be assumed by the connector on top of
+the implicit credentials. This is only possible when the implicit credentials
+have permissions to assume the target IAM role. Configuring an IAM role has all
+the advantages of the AWS IAM Role authentication method plus the added benefit
+of not requiring any explicit credentials to be configured and stored:
+
+* The connector will generate temporary STS tokens upon request by
+[calling the AssumeRole STS API](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp_request.html#api_assumerole).
+* Allows implementing a two layer authentication scheme that keeps the set of
+permissions associated with implicit credentials down to the bare minimum and
+grants permissions to the privilege-bearing IAM role instead.
+* One or more optional [IAM session policies](https://docs.aws.amazon.com/IAM/latest/UserGuide/access_policies.html#policies_session)
+may also be configured to further restrict the permissions of the generated
+STS tokens. If not specified, IAM session policies are automatically configured
+for the generated STS tokens to restrict them to the minimum set of permissions
+required to access the target resource. Refer to the documentation for each
+supported Resource Type for the complete list of AWS permissions automatically
+granted to the generated STS tokens.
+* The default expiration period for generated STS tokens is 1 hour with a
+minimum of 15 minutes up to the maximum session duration setting configured for
+the IAM role (default is 1 hour). If you need longer-lived tokens, you can
+configure the IAM role to use a higher maximum expiration value (up to 12 hours)
+or use the AWS Federation Token or AWS Session Token authentication methods.
+
 Note that the discovered credentials inherit the full set of permissions of the
-local AWS client configuration, environment variables or remote AWS IAM role.
+local AWS client configuration, environment variables or attached AWS IAM role.
 Depending on the extent of those permissions, this authentication method might
 not be recommended for production use, as it can lead to accidental privilege
-escalation. Instead, it is recommended to use the AWS IAM Role, AWS Session
-Token or AWS Federation Token authentication methods to limit the validity
-and/or permissions of the credentials being issued to connector clients.
+escalation. It is recommended to also configure an IAM role when using
+the implicit authentication method, or to use the AWS IAM Role, AWS Session
+Token or AWS Federation Token authentication methods instead to limit the
+validity and/or permissions of the credentials being issued to connector
+clients.
 
 If you need to access an EKS kubernetes cluster with this authentication method,
 please be advised that the EKS cluster's aws-auth ConfigMap may need to be
@@ -256,9 +284,10 @@ picked up by the Service Connector. For more information,
 [see this documentation](https://docs.aws.amazon.com/eks/latest/userguide/add-user-role.html).
 
 An AWS region is required and the connector may only be used to access AWS
-resources in the specified region. When used with a remote IAM role, the region
-has to be the same as the region where the IAM role is configured.
+resources in the specified region.
 """,
+            min_expiration_seconds=900,  # 15 minutes
+            default_expiration_seconds=DEFAULT_IAM_ROLE_TOKEN_EXPIRATION,  # 1 hour
             config_class=AWSImplicitConfig,
         ),
         AuthenticationMethodModel(
@@ -315,6 +344,12 @@ another IAM role. The IAM user or IAM role must have permissions to assume the
 target IAM role. The connector will generate temporary STS tokens upon
 request by [calling the AssumeRole STS API](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp_request.html#api_assumerole).
 
+This authentication method still requires credentials to be explicitly
+configured. If your ZenML server is running in AWS and you're looking for an
+alternative that uses implicit credentials while at the same time benefits from
+all the security advantages of assuming an IAM role, you should use the implicit
+authentication method with a configured IAM role instead.
+
 The best practice implemented with this authentication scheme is to keep the set
 of permissions associated with the primary IAM user or IAM role down to the bare
 minimum and grant permissions to the privilege bearing IAM role instead.
@@ -487,6 +522,9 @@ that the connector will be allowed to access (e.g. `arn:aws:s3:::*` and
 - `s3:PutObject`
 - `s3:DeleteObject`
 - `s3:ListAllMyBuckets`
+- `s3:GetBucketVersioning`
+- `s3:ListBucketVersions`
+- `s3:DeleteObjectVersion`
 
 If set, the resource name must identify an S3 bucket using one of the following
 formats:
@@ -675,7 +713,10 @@ class AWSServiceConnector(ServiceConnector):
             # Refresh expired sessions
             now = datetime.datetime.now(datetime.timezone.utc)
             expires_at = expires_at.replace(tzinfo=datetime.timezone.utc)
-            if expires_at > now:
+            # check if the token expires in the near future
+            if expires_at > now + datetime.timedelta(
+                minutes=BOTO3_SESSION_EXPIRATION_BUFFER
+            ):
                 return session, expires_at
 
         logger.debug(
@@ -689,6 +730,37 @@ class AWSServiceConnector(ServiceConnector):
         self._session_cache[key] = (session, expires_at)
         return session, expires_at
 
+    def get_ecr_client(self) -> BaseClient:
+        """Get an ECR client.
+
+        Raises:
+            ValueError: If the service connector is not able to instantiate an
+                ECR client.
+
+        Returns:
+            An ECR client.
+        """
+        if self.resource_type and self.resource_type not in {
+            AWS_RESOURCE_TYPE,
+            DOCKER_REGISTRY_RESOURCE_TYPE,
+        }:
+            raise ValueError(
+                f"Unable to instantiate ECR client for a connector that is "
+                f"configured to provide access to a '{self.resource_type}' "
+                "resource type."
+            )
+
+        session, _ = self.get_boto3_session(
+            auth_method=self.auth_method,
+            resource_type=DOCKER_REGISTRY_RESOURCE_TYPE,
+            resource_id=self.config.region,
+        )
+        return session.client(
+            "ecr",
+            region_name=self.config.region,
+            endpoint_url=self.config.endpoint_url,
+        )
+
     def _get_iam_policy(
         self,
         region_id: str,
@@ -726,6 +798,9 @@ class AWSServiceConnector(ServiceConnector):
                             "s3:PutObject",
                             "s3:DeleteObject",
                             "s3:ListAllMyBuckets",
+                            "s3:GetBucketVersioning",
+                            "s3:ListBucketVersions",
+                            "s3:DeleteObjectVersion",
                         ],
                         "Resource": resource,
                     },
@@ -818,6 +893,8 @@ class AWSServiceConnector(ServiceConnector):
             NotImplementedError: If the authentication method is not supported.
         """
         cfg = self.config
+        policy_kwargs: Dict[str, Any] = {}
+
         if auth_method == AWSAuthenticationMethods.IMPLICIT:
             self._check_implicit_auth_method_allowed()
 
@@ -827,6 +904,50 @@ class AWSServiceConnector(ServiceConnector):
                 profile_name=cfg.profile_name, region_name=cfg.region
             )
 
+            if cfg.role_arn:
+                # If an IAM role is configured, assume it
+                policy = cfg.policy
+                if not cfg.policy and not cfg.policy_arns:
+                    policy = self._get_iam_policy(
+                        region_id=cfg.region,
+                        resource_type=resource_type,
+                        resource_id=resource_id,
+                    )
+                if policy:
+                    policy_kwargs["Policy"] = policy
+                elif cfg.policy_arns:
+                    policy_kwargs["PolicyArns"] = cfg.policy_arns
+
+                sts = session.client("sts", region_name=cfg.region)
+                session_name = "zenml-connector"
+                if self.id:
+                    session_name += f"-{self.id}"
+
+                try:
+                    response = sts.assume_role(
+                        RoleArn=cfg.role_arn,
+                        RoleSessionName=session_name,
+                        DurationSeconds=self.expiration_seconds,
+                        **policy_kwargs,
+                    )
+                except (ClientError, BotoCoreError) as e:
+                    raise AuthorizationException(
+                        f"Failed to assume IAM role {cfg.role_arn} "
+                        f"using the implicit AWS credentials: {e}"
+                    ) from e
+
+                session = boto3.Session(
+                    aws_access_key_id=response["Credentials"]["AccessKeyId"],
+                    aws_secret_access_key=response["Credentials"][
+                        "SecretAccessKey"
+                    ],
+                    aws_session_token=response["Credentials"]["SessionToken"],
+                )
+                expiration = response["Credentials"]["Expiration"]
+                # Add the UTC timezone to the expiration time
+                expiration = expiration.replace(tzinfo=datetime.timezone.utc)
+                return session, expiration
+
             credentials = session.get_credentials()
             if not credentials:
                 raise AuthorizationException(
@@ -864,7 +985,7 @@ class AWSServiceConnector(ServiceConnector):
                 aws_session_token=cfg.aws_session_token.get_secret_value(),
                 region_name=cfg.region,
             )
-            return session, cfg.expires_at
+            return session, self.expires_at
         elif auth_method in [
             AWSAuthenticationMethods.IAM_ROLE,
             AWSAuthenticationMethods.SESSION_TOKEN,
@@ -895,7 +1016,6 @@ class AWSServiceConnector(ServiceConnector):
                 AWSAuthenticationMethods.FEDERATION_TOKEN,
             ]:
                 assert isinstance(cfg, AWSSessionPolicy)
-                policy_kwargs = {}
                 policy = cfg.policy
                 if not cfg.policy and not cfg.policy_arns:
                     policy = self._get_iam_policy(
@@ -991,6 +1111,8 @@ class AWSServiceConnector(ServiceConnector):
         Returns:
             A bearer token for authenticating to the EKS API server.
         """
+        STS_TOKEN_EXPIRES_IN = 60
+
         client = session.client("sts", region_name=region)
         service_id = client.meta.service_model.service_id
 
@@ -1014,7 +1136,7 @@ class AWSServiceConnector(ServiceConnector):
         signed_url = signer.generate_presigned_url(
             params,
             region_name=region,
-            expires_in=EKS_KUBE_API_TOKEN_EXPIRATION,
+            expires_in=STS_TOKEN_EXPIRES_IN,
             operation_name="",
         )
 
@@ -1047,19 +1169,19 @@ class AWSServiceConnector(ServiceConnector):
         # We need to extract the bucket name from the provided resource ID
         bucket_name: Optional[str] = None
         if re.match(
-            r"^arn:aws:s3:::[a-z0-9-]+(/.*)*$",
+            r"^arn:aws:s3:::[a-z0-9][a-z0-9\-\.]{1,61}[a-z0-9](/.*)*$",
             resource_id,
         ):
             # The resource ID is an S3 bucket ARN
             bucket_name = resource_id.split(":")[-1].split("/")[0]
         elif re.match(
-            r"^s3://[a-z0-9-]+(/.*)*$",
+            r"^s3://[a-z0-9][a-z0-9\-\.]{1,61}[a-z0-9](/.*)*$",
             resource_id,
         ):
             # The resource ID is an S3 bucket URI
             bucket_name = resource_id.split("/")[2]
         elif re.match(
-            r"^[a-z0-9][a-z0-9-]{1,61}[a-z0-9]$",
+            r"^[a-z0-9][a-z0-9\-\.]{1,61}[a-z0-9]$",
             resource_id,
         ):
             # The resource ID is the S3 bucket name
@@ -1158,14 +1280,14 @@ class AWSServiceConnector(ServiceConnector):
         cluster_name: Optional[str] = None
         region_id: Optional[str] = None
         if re.match(
-            r"^arn:aws:eks:[a-z0-9-]+:\d{12}:cluster/.+$",
+            r"^arn:aws:eks:[a-z0-9-]+:\d{12}:cluster/[0-9A-Za-z][A-Za-z0-9\-_]*$",
             resource_id,
         ):
             # The resource ID is an EKS cluster ARN
             cluster_name = resource_id.split("/")[-1]
             region_id = resource_id.split(":")[3]
         elif re.match(
-            r"^[a-z0-9]+[a-z0-9_-]*$",
+            r"^[0-9A-Za-z][A-Za-z0-9\-_]*$",
             resource_id,
         ):
             # Assume the resource ID is an EKS cluster name
@@ -1368,9 +1490,9 @@ class AWSServiceConnector(ServiceConnector):
             }
 
             if credentials.token:
-                all_profiles[aws_profile_name][
-                    "aws_session_token"
-                ] = credentials.token
+                all_profiles[aws_profile_name]["aws_session_token"] = (
+                    credentials.token
+                )
 
             aws_credentials_path = os.path.join(
                 users_home, ".aws", "credentials"
@@ -1453,8 +1575,6 @@ class AWSServiceConnector(ServiceConnector):
         expiration_seconds: Optional[int] = None
         expires_at: Optional[datetime.datetime] = None
         if auth_method == AWSAuthenticationMethods.IMPLICIT:
-            cls._check_implicit_auth_method_allowed()
-
             if region_name is None:
                 raise ValueError(
                     "The AWS region name must be specified when using the "
@@ -1971,7 +2091,7 @@ class AWSServiceConnector(ServiceConnector):
             assert resource_id is not None
 
             # Get an authenticated boto3 session
-            session, expires_at = self.get_boto3_session(
+            session, _ = self.get_boto3_session(
                 self.auth_method,
                 resource_type=resource_type,
                 resource_id=resource_id,
@@ -2007,6 +2127,13 @@ class AWSServiceConnector(ServiceConnector):
                     f"Failed to get EKS bearer token: {e}"
                 ) from e
 
+            # Kubernetes authentication tokens issued by AWS EKS have a fixed
+            # expiration time of 15 minutes
+            # source: https://aws.github.io/aws-eks-best-practices/security/docs/iam/#controlling-access-to-eks-clusters
+            expires_at = datetime.datetime.now(
+                tz=datetime.timezone.utc
+            ) + datetime.timedelta(minutes=EKS_KUBE_API_TOKEN_EXPIRATION)
+
             # get cluster details
             cluster_arn = cluster["cluster"]["arn"]
             cluster_ca_cert = cluster["cluster"]["certificateAuthority"][

zenml/integrations/aws/step_operators/sagemaker_step_operator.py

@@ -261,7 +261,7 @@ class SagemakerStepOperator(BaseStepOperator):
                 "ExperimentName": settings.experiment_name,
                 "TrialName": sanitized_training_job_name,
             }
-
+        info.force_write_logs()
         estimator.fit(
             wait=True,
             inputs=inputs,

zenml/integrations/azure/__init__.py

@@ -27,6 +27,7 @@ from zenml.stack import Flavor
 
 AZURE_ARTIFACT_STORE_FLAVOR = "azure"
 AZUREML_STEP_OPERATOR_FLAVOR = "azureml"
+AZUREML_ORCHESTRATOR_FLAVOR = "azureml"
 
 # Service connector constants
 AZURE_CONNECTOR_TYPE = "azure"
@@ -39,18 +40,24 @@ class AzureIntegration(Integration):
 
     NAME = AZURE
     REQUIREMENTS = [
-        "adlfs==2021.10.0",
+        "adlfs>=2021.10.0",
         "azure-keyvault-keys",
         "azure-keyvault-secrets",
-        "azure-identity==1.10.0",
-        "azureml-core==1.48.0",
+        "azure-identity",
+        "azureml-core==1.56.0",
         "azure-mgmt-containerservice>=20.0.0",
         "azure-storage-blob==12.17.0",  # temporary fix for https://github.com/Azure/azure-sdk-for-python/issues/32056
         "kubernetes",
+        "azure-ai-ml==1.23.1",
+        # In azureml/core/_metrics.py:212 of azureml-core 1.56.0, they use 
+        # an attribute that was removed in Numpy 2.0. However, AzureML itself
+        # does not have a limitation on numpy.
+        "numpy<2.0",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes", "numpy"]
 
-    @staticmethod
-    def activate() -> None:
+    @classmethod
+    def activate(cls) -> None:
         """Activate the Azure integration."""
         from zenml.integrations.azure import service_connectors  # noqa
 
@@ -64,11 +71,13 @@ class AzureIntegration(Integration):
         from zenml.integrations.azure.flavors import (
             AzureArtifactStoreFlavor,
             AzureMLStepOperatorFlavor,
+            AzureMLOrchestratorFlavor,
         )
 
         return [
             AzureArtifactStoreFlavor,
             AzureMLStepOperatorFlavor,
+            AzureMLOrchestratorFlavor,
         ]
 
 

zenml/integrations/azure/azureml_utils.py

@@ -0,0 +1,201 @@
+#  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 azure.ai.ml import MLClient
+from azure.ai.ml.entities import Compute
+from azure.core.exceptions import (
+    ResourceNotFoundError,
+)
+
+from zenml.integrations.azure.flavors.azureml import (
+    AzureMLComputeSettings,
+    AzureMLComputeTypes,
+)
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def check_settings_and_compute_configuration(
+    parameter: str,
+    settings: AzureMLComputeSettings,
+    compute: Compute,
+) -> None:
+    """Utility function comparing a parameter between settings and compute.
+
+    Args:
+        parameter: the name of the parameter.
+        settings: The AzureML orchestrator settings.
+        compute: The compute instance or cluster from AzureML.
+    """
+    # Check the compute size
+    compute_value = getattr(compute, parameter)
+    settings_value = getattr(settings, parameter)
+
+    if settings_value is not None and settings_value != compute_value:
+        logger.warning(
+            f"The '{parameter}' defined in the settings '{settings_value}' "
+            "does not match the actual parameter of the instance: "
+            f"'{compute_value}'. Will ignore this setting for now."
+        )
+
+
+def create_or_get_compute(
+    client: MLClient,
+    settings: AzureMLComputeSettings,
+    default_compute_name: str,
+) -> Optional[str]:
+    """Creates or fetches the compute target if defined in the settings.
+
+    Args:
+        client: the AzureML client.
+        settings: the settings for the orchestrator.
+        default_compute_name: the default name for the compute target, if one
+            is not provided in the settings.
+
+    Returns:
+        None, if the orchestrator is using serverless compute or
+        str, the name of the compute target (instance or cluster).
+
+    Raises:
+        RuntimeError: if the fetched compute target is unsupported or the
+            mode defined in the setting does not match the type of the
+            compute target.
+    """
+    # If the mode is serverless, we can not fetch anything anyhow
+    if settings.mode == AzureMLComputeTypes.SERVERLESS:
+        return None
+
+    # If a name is not provided, generate one based on the orchestrator id
+    compute_name = settings.compute_name or default_compute_name
+
+    # Try to fetch the compute target
+    try:
+        compute = client.compute.get(compute_name)
+
+        logger.info(f"Using existing compute target: '{compute_name}'.")
+
+        # Check if compute size matches with the settings
+        check_settings_and_compute_configuration(
+            parameter="size", settings=settings, compute=compute
+        )
+
+        compute_type = compute.type
+
+        # Check the type and matches the settings
+        if compute_type == "computeinstance":  # Compute Instance
+            if settings.mode != AzureMLComputeTypes.COMPUTE_INSTANCE:
+                raise RuntimeError(
+                    "The mode of operation for the compute target defined"
+                    f"in the settings '{settings.mode}' does not match "
+                    f"the type of the compute target: `{compute_name}` "
+                    "which is a 'compute-instance'. Please make sure that "
+                    "the settings are adjusted properly."
+                )
+
+            if compute.state != "Running":
+                raise RuntimeError(
+                    f"The compute instance `{compute_name}` is not in a "
+                    "running state at the moment. Please make sure that "
+                    "the compute target is running, before executing the "
+                    "pipeline."
+                )
+
+            # Idle time before shutdown
+            check_settings_and_compute_configuration(
+                parameter="idle_time_before_shutdown_minutes",
+                settings=settings,
+                compute=compute,
+            )
+
+        elif compute_type == "amIcompute":  # Compute Cluster
+            if settings.mode != AzureMLComputeTypes.COMPUTE_CLUSTER:
+                raise RuntimeError(
+                    "The mode of operation for the compute target defined "
+                    f"in the settings '{settings.mode}' does not match "
+                    f"the type of the compute target: `{compute_name}` "
+                    "which is a 'compute-cluster'. Please make sure that "
+                    "the settings are adjusted properly."
+                )
+
+            if compute.provisioning_state != "Succeeded":
+                raise RuntimeError(
+                    f"The provisioning state '{compute.provisioning_state}'"
+                    f"of the compute cluster `{compute_name}` is not "
+                    "successful. Please make sure that the compute cluster "
+                    "is provisioned properly, before executing the "
+                    "pipeline."
+                )
+
+            for parameter in [
+                "idle_time_before_scale_down",
+                "max_instances",
+                "min_instances",
+                "tier",
+                "location",
+            ]:
+                # Check all possible configurations
+                check_settings_and_compute_configuration(
+                    parameter=parameter, settings=settings, compute=compute
+                )
+        else:
+            raise RuntimeError(f"Unsupported compute type: {compute_type}")
+        return compute_name
+
+    # If the compute target does not exist create it
+    except ResourceNotFoundError:
+        logger.info(
+            f"Can not find the compute target with name: '{compute_name}':"
+        )
+
+        if settings.mode == AzureMLComputeTypes.COMPUTE_INSTANCE:
+            logger.info(
+                "Creating a new compute instance. This might take a "
+                "few minutes."
+            )
+
+            from azure.ai.ml.entities import ComputeInstance
+
+            compute_instance = ComputeInstance(
+                name=compute_name,
+                size=settings.size,
+                idle_time_before_shutdown_minutes=settings.idle_time_before_shutdown_minutes,
+            )
+            client.begin_create_or_update(compute_instance).result()
+            return compute_name
+
+        elif settings.mode == AzureMLComputeTypes.COMPUTE_CLUSTER:
+            logger.info(
+                "Creating a new compute cluster. This might take a "
+                "few minutes."
+            )
+
+            from azure.ai.ml.entities import AmlCompute
+
+            compute_cluster = AmlCompute(
+                name=compute_name,
+                size=settings.size,
+                location=settings.location,
+                min_instances=settings.min_instances,
+                max_instances=settings.max_instances,
+                idle_time_before_scale_down=settings.idle_time_before_scaledown_down,
+                tier=settings.tier,
+            )
+            client.begin_create_or_update(compute_cluster).result()
+            return compute_name
+
+    return None