zenml-nightly 0.64.0.dev20240811__py3-none-any.whl → 0.66.0.dev20240910__py3-none-any.whl

zenml/integrations/azure/orchestrators/azureml_orchestrator.py

@@ -26,6 +26,7 @@ from typing import (
     Union,
     cast,
 )
+from uuid import UUID
 
 from azure.ai.ml import Input, MLClient, Output
 from azure.ai.ml.constants import TimeZone
@@ -40,15 +41,16 @@ from azure.ai.ml.entities import (
 from azure.core.exceptions import (
     HttpResponseError,
     ResourceExistsError,
-    ResourceNotFoundError,
 )
 from azure.identity import DefaultAzureCredential
 
 from zenml.config.base_settings import BaseSettings
 from zenml.config.step_configurations import Step
+from zenml.constants import METADATA_ORCHESTRATOR_URL
 from zenml.enums import StackComponentType
+from zenml.integrations.azure.azureml_utils import create_or_get_compute
+from zenml.integrations.azure.flavors.azureml import AzureMLComputeTypes
 from zenml.integrations.azure.flavors.azureml_orchestrator_flavor import (
-    AzureMLComputeTypes,
     AzureMLOrchestratorConfig,
     AzureMLOrchestratorSettings,
 )
@@ -56,14 +58,13 @@ from zenml.integrations.azure.orchestrators.azureml_orchestrator_entrypoint_conf
     AzureMLEntrypointConfiguration,
 )
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType, Uri
 from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 from zenml.utils.string_utils import b64_encode
 
 if TYPE_CHECKING:
-    from azure.ai.ml.entities import AmlCompute, ComputeInstance
-
     from zenml.models import PipelineDeploymentResponse
     from zenml.stack import Stack
 
@@ -193,172 +194,6 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
             command=" ".join(command + arguments),
         )
 
