Lightning

zenml.integrations.lightning

Initialization of the Lightning integration for ZenML.

Attributes

LIGHTNING = 'lightning' module-attribute

LIGHTNING_ORCHESTRATOR_FLAVOR = 'lightning' module-attribute

Classes

Flavor

Class for ZenML Flavors.

Attributes

config_class: Type[StackComponentConfig] abstractmethod property

Returns StackComponentConfig config class.

Returns:

Type	Description
Type[StackComponentConfig]	The config class.

config_schema: Dict[str, Any] property

The config schema for a flavor.

Returns:

Type	Description
Dict[str, Any]	The config schema.

docs_url: Optional[str] property

A url to point at docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor docs url.

implementation_class: Type[StackComponent] abstractmethod property

Implementation class for this flavor.

Returns:

Type	Description
Type[StackComponent]	The implementation class for this flavor.

logo_url: Optional[str] property

A url to represent the flavor in the dashboard.

Returns:

Type	Description
Optional[str]	The flavor logo.

name: str abstractmethod property

The flavor name.

Returns:

Type	Description
str	The flavor name.

sdk_docs_url: Optional[str] property

A url to point at SDK docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor SDK docs url.

service_connector_requirements: Optional[ServiceConnectorRequirements] property

Service connector resource requirements for service connectors.

Specifies resource requirements that are used to filter the available service connector types that are compatible with this flavor.

Returns:

Type	Description
Optional[ServiceConnectorRequirements]	Requirements for compatible service connectors, if a service
Optional[ServiceConnectorRequirements]	connector is required for this flavor.

type: StackComponentType abstractmethod property

The stack component type.

Returns:

Type	Description
StackComponentType	The stack component type.

Functions

from_model(flavor_model: FlavorResponse) -> Flavor classmethod

Loads a flavor from a model.

Parameters:

Name	Type	Description	Default
flavor_model	FlavorResponse	The model to load from.	required

Raises:

Type	Description
CustomFlavorImportError	If the custom flavor can't be imported.
ImportError	If the flavor can't be imported.

Returns:

Type	Description
Flavor	The loaded flavor.

Source code in src/zenml/stack/flavor.py

@classmethod
def from_model(cls, flavor_model: FlavorResponse) -> "Flavor":
    """Loads a flavor from a model.

    Args:
        flavor_model: The model to load from.

    Raises:
        CustomFlavorImportError: If the custom flavor can't be imported.
        ImportError: If the flavor can't be imported.

    Returns:
        The loaded flavor.
    """
    try:
        flavor = source_utils.load(flavor_model.source)()
    except (ModuleNotFoundError, ImportError, NotImplementedError) as err:
        if flavor_model.is_custom:
            flavor_module, _ = flavor_model.source.rsplit(".", maxsplit=1)
            expected_file_path = os.path.join(
                source_utils.get_source_root(),
                flavor_module.replace(".", os.path.sep),
            )
            raise CustomFlavorImportError(
                f"Couldn't import custom flavor {flavor_model.name}: "
                f"{err}. Make sure the custom flavor class "
                f"`{flavor_model.source}` is importable. If it is part of "
                "a library, make sure it is installed. If "
                "it is a local code file, make sure it exists at "
                f"`{expected_file_path}.py`."
            )
        else:
            raise ImportError(
                f"Couldn't import flavor {flavor_model.name}: {err}"
            )
    return cast(Flavor, flavor)

generate_default_docs_url() -> str

Generate the doc urls for all inbuilt and integration flavors.

Note that this method is not going to be useful for custom flavors, which do not have any docs in the main zenml docs.

Returns:

Type	Description
str	The complete url to the zenml documentation

Source code in src/zenml/stack/flavor.py

def generate_default_docs_url(self) -> str:
    """Generate the doc urls for all inbuilt and integration flavors.

    Note that this method is not going to be useful for custom flavors,
    which do not have any docs in the main zenml docs.

    Returns:
        The complete url to the zenml documentation
    """
    from zenml import __version__

    component_type = self.type.plural.replace("_", "-")
    name = self.name.replace("_", "-")

    try:
        is_latest = is_latest_zenml_version()
    except RuntimeError:
        # We assume in error cases that we are on the latest version
        is_latest = True

    if is_latest:
        base = "https://docs.zenml.io"
    else:
        base = f"https://zenml-io.gitbook.io/zenml-legacy-documentation/v/{__version__}"
    return f"{base}/stack-components/{component_type}/{name}"

generate_default_sdk_docs_url() -> str

Generate SDK docs url for a flavor.

Returns:

Type	Description
str	The complete url to the zenml SDK docs

Source code in src/zenml/stack/flavor.py

def generate_default_sdk_docs_url(self) -> str:
    """Generate SDK docs url for a flavor.

    Returns:
        The complete url to the zenml SDK docs
    """
    from zenml import __version__

    base = f"https://sdkdocs.zenml.io/{__version__}"

    component_type = self.type.plural

    if "zenml.integrations" in self.__module__:
        # Get integration name out of module path which will look something
        #  like this "zenml.integrations.<integration>....
        integration = self.__module__.split(
            "zenml.integrations.", maxsplit=1
        )[1].split(".")[0]

        return (
            f"{base}/integration_code_docs"
            f"/integrations-{integration}/#{self.__module__}"
        )

    else:
        return (
            f"{base}/core_code_docs/core-{component_type}/"
            f"#{self.__module__}"
        )

to_model(integration: Optional[str] = None, is_custom: bool = True) -> FlavorRequest

Converts a flavor to a model.

Parameters:

Name	Type	Description	Default
integration	Optional[str]	The integration to use for the model.	None
is_custom	bool	Whether the flavor is a custom flavor.	True

Returns:

Type	Description
FlavorRequest	The model.

Source code in src/zenml/stack/flavor.py

def to_model(
    self,
    integration: Optional[str] = None,
    is_custom: bool = True,
) -> FlavorRequest:
    """Converts a flavor to a model.

    Args:
        integration: The integration to use for the model.
        is_custom: Whether the flavor is a custom flavor.

    Returns:
        The model.
    """
    connector_requirements = self.service_connector_requirements
    connector_type = (
        connector_requirements.connector_type
        if connector_requirements
        else None
    )
    resource_type = (
        connector_requirements.resource_type
        if connector_requirements
        else None
    )
    resource_id_attr = (
        connector_requirements.resource_id_attr
        if connector_requirements
        else None
    )

    model = FlavorRequest(
        name=self.name,
        type=self.type,
        source=source_utils.resolve(self.__class__).import_path,
        config_schema=self.config_schema,
        connector_type=connector_type,
        connector_resource_type=resource_type,
        connector_resource_id_attr=resource_id_attr,
        integration=integration,
        logo_url=self.logo_url,
        docs_url=self.docs_url,
        sdk_docs_url=self.sdk_docs_url,
        is_custom=is_custom,
    )
    return model

