skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/dag.py

@@ -1,8 +1,12 @@
 """DAGs: user applications to be run."""
 import pprint
 import threading
+import typing
 from typing import List, Optional
 
+if typing.TYPE_CHECKING:
+    from sky import task
+
 
 class Dag:
     """Dag: a user application, represented as a DAG of Tasks.
@@ -13,37 +17,38 @@ class Dag:
         >>>     task = sky.Task(...)
     """
 
-    def __init__(self):
-        self.tasks = []
+    def __init__(self) -> None:
+        self.tasks: List['task.Task'] = []
         import networkx as nx  # pylint: disable=import-outside-toplevel
 
         self.graph = nx.DiGraph()
-        self.name = None
+        self.name: Optional[str] = None
+        self.policy_applied: bool = False
 
-    def add(self, task):
+    def add(self, task: 'task.Task') -> None:
         self.graph.add_node(task)
         self.tasks.append(task)
 
-    def remove(self, task):
+    def remove(self, task: 'task.Task') -> None:
         self.tasks.remove(task)
         self.graph.remove_node(task)
 
-    def add_edge(self, op1, op2):
+    def add_edge(self, op1: 'task.Task', op2: 'task.Task') -> None:
         assert op1 in self.graph.nodes
         assert op2 in self.graph.nodes
         self.graph.add_edge(op1, op2)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.tasks)
 
-    def __enter__(self):
+    def __enter__(self) -> 'Dag':
         push_dag(self)
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
         pop_dag()
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         pformat = pprint.pformat(self.tasks)
         return f'DAG:\n{pformat}'
 
@@ -51,34 +56,42 @@ class Dag:
         return self.graph
 
     def is_chain(self) -> bool:
-        # NOTE: this method assumes that the graph has no cycle.
-        is_chain = True
-        visited_zero_out_degree = False
-        for node in self.graph.nodes:
-            out_degree = self.graph.out_degree(node)
-            if out_degree > 1:
-                is_chain = False
-                break
-            elif out_degree == 0:
-                if visited_zero_out_degree:
-                    is_chain = False
-                    break
-                else:
-                    visited_zero_out_degree = True
-        return is_chain
+        """Check if the DAG is a linear chain of tasks."""
+
+        nodes = list(self.graph.nodes)
+
+        if len(nodes) == 0:
+            return True
+
+        in_degrees = [self.graph.in_degree(node) for node in nodes]
+        out_degrees = [self.graph.out_degree(node) for node in nodes]
+
+        # Check out-degrees: all <= 1 and exactly one node has out_degree == 0
+        out_degree_condition = (all(degree <= 1 for degree in out_degrees) and
+                                sum(degree == 0 for degree in out_degrees) == 1)
+
+        # Check in-degrees: all <= 1 and exactly one node has in_degree == 0
+        in_degree_condition = (all(degree <= 1 for degree in in_degrees) and
+                               sum(degree == 0 for degree in in_degrees) == 1)
+
+        return out_degree_condition and in_degree_condition
+
+    def validate(self, workdir_only: bool = False):
+        for task in self.tasks:
+            task.validate(workdir_only=workdir_only)
 
 
 class _DagContext(threading.local):
     """A thread-local stack of Dags."""
-    _current_dag = None
+    _current_dag: Optional[Dag] = None
     _previous_dags: List[Dag] = []
 
-    def push_dag(self, dag):
+    def push_dag(self, dag: Dag):
         if self._current_dag is not None:
             self._previous_dags.append(self._current_dag)
         self._current_dag = dag
 
-    def pop_dag(self):
+    def pop_dag(self) -> Optional[Dag]:
         old_dag = self._current_dag
         if self._previous_dags:
             self._current_dag = self._previous_dags.pop()

sky/data/data_transfer.py

@@ -200,3 +200,40 @@ def _add_bucket_iam_member(bucket_name: str, role: str, member: str) -> None:
     bucket.set_iam_policy(policy)
 
     logger.debug(f'Added {member} with role {role} to {bucket_name}.')
