Skypilot Lambda

zenml.integrations.skypilot_lambda special

Initialization of the Skypilot Lambda integration for ZenML.

The Skypilot integration sub-module powers an alternative to the local orchestrator for a remote orchestration of ZenML pipelines on VMs.

SkypilotLambdaIntegration (Integration)

Definition of Skypilot Lambda Integration for ZenML.

Source code in zenml/integrations/skypilot_lambda/__init__.py

class SkypilotLambdaIntegration(Integration):
    """Definition of Skypilot Lambda Integration for ZenML."""

    NAME = SKYPILOT_LAMBDA
    REQUIREMENTS = ["skypilot[lambda]~=0.6.0"]

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

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

        return [SkypilotLambdaOrchestratorFlavor]

flavors() classmethod

Declare the stack component flavors for the Skypilot Lambda integration.

Returns:

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

Source code in zenml/integrations/skypilot_lambda/__init__.py

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

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

    return [SkypilotLambdaOrchestratorFlavor]

flavors special

Skypilot integration flavor for Skypilot Lambda orchestrator.

skypilot_orchestrator_lambda_vm_flavor

Skypilot orchestrator Lambda flavor.

SkypilotLambdaOrchestratorConfig (SkypilotBaseOrchestratorConfig, SkypilotLambdaOrchestratorSettings)

Skypilot orchestrator config.

Source code in zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

class SkypilotLambdaOrchestratorConfig(
    SkypilotBaseOrchestratorConfig, SkypilotLambdaOrchestratorSettings
):
    """Skypilot orchestrator config."""

    api_key: Optional[str] = SecretField(default=None)

model_post_init(/, self, context)

We need to both initialize private attributes and call the user-defined model_post_init method.

Source code in zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

def wrapped_model_post_init(self: BaseModel, context: Any, /) -> None:
    """We need to both initialize private attributes and call the user-defined model_post_init
    method.
    """
    init_private_attributes(self, context)
    original_model_post_init(self, context)

SkypilotLambdaOrchestratorFlavor (BaseOrchestratorFlavor)

Flavor for the Skypilot Lambda orchestrator.

Source code in zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

class SkypilotLambdaOrchestratorFlavor(BaseOrchestratorFlavor):
    """Flavor for the Skypilot Lambda orchestrator."""

    @property
    def name(self) -> str:
        """Name of the orchestrator flavor.

        Returns:
            Name of the orchestrator flavor.
        """
        return SKYPILOT_LAMBDA_ORCHESTRATOR_FLAVOR

    @property
    def docs_url(self) -> Optional[str]:
        """A url to point at docs explaining this flavor.

        Returns:
            A flavor docs url.
        """
        return self.generate_default_docs_url()

    @property
    def sdk_docs_url(self) -> Optional[str]:
        """A url to point at SDK docs explaining this flavor.

        Returns:
            A flavor SDK docs url.
        """
        return self.generate_default_sdk_docs_url()

    @property
    def logo_url(self) -> str:
        """A url to represent the flavor in the dashboard.

        Returns:
            The flavor logo.
        """
        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/orchestrator/lambda.png"

    @property
    def config_class(self) -> Type[BaseOrchestratorConfig]:
        """Config class for the base orchestrator flavor.

        Returns:
            The config class.
        """
        return SkypilotLambdaOrchestratorConfig

    @property
    def implementation_class(self) -> Type["SkypilotLambdaOrchestrator"]:
        """Implementation class for this flavor.

        Returns:
            Implementation class for this flavor.
        """
        from zenml.integrations.skypilot_lambda.orchestrators import (
            SkypilotLambdaOrchestrator,
        )

        return SkypilotLambdaOrchestrator

config_class: Type[zenml.orchestrators.base_orchestrator.BaseOrchestratorConfig] property readonly

Config class for the base orchestrator flavor.

Returns:

Type	Description
Type[zenml.orchestrators.base_orchestrator.BaseOrchestratorConfig]	The config class.

docs_url: Optional[str] property readonly

A url to point at docs explaining this flavor.

Returns:

Type	Description
Optional[str]	A flavor docs url.

implementation_class: Type[SkypilotLambdaOrchestrator] property readonly

Implementation class for this flavor.

Returns:

Type	Description
Type[SkypilotLambdaOrchestrator]	Implementation class for this flavor.

logo_url: str property readonly

A url to represent the flavor in the dashboard.

Returns:

Type	Description
str	The flavor logo.

name: str property readonly

Name of the orchestrator flavor.

Returns:

Type	Description
str	Name of the orchestrator flavor.

sdk_docs_url: Optional[str] property readonly

A url to point at SDK docs explaining this flavor.

Returns:

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

SkypilotLambdaOrchestratorSettings (SkypilotBaseOrchestratorSettings)

Skypilot orchestrator settings.

Source code in zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

class SkypilotLambdaOrchestratorSettings(SkypilotBaseOrchestratorSettings):
    """Skypilot orchestrator settings."""

    _UNSUPPORTED_FEATURES = {
        "use_spot": "Spot instances not supported for Lambda orchestrator.",
        "job_recovery": "Job recovery not supported for Lambda orchestrator.",
        "image_id": "Custom image IDs not supported for Lambda orchestrator.",
        # Add other unsupported features as needed
    }

    def __setattr__(self, name: str, value: Any) -> None:
        """Set attribute.

        Args:
            name: Name of the attribute.
            value: Value of the attribute.

        Raises:
            AttributeError: If the attribute is not supported.
        """
        if name in self._UNSUPPORTED_FEATURES:
            raise AttributeError(f"{name} is not supported on Lambda.")
        super().__setattr__(name, value)