Integration

Base class for integration in ZenML.

Functions

activate() -> None classmethod

Abstract method to activate the integration.

Source code in src/zenml/integrations/integration.py

@classmethod
def activate(cls) -> None:
    """Abstract method to activate the integration."""

check_installation() -> bool classmethod

Method to check whether the required packages are installed.

Returns:

Type	Description
bool	True if all required packages are installed, False otherwise.

Source code in src/zenml/integrations/integration.py

@classmethod
def check_installation(cls) -> bool:
    """Method to check whether the required packages are installed.

    Returns:
        True if all required packages are installed, False otherwise.
    """
    for r in cls.get_requirements():
        try:
            # First check if the base package is installed
            dist = pkg_resources.get_distribution(r)

            # Next, check if the dependencies (including extras) are
            # installed
            deps: List[Requirement] = []

            _, extras = parse_requirement(r)
            if extras:
                extra_list = extras[1:-1].split(",")
                for extra in extra_list:
                    try:
                        requirements = dist.requires(extras=[extra])  # type: ignore[arg-type]
                    except pkg_resources.UnknownExtra as e:
                        logger.debug(f"Unknown extra: {str(e)}")
                        return False
                    deps.extend(requirements)
            else:
                deps = dist.requires()

            for ri in deps:
                try:
                    # Remove the "extra == ..." part from the requirement string
                    cleaned_req = re.sub(
                        r"; extra == \"\w+\"", "", str(ri)
                    )
                    pkg_resources.get_distribution(cleaned_req)
                except pkg_resources.DistributionNotFound as e:
                    logger.debug(
                        f"Unable to find required dependency "
                        f"'{e.req}' for requirement '{r}' "
                        f"necessary for integration '{cls.NAME}'."
                    )
                    return False
                except pkg_resources.VersionConflict as e:
                    logger.debug(
                        f"Package version '{e.dist}' does not match "
                        f"version '{e.req}' required by '{r}' "
                        f"necessary for integration '{cls.NAME}'."
                    )
                    return False

        except pkg_resources.DistributionNotFound as e:
            logger.debug(
                f"Unable to find required package '{e.req}' for "
                f"integration {cls.NAME}."
            )
            return False
        except pkg_resources.VersionConflict as e:
            logger.debug(
                f"Package version '{e.dist}' does not match version "
                f"'{e.req}' necessary for integration {cls.NAME}."
            )
            return False

    logger.debug(
        f"Integration {cls.NAME} is installed correctly with "
        f"requirements {cls.get_requirements()}."
    )
    return True

flavors() -> List[Type[Flavor]] classmethod

Abstract method to declare new stack component flavors.

Returns:

Type	Description
List[Type[Flavor]]	A list of new stack component flavors.

Source code in src/zenml/integrations/integration.py

@classmethod
def flavors(cls) -> List[Type[Flavor]]:
    """Abstract method to declare new stack component flavors.

    Returns:
        A list of new stack component flavors.
    """
    return []

get_requirements(target_os: Optional[str] = None, python_version: Optional[str] = None) -> List[str] classmethod

Method to get the requirements for the integration.

Parameters:

Name	Type	Description	Default
target_os	Optional[str]	The target operating system to get the requirements for.	None
python_version	Optional[str]	The Python version to use for the requirements.	None

Returns:

Type	Description
List[str]	A list of requirements.

Source code in src/zenml/integrations/integration.py