+
+
+def s3_to_oci(s3_bucket_name: str, oci_bucket_name: str) -> None:
+    """Creates a one-time transfer from Amazon S3 to OCI Object Storage.
+    Args:
+      s3_bucket_name: str; Name of the Amazon S3 Bucket
+      oci_bucket_name: str; Name of the OCI Bucket
+    """
+    # TODO(HysunHe): Implement sync with other clouds (s3, gs)
+    raise NotImplementedError('Moving data directly from S3 to OCI bucket '
+                              'is currently not supported. Please specify '
+                              'a local source for the storage object.')
+
+
+def gcs_to_oci(gs_bucket_name: str, oci_bucket_name: str) -> None:
+    """Creates a one-time transfer from Google Cloud Storage to
+    OCI Object Storage.
+    Args:
+      gs_bucket_name: str; Name of the Google Cloud Storage Bucket
+      oci_bucket_name: str; Name of the OCI Bucket
+    """
+    # TODO(HysunHe): Implement sync with other clouds (s3, gs)
+    raise NotImplementedError('Moving data directly from GCS to OCI bucket '
+                              'is currently not supported. Please specify '
+                              'a local source for the storage object.')
+
+
+def r2_to_oci(r2_bucket_name: str, oci_bucket_name: str) -> None:
+    """Creates a one-time transfer from Cloudflare R2 to OCI Bucket.
+    Args:
+      r2_bucket_name: str; Name of the Cloudflare R2 Bucket
+      oci_bucket_name: str; Name of the OCI Bucket
+    """
+    raise NotImplementedError(
+        'Moving data directly from Cloudflare R2 to OCI '
+        'bucket is currently not supported. Please specify '
+        'a local source for the storage object.')

sky/data/data_utils.py

@@ -7,6 +7,7 @@ import os
 import re
 import subprocess
 import textwrap
+import time
 from typing import Any, Callable, Dict, List, Optional, Tuple
 import urllib.parse
 
@@ -15,15 +16,25 @@ from filelock import FileLock
 from sky import exceptions
 from sky import sky_logging
 from sky.adaptors import aws
+from sky.adaptors import azure
 from sky.adaptors import cloudflare
 from sky.adaptors import gcp
 from sky.adaptors import ibm
+from sky.skylet import log_lib
+from sky.utils import common_utils
 from sky.utils import ux_utils
 
 Client = Any
 
 logger = sky_logging.init_logger(__name__)
 
+AZURE_CONTAINER_URL = (
+    'https://{storage_account_name}.blob.core.windows.net/{container_name}')
+
+# Retry 5 times by default for delayed propagation to Azure system
+# when creating Storage Account.
+_STORAGE_ACCOUNT_KEY_RETRIEVE_MAX_ATTEMPT = 5
+
 
 def split_s3_path(s3_path: str) -> Tuple[str, str]:
     """Splits S3 Path into Bucket name and Relative Path to Bucket
@@ -49,6 +60,28 @@ def split_gcs_path(gcs_path: str) -> Tuple[str, str]:
     return bucket, key
 
 
+def split_az_path(az_path: str) -> Tuple[str, str, str]:
+    """Splits Path into Storage account and Container names and Relative Path
+
+    Args:
+        az_path: Container Path,
+          e.g. https://azureopendatastorage.blob.core.windows.net/nyctlc
+
+    Returns:
+        str: Name of the storage account
+        str: Name of the container
+        str: Paths of the file/directory defined within the container
+    """
+    path_parts = az_path.replace('https://', '').split('/')
+    service_endpoint = path_parts.pop(0)
+    service_endpoint_parts = service_endpoint.split('.')
+    storage_account_name = service_endpoint_parts[0]
+    container_name = path_parts.pop(0)
+    path = '/'.join(path_parts)
+
+    return storage_account_name, container_name, path
+
+
 def split_r2_path(r2_path: str) -> Tuple[str, str]:
     """Splits R2 Path into Bucket name and Relative Path to Bucket
 
@@ -126,6 +159,145 @@ def verify_gcs_bucket(name: str) -> bool:
         return False
 
 
