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

zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint.py

@@ -0,0 +1,307 @@
+#  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.
+"""Entrypoint of the Lightning master/orchestrator STUDIO."""
+
+import argparse
+import os
+from typing import Dict, cast
+
+from lightning_sdk import Machine, Studio
+
+from zenml.client import Client
+from zenml.entrypoints.step_entrypoint_configuration import (
+    StepEntrypointConfiguration,
+)
+from zenml.enums import ExecutionStatus
+from zenml.integrations.lightning.flavors.lightning_orchestrator_flavor import (
+    LightningOrchestratorSettings,
+)
+from zenml.integrations.lightning.orchestrators.lightning_orchestrator import (
+    ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID,
+    LightningOrchestrator,
+)
+from zenml.integrations.lightning.orchestrators.utils import (
+    gather_requirements,
+    sanitize_studio_name,
+)
+from zenml.logger import get_logger
+from zenml.orchestrators.dag_runner import ThreadedDagRunner
+
+logger = get_logger(__name__)
+
+
+LIGHTNING_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH = "."
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse entrypoint arguments.
+
+    Returns:
+        Parsed args.
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--run_name", type=str, required=True)
+    parser.add_argument("--deployment_id", type=str, required=True)
+    return parser.parse_args()
+
+
+def main() -> None:
+    """Entrypoint of the Lightning master/orchestrator STUDIO.
+
+    This is the entrypoint of the Lightning master/orchestrator STUDIO. It is
+    responsible for provisioning the STUDIO and running the pipeline steps in
+    separate STUDIO.
+
+    Raises:
+        TypeError: If the active stack's orchestrator is not an instance of
+            LightningOrchestrator.
+        ValueError: If the active stack's container registry is None.
+    """
+    # Log to the container's stdout so it can be streamed by the client.
+    logger.info("Lightning orchestrator STUDIO started.")
+
+    # Parse / extract args.
+    args = parse_args()
+
+    orchestrator_run_id = os.environ.get(
+        ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID
+    )
+    if not orchestrator_run_id:
+        raise ValueError(
+            f"Environment variable '{ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID}' is not set."
+        )
+
+    logger.info(f"Orchestrator run id: {orchestrator_run_id}")
+
+    deployment = Client().get_deployment(args.deployment_id)
+    filename = f"{args.run_name}.tar.gz"
+
+    pipeline_dag = {
+        step_name: step.spec.upstream_steps
+        for step_name, step in deployment.step_configurations.items()
+    }
+    entrypoint_command = StepEntrypointConfiguration.get_entrypoint_command()
+
+    active_stack = Client().active_stack
+
+    orchestrator = active_stack.orchestrator
+    if not isinstance(orchestrator, LightningOrchestrator):
+        raise TypeError(
+            "The active stack's orchestrator is not an instance of LightningOrchestrator."
+        )
+
+    # Set up credentials
+    orchestrator._set_lightning_env_vars(deployment)
+
+    pipeline_settings = cast(
+        LightningOrchestratorSettings, orchestrator.get_settings(deployment)
+    )
+
+    # Gather the requirements
+    pipeline_docker_settings = (
+        deployment.pipeline_configuration.docker_settings
+    )
+    pipeline_requirements = gather_requirements(pipeline_docker_settings)
+    pipeline_requirements_to_string = " ".join(
+        f'"{req}"' for req in pipeline_requirements
+    )
+
+    unique_resource_configs: Dict[str, str] = {}
+    main_studio_name = sanitize_studio_name(
+        f"zenml_{orchestrator_run_id}_pipeline"
+    )
+    for step_name, step in deployment.step_configurations.items():
+        step_settings = cast(
+            LightningOrchestratorSettings,
+            orchestrator.get_settings(step),
+        )
+        unique_resource_configs[step_name] = main_studio_name
+        if pipeline_settings.machine_type != step_settings.machine_type:
+            unique_resource_configs[step_name] = (
+                f"zenml-{orchestrator_run_id}_{step_name}"
+            )
+
+    logger.info(f"Creating main studio: {main_studio_name}")
+    main_studio = Studio(name=main_studio_name)
+    if pipeline_settings.machine_type:
+        main_studio.start(Machine(pipeline_settings.machine_type))
+    else:
+        main_studio.start()
+
+    logger.info("Main studio started.")
+    logger.info("Uploading code to main studio the code path: %s", filename)
+    main_studio.upload_file(
+        "/teamspace/studios/this_studio/.lightning_studio/.studiorc",
+        remote_path="/teamspace/studios/this_studio/.lightning_studio/.studiorc",
+    )
+    output = main_studio.run(
+        f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
+    )
+    logger.info(output)
+    main_studio.upload_file(
+        f"/teamspace/studios/this_studio/zenml_codes/{filename}",
+        remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
+    )
+    logger.info("Extracting code... ")
+    output = main_studio.run(
+        f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
+    )
+    logger.info(f"Code extraction output: {output}")
+    logger.info("Installing requirements... ")
+
+    output = main_studio.run("pip install uv")
+    logger.info(output)
+    output = main_studio.run(
+        f"uv pip install {pipeline_requirements_to_string}"
+    )
+    logger.info(output)
+    output = main_studio.run(
+        "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
+    )
+    logger.info(output)
+
+    for command in pipeline_settings.custom_commands or []:
+        output = main_studio.run(
+            f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
+        )
+        logger.info(f"Custom command output: {output}")
+
+    run = Client().list_pipeline_runs(
+        sort_by="asc:created",
+        size=1,
+        deployment_id=args.deployment_id,
+        status=ExecutionStatus.INITIALIZING,
+    )[0]
+
+    logger.info("Fetching pipeline run: %s", run.id)
+
+    def run_step_on_lightning_studio(step_name: str) -> None:
+        """Run a pipeline step in a separate Lightning STUDIO.
+
+        Args:
+            step_name: Name of the step.
+
+        Raises:
+            Exception: If an error occurs while running the step on the STUDIO.
+        """
+        step_args = StepEntrypointConfiguration.get_entrypoint_arguments(
+            step_name=step_name,
+            deployment_id=args.deployment_id,
+        )
+
+        entrypoint = entrypoint_command + step_args
+        entrypoint_string = " ".join(entrypoint)
+        run_command = f"{entrypoint_string}"
+
+        step = deployment.step_configurations[step_name]
+        if unique_resource_configs[step_name] != main_studio_name:
+            logger.info(
+                f"Creating separate studio for step: {unique_resource_configs[step_name]}"
+            )
+            # Get step settings
+            step_settings = cast(
+                LightningOrchestratorSettings,
+                orchestrator.get_settings(step),
+            )
+            # Gather the requirements
+            step_docker_settings = step.config.docker_settings
+            step_requirements = gather_requirements(step_docker_settings)
+            step_requirements_to_string = " ".join(
+                f'"{req}"' for req in step_requirements
+            )
+            run_command = f"{entrypoint_string}"
+
+            logger.info(
+                f"Creating separate studio for step: {unique_resource_configs[step_name]}"
+            )
+            studio = Studio(name=unique_resource_configs[step_name])
+            try:
+                studio.start(Machine(step_settings.machine_type))
+                output = studio.run(
+                    f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
+                )
+                logger.info(output)
+                studio.upload_file(
+                    f"/teamspace/studios/this_studio/zenml_codes/{filename}",
+                    remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
+                )
+                output = studio.run(
+                    f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
+                )
+                logger.info(output)
+                studio.upload_file(
+                    "/teamspace/studios/this_studio/.lightning_studio/.studiorc",
+                    remote_path="/teamspace/studios/this_studio/.lightning_studio/.studiorc",
+                )
+                output = studio.run("pip install uv")
+                logger.info(output)
+                output = studio.run(
+                    f"uv pip install {step_requirements_to_string}"
+                )
+                logger.info(output)
+                output = studio.run(
+                    "pip uninstall zenml -y && pip install git+https://github.com/zenml-io/zenml.git@feature/lightening-studio-orchestrator"
+                )
+                logger.info(output)
+                for command in step_settings.custom_commands or []:
+                    output = studio.run(
+                        f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
+                    )
+                    logger.info(f"Custom command output: {output}")
+                output = studio.run(
+                    f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {run_command}"
+                )
+                logger.info(output)
+            except Exception as e:
+                logger.error(
+                    f"Error running step {step_name} on studio {unique_resource_configs[step_name]}: {e}"
+                )
+                raise e
+            finally:
+                studio.delete()
+                studio.delete()
+        else:
+            output = main_studio.run(
+                f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {run_command}"
+            )
+            logger.info(output)
+
+            # Pop the resource configuration for this step
+        unique_resource_configs.pop(step_name)
+
+        if main_studio_name in unique_resource_configs.values():
+            # If there are more steps using this configuration, skip deprovisioning the cluster
+            logger.info(
+                f"Resource configuration for studio '{main_studio_name}' "
+                "is used by subsequent steps. Skipping the deprovisioning of "
+                "the studio."
+            )
+        else:
+            # If there are no more steps using this configuration, down the cluster
+            logger.info(
+                f"Resource configuration for cluster '{main_studio_name}' "
+                "is not used by subsequent steps. deprovisioning the cluster."
+            )
+            main_studio.delete()
+        logger.info(f"Running step `{step_name}` on a Studio is completed.")
+
+    ThreadedDagRunner(
+        dag=pipeline_dag, run_fn=run_step_on_lightning_studio
+    ).run()
+
+    logger.info("Orchestration STUDIO provisioned.")
+
+
+if __name__ == "__main__":
+    main()

zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint_configuration.py

@@ -0,0 +1,77 @@
+#  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.
+"""Entrypoint configuration for the Lightning master/orchestrator VM."""
+
+from typing import TYPE_CHECKING, List, Set
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+RUN_NAME_OPTION = "run_name"
+DEPLOYMENT_ID_OPTION = "deployment_id"
+
+
+class LightningOrchestratorEntrypointConfiguration:
+    """Entrypoint configuration for the Lightning master/orchestrator VM."""
+
+    @classmethod
+    def get_entrypoint_options(cls) -> Set[str]:
+        """Gets all the options required for running this entrypoint.
+
+        Returns:
+            Entrypoint options.
+        """
+        options = {
+            RUN_NAME_OPTION,
+            DEPLOYMENT_ID_OPTION,
+        }
+        return options
+
+    @classmethod
+    def get_entrypoint_command(cls) -> List[str]:
+        """Returns a command that runs the entrypoint module.
+
+        Returns:
+            Entrypoint command.
+        """
+        command = [
+            "python",
+            "-m",
+            "zenml.integrations.lightning.orchestrators.lightning_orchestrator_entrypoint",
+        ]
+        return command
+
+    @classmethod
+    def get_entrypoint_arguments(
+        cls,
+        run_name: str,
+        deployment_id: "UUID",
+    ) -> List[str]:
+        """Gets all arguments that the entrypoint command should be called with.
+
+        Args:
+            run_name: Name of the ZenML run.
+            deployment_id: ID of the deployment.
+
+        Returns:
+            List of entrypoint arguments.
+        """
+        args = [
+            f"--{RUN_NAME_OPTION}",
+            run_name,
+            f"--{DEPLOYMENT_ID_OPTION}",
+            str(deployment_id),
+        ]
+
+        return args