-    @staticmethod
-    def _check_settings_and_compute_configuration(
-        parameter: str,
-        settings: AzureMLOrchestratorSettings,
-        compute: Union["ComputeInstance", "AmlCompute"],
-    ) -> 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(
-        self, client: MLClient, settings: AzureMLOrchestratorSettings
-    ) -> Optional[str]:
-        """Creates or fetches the compute target if defined in the settings.
-
-        Args:
-            client: the AzureML client.
-            settings: the settings for the orchestrator.
-
-        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 f"compute_{self.id}"
-        # 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
-            self._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
-                self._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
-                    self._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(
-                "Can not find the compute target with name: "
-                f"'{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
-
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
@@ -429,7 +264,9 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
         )
         pipeline_args["name"] = run_name
 
-        if compute_target := self._create_or_get_compute(ml_client, settings):
+        if compute_target := create_or_get_compute(
+            ml_client, settings, default_compute_name=f"zenml_{self.id}"
+        ):
             pipeline_args["compute"] = compute_target
 
         @pipeline(force_rerun=True, **pipeline_args)  # type: ignore[call-overload, misc]
@@ -540,5 +377,54 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
                 )
 
         else:
-            ml_client.jobs.create_or_update(pipeline_job)
+            job = ml_client.jobs.create_or_update(pipeline_job)
             logger.info(f"Pipeline {run_name} has been started.")
+            assert job.services is not None
+            assert job.name is not None
+
+            logger.info(
+                f"Pipeline {run_name} is running. "
+                "You can view the pipeline in the AzureML portal at "
+                f"{job.services['Studio'].endpoint}"
+            )
+
+            if self.config.synchronous:
+                logger.info("Waiting for pipeline to finish...")
+                ml_client.jobs.stream(job.name)
+
+    def get_pipeline_run_metadata(
+        self, run_id: UUID
+    ) -> Dict[str, "MetadataType"]:
+        """Get general component-specific metadata for a pipeline run.
+
+        Args:
+            run_id: The ID of the pipeline run.
+
+        Returns:
+            A dictionary of metadata.
+        """
+        try:
+            if connector := self.get_connector():
+                credentials = connector.connect()
+            else:
+                credentials = DefaultAzureCredential()
+
+            ml_client = MLClient(
+                credential=credentials,
+                subscription_id=self.config.subscription_id,
+                resource_group_name=self.config.resource_group,
+                workspace_name=self.config.workspace,
+            )
+
+            azureml_root_run_id = os.environ[ENV_ZENML_AZUREML_RUN_ID]
+            azureml_job = ml_client.jobs.get(azureml_root_run_id)
+
+            return {
+                METADATA_ORCHESTRATOR_URL: Uri(azureml_job.studio_url),
+            }
+        except Exception as e:
+            logger.warning(
+                f"Failed to fetch the Studio URL of the AzureML pipeline "
+                f"job: {e}"
+            )
+            return {}

zenml/integrations/azure/orchestrators/azureml_orchestrator_entrypoint_config.py

@@ -69,6 +69,7 @@ class AzureMLEntrypointConfiguration(StepEntrypointConfiguration):
 
         # Azure automatically changes the working directory, we have to set it
         # back to /app before running the step.
+        os.makedirs("/app", exist_ok=True)
         os.chdir("/app")
 
         # Run the step

zenml/integrations/azure/service_connectors/azure_service_connector.py

@@ -14,6 +14,7 @@
 """Azure Service Connector."""
 
 import datetime
+import logging
 import re
 import subprocess
 from typing import Any, Dict, List, Optional, Tuple
@@ -58,6 +59,9 @@ from zenml.service_connectors.service_connector import (
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.secret_utils import PlainSerializedSecretStr
 
+# Configure the logging level for azure.identity
+logging.getLogger("azure.identity").setLevel(logging.WARNING)
+
 logger = get_logger(__name__)
 
 

zenml/integrations/azure/step_operators/azureml_step_operator.py

@@ -13,29 +13,16 @@
 #  permissions and limitations under the License.
 """Implementation of the ZenML AzureML Step Operator."""
 
-import itertools
-import os
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, cast
 
-from azureml.core import (
-    ComputeTarget,
-    Environment,
-    Experiment,
-    ScriptRunConfig,
-    Workspace,
-)
-from azureml.core.authentication import (
-    AbstractAuthentication,
-    ServicePrincipalAuthentication,
-)
-from azureml.core.conda_dependencies import CondaDependencies
+from azure.ai.ml import MLClient, command
+from azure.ai.ml.entities import Environment
+from azure.core.credentials import TokenCredential
+from azure.identity import ClientSecretCredential, DefaultAzureCredential
 
-import zenml
-from zenml.client import Client
-from zenml.constants import (
-    ENV_ZENML_CONFIG_PATH,
-)
-from zenml.environment import Environment as ZenMLEnvironment
+from zenml.config.build_configuration import BuildConfiguration
+from zenml.enums import StackComponentType
+from zenml.integrations.azure.azureml_utils import create_or_get_compute
 from zenml.integrations.azure.flavors.azureml_step_operator_flavor import (
     AzureMLStepOperatorConfig,
     AzureMLStepOperatorSettings,
@@ -43,20 +30,18 @@ from zenml.integrations.azure.flavors.azureml_step_operator_flavor import (
 from zenml.logger import get_logger
 from zenml.stack import Stack, StackValidator
 from zenml.step_operators import BaseStepOperator
-from zenml.utils import source_utils
-from zenml.utils.pipeline_docker_image_builder import (
-    DOCKER_IMAGE_ZENML_CONFIG_DIR,
-    PipelineDockerImageBuilder,
-)
 
 if TYPE_CHECKING:
-    from zenml.config import DockerSettings
     from zenml.config.base_settings import BaseSettings
     from zenml.config.step_run_info import StepRunInfo
+    from zenml.models import PipelineDeploymentBase
 
 logger = get_logger(__name__)
 
 
+AZUREML_STEP_OPERATOR_DOCKER_IMAGE_KEY = "azureml_step_operator"
+
+
 class AzureMLStepOperator(BaseStepOperator):
     """Step operator to run a step on AzureML.
 
@@ -87,11 +72,11 @@ class AzureMLStepOperator(BaseStepOperator):
         """Validates the stack.
 
         Returns:
-            A validator that checks that the stack contains a remote artifact
-            store.
+            A validator that checks that the stack contains a remote container
+            registry and a remote artifact store.
         """
 
-        def _validate_remote_artifact_store(
+        def _validate_remote_components(
             stack: "Stack",
         ) -> Tuple[bool, str]:
             if stack.artifact_store.config.is_local:
@@ -104,113 +89,75 @@ class AzureMLStepOperator(BaseStepOperator):
                     "step operator."
                 )
 