+def create_az_client(client_type: str, **kwargs: Any) -> Client:
+    """Helper method that connects to AZ client for diverse Resources.
+
+    Args:
+      client_type: str; specify client type, e.g. storage, resource, container
+
+    Returns:
+        Client object facing AZ Resource of the 'client_type'.
+    """
+    resource_group_name = kwargs.pop('resource_group_name', None)
+    container_url = kwargs.pop('container_url', None)
+    storage_account_name = kwargs.pop('storage_account_name', None)
+    refresh_client = kwargs.pop('refresh_client', False)
+    if client_type == 'container':
+        # We do not assert on resource_group_name as it is set to None when the
+        # container_url is for public container with user access.
+        assert container_url is not None, ('container_url must be provided for '
+                                           'container client')
+        assert storage_account_name is not None, ('storage_account_name must '
+                                                  'be provided for container '
+                                                  'client')
+
+    if refresh_client:
+        azure.get_client.cache_clear()
+
+    subscription_id = azure.get_subscription_id()
+    client = azure.get_client(client_type,
+                              subscription_id,
+                              container_url=container_url,
+                              storage_account_name=storage_account_name,
+                              resource_group_name=resource_group_name)
+    return client
+
+
+def verify_az_bucket(storage_account_name: str, container_name: str) -> bool:
+    """Helper method that checks if the AZ Container exists
+
+    Args:
+      storage_account_name: str; Name of the storage account
+      container_name: str; Name of the container
+
+    Returns:
+      True if the container exists, False otherwise.
+    """
+    container_url = AZURE_CONTAINER_URL.format(
+        storage_account_name=storage_account_name,
+        container_name=container_name)
+    resource_group_name = azure.get_az_resource_group(storage_account_name)
+    container_client = create_az_client(
+        client_type='container',
+        container_url=container_url,
+        storage_account_name=storage_account_name,
+        resource_group_name=resource_group_name)
+    return container_client.exists()
+
+
+def get_az_storage_account_key(
+    storage_account_name: str,
+    resource_group_name: Optional[str] = None,
+    storage_client: Optional[Client] = None,
+    resource_client: Optional[Client] = None,
+) -> Optional[str]:
+    """Returns access key of the given name of storage account.
+
+    Args:
+        storage_account_name: Name of the storage account
+        resource_group_name: Name of the resource group the
+            passed storage account belongs to.
+        storage_clent: Client object facing Storage
+        resource_client: Client object facing Resource
+
+    Returns:
+        One of the two access keys to the given storage account, or None if
+        the account is not found.
+    """
+    if resource_client is None:
+        resource_client = create_az_client('resource')
+    if storage_client is None:
+        storage_client = create_az_client('storage')
+    if resource_group_name is None:
+        resource_group_name = azure.get_az_resource_group(
+            storage_account_name, storage_client)
+    # resource_group_name is None when using a public container or
+    # a private container not belonging to the user.
+    if resource_group_name is None:
+        return None
+
+    attempt = 0
+    backoff = common_utils.Backoff()
+    while True:
+        storage_account_keys = None
+        resources = resource_client.resources.list_by_resource_group(
+            resource_group_name)
+        # resource group is either created or read when Storage initializes.
+        assert resources is not None
+        for resource in resources:
+            if (resource.type == 'Microsoft.Storage/storageAccounts' and
+                    resource.name == storage_account_name):
+                assert storage_account_keys is None
+                keys = storage_client.storage_accounts.list_keys(
+                    resource_group_name, storage_account_name)
+                storage_account_keys = [key.value for key in keys.keys]
+        # If storage account was created right before call to this method,
+        # it is possible to fail to retrieve the key as the creation did not
+        # propagate to Azure yet. We retry several times.
+        if storage_account_keys is None:
+            attempt += 1
+            time.sleep(backoff.current_backoff())
+            if attempt > _STORAGE_ACCOUNT_KEY_RETRIEVE_MAX_ATTEMPT:
+                raise RuntimeError('Failed to obtain key value of storage '
+                                   f'account {storage_account_name!r}. '
+                                   'Check if the storage account was created.')
+            continue
+        # Azure provides two sets of working storage account keys and we use
+        # one of it.
+        storage_account_key = storage_account_keys[0]
+        return storage_account_key
+
+
+def is_az_container_endpoint(endpoint_url: str) -> bool:
+    """Checks if provided url follows a valid container endpoint naming format.
+
+    Args:
+      endpoint_url: Url of container endpoint.
+        e.g. https://azureopendatastorage.blob.core.windows.net/nyctlc
+
+    Returns:
+      bool: True if the endpoint is valid, False otherwise.
+    """
+    # Storage account must be length of 3-24
+    # Reference: https://learn.microsoft.com/en-us/azure/azure-resource-manager/management/resource-name-rules#microsoftstorage # pylint: disable=line-too-long
+    pattern = re.compile(
+        r'^https://([a-z0-9]{3,24})\.blob\.core\.windows\.net(/[^/]+)*$')
+    match = pattern.match(endpoint_url)
+    if match is None:
+        return False
+    return True
+
+
 def create_r2_client(region: str = 'auto') -> Client:
     """Helper method that connects to Boto3 client for R2 Bucket
 
@@ -259,6 +431,7 @@ def _group_files_by_dir(
 def parallel_upload(source_path_list: List[str],
                     filesync_command_generator: Callable[[str, List[str]], str],
                     dirsync_command_generator: Callable[[str, str], str],
+                    log_path: str,
                     bucket_name: str,
                     access_denied_message: str,
                     create_dirs: bool = False,
@@ -274,6 +447,7 @@ def parallel_upload(source_path_list: List[str],
             for a list of files belonging to the same dir.
         dirsync_command_generator: Callable that generates rsync command
             for a directory.
+        log_path: Path to the log file.
         access_denied_message: Message to intercept from the underlying
             upload utility when permissions are insufficient. Used in
             exception handling.
@@ -306,7 +480,7 @@ def parallel_upload(source_path_list: List[str],
         p.starmap(
             run_upload_cli,
             zip(commands, [access_denied_message] * len(commands),
-                [bucket_name] * len(commands)))
+                [bucket_name] * len(commands), [log_path] * len(commands)))
 
 
 def get_gsutil_command() -> Tuple[str, str]:
@@ -347,37 +521,31 @@ def get_gsutil_command() -> Tuple[str, str]:
     return gsutil_alias, alias_gen
 
 
-def run_upload_cli(command: str, access_denied_message: str, bucket_name: str):
-    # TODO(zhwu): Use log_lib.run_with_log() and redirect the output
-    # to a log file.
-    with subprocess.Popen(command,
-                          stderr=subprocess.PIPE,
-                          stdout=subprocess.DEVNULL,
-                          shell=True) as process:
-        stderr = []
-        assert process.stderr is not None  # for mypy
-        while True:
-            line = process.stderr.readline()
-            if not line:
-                break
-            str_line = line.decode('utf-8')
-            stderr.append(str_line)
-            if access_denied_message in str_line:
-                process.kill()
-                with ux_utils.print_exception_no_traceback():
-                    raise PermissionError(
-                        'Failed to upload files to '
-                        'the remote bucket. The bucket does not have '
-                        'write permissions. It is possible that '
-                        'the bucket is public.')
-        returncode = process.wait()
-        if returncode != 0:
-            stderr_str = '\n'.join(stderr)
-            with ux_utils.print_exception_no_traceback():
-                logger.error(stderr_str)
-                raise exceptions.StorageUploadError(
-                    f'Upload to bucket failed for store {bucket_name}. '
-                    'Please check the logs.')
+def run_upload_cli(command: str, access_denied_message: str, bucket_name: str,
+                   log_path: str):
+    returncode, stdout, stderr = log_lib.run_with_log(
+        command,
+        log_path,
+        shell=True,
+        require_outputs=True,
+        # We need to use bash as some of the cloud commands uses bash syntax,
+        # such as [[ ... ]]
+        executable='/bin/bash')
+    if access_denied_message in stderr:
+        with ux_utils.print_exception_no_traceback():
+            raise PermissionError('Failed to upload files to '
+                                  'the remote bucket. The bucket does not have '
+                                  'write permissions. It is possible that '
+                                  'the bucket is public.')
+    if returncode != 0:
+        with ux_utils.print_exception_no_traceback():
+            logger.error(stderr)
+            raise exceptions.StorageUploadError(
+                f'Upload to bucket failed for store {bucket_name}. '
+                f'Please check the logs: {log_path}')
+    if not stdout:
+        logger.debug('No file uploaded. This could be due to an error or '
+                     'because all files already exist on the cloud.')
 
 
 def get_cos_regions() -> List[str]:
@@ -566,3 +734,14 @@ class Rclone():
                 lines_to_keep.append(line)
 
         return lines_to_keep
+
+
+def split_oci_path(oci_path: str) -> Tuple[str, str]:
+    """Splits OCI Path into Bucket name and Relative Path to Bucket
+    Args:
+      oci_path: str; OCI Path, e.g. oci://imagenet/train/
+    """
+    path_parts = oci_path.replace('oci://', '').split('/')
+    bucket = path_parts.pop(0)
+    key = '/'.join(path_parts)
+    return bucket, key