__setattr__(self, name, value) special

Set attribute.

Parameters:

Name	Type	Description	Default
name	str	Name of the attribute.	required
value	Any	Value of the attribute.	required

Exceptions:

Type	Description
AttributeError	If the attribute is not supported.

Source code in zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

def __setattr__(self, name: str, value: Any) -> None:
    """Set attribute.

    Args:
        name: Name of the attribute.
        value: Value of the attribute.

    Raises:
        AttributeError: If the attribute is not supported.
    """
    if name in self._UNSUPPORTED_FEATURES:
        raise AttributeError(f"{name} is not supported on Lambda.")
    super().__setattr__(name, value)

model_post_init(/, self, context)

This function is meant to behave like a BaseModel method to initialise private attributes.

It takes context as an argument since that's what pydantic-core passes when calling it.

Parameters:

Name	Type	Description	Default
self	BaseModel	The BaseModel instance.	required
context	Any	The context.	required

Source code in zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

def init_private_attributes(self: BaseModel, context: Any, /) -> None:
    """This function is meant to behave like a BaseModel method to initialise private attributes.

    It takes context as an argument since that's what pydantic-core passes when calling it.

    Args:
        self: The BaseModel instance.
        context: The context.
    """
    if getattr(self, '__pydantic_private__', None) is None:
        pydantic_private = {}
        for name, private_attr in self.__private_attributes__.items():
            default = private_attr.get_default()
            if default is not PydanticUndefined:
                pydantic_private[name] = default
        object_setattr(self, '__pydantic_private__', pydantic_private)

orchestrators special

Initialization of the Skypilot Lambda ZenML orchestrator.

skypilot_lambda_vm_orchestrator

Implementation of the a Skypilot based Lambda VM orchestrator.

SkypilotLambdaOrchestrator (SkypilotBaseOrchestrator)

Orchestrator responsible for running pipelines remotely in a VM on Lambda.

This orchestrator does not support running on a schedule.

Source code in zenml/integrations/skypilot_lambda/orchestrators/skypilot_lambda_vm_orchestrator.py

class SkypilotLambdaOrchestrator(SkypilotBaseOrchestrator):
    """Orchestrator responsible for running pipelines remotely in a VM on Lambda.

    This orchestrator does not support running on a schedule.
    """

    DEFAULT_INSTANCE_TYPE: str = "gpu_1x_a10"

    @property
    def cloud(self) -> sky.clouds.Cloud:
        """The type of sky cloud to use.

        Returns:
            A `sky.clouds.Cloud` instance.
        """
        return sky.clouds.Lambda()

    @property
    def config(self) -> SkypilotLambdaOrchestratorConfig:
        """Returns the `SkypilotLambdaOrchestratorConfig` config.

        Returns:
            The configuration.
        """
        return cast(SkypilotLambdaOrchestratorConfig, self._config)

    @property
    def settings_class(self) -> Optional[Type["BaseSettings"]]:
        """Settings class for the Skypilot orchestrator.

        Returns:
            The settings class.
        """
        return SkypilotLambdaOrchestratorSettings

    def prepare_environment_variable(self, set: bool = True) -> None:
        """Set up Environment variables that are required for the orchestrator.

        Args:
            set: Whether to set the environment variables or not.
        """
        pass

    def setup_credentials(self) -> None:
        """Set up credentials for the orchestrator."""
        # Define the directory and file paths
        directory = os.path.expanduser("~/.lambda_cloud")
        file_path = os.path.join(directory, "lambda_keys")

        # Check if the directory exists, and create it if it doesn't
        if not os.path.exists(directory):
            os.makedirs(directory)

        # Write the API key to the file, creating or overwriting it
        with fileio.open(file_path, "w") as file:
            file.write(f"api_key = {self.config.api_key}")

cloud: sky.clouds.Cloud property readonly

The type of sky cloud to use.

Returns:

Type	Description
sky.clouds.Cloud	A sky.clouds.Cloud instance.

config: SkypilotLambdaOrchestratorConfig property readonly

Returns the SkypilotLambdaOrchestratorConfig config.

Returns:

Type	Description
SkypilotLambdaOrchestratorConfig	The configuration.

settings_class: Optional[Type[BaseSettings]] property readonly

Settings class for the Skypilot orchestrator.

Returns:

Type	Description
Optional[Type[BaseSettings]]	The settings class.

prepare_environment_variable(self, set=True)

Set up Environment variables that are required for the orchestrator.

Parameters:

Name	Type	Description	Default
set	bool	Whether to set the environment variables or not.	True

Source code in zenml/integrations/skypilot_lambda/orchestrators/skypilot_lambda_vm_orchestrator.py

def prepare_environment_variable(self, set: bool = True) -> None:
    """Set up Environment variables that are required for the orchestrator.

    Args:
        set: Whether to set the environment variables or not.
    """
    pass

setup_credentials(self)

Set up credentials for the orchestrator.

Source code in zenml/integrations/skypilot_lambda/orchestrators/skypilot_lambda_vm_orchestrator.py

def setup_credentials(self) -> None:
    """Set up credentials for the orchestrator."""
    # Define the directory and file paths
    directory = os.path.expanduser("~/.lambda_cloud")
    file_path = os.path.join(directory, "lambda_keys")

    # Check if the directory exists, and create it if it doesn't
    if not os.path.exists(directory):
        os.makedirs(directory)

    # Write the API key to the file, creating or overwriting it
    with fileio.open(file_path, "w") as file:
        file.write(f"api_key = {self.config.api_key}")