+            container_registry = stack.container_registry
+            assert container_registry is not None
+
+            if container_registry.config.is_local:
+                return False, (
+                    "The AzureML step operator runs code remotely and "
+                    "needs to push/pull Docker images, but the "
+                    f"container registry `{container_registry.name}` of the "
+                    "active stack is local. Please ensure that your stack "
+                    "contains a remote container registry when using the "
+                    "AzureML step operator."
+                )
+
             return True, ""
 
         return StackValidator(
-            custom_validation_function=_validate_remote_artifact_store,
+            required_components={
+                StackComponentType.CONTAINER_REGISTRY,
+                StackComponentType.IMAGE_BUILDER,
+            },
+            custom_validation_function=_validate_remote_components,
         )
 
-    def _get_authentication(self) -> Optional[AbstractAuthentication]:
+    def _get_credentials(self) -> TokenCredential:
         """Returns the authentication object for the AzureML environment.
 
         Returns:
             The authentication object for the AzureML environment.
         """
-        if (
+        # Authentication
+        if connector := self.get_connector():
+            credentials = connector.connect()
+            assert isinstance(credentials, TokenCredential)
+            return credentials
+        elif (
             self.config.tenant_id
             and self.config.service_principal_id
             and self.config.service_principal_password
         ):
-            return ServicePrincipalAuthentication(
+            return ClientSecretCredential(
                 tenant_id=self.config.tenant_id,
-                service_principal_id=self.config.service_principal_id,
-                service_principal_password=self.config.service_principal_password,
+                client_id=self.config.service_principal_id,
+                client_secret=self.config.service_principal_password,
             )
-        return None
+        else:
+            return DefaultAzureCredential()
 
-    def _prepare_environment(
-        self,
-        workspace: Workspace,
-        docker_settings: "DockerSettings",
-        run_name: str,
-        environment_variables: Dict[str, str],
-        environment_name: Optional[str] = None,
-    ) -> Environment:
-        """Prepares the environment in which Azure will run all jobs.
+    def get_docker_builds(
+        self, deployment: "PipelineDeploymentBase"
+    ) -> List["BuildConfiguration"]:
+        """Gets the Docker builds required for the component.
 
         Args:
-            workspace: The AzureML Workspace that has configuration
-                for a storage account, container registry among other
-                things.
-            docker_settings: The Docker settings for this step.
-            run_name: The name of the pipeline run that can be used
-                for naming environments and runs.
-            environment_variables: Environment variables to set in the
-                environment.
-            environment_name: Optional name of an existing environment to use.
+            deployment: The pipeline deployment for which to get the builds.
 
         Returns:
-            The AzureML Environment object.
+            The required Docker builds.
         """
-        docker_image_builder = PipelineDockerImageBuilder()
-        requirements_files = docker_image_builder.gather_requirements_files(
-            docker_settings=docker_settings,
-            stack=Client().active_stack,
-            log=False,
-        )
-        requirements = list(
-            itertools.chain.from_iterable(
-                r[1].split("\n") for r in requirements_files
-            )
-        )
-        requirements.append(f"zenml=={zenml.__version__}")
-        logger.info(
-            "Using requirements for AzureML step operator environment: %s",
-            requirements,
-        )
-        if environment_name:
-            environment = Environment.get(
-                workspace=workspace, name=environment_name
-            )
-            if not environment.python.conda_dependencies:
-                environment.python.conda_dependencies = (
-                    CondaDependencies.create(
-                        python_version=ZenMLEnvironment.python_version()
-                    )
+        builds = []
+        for step_name, step in deployment.step_configurations.items():
+            if step.config.step_operator == self.name:
+                build = BuildConfiguration(
+                    key=AZUREML_STEP_OPERATOR_DOCKER_IMAGE_KEY,
+                    settings=step.config.docker_settings,
+                    step_name=step_name,
                 )
+                builds.append(build)
 
-            for requirement in requirements:
-                environment.python.conda_dependencies.add_pip_package(
-                    requirement
-                )
-        else:
-            environment = Environment(name=f"zenml-{run_name}")
-            environment.python.conda_dependencies = CondaDependencies.create(
-                pip_packages=requirements,
-                python_version=ZenMLEnvironment.python_version(),
-            )
-
-            if docker_settings.parent_image:
-                # replace the default azure base image
-                environment.docker.base_image = docker_settings.parent_image
-
-        # set credentials to access azure storage
-        for key in [
-            "AZURE_STORAGE_ACCOUNT_KEY",
-            "AZURE_STORAGE_ACCOUNT_NAME",
-            "AZURE_STORAGE_CONNECTION_STRING",
-            "AZURE_STORAGE_SAS_TOKEN",
-        ]:
-            value = os.getenv(key)
-            if value:
-                environment_variables[key] = value
-
-        environment_variables[ENV_ZENML_CONFIG_PATH] = (
-            f"./{DOCKER_IMAGE_ZENML_CONFIG_DIR}"
-        )
-        environment_variables.update(docker_settings.environment)
-        environment.environment_variables = environment_variables
-        return environment
+        return builds
 
     def launch(
         self,
@@ -226,79 +173,33 @@ class AzureMLStepOperator(BaseStepOperator):
             environment: Environment variables to set in the step operator
                 environment.
         """
-        if not info.config.resource_settings.empty:
-            logger.warning(
-                "Specifying custom step resources is not supported for "
-                "the AzureML step operator. If you want to run this step "
-                "operator on specific resources, you can do so by creating an "
-                "Azure compute target (https://docs.microsoft.com/en-us/azure/machine-learning/concept-compute-target) "
-                "with a specific machine type and then updating this step "
-                "operator: `zenml step-operator update %s "
-                "--compute_target_name=<COMPUTE_TARGET_NAME>`",
-                self.name,
-            )
-
-        unused_docker_fields = [
-            "dockerfile",
-            "build_context_root",
-            "build_options",
-            "skip_build",
-            "target_repository",
-            "dockerignore",
-            "copy_files",
-            "copy_global_config",
-            "apt_packages",
-            "user",
-            "source_files",
-            "allow_including_files_in_images",
-            "allow_download_from_code_repository",
-            "allow_download_from_artifact_store",
-        ]
-        docker_settings = info.config.docker_settings
-        ignored_docker_fields = docker_settings.model_fields_set.intersection(
-            unused_docker_fields
-        )
-
-        if ignored_docker_fields:
-            logger.warning(
-                "The AzureML step operator currently does not support all "
-                "options defined in your Docker settings. Ignoring all "
-                "values set for the attributes: %s",
-                ignored_docker_fields,
-            )
-
         settings = cast(AzureMLStepOperatorSettings, self.get_settings(info))
+        image_name = info.get_image(key=AZUREML_STEP_OPERATOR_DOCKER_IMAGE_KEY)
 
-        workspace = Workspace.get(
+        # Client creation
+        ml_client = MLClient(
+            credential=self._get_credentials(),
             subscription_id=self.config.subscription_id,
-            resource_group=self.config.resource_group,
-            name=self.config.workspace_name,
-            auth=self._get_authentication(),
+            resource_group_name=self.config.resource_group,
+            workspace_name=self.config.workspace_name,
         )
 
-        source_directory = source_utils.get_source_root()
+        env = Environment(name=f"zenml-{info.run_name}", image=image_name)
 
-        environment = self._prepare_environment(
-            workspace=workspace,
-            docker_settings=docker_settings,
-            run_name=info.run_name,
-            environment_variables=environment,
-            environment_name=settings.environment_name,
-        )
-        compute_target = ComputeTarget(
-            workspace=workspace, name=self.config.compute_target_name
+        compute_target = create_or_get_compute(
+            ml_client, settings, default_compute_name=f"zenml_{self.id}"
         )
 
-        run_config = ScriptRunConfig(
-            source_directory=source_directory,
-            environment=environment,
-            compute_target=compute_target,
-            command=entrypoint_command,
+        command_job = command(
+            name=info.run_name,
+            command=" ".join(entrypoint_command),
+            environment=env,
+            environment_variables=environment,
+            compute=compute_target,
+            experiment_name=info.pipeline.name,
         )
 
-        experiment = Experiment(workspace=workspace, name=info.pipeline.name)
-        run = experiment.submit(config=run_config)
+        job = ml_client.jobs.create_or_update(command_job)
 
-        run.display_name = info.run_name
-        info.force_write_logs()
-        run.wait_for_completion(show_output=True)
+        logger.info(f"AzureML job created with id: {job.id}")
+        ml_client.jobs.stream(info.run_name)

zenml/integrations/constants.py

@@ -44,7 +44,9 @@ LIGHTGBM = "lightgbm"
 MLFLOW = "mlflow"
 NEPTUNE = "neptune"
 NEURAL_PROPHET = "neural_prophet"
+NUMPY = "numpy"
 OPEN_AI = "openai"
+PANDAS = "pandas"
 PIGEON = "pigeon"
 PILLOW = "pillow"
 PLOTLY = "plotly"
@@ -73,3 +75,4 @@ WANDB = "wandb"
 VERTEX = "vertex"
 XGBOOST = "xgboost"
 VAULT = "vault"
+LIGHTNING = "lightning"