zenml/integrations/lightning/orchestrators/utils.py

@@ -0,0 +1,67 @@
+#  Copyright (c) ZenML GmbH 2023. 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.
+"""Utility functions for the Lightning orchestrator."""
+
+import itertools
+import re
+from typing import List
+
+from zenml.client import Client
+from zenml.config import DockerSettings
+from zenml.utils.pipeline_docker_image_builder import (
+    PipelineDockerImageBuilder,
+)
+
+
+def sanitize_studio_name(studio_name: str) -> str:
+    """Sanitize studio_names so they conform to Kubernetes studio naming convention.
+
+    Args:
+        studio_name: Arbitrary input studio_name.
+
+    Returns:
+        Sanitized pod name.
+    """
+    studio_name = re.sub(r"[^a-z0-9-]", "-", studio_name.lower())
+    studio_name = re.sub(r"^[-]+", "", studio_name)
+    return re.sub(r"[-]+", "-", studio_name)
+
+
+def gather_requirements(docker_settings: "DockerSettings") -> List[str]:
+    """Gather the requirements files.
+
+    Args:
+        docker_settings: Docker settings.
+
+    Returns:
+        List of requirements.
+    """
+    docker_image_builder = PipelineDockerImageBuilder()
+    requirements_files = docker_image_builder.gather_requirements_files(
+        docker_settings=docker_settings,
+        stack=Client().active_stack,
+        log=False,
+    )
+
+    # Extract and clean the requirements
+    requirements = list(
+        itertools.chain.from_iterable(
+            r[1].strip().split("\n") for r in requirements_files
+        )
+    )
+
+    # Remove empty items and duplicates
+    requirements = sorted(set(filter(None, requirements)))
+
+    return requirements

zenml/integrations/mlflow/__init__.py