@classmethod
def get_requirements(
    cls,
    target_os: Optional[str] = None,
    python_version: 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.
        python_version: The Python version to use for the requirements.

    Returns:
        A list of requirements.
    """
    return cls.REQUIREMENTS

get_uninstall_requirements(target_os: Optional[str] = None) -> List[str] classmethod

Method to get the uninstall requirements for the integration.

Parameters:

Name	Type	Description	Default
target_os	Optional[str]	The target operating system to get the requirements for.	None

Returns:

Type	Description
List[str]	A list of requirements.

Source code in src/zenml/integrations/integration.py

@classmethod
def get_uninstall_requirements(
    cls, target_os: Optional[str] = None
) -> List[str]:
    """Method to get the uninstall requirements for the integration.

    Args:
        target_os: The target operating system to get the requirements for.

    Returns:
        A list of requirements.
    """
    ret = []
    for each in cls.get_requirements(target_os=target_os):
        is_ignored = False
        for ignored in cls.REQUIREMENTS_IGNORED_ON_UNINSTALL:
            if each.startswith(ignored):
                is_ignored = True
                break
        if not is_ignored:
            ret.append(each)
    return ret

plugin_flavors() -> List[Type[BasePluginFlavor]] classmethod

Abstract method to declare new plugin flavors.

Returns:

Type	Description
List[Type[BasePluginFlavor]]	A list of new plugin flavors.

Source code in src/zenml/integrations/integration.py

@classmethod
def plugin_flavors(cls) -> List[Type["BasePluginFlavor"]]:
    """Abstract method to declare new plugin flavors.

    Returns:
        A list of new plugin flavors.
    """
    return []

LightningIntegration

Bases: Integration

Definition of Lightning Integration for ZenML.

Functions

flavors() -> List[Type[Flavor]] classmethod

Declare the stack component flavors for the Lightning integration.

Returns:

Type	Description
List[Type[Flavor]]	List of stack component flavors for this integration.

Source code in src/zenml/integrations/lightning/__init__.py

@classmethod
def flavors(cls) -> List[Type[Flavor]]:
    """Declare the stack component flavors for the Lightning integration.

    Returns:
        List of stack component flavors for this integration.
    """
    from zenml.integrations.lightning.flavors import (
        LightningOrchestratorFlavor,
    )

    return [
        LightningOrchestratorFlavor,
    ]

Modules

flavors

Lightning integration flavors.

Classes

LightningOrchestratorConfig(warn_about_plain_text_secrets: bool = False, **kwargs: Any)

Bases: BaseOrchestratorConfig, LightningOrchestratorSettings

Lightning orchestrator base config.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self, warn_about_plain_text_secrets: bool = False, **kwargs: Any
) -> None:
    """Ensures that secret references don't clash with pydantic validation.

    StackComponents allow the specification of all their string attributes
    using secret references of the form `{{secret_name.key}}`. This however
    is only possible when the stack component does not perform any explicit
    validation of this attribute using pydantic validators. If this were
    the case, the validation would run on the secret reference and would
    fail or in the worst case, modify the secret reference and lead to
    unexpected behavior. This method ensures that no attributes that require
    custom pydantic validation are set as secret references.

    Args:
        warn_about_plain_text_secrets: If true, then warns about using
            plain-text secrets.
        **kwargs: Arguments to initialize this stack component.

    Raises:
        ValueError: If an attribute that requires custom pydantic validation
            is passed as a secret reference, or if the `name` attribute
            was passed as a secret reference.
    """
    for key, value in kwargs.items():
        try:
            field = self.__class__.model_fields[key]
        except KeyError:
            # Value for a private attribute or non-existing field, this
            # will fail during the upcoming pydantic validation
            continue

        if value is None:
            continue

        if not secret_utils.is_secret_reference(value):
            if (
                secret_utils.is_secret_field(field)
                and warn_about_plain_text_secrets
            ):
                logger.warning(
                    "You specified a plain-text value for the sensitive "
                    f"attribute `{key}` for a `{self.__class__.__name__}` "
                    "stack component. This is currently only a warning, "
                    "but future versions of ZenML will require you to pass "
                    "in sensitive information as secrets. Check out the "
                    "documentation on how to configure your stack "
                    "components with secrets here: "
                    "https://docs.zenml.io/getting-started/deploying-zenml/secret-management"
                )
            continue

        if pydantic_utils.has_validators(
            pydantic_class=self.__class__, field_name=key
        ):
            raise ValueError(
                f"Passing the stack component attribute `{key}` as a "
                "secret reference is not allowed as additional validation "
                "is required for this attribute."
            )

    super().__init__(**kwargs)

Attributes

is_local: bool property

Checks if this stack component is running locally.

Returns:

Type	Description
bool	True if this config is for a local component, False otherwise.

is_schedulable: bool property

Whether the orchestrator is schedulable or not.

Returns:

Type	Description
bool	Whether the orchestrator is schedulable or not.

is_synchronous: bool property

Whether the orchestrator runs synchronous or not.

Returns:

Type	Description
bool	Whether the orchestrator runs synchronous or not.

supports_client_side_caching: bool property

Whether the orchestrator supports client side caching.

Returns:

Type	Description
bool	Whether the orchestrator supports client side caching.

LightningOrchestratorFlavor

Bases: BaseOrchestratorFlavor

Lightning orchestrator flavor.

Attributes

config_class: Type[LightningOrchestratorConfig] property

Returns KubeflowOrchestratorConfig config class.

Returns:

Type	Description
Type[LightningOrchestratorConfig]	The config class.

docs_url: Optional[str] property

A url to point at docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor docs url.

implementation_class: Type[LightningOrchestrator] property

Implementation class for this flavor.

Returns:

Type	Description
Type[LightningOrchestrator]	The implementation class.

logo_url: str property

A url to represent the flavor in the dashboard.

Returns:

Type	Description
str	The flavor logo.

name: str property

Name of the flavor.

Returns:

Type	Description
str	The name of the flavor.

sdk_docs_url: Optional[str] property

A url to point at SDK docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor SDK docs url.

Modules

lightning_orchestrator_flavor

Lightning orchestrator base config and settings.

Classes

LightningOrchestratorConfig(warn_about_plain_text_secrets: bool = False, **kwargs: Any)

Bases: BaseOrchestratorConfig, LightningOrchestratorSettings

Lightning orchestrator base config.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self, warn_about_plain_text_secrets: bool = False, **kwargs: Any
) -> None:
    """Ensures that secret references don't clash with pydantic validation.

    StackComponents allow the specification of all their string attributes
    using secret references of the form `{{secret_name.key}}`. This however
    is only possible when the stack component does not perform any explicit
    validation of this attribute using pydantic validators. If this were
    the case, the validation would run on the secret reference and would
    fail or in the worst case, modify the secret reference and lead to
    unexpected behavior. This method ensures that no attributes that require
    custom pydantic validation are set as secret references.

    Args:
        warn_about_plain_text_secrets: If true, then warns about using
            plain-text secrets.
        **kwargs: Arguments to initialize this stack component.

    Raises:
        ValueError: If an attribute that requires custom pydantic validation
            is passed as a secret reference, or if the `name` attribute
            was passed as a secret reference.
    """
    for key, value in kwargs.items():
        try:
            field = self.__class__.model_fields[key]
        except KeyError:
            # Value for a private attribute or non-existing field, this
            # will fail during the upcoming pydantic validation
            continue

        if value is None:
            continue

        if not secret_utils.is_secret_reference(value):
            if (
                secret_utils.is_secret_field(field)
                and warn_about_plain_text_secrets
            ):
                logger.warning(
                    "You specified a plain-text value for the sensitive "
                    f"attribute `{key}` for a `{self.__class__.__name__}` "
                    "stack component. This is currently only a warning, "
                    "but future versions of ZenML will require you to pass "
                    "in sensitive information as secrets. Check out the "
                    "documentation on how to configure your stack "
                    "components with secrets here: "
                    "https://docs.zenml.io/getting-started/deploying-zenml/secret-management"
                )
            continue

        if pydantic_utils.has_validators(
            pydantic_class=self.__class__, field_name=key
        ):
            raise ValueError(
                f"Passing the stack component attribute `{key}` as a "
                "secret reference is not allowed as additional validation "
                "is required for this attribute."
            )

    super().__init__(**kwargs)

Attributes

is_local: bool property

Checks if this stack component is running locally.

Returns:

Type	Description
bool	True if this config is for a local component, False otherwise.

is_schedulable: bool property

Whether the orchestrator is schedulable or not.

Returns:

Type	Description
bool	Whether the orchestrator is schedulable or not.

is_synchronous: bool property

Whether the orchestrator runs synchronous or not.

Returns:

Type	Description
bool	Whether the orchestrator runs synchronous or not.

supports_client_side_caching: bool property

Whether the orchestrator supports client side caching.

Returns:

Type	Description
bool	Whether the orchestrator supports client side caching.

LightningOrchestratorFlavor

Bases: BaseOrchestratorFlavor

Lightning orchestrator flavor.

Attributes

config_class: Type[LightningOrchestratorConfig] property

Returns KubeflowOrchestratorConfig config class.

Returns:

Type	Description
Type[LightningOrchestratorConfig]	The config class.

docs_url: Optional[str] property

A url to point at docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor docs url.

implementation_class: Type[LightningOrchestrator] property

Implementation class for this flavor.

Returns:

Type	Description
Type[LightningOrchestrator]	The implementation class.

logo_url: str property

A url to represent the flavor in the dashboard.

Returns:

Type	Description
str	The flavor logo.

name: str property

Name of the flavor.

Returns:

Type	Description
str	The name of the flavor.

sdk_docs_url: Optional[str] property

A url to point at SDK docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor SDK docs url.

LightningOrchestratorSettings(warn_about_plain_text_secrets: bool = False, **kwargs: Any)

Bases: BaseSettings

Lightning orchestrator base settings.

Attributes:

Name	Type	Description
main_studio_name	Optional[str]	Main studio name.
machine_type	Optional[str]	Machine type.
user_id	Optional[str]	User id.
api_key	Optional[str]	api_key.
username	Optional[str]	Username.
teamspace	Optional[str]	Teamspace.
organization	Optional[str]	Organization.
custom_commands	Optional[List[str]]	Custom commands to run.
synchronous	bool	If True, the client running a pipeline using this orchestrator waits until all steps finish running. If False, the client returns immediately and the pipeline is executed asynchronously. Defaults to True. This setting only has an effect when specified on the pipeline and will be ignored if specified on steps.

Source code in src/zenml/config/secret_reference_mixin.py

def __init__(
    self, warn_about_plain_text_secrets: bool = False, **kwargs: Any
) -> None:
    """Ensures that secret references are only passed for valid fields.

    This method ensures that secret references are not passed for fields
    that explicitly prevent them or require pydantic validation.

    Args:
        warn_about_plain_text_secrets: If true, then warns about using plain-text secrets.
        **kwargs: Arguments to initialize this object.

    Raises:
        ValueError: If an attribute that requires custom pydantic validation
            or an attribute which explicitly disallows secret references
            is passed as a secret reference.
    """
    for key, value in kwargs.items():
        try:
            field = self.__class__.model_fields[key]
        except KeyError:
            # Value for a private attribute or non-existing field, this
            # will fail during the upcoming pydantic validation
            continue

        if value is None:
            continue

        if not secret_utils.is_secret_reference(value):
            if (
                secret_utils.is_secret_field(field)
                and warn_about_plain_text_secrets
            ):
                logger.warning(
                    "You specified a plain-text value for the sensitive "
                    f"attribute `{key}`. This is currently only a warning, "
                    "but future versions of ZenML will require you to pass "
                    "in sensitive information as secrets. Check out the "
                    "documentation on how to configure values with secrets "
                    "here: https://docs.zenml.io/getting-started/deploying-zenml/secret-management"
                )
            continue

        if secret_utils.is_clear_text_field(field):
            raise ValueError(
                f"Passing the `{key}` attribute as a secret reference is "
                "not allowed."
            )

        requires_validation = has_validators(
            pydantic_class=self.__class__, field_name=key
        )
        if requires_validation:
            raise ValueError(
                f"Passing the attribute `{key}` as a secret reference is "
                "not allowed as additional validation is required for "
                "this attribute."
            )

    super().__init__(**kwargs)

Functions

orchestrators

Initialization of the Lightning ZenML orchestrator.

Classes

LightningOrchestrator(name: str, id: UUID, config: StackComponentConfig, flavor: str, type: StackComponentType, user: Optional[UUID], created: datetime, updated: datetime, labels: Optional[Dict[str, Any]] = None, connector_requirements: Optional[ServiceConnectorRequirements] = None, connector: Optional[UUID] = None, connector_resource_id: Optional[str] = None, *args: Any, **kwargs: Any)

Bases: WheeledOrchestrator

Base class for Orchestrator responsible for running pipelines remotely in a VM.

This orchestrator does not support running on a schedule.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self,
    name: str,
    id: UUID,
    config: StackComponentConfig,
    flavor: str,
    type: StackComponentType,
    user: Optional[UUID],
    created: datetime,
    updated: datetime,
    labels: Optional[Dict[str, Any]] = None,
    connector_requirements: Optional[ServiceConnectorRequirements] = None,
    connector: Optional[UUID] = None,
    connector_resource_id: Optional[str] = None,
    *args: Any,
    **kwargs: Any,
):
    """Initializes a StackComponent.

    Args:
        name: The name of the component.
        id: The unique ID of the component.
        config: The config of the component.
        flavor: The flavor of the component.
        type: The type of the component.
        user: The ID of the user who created the component.
        created: The creation time of the component.
        updated: The last update time of the component.
        labels: The labels of the component.
        connector_requirements: The requirements for the connector.
        connector: The ID of a connector linked to the component.
        connector_resource_id: The custom resource ID to access through
            the connector.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Raises:
        ValueError: If a secret reference is passed as name.
    """
    if secret_utils.is_secret_reference(name):
        raise ValueError(
            "Passing the `name` attribute of a stack component as a "
            "secret reference is not allowed."
        )

    self.id = id
    self.name = name
    self._config = config
    self.flavor = flavor
    self.type = type
    self.user = user
    self.created = created
    self.updated = updated
    self.labels = labels
    self.connector_requirements = connector_requirements
    self.connector = connector
    self.connector_resource_id = connector_resource_id
    self._connector_instance: Optional[ServiceConnector] = None

Attributes

config: LightningOrchestratorConfig property

Returns the LightningOrchestratorConfig config.

Returns:

Type	Description
LightningOrchestratorConfig	The configuration.

pipeline_directory: str property

Returns path to a directory in which the kubeflow pipeline files are stored.

Returns:

Type	Description
str	Path to the pipeline directory.

root_directory: str property

Path to the root directory for all files concerning this orchestrator.

Returns:

Type	Description
str	Path to the root directory.

settings_class: Type[LightningOrchestratorSettings] property

Settings class for the Lightning orchestrator.

Returns:

Type	Description
Type[LightningOrchestratorSettings]	The settings class.

validator: Optional[StackValidator] property

Validates the stack.

In the remote case, checks that the stack contains a container registry, image builder and only remote components.

Returns:

Type	Description
Optional[StackValidator]	A StackValidator instance.

Functions

get_orchestrator_run_id() -> str

Returns the active orchestrator run id.

Raises:

Type	Description
RuntimeError	If no run id exists. This happens when this method gets called while the orchestrator is not running a pipeline.

Returns:

Type	Description
str	The orchestrator run id.

Raises:

Type	Description
RuntimeError	If the run id cannot be read from the environment.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

def get_orchestrator_run_id(self) -> str:
    """Returns the active orchestrator run id.

    Raises:
        RuntimeError: If no run id exists. This happens when this method
            gets called while the orchestrator is not running a pipeline.

    Returns:
        The orchestrator run id.

    Raises:
        RuntimeError: If the run id cannot be read from the environment.
    """
    try:
        return os.environ[ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID]
    except KeyError:
        raise RuntimeError(
            "Unable to read run id from environment variable "
            f"{ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID}."
        )

prepare_or_run_pipeline(deployment: PipelineDeploymentResponse, stack: Stack, environment: Dict[str, str], placeholder_run: Optional[PipelineRunResponse] = None) -> Any

Creates a wheel and uploads the pipeline to Lightning.

This functions as an intermediary representation of the pipeline which is then deployed to the kubeflow pipelines instance.

How it works:

Before this method is called the prepare_pipeline_deployment() method builds a docker image that contains the code for the pipeline, all steps the context around these files.

Based on this docker image a callable is created which builds task for each step (_construct_lightning_pipeline). To do this the entrypoint of the docker image is configured to run the correct step within the docker image. The dependencies between these task are then also configured onto each task by pointing at the downstream steps.

Parameters:

Name	Type	Description	Default
deployment	PipelineDeploymentResponse	The pipeline deployment to prepare or run.	required
stack	Stack	The stack the pipeline will run on.	required
environment	Dict[str, str]	Environment variables to set in the orchestration environment.	required
placeholder_run	Optional[PipelineRunResponse]	An optional placeholder run for the deployment.	None

Raises:

Type	Description
ValueError	If the schedule is not set or if the cron expression is not set.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

def prepare_or_run_pipeline(
    self,
    deployment: "PipelineDeploymentResponse",
    stack: "Stack",
    environment: Dict[str, str],
    placeholder_run: Optional["PipelineRunResponse"] = None,
) -> Any:
    """Creates a wheel and uploads the pipeline to Lightning.

    This functions as an intermediary representation of the pipeline which
    is then deployed to the kubeflow pipelines instance.

    How it works:
    -------------
    Before this method is called the `prepare_pipeline_deployment()`
    method builds a docker image that contains the code for the
    pipeline, all steps the context around these files.

    Based on this docker image a callable is created which builds
    task for each step (`_construct_lightning_pipeline`).
    To do this the entrypoint of the docker image is configured to
    run the correct step within the docker image. The dependencies
    between these task are then also configured onto each
    task by pointing at the downstream steps.

    Args:
        deployment: The pipeline deployment to prepare or run.
        stack: The stack the pipeline will run on.
        environment: Environment variables to set in the orchestration
            environment.
        placeholder_run: An optional placeholder run for the deployment.

    Raises:
        ValueError: If the schedule is not set or if the cron expression
            is not set.
    """
    settings = cast(
        LightningOrchestratorSettings, self.get_settings(deployment)
    )
    if deployment.schedule:
        if (
            deployment.schedule.catchup
            or deployment.schedule.interval_second
        ):
            logger.warning(
                "Lightning orchestrator only uses schedules with the "
                "`cron_expression` property, with optional `start_time` and/or `end_time`. "
                "All other properties are ignored."
            )
        if deployment.schedule.cron_expression is None:
            raise ValueError(
                "Property `cron_expression` must be set when passing "
                "schedule to a Lightning orchestrator."
            )
        if deployment.schedule.cron_expression:
            raise ValueError(
                "Property `schedule_timezone` must be set when passing "
                "`cron_expression` to a Lightning orchestrator."
                "Lightning orchestrator requires a Java Timezone ID to run the pipeline on schedule."
                "Please refer to https://docs.oracle.com/middleware/1221/wcs/tag-ref/MISC/TimeZones.html for more information."
            )

    # Get deployment id
    deployment_id = deployment.id

    pipeline_name = deployment.pipeline_configuration.name
    orchestrator_run_name = get_orchestrator_run_name(pipeline_name)

    # Copy the repository to a temporary directory and add a setup.py file
    # repository_temp_dir = (
    #    self.copy_repository_to_temp_dir_and_add_setup_py()
    # )

    # Create a wheel for the package in the temporary directory
    # wheel_path = self.create_wheel(temp_dir=repository_temp_dir)
    code_archive = code_utils.CodeArchive(
        root=source_utils.get_source_root()
    )
    logger.info("Archiving pipeline code...")
    with tempfile.NamedTemporaryFile(
        mode="w+b", delete=False, suffix=".tar.gz"
    ) as code_file:
        code_archive.write_archive(code_file)
        code_path = code_file.name
    filename = f"{orchestrator_run_name}.tar.gz"
    # Construct the env variables for the pipeline
    env_vars = environment.copy()
    orchestrator_run_id = str(uuid4())
    env_vars[ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID] = orchestrator_run_id
    # Set up some variables for configuration
    env_vars[ENV_ZENML_CUSTOM_SOURCE_ROOT] = (
        LIGHTNING_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH
    )
    env_vars[ENV_ZENML_WHEEL_PACKAGE_NAME] = self.package_name

    # Create a line-by-line export of environment variables
    env_exports = "\n".join(
        [f"export {key}='{value}'" for key, value in env_vars.items()]
    )

    # Write the environment variables to a temporary file
    with tempfile.NamedTemporaryFile(
        mode="w", delete=False, suffix=".studiorc"
    ) as temp_file:
        temp_file.write(env_exports)
        env_file_path = temp_file.name

    # 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
    )

    def _construct_lightning_steps(
        deployment: "PipelineDeploymentResponse",
    ) -> Dict[str, Dict[str, Any]]:
        """Construct the steps for the pipeline.

        Args:
            deployment: The pipeline deployment to prepare or run.

        Returns:
            The steps for the pipeline.
        """
        steps = {}
        for step_name, step in deployment.step_configurations.items():
            # The arguments are passed to configure the entrypoint of the
            # docker container when the step is called.
            entrypoint_command = (
                StepEntrypointConfiguration.get_entrypoint_command()
            )
            entrypoint_arguments = (
                StepEntrypointConfiguration.get_entrypoint_arguments(
                    step_name=step_name,
                    deployment_id=deployment_id,
                )
            )
            entrypoint = entrypoint_command + entrypoint_arguments
            entrypoint_string = " ".join(entrypoint)

            step_settings = cast(
                LightningOrchestratorSettings, self.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
            )

            # Construct the command to run the step
            run_command = f"{entrypoint_string}"
            commands = [run_command]
            steps[step_name] = {
                "commands": commands,
                "requirements": step_requirements_to_string,
                "machine": step_settings.machine_type
                if step_settings != settings
                else None,
            }
        return steps

    if not settings.synchronous:
        entrypoint_command = LightningOrchestratorEntrypointConfiguration.get_entrypoint_command()
        entrypoint_arguments = LightningOrchestratorEntrypointConfiguration.get_entrypoint_arguments(
            run_name=orchestrator_run_name,
            deployment_id=deployment.id,
        )
        entrypoint = entrypoint_command + entrypoint_arguments
        entrypoint_string = " ".join(entrypoint)
        logger.info("Setting up Lightning AI client")
        self._set_lightning_env_vars(deployment)

        studio_name = sanitize_studio_name(
            "zenml_async_orchestrator_studio"
        )
        logger.info(f"Creating main studio: {studio_name}")
        studio = Studio(name=studio_name)
        studio.start()

        logger.info(
            "Uploading wheel package and installing dependencies on main studio"
        )
        studio.run(
            f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
        )
        studio.upload_file(
            code_path,
            remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
        )
        time.sleep(10)
        studio.run(
            f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
        )
        studio.upload_file(env_file_path)
        time.sleep(6)
        studio.run(
            f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
        )
        studio.run(f"rm {env_file_path.split('/')[-1]}")

        studio.run("pip install uv")
        logger.info(
            f"Installing requirements: {pipeline_requirements_to_string}"
        )
        studio.run(f"uv pip install {pipeline_requirements_to_string}")
        studio.run("pip install zenml")

        for custom_command in settings.custom_commands or []:
            studio.run(
                f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {custom_command}"
            )
        # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
        logger.info("Running pipeline in async mode")
        studio.run(
            f"nohup bash -c 'cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {entrypoint_string}' > log_{filename.rsplit('.', 2)[0]}.txt 2>&1 &"
        )
        logger.info(
            f"The pipeline is running in async mode, you can keep checking the logs by running the following command: `lightning download -s vision-model/zenml-async-orchestrator-studio -p /teamspace/studios/this_studio/log_{filename.rsplit('.', 2)[0]}.txt && cat log_{filename.rsplit('.', 2)[0]}.txt`"
        )
    else:
        self._upload_and_run_pipeline(
            deployment,
            orchestrator_run_id,
            pipeline_requirements_to_string,
            settings,
            _construct_lightning_steps(deployment),
            code_path,
            filename,
            env_file_path,
        )
    os.unlink(env_file_path)

setup_credentials() -> None

Set up credentials for the orchestrator.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

def setup_credentials(self) -> None:
    """Set up credentials for the orchestrator."""
    connector = self.get_connector()
    assert connector is not None
    connector.configure_local_client()

LightningOrchestratorEntrypointConfiguration

Entrypoint configuration for the Lightning master/orchestrator VM.

Functions

get_entrypoint_arguments(run_name: str, deployment_id: UUID) -> List[str] classmethod

Gets all arguments that the entrypoint command should be called with.

Parameters:

Name	Type	Description	Default
run_name	str	Name of the ZenML run.	required
deployment_id	UUID	ID of the deployment.	required

Returns:

Type	Description
List[str]	List of entrypoint arguments.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint_configuration.py

@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

get_entrypoint_command() -> List[str] classmethod

Returns a command that runs the entrypoint module.

Returns:

Type	Description
List[str]	Entrypoint command.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint_configuration.py

@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

get_entrypoint_options() -> Set[str] classmethod

Gets all the options required for running this entrypoint.

Returns:

Type	Description
Set[str]	Entrypoint options.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint_configuration.py

@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

Modules

lightning_orchestrator

Implementation of the Lightning orchestrator.

Classes

LightningOrchestrator(name: str, id: UUID, config: StackComponentConfig, flavor: str, type: StackComponentType, user: Optional[UUID], created: datetime, updated: datetime, labels: Optional[Dict[str, Any]] = None, connector_requirements: Optional[ServiceConnectorRequirements] = None, connector: Optional[UUID] = None, connector_resource_id: Optional[str] = None, *args: Any, **kwargs: Any)

Bases: WheeledOrchestrator

Base class for Orchestrator responsible for running pipelines remotely in a VM.

This orchestrator does not support running on a schedule.

Source code in src/zenml/stack/stack_component.py

def __init__(
    self,
    name: str,
    id: UUID,
    config: StackComponentConfig,
    flavor: str,
    type: StackComponentType,
    user: Optional[UUID],
    created: datetime,
    updated: datetime,
    labels: Optional[Dict[str, Any]] = None,
    connector_requirements: Optional[ServiceConnectorRequirements] = None,
    connector: Optional[UUID] = None,
    connector_resource_id: Optional[str] = None,
    *args: Any,
    **kwargs: Any,
):
    """Initializes a StackComponent.

    Args:
        name: The name of the component.
        id: The unique ID of the component.
        config: The config of the component.
        flavor: The flavor of the component.
        type: The type of the component.
        user: The ID of the user who created the component.
        created: The creation time of the component.
        updated: The last update time of the component.
        labels: The labels of the component.
        connector_requirements: The requirements for the connector.
        connector: The ID of a connector linked to the component.
        connector_resource_id: The custom resource ID to access through
            the connector.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.

    Raises:
        ValueError: If a secret reference is passed as name.
    """
    if secret_utils.is_secret_reference(name):
        raise ValueError(
            "Passing the `name` attribute of a stack component as a "
            "secret reference is not allowed."
        )

    self.id = id
    self.name = name
    self._config = config
    self.flavor = flavor
    self.type = type
    self.user = user
    self.created = created
    self.updated = updated
    self.labels = labels
    self.connector_requirements = connector_requirements
    self.connector = connector
    self.connector_resource_id = connector_resource_id
    self._connector_instance: Optional[ServiceConnector] = None

Attributes

config: LightningOrchestratorConfig property

Returns the LightningOrchestratorConfig config.

Returns:

Type	Description
LightningOrchestratorConfig	The configuration.

pipeline_directory: str property

Returns path to a directory in which the kubeflow pipeline files are stored.

Returns:

Type	Description
str	Path to the pipeline directory.

root_directory: str property

Path to the root directory for all files concerning this orchestrator.

Returns:

Type	Description
str	Path to the root directory.

settings_class: Type[LightningOrchestratorSettings] property

Settings class for the Lightning orchestrator.

Returns:

Type	Description
Type[LightningOrchestratorSettings]	The settings class.

validator: Optional[StackValidator] property

Validates the stack.

In the remote case, checks that the stack contains a container registry, image builder and only remote components.

Returns:

Type	Description
Optional[StackValidator]	A StackValidator instance.

Functions

get_orchestrator_run_id() -> str

Returns the active orchestrator run id.

Raises:

Type	Description
RuntimeError	If no run id exists. This happens when this method gets called while the orchestrator is not running a pipeline.

Returns:

Type	Description
str	The orchestrator run id.

Raises:

Type	Description
RuntimeError	If the run id cannot be read from the environment.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

def get_orchestrator_run_id(self) -> str:
    """Returns the active orchestrator run id.

    Raises:
        RuntimeError: If no run id exists. This happens when this method
            gets called while the orchestrator is not running a pipeline.

    Returns:
        The orchestrator run id.

    Raises:
        RuntimeError: If the run id cannot be read from the environment.
    """
    try:
        return os.environ[ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID]
    except KeyError:
        raise RuntimeError(
            "Unable to read run id from environment variable "
            f"{ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID}."
        )

prepare_or_run_pipeline(deployment: PipelineDeploymentResponse, stack: Stack, environment: Dict[str, str], placeholder_run: Optional[PipelineRunResponse] = None) -> Any

Creates a wheel and uploads the pipeline to Lightning.

This functions as an intermediary representation of the pipeline which is then deployed to the kubeflow pipelines instance.

How it works:

Before this method is called the prepare_pipeline_deployment() method builds a docker image that contains the code for the pipeline, all steps the context around these files.

Based on this docker image a callable is created which builds task for each step (_construct_lightning_pipeline). To do this the entrypoint of the docker image is configured to run the correct step within the docker image. The dependencies between these task are then also configured onto each task by pointing at the downstream steps.

Parameters:

Name	Type	Description	Default
deployment	PipelineDeploymentResponse	The pipeline deployment to prepare or run.	required
stack	Stack	The stack the pipeline will run on.	required
environment	Dict[str, str]	Environment variables to set in the orchestration environment.	required
placeholder_run	Optional[PipelineRunResponse]	An optional placeholder run for the deployment.	None

Raises:

Type	Description
ValueError	If the schedule is not set or if the cron expression is not set.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

def prepare_or_run_pipeline(
    self,
    deployment: "PipelineDeploymentResponse",
    stack: "Stack",
    environment: Dict[str, str],
    placeholder_run: Optional["PipelineRunResponse"] = None,
) -> Any:
    """Creates a wheel and uploads the pipeline to Lightning.

    This functions as an intermediary representation of the pipeline which
    is then deployed to the kubeflow pipelines instance.

    How it works:
    -------------
    Before this method is called the `prepare_pipeline_deployment()`
    method builds a docker image that contains the code for the
    pipeline, all steps the context around these files.

    Based on this docker image a callable is created which builds
    task for each step (`_construct_lightning_pipeline`).
    To do this the entrypoint of the docker image is configured to
    run the correct step within the docker image. The dependencies
    between these task are then also configured onto each
    task by pointing at the downstream steps.

    Args:
        deployment: The pipeline deployment to prepare or run.
        stack: The stack the pipeline will run on.
        environment: Environment variables to set in the orchestration
            environment.
        placeholder_run: An optional placeholder run for the deployment.

    Raises:
        ValueError: If the schedule is not set or if the cron expression
            is not set.
    """
    settings = cast(
        LightningOrchestratorSettings, self.get_settings(deployment)
    )
    if deployment.schedule:
        if (
            deployment.schedule.catchup
            or deployment.schedule.interval_second
        ):
            logger.warning(
                "Lightning orchestrator only uses schedules with the "
                "`cron_expression` property, with optional `start_time` and/or `end_time`. "
                "All other properties are ignored."
            )
        if deployment.schedule.cron_expression is None:
            raise ValueError(
                "Property `cron_expression` must be set when passing "
                "schedule to a Lightning orchestrator."
            )
        if deployment.schedule.cron_expression:
            raise ValueError(
                "Property `schedule_timezone` must be set when passing "
                "`cron_expression` to a Lightning orchestrator."
                "Lightning orchestrator requires a Java Timezone ID to run the pipeline on schedule."
                "Please refer to https://docs.oracle.com/middleware/1221/wcs/tag-ref/MISC/TimeZones.html for more information."
            )

    # Get deployment id
    deployment_id = deployment.id

    pipeline_name = deployment.pipeline_configuration.name
    orchestrator_run_name = get_orchestrator_run_name(pipeline_name)

    # Copy the repository to a temporary directory and add a setup.py file
    # repository_temp_dir = (
    #    self.copy_repository_to_temp_dir_and_add_setup_py()
    # )

    # Create a wheel for the package in the temporary directory
    # wheel_path = self.create_wheel(temp_dir=repository_temp_dir)
    code_archive = code_utils.CodeArchive(
        root=source_utils.get_source_root()
    )
    logger.info("Archiving pipeline code...")
    with tempfile.NamedTemporaryFile(
        mode="w+b", delete=False, suffix=".tar.gz"
    ) as code_file:
        code_archive.write_archive(code_file)
        code_path = code_file.name
    filename = f"{orchestrator_run_name}.tar.gz"
    # Construct the env variables for the pipeline
    env_vars = environment.copy()
    orchestrator_run_id = str(uuid4())
    env_vars[ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID] = orchestrator_run_id
    # Set up some variables for configuration
    env_vars[ENV_ZENML_CUSTOM_SOURCE_ROOT] = (
        LIGHTNING_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH
    )
    env_vars[ENV_ZENML_WHEEL_PACKAGE_NAME] = self.package_name

    # Create a line-by-line export of environment variables
    env_exports = "\n".join(
        [f"export {key}='{value}'" for key, value in env_vars.items()]
    )

    # Write the environment variables to a temporary file
    with tempfile.NamedTemporaryFile(
        mode="w", delete=False, suffix=".studiorc"
    ) as temp_file:
        temp_file.write(env_exports)
        env_file_path = temp_file.name

    # 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
    )

    def _construct_lightning_steps(
        deployment: "PipelineDeploymentResponse",
    ) -> Dict[str, Dict[str, Any]]:
        """Construct the steps for the pipeline.

        Args:
            deployment: The pipeline deployment to prepare or run.

        Returns:
            The steps for the pipeline.
        """
        steps = {}
        for step_name, step in deployment.step_configurations.items():
            # The arguments are passed to configure the entrypoint of the
            # docker container when the step is called.
            entrypoint_command = (
                StepEntrypointConfiguration.get_entrypoint_command()
            )
            entrypoint_arguments = (
                StepEntrypointConfiguration.get_entrypoint_arguments(
                    step_name=step_name,
                    deployment_id=deployment_id,
                )
            )
            entrypoint = entrypoint_command + entrypoint_arguments
            entrypoint_string = " ".join(entrypoint)

            step_settings = cast(
                LightningOrchestratorSettings, self.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
            )

            # Construct the command to run the step
            run_command = f"{entrypoint_string}"
            commands = [run_command]
            steps[step_name] = {
                "commands": commands,
                "requirements": step_requirements_to_string,
                "machine": step_settings.machine_type
                if step_settings != settings
                else None,
            }
        return steps

    if not settings.synchronous:
        entrypoint_command = LightningOrchestratorEntrypointConfiguration.get_entrypoint_command()
        entrypoint_arguments = LightningOrchestratorEntrypointConfiguration.get_entrypoint_arguments(
            run_name=orchestrator_run_name,
            deployment_id=deployment.id,
        )
        entrypoint = entrypoint_command + entrypoint_arguments
        entrypoint_string = " ".join(entrypoint)
        logger.info("Setting up Lightning AI client")
        self._set_lightning_env_vars(deployment)

        studio_name = sanitize_studio_name(
            "zenml_async_orchestrator_studio"
        )
        logger.info(f"Creating main studio: {studio_name}")
        studio = Studio(name=studio_name)
        studio.start()

        logger.info(
            "Uploading wheel package and installing dependencies on main studio"
        )
        studio.run(
            f"mkdir -p /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
        )
        studio.upload_file(
            code_path,
            remote_path=f"/teamspace/studios/this_studio/zenml_codes/{filename}",
        )
        time.sleep(10)
        studio.run(
            f"tar -xvzf /teamspace/studios/this_studio/zenml_codes/{filename} -C /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]}"
        )
        studio.upload_file(env_file_path)
        time.sleep(6)
        studio.run(
            f"cp {env_file_path.split('/')[-1]} ./.lightning_studio/.studiorc"
        )
        studio.run(f"rm {env_file_path.split('/')[-1]}")

        studio.run("pip install uv")
        logger.info(
            f"Installing requirements: {pipeline_requirements_to_string}"
        )
        studio.run(f"uv pip install {pipeline_requirements_to_string}")
        studio.run("pip install zenml")

        for custom_command in settings.custom_commands or []:
            studio.run(
                f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {custom_command}"
            )
        # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
        logger.info("Running pipeline in async mode")
        studio.run(
            f"nohup bash -c 'cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {entrypoint_string}' > log_{filename.rsplit('.', 2)[0]}.txt 2>&1 &"
        )
        logger.info(
            f"The pipeline is running in async mode, you can keep checking the logs by running the following command: `lightning download -s vision-model/zenml-async-orchestrator-studio -p /teamspace/studios/this_studio/log_{filename.rsplit('.', 2)[0]}.txt && cat log_{filename.rsplit('.', 2)[0]}.txt`"
        )
    else:
        self._upload_and_run_pipeline(
            deployment,
            orchestrator_run_id,
            pipeline_requirements_to_string,
            settings,
            _construct_lightning_steps(deployment),
            code_path,
            filename,
            env_file_path,
        )
    os.unlink(env_file_path)

setup_credentials() -> None

Set up credentials for the orchestrator.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

def setup_credentials(self) -> None:
    """Set up credentials for the orchestrator."""
    connector = self.get_connector()
    assert connector is not None
    connector.configure_local_client()

Functions Modules

lightning_orchestrator_entrypoint

Entrypoint of the Lightning master/orchestrator STUDIO.

Classes Functions

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:

Type	Description
TypeError	If the active stack's orchestrator is not an instance of LightningOrchestrator.
ValueError	If the active stack's container registry is None.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint.py

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 install zenml")
    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 install zenml")
                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.")

parse_args() -> argparse.Namespace

Parse entrypoint arguments.

Returns:

Type	Description
Namespace	Parsed args.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint.py

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()

lightning_orchestrator_entrypoint_configuration

Entrypoint configuration for the Lightning master/orchestrator VM.

Classes

LightningOrchestratorEntrypointConfiguration

Entrypoint configuration for the Lightning master/orchestrator VM.

Functions

get_entrypoint_arguments(run_name: str, deployment_id: UUID) -> List[str] classmethod

Gets all arguments that the entrypoint command should be called with.

Parameters:

Name	Type	Description	Default
run_name	str	Name of the ZenML run.	required
deployment_id	UUID	ID of the deployment.	required

Returns:

Type	Description
List[str]	List of entrypoint arguments.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint_configuration.py

@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

get_entrypoint_command() -> List[str] classmethod

Returns a command that runs the entrypoint module.

Returns:

Type	Description
List[str]	Entrypoint command.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint_configuration.py

@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

get_entrypoint_options() -> Set[str] classmethod

Gets all the options required for running this entrypoint.

Returns:

Type	Description
Set[str]	Entrypoint options.

Source code in src/zenml/integrations/lightning/orchestrators/lightning_orchestrator_entrypoint_configuration.py

@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

utils

Utility functions for the Lightning orchestrator.

Classes Functions

gather_requirements(docker_settings: DockerSettings) -> List[str]

Gather the requirements files.

Parameters:

Name	Type	Description	Default
docker_settings	DockerSettings	Docker settings.	required

Returns:

Type	Description
List[str]	List of requirements.

Source code in src/zenml/integrations/lightning/orchestrators/utils.py

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

sanitize_studio_name(studio_name: str) -> str

Sanitize studio_names so they conform to Kubernetes studio naming convention.

Parameters:

Name	Type	Description	Default
studio_name	str	Arbitrary input studio_name.	required

Returns:

Type	Description
str	Sanitized pod name.

Source code in src/zenml/integrations/lightning/orchestrators/utils.py

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)