Source code for oumi.launcher.clouds.sky_cloud

# Copyright 2025 - Oumi
#
# 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
#
#     http://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.

from typing import TypeVar

from oumi.core.configs import JobConfig
from oumi.core.launcher import BaseCloud, BaseCluster, JobStatus
from oumi.core.registry import register_cloud_builder
from oumi.launcher.clients.sky_client import SkyClient, SkyRequestStatus
from oumi.launcher.clusters.sky_cluster import SkyCluster

T = TypeVar("T")




class SkyCloud(BaseCloud):
    """A resource pool capable of creating clusters using Sky Pilot."""

    @property
    def _client(self) -> SkyClient:
        """Returns the SkyClient instance."""
        # Instantiating a SkyClient imports sky.
        # Delay sky import: https://github.com/oumi-ai/oumi/issues/1605
        if not self._sky_client:
            self._sky_client = SkyClient()
        return self._sky_client

    def __init__(self, cloud_name: str):
        """Initializes a new instance of the SkyCloud class."""
        self._cloud_name = cloud_name
        self._sky_client: SkyClient | None = None

    def _get_clusters_by_class(self, cloud_class: type[T]) -> list[BaseCluster]:
        """Gets the appropriate clusters of type T."""
        # Delay sky import: https://github.com/oumi-ai/oumi/issues/1605
        import sky

        return [
            SkyCluster(cluster["name"], self._client)
            for cluster in self._client.status()
            if (
                isinstance(cluster["handle"].launched_resources.cloud, cloud_class)
                and cluster["status"] in (sky.ClusterStatus.UP, sky.ClusterStatus.INIT)
            )
        ]



    def up_cluster(self, job: JobConfig, name: str | None, **kwargs) -> JobStatus:
        """Creates a cluster and starts the provided Job."""
        return self._started_job(self._client.launch(job, name, **kwargs))




    def submit_job(self, job: JobConfig, name: str | None, **kwargs) -> str:
        """Submit a cluster launch without waiting for the job to start.

        The returned request ID can be passed to :meth:`wait_for_job`,
        :meth:`cancel_request`, or :meth:`request_status`.

        Args:
            job: Job to run on the cluster.
            name: Name of the cluster to create.
            kwargs: Launch options passed to the Sky Pilot client.

        Returns:
            ID of the launch request on the Sky Pilot server.
        """
        return self._client.submit(job, name, **kwargs)




    def wait_for_job(self, request_id: str) -> JobStatus:
        """Wait for a submitted launch and read its job from the cluster."""
        return self._started_job(self._client.wait(request_id))




    def cancel_request(self, request_id: str) -> bool:
        """Cancel a launch request submitted with :meth:`submit_job`.

        Stopping the request does not remove a cluster it had started to
        build; tear that down by name. Returns whether anything was cancelled.
        """
        return self._client.cancel_request(request_id)




    def request_status(self, request_id: str) -> SkyRequestStatus:
        """Get the status of a request submitted with :meth:`submit_job`."""
        return self._client.request_status(request_id)


    def _started_job(self, launch_status: JobStatus) -> JobStatus:
        """Read the started job's current status from its cluster."""
        cluster = self.get_cluster(launch_status.cluster)
        if not cluster:
            raise RuntimeError(f"Cluster {launch_status.cluster} not found.")
        final_status = cluster.get_job(launch_status.id)
        if not final_status:
            raise RuntimeError(
                f"Job {launch_status.id} not found on cluster {launch_status.cluster}."
            )
        return final_status



    def get_cluster(self, name) -> BaseCluster | None:
        """Gets the cluster with the specified name, or None if not found."""
        clusters = self.list_clusters()
        for cluster in clusters:
            if cluster.name() == name:
                return cluster
        return None




    def list_clusters(self) -> list[BaseCluster]:
        """Lists the active clusters on this cloud."""
        # Delay sky import: https://github.com/oumi-ai/oumi/issues/1605
        import sky

        if self._cloud_name == SkyClient.SupportedClouds.GCP.value:
            return self._get_clusters_by_class(sky.clouds.GCP)
        elif self._cloud_name == SkyClient.SupportedClouds.RUNPOD.value:
            return self._get_clusters_by_class(sky.clouds.RunPod)
        elif self._cloud_name == SkyClient.SupportedClouds.LAMBDA.value:
            return self._get_clusters_by_class(sky.clouds.Lambda)
        elif self._cloud_name == SkyClient.SupportedClouds.AWS.value:
            return self._get_clusters_by_class(sky.clouds.AWS)
        elif self._cloud_name == SkyClient.SupportedClouds.AZURE.value:
            return self._get_clusters_by_class(sky.clouds.Azure)
        elif self._cloud_name == SkyClient.SupportedClouds.K8S.value:
            return self._get_clusters_by_class(sky.clouds.Kubernetes)
        elif self._cloud_name == SkyClient.SupportedClouds.NEBIUS.value:
            return self._get_clusters_by_class(sky.clouds.Nebius)
        elif self._cloud_name == SkyClient.SupportedClouds.SLURM.value:
            return self._get_clusters_by_class(sky.clouds.Slurm)
        raise ValueError(f"Unsupported cloud: {self._cloud_name}")




@register_cloud_builder("runpod")
def runpod_cloud_builder() -> SkyCloud:
    """Builds a SkyCloud instance for runpod."""
    return SkyCloud(SkyClient.SupportedClouds.RUNPOD.value)


@register_cloud_builder("gcp")
def gcp_cloud_builder() -> SkyCloud:
    """Builds a SkyCloud instance for Google Cloud Platform."""
    return SkyCloud(SkyClient.SupportedClouds.GCP.value)


@register_cloud_builder("lambda")
def lambda_cloud_builder() -> SkyCloud:
    """Builds a SkyCloud instance for Lambda."""
    return SkyCloud(SkyClient.SupportedClouds.LAMBDA.value)


@register_cloud_builder("aws")
def aws_cloud_builder() -> SkyCloud:
    """Builds a SkyCloud instance for AWS."""
    return SkyCloud(SkyClient.SupportedClouds.AWS.value)


@register_cloud_builder("azure")
def azure_cloud_builder() -> SkyCloud:
    """Builds a SkyCloud instance for Azure."""
    return SkyCloud(SkyClient.SupportedClouds.AZURE.value)


@register_cloud_builder("k8s")
def k8s_cloud_builder() -> SkyCloud:
    """Builds a SkyCloud instance for Kubernetes."""
    return SkyCloud(SkyClient.SupportedClouds.K8S.value)


@register_cloud_builder("nebius")
def nebius_cloud_builder() -> SkyCloud:
    """Builds a SkyCloud instance for Nebius."""
    return SkyCloud(SkyClient.SupportedClouds.NEBIUS.value)


@register_cloud_builder("sky-slurm")
def sky_slurm_cloud_builder() -> SkyCloud:
    """Builds a SkyCloud instance for SkyPilot-managed Slurm."""
    return SkyCloud(SkyClient.SupportedClouds.SLURM.value)