@@ -16,16 +16,21 @@
 The MLflow integrations currently enables you to use MLflow tracking as a
 convenient way to visualize your experiment runs within the MLflow UI.
 """
-from typing import List, Type
+from typing import List, Type, Optional
 
 from zenml.integrations.constants import MLFLOW
 from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
+import sys
+
+from zenml.logger import get_logger
 
 MLFLOW_MODEL_DEPLOYER_FLAVOR = "mlflow"
 MLFLOW_MODEL_EXPERIMENT_TRACKER_FLAVOR = "mlflow"
 MLFLOW_MODEL_REGISTRY_FLAVOR = "mlflow"
 
+logger = get_logger(__name__)
+
 
 class MlflowIntegration(Integration):
     """Definition of MLflow integration for ZenML."""
@@ -34,8 +39,6 @@ class MlflowIntegration(Integration):
 
     REQUIREMENTS = [
         "mlflow>=2.1.1,<3",
-        "mlserver>=1.3.3",
-        "mlserver-mlflow>=1.3.3",
         # TODO: remove this requirement once rapidjson is fixed
         "python-rapidjson<1.15",
         # When you do:
@@ -45,10 +48,45 @@ class MlflowIntegration(Integration):
         # This downgrades pydantic to v1 even though mlflow does not have
         # any issues with v2. This is why we have to pin it here so a downgrade
         # will not happen.
-        "pydantic>=2.7.0,<2.8.0",
+        "pydantic>=2.8.0,<2.9.0",
+    ]
+
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "python-rapidjson",
+        "pydantic",
+        "numpy",
+        "pandas",
     ]
 
-    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["python-rapidjson", "pydantic"]
+    @classmethod
+    def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
+        """Method to get the requirements for the integration.
+
+        Args:
+            target_os: The target operating system to get the requirements for.
+
+        Returns:
+            A list of requirements.
+        """
+        from zenml.integrations.numpy import NumpyIntegration
+        from zenml.integrations.pandas import PandasIntegration
+
+        reqs = cls.REQUIREMENTS
+        if sys.version_info.minor >= 12:
+            logger.debug(
+                "The MLflow integration on Python 3.12 and above is not yet "
+                "fully supported: The extra dependencies 'mlserver' and "
+                "'mlserver-mlflow' will be skipped."
+            )
+        else:
+            reqs.extend([
+                "mlserver>=1.3.3",
+                "mlserver-mlflow>=1.3.3",
+            ])
+
+        reqs.extend(NumpyIntegration.get_requirements(target_os=target_os))
+        reqs.extend(PandasIntegration.get_requirements(target_os=target_os))
+        return reqs
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/mlflow/services/mlflow_deployment.py

@@ -14,6 +14,7 @@
 """Implementation of the MLflow deployment functionality."""
 
 import os
+import sys
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import numpy as np
@@ -21,6 +22,7 @@ import pandas as pd
 import requests
 from mlflow.pyfunc.backend import PyFuncBackend
 from mlflow.version import VERSION as MLFLOW_VERSION
+from pydantic import field_validator
 
 from zenml.client import Client
 from zenml.constants import DEFAULT_SERVICE_START_STOP_TIMEOUT
@@ -110,6 +112,30 @@ class MLFlowDeploymentConfig(LocalDaemonServiceConfig):
     mlserver: bool = False
     timeout: int = DEFAULT_SERVICE_START_STOP_TIMEOUT
 
+    @field_validator("mlserver")
+    @classmethod
+    def validate_mlserver_python_version(cls, mlserver: bool) -> bool:
+        """Validates the Python version if mlserver is used.
+
+        Args:
+            mlserver: set to True if the MLflow MLServer backend is used,
+                else set to False and MLflow built-in scoring server will be
+                used.
+
+        Returns:
+            the validated value
+
+        Raises:
+            ValueError: if mlserver is set to true on Python 3.12 as it is not
+                yet supported.
+        """
+        if mlserver is True and sys.version_info.minor >= 12:
+            raise ValueError(
+                "The mlserver deployment is not yet supported on Python 3.12 "
+                "or above."
+            )
+        return mlserver
+
 
 class MLFlowDeploymentService(LocalDaemonService, BaseDeploymentService):
     """MLflow deployment service used to start a local prediction server for MLflow models.

zenml/integrations/numpy/__init__.py

@@ -0,0 +1,32 @@
+#  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.
+"""Initialization of the Numpy integration."""
+
+from zenml.integrations.constants import NUMPY
+from zenml.integrations.integration import Integration
+
+
+class NumpyIntegration(Integration):
+    """Definition of Numpy integration for ZenML."""
+
+    NAME = NUMPY
+    REQUIREMENTS = ["numpy<2.0.0"]
+
+    @classmethod
+    def activate(cls) -> None:
+        """Activates the integration."""
+        from zenml.integrations.numpy import materializers  # noqa
+
+
+NumpyIntegration.check_installation()

zenml/integrations/numpy/materializers/__init__.py

@@ -0,0 +1,18 @@
+#  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.
+"""Initialization of the Numpy materializer."""
+
+from zenml.integrations.numpy.materializers.numpy_materializer import (  # noqa
+    NumpyMaterializer,
+)