skypilot-nightly 1.0.0.dev20241111__py3-none-any.whl → 1.0.0.dev20241113__py3-none-any.whl

sky/resources.py

@@ -14,6 +14,7 @@ from sky import sky_logging
 from sky import skypilot_config
 from sky.clouds import service_catalog
 from sky.provision import docker_utils
+from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.skylet import constants
 from sky.utils import accelerator_registry
 from sky.utils import common_utils
@@ -582,36 +583,46 @@ class Resources:
             acc, _ = list(accelerators.items())[0]
             if 'tpu' in acc.lower():
                 if self.cloud is None:
-                    self._cloud = clouds.GCP()
-                assert self.cloud.is_same_cloud(
-                    clouds.GCP()), 'Cloud must be GCP.'
+                    if kubernetes_utils.is_tpu_on_gke(acc):
+                        self._cloud = clouds.Kubernetes()
+                    else:
+                        self._cloud = clouds.GCP()
+                assert (self.cloud.is_same_cloud(clouds.GCP()) or
+                        self.cloud.is_same_cloud(clouds.Kubernetes())), (
+                            'Cloud must be GCP or Kubernetes for TPU '
+                            'accelerators.')
+
                 if accelerator_args is None:
                     accelerator_args = {}
+
                 use_tpu_vm = accelerator_args.get('tpu_vm', True)
-                if self.instance_type is not None and use_tpu_vm:
-                    if self.instance_type != 'TPU-VM':
-                        with ux_utils.print_exception_no_traceback():
-                            raise ValueError(
-                                'Cannot specify instance type'
-                                f' (got "{self.instance_type}") for TPU VM.')
-                if 'runtime_version' not in accelerator_args:
-
-                    def _get_default_runtime_version() -> str:
-                        if not use_tpu_vm:
-                            return '2.12.0'
-                        # TPU V5 requires a newer runtime version.
-                        if acc.startswith('tpu-v5'):
-                            return 'v2-alpha-tpuv5'
-                        # TPU V6e requires a newer runtime version.
-                        if acc.startswith('tpu-v6e'):
-                            return 'v2-alpha-tpuv6e'
-                        return 'tpu-vm-base'
-
-                    accelerator_args['runtime_version'] = (
-                        _get_default_runtime_version())
-                    logger.info(
-                        'Missing runtime_version in accelerator_args, using'
-                        f' default ({accelerator_args["runtime_version"]})')
+                if (self.cloud.is_same_cloud(clouds.GCP()) and
+                        not kubernetes_utils.is_tpu_on_gke(acc)):
+                    if 'runtime_version' not in accelerator_args:
+
+                        def _get_default_runtime_version() -> str:
+                            if not use_tpu_vm:
+                                return '2.12.0'
+                            # TPU V5 requires a newer runtime version.
+                            if acc.startswith('tpu-v5'):
+                                return 'v2-alpha-tpuv5'
+                            # TPU V6e requires a newer runtime version.
+                            elif acc.startswith('tpu-v6e'):
+                                return 'v2-alpha-tpuv6e'
+                            return 'tpu-vm-base'
+
+                        accelerator_args['runtime_version'] = (
+                            _get_default_runtime_version())
+                        logger.info(
+                            'Missing runtime_version in accelerator_args, using'
+                            f' default ({accelerator_args["runtime_version"]})')
+
+                    if self.instance_type is not None and use_tpu_vm:
+                        if self.instance_type != 'TPU-VM':
+                            with ux_utils.print_exception_no_traceback():
+                                raise ValueError(
+                                    'Cannot specify instance type (got '
+                                    f'{self.instance_type!r}) for TPU VM.')
 
         self._accelerators = accelerators
         self._accelerator_args = accelerator_args

sky/serve/serve_utils.py

@@ -592,15 +592,26 @@ def get_latest_version_with_min_replicas(
 
 
 def _follow_replica_logs(
-        file: TextIO,
-        cluster_name: str,
-        *,
-        finish_stream: Callable[[], bool],
-        exit_if_stream_end: bool = False,
-        no_new_content_timeout: Optional[int] = None) -> Iterator[str]:
-    line = ''
-    log_file = None
-    no_new_content_cnt = 0
+    file: TextIO,
+    cluster_name: str,
+    *,
+    should_stop: Callable[[], bool],
+    stop_on_eof: bool = False,
+    idle_timeout_seconds: Optional[int] = None,
+) -> Iterator[str]:
+    """Follows logs for a replica, handling nested log files.
+
+    Args:
+        file: Log file to read from.
+        cluster_name: Name of the cluster being launched.
+        should_stop: Callback that returns True when streaming should stop.
+        stop_on_eof: If True, stop when reaching end of file.
+        idle_timeout_seconds: If set, stop after these many seconds without
+            new content.
+
+    Yields:
+        Log lines from the main file and any nested log files.
+    """
 
     def cluster_is_up() -> bool:
         cluster_record = global_user_state.get_cluster_from_name(cluster_name)
@@ -608,51 +619,52 @@ def _follow_replica_logs(
             return False
         return cluster_record['status'] == status_lib.ClusterStatus.UP
 
-    while True:
-        tmp = file.readline()
-        if tmp is not None and tmp != '':
-            no_new_content_cnt = 0
-            line += tmp
-            if '\n' in line or '\r' in line:
-                # Tailing detailed progress for user. All logs in skypilot is
-                # of format `To view detailed progress: tail -n100 -f *.log`.
-                x = re.match(_SKYPILOT_PROVISION_LOG_PATTERN, line)
-                if x is not None:
-                    log_file = os.path.expanduser(x.group(1))
-                elif re.match(_SKYPILOT_LOG_PATTERN, line) is None:
-                    # Not print other logs (file sync logs) since we lack
-                    # utility to determine when these log files are finished
-                    # writing.
-                    # TODO(tian): Not skip these logs since there are small
-                    # chance that error will happen in file sync. Need to find
-                    # a better way to do this.
-                    yield line
-                    # Output next line first since it indicates the process is
-                    # starting. For our launching logs, it's always:
-                    # Launching on <cloud> <region> (<zone>)
-                    if log_file is not None:
-                        with open(log_file, 'r', newline='',
-                                  encoding='utf-8') as f:
-                            # We still exit if more than 10 seconds without new
-                            # content to avoid any internal bug that causes
-                            # the launch failed and cluster status remains INIT.
-                            for l in _follow_replica_logs(
-                                    f,
-                                    cluster_name,
-                                    finish_stream=cluster_is_up,
-                                    exit_if_stream_end=exit_if_stream_end,
-                                    no_new_content_timeout=10):
-                                yield l
-                        log_file = None
-                line = ''
-        else:
-            if exit_if_stream_end or finish_stream():
-                break
-            if no_new_content_timeout is not None:
-                if no_new_content_cnt >= no_new_content_timeout:
-                    break
-                no_new_content_cnt += 1
-            time.sleep(1)
+    def process_line(line: str) -> Iterator[str]:
+        # Tailing detailed progress for user. All logs in skypilot is
+        # of format `To view detailed progress: tail -n100 -f *.log`.
+        # Check if the line is directing users to view logs
+        provision_log_prompt = re.match(_SKYPILOT_PROVISION_LOG_PATTERN, line)
+        other_log_prompt = re.match(_SKYPILOT_LOG_PATTERN, line)
+
+        if provision_log_prompt is not None:
+            nested_log_path = os.path.expanduser(provision_log_prompt.group(1))
+            with open(nested_log_path, 'r', newline='', encoding='utf-8') as f:
+                # We still exit if more than 10 seconds without new content
+                # to avoid any internal bug that causes the launch to fail
+                # while cluster status remains INIT.
+                # Originally, we output the next line first before printing
+                # the launching logs. Since the next line is always
+                # `Launching on <cloud> <region> (<zone>)`, we output it first
+                # to indicate the process is starting.
+                # TODO(andyl): After refactor #4323, the above logic is broken,
+                # but coincidentally with the new UX 3.0, the `Cluster launched`
+                # message is printed first, making the output appear correct.
+                # Explaining this since it's technically a breaking change
+                # for this refactor PR #4323. Will remove soon in a fix PR
+                # for adapting the serve.follow_logs to the new UX.
+                yield from _follow_replica_logs(f,
+                                                cluster_name,
+                                                should_stop=cluster_is_up,
+                                                stop_on_eof=stop_on_eof,
+                                                idle_timeout_seconds=10)
+            return
+
+        if other_log_prompt is not None:
+            # Now we skip other logs (file sync logs) since we lack
+            # utility to determine when these log files are finished
+            # writing.
+            # TODO(tian): We should not skip these logs since there are
+            # small chance that error will happen in file sync. Need to
+            # find a better way to do this.
+            return
+
+        yield line
+
+    return log_utils.follow_logs(file,
+                                 should_stop=should_stop,
+                                 stop_on_eof=stop_on_eof,
+                                 process_line=process_line,
+                                 idle_timeout_seconds=idle_timeout_seconds)
 
 
 def stream_replica_logs(service_name: str, replica_id: int,
@@ -687,14 +699,17 @@ def stream_replica_logs(service_name: str, replica_id: int,
             raise ValueError(
                 _FAILED_TO_FIND_REPLICA_MSG.format(replica_id=replica_id))
 
-    finish_stream = (
+    replica_provisioned = (
         lambda: _get_replica_status() != serve_state.ReplicaStatus.PROVISIONING)
     with open(launch_log_file_name, 'r', newline='', encoding='utf-8') as f:
-        for line in _follow_replica_logs(f,
-                                         replica_cluster_name,
-                                         finish_stream=finish_stream,
-                                         exit_if_stream_end=not follow):
+        for line in _follow_replica_logs(
+                f,
+                replica_cluster_name,
+                should_stop=replica_provisioned,
+                stop_on_eof=not follow,
+        ):
             print(line, end='', flush=True)
+
     if (not follow and
             _get_replica_status() == serve_state.ReplicaStatus.PROVISIONING):
         # Early exit if not following the logs.
@@ -719,22 +734,6 @@ def stream_replica_logs(service_name: str, replica_id: int,
     return ''
 
 
-def _follow_logs(file: TextIO, *, finish_stream: Callable[[], bool],
-                 exit_if_stream_end: bool) -> Iterator[str]:
-    line = ''
-    while True:
-        tmp = file.readline()
-        if tmp is not None and tmp != '':
-            line += tmp
-            if '\n' in line or '\r' in line:
-                yield line
-                line = ''
-        else:
-            if exit_if_stream_end or finish_stream():
-                break
-            time.sleep(1)
-
-
 def stream_serve_process_logs(service_name: str, stream_controller: bool,
                               follow: bool) -> str:
     msg = check_service_status_healthy(service_name)
@@ -753,9 +752,11 @@ def stream_serve_process_logs(service_name: str, stream_controller: bool,
 
     with open(os.path.expanduser(log_file), 'r', newline='',
               encoding='utf-8') as f:
-        for line in _follow_logs(f,
-                                 finish_stream=_service_is_terminal,
-                                 exit_if_stream_end=not follow):
+        for line in log_utils.follow_logs(
+                f,
+                should_stop=_service_is_terminal,
+                stop_on_eof=not follow,
+        ):
             print(line, end='', flush=True)
     return ''
 

sky/skylet/log_lib.py

@@ -320,11 +320,8 @@ def run_bash_command_with_log(bash_command: str,
         # Need this `-i` option to make sure `source ~/.bashrc` work.
         inner_command = f'/bin/bash -i {script_path}'
 
-        subprocess_cmd: Union[str, List[str]]
-        subprocess_cmd = inner_command
-
         return run_with_log(
-            subprocess_cmd,
+            inner_command,
             log_path,
             stream_logs=stream_logs,
             with_ray=with_ray,

sky/templates/kubernetes-ray.yml.j2

@@ -283,12 +283,15 @@ available_node_types:
 
         restartPolicy: Never
 
-        # Add node selector if GPUs are requested:
+        # Add node selector if GPU/TPUs are requested:
         {% if (k8s_acc_label_key is not none and k8s_acc_label_value is not none) or (k8s_spot_label_key is not none) %}
         nodeSelector:
             {% if k8s_acc_label_key is not none and k8s_acc_label_value is not none %}
             {{k8s_acc_label_key}}: {{k8s_acc_label_value}}
             {% endif %}
+            {% if k8s_topology_label_key is not none and k8s_topology_label_value is not none %}
+            {{k8s_topology_label_key}}: {{k8s_topology_label_value}}
+            {% endif %}
             {% if k8s_spot_label_key is not none %}
             {{k8s_spot_label_key}}: {{k8s_spot_label_value|tojson}}
             {% endif %}
@@ -409,14 +412,24 @@ available_node_types:
             requests:
               cpu: {{cpus}}
               memory: {{memory}}G
-              nvidia.com/gpu: {{accelerator_count}}
+              {% if k8s_resource_key is not none %}
+              # Number of requested google.com/tpu must be equal to the total
+              # number of available TPU chips on the TPU slice node either it
+              # being a node from multi-host TPU slice or single-host TPU
+              # slice. Example reference:
+              # https://cloud.google.com/kubernetes-engine/docs/concepts/tpus#how_tpus_work
+              {{k8s_resource_key}}: {{accelerator_count}}
+              {% endif %}
               {% if k8s_fuse_device_required %}
               # Kubernetes resource exposed by the fuse device manager
               # https://gitlab.com/arm-research/smarter/smarter-device-manager
               smarter-devices/fuse: "1"
               {% endif %}
             limits:
-              nvidia.com/gpu: {{accelerator_count}} # Limits need to be defined for GPU requests
+              # Limits need to be defined for GPU/TPU requests
+              {% if k8s_resource_key is not none %}
+              {{k8s_resource_key}}: {{accelerator_count}}
+              {% endif %}
               {% if k8s_fuse_device_required %}
               smarter-devices/fuse: "1"
               {% endif %}
@@ -451,6 +464,19 @@ setup_commands:
     sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
     mkdir -p ~/.ssh; (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n" >> ~/.ssh/config;
     [ -f /etc/fuse.conf ] && sudo sed -i 's/#user_allow_other/user_allow_other/g' /etc/fuse.conf || (sudo sh -c 'echo "user_allow_other" > /etc/fuse.conf'); # This is needed for `-o allow_other` option for `goofys`;
+  {% if tpu_requested %}
+  # The /tmp/tpu_logs directory is where TPU-related logs, such as logs from
+  # the TPU runtime, are written. These capture runtime information about the
+  # TPU execution, including any warnings, errors, or general activity of
+  # the TPU driver. By default, the /tmp/tpu_logs directory is created with
+  # 755 permissions, and the user of the provisioned pod is not necessarily
+  # a root. Hence, we need to update the write permission so the logs can be
+  # properly written.
+  # TODO(Doyoung): Investigate to see why TPU workload fails to run without
+  # execution permission, such as granting 766 to log file. Check if it's a
+  # must and see if there's a workaround to grant minimum permission. 
+  - sudo chmod 777 /tmp/tpu_logs;
+  {% endif %}
 
 # Format: `REMOTE_PATH : LOCAL_PATH`
 file_mounts: {

sky/utils/kubernetes/generate_kubeconfig.sh

@@ -112,6 +112,9 @@ rules:
   - apiGroups: ["networking.k8s.io"]   # Required for exposing services through ingresses
     resources: ["ingressclasses"]
     verbs: ["get", "list", "watch"]
+  - apiGroups: [""]                 # Required for sky show-gpus command
+    resources: ["pods"]
+    verbs: ["get", "list"]
 ---
 # ClusterRoleBinding for the service account
 apiVersion: rbac.authorization.k8s.io/v1

sky/utils/kubernetes/gpu_labeler.py

@@ -101,7 +101,7 @@ def label():
         # Get the list of nodes with GPUs
         gpu_nodes = []
         for node in nodes:
-            if 'nvidia.com/gpu' in node.status.capacity:
+            if kubernetes_utils.GPU_RESOURCE_KEY in node.status.capacity:
                 gpu_nodes.append(node)
 
         print(f'Found {len(gpu_nodes)} GPU nodes in the cluster')
@@ -142,7 +142,7 @@ def label():
     if len(gpu_nodes) == 0:
         print('No GPU nodes found in the cluster. If you have GPU nodes, '
               'please ensure that they have the label '
-              '`nvidia.com/gpu: <number of GPUs>`')
+              f'`{kubernetes_utils.GPU_RESOURCE_KEY}: <number of GPUs>`')
     else:
         print('GPU labeling started - this may take 10 min or more to complete.'
               '\nTo check the status of GPU labeling jobs, run '

sky/utils/log_utils.py

@@ -1,7 +1,8 @@
 """Logging utils."""
 import enum
+import time
 import types
-from typing import List, Optional, Type
+from typing import Callable, Iterator, List, Optional, TextIO, Type
 
 import colorama
 import pendulum
@@ -284,3 +285,53 @@ def readable_time_duration(start: Optional[float],
         diff = diff.replace('hour', 'hr')
 
     return diff
+
+
+def follow_logs(
+    file: TextIO,
+    *,
+    should_stop: Callable[[], bool],
+    stop_on_eof: bool = False,
+    process_line: Optional[Callable[[str], Iterator[str]]] = None,
+    idle_timeout_seconds: Optional[int] = None,
+) -> Iterator[str]:
+    """Streams and processes logs line by line from a file.
+
+    Args:
+        file: File object to read logs from.
+        should_stop: Callback that returns True when streaming should stop.
+        stop_on_eof: If True, stop when reaching end of file.
+        process_line: Optional callback to transform/filter each line.
+        idle_timeout_seconds: If set, stop after these many seconds without
+            new content.
+
+    Yields:
+        Log lines, possibly transformed by process_line if provided.
+    """
+    current_line: str = ''
+    seconds_without_content: int = 0
+
+    while True:
+        content = file.readline()
+
+        if not content:
+            if stop_on_eof or should_stop():
+                break
+
+            if idle_timeout_seconds is not None:
+                if seconds_without_content >= idle_timeout_seconds:
+                    break
+                seconds_without_content += 1
+
+            time.sleep(1)
+            continue
+
+        seconds_without_content = 0
+        current_line += content
+
+        if '\n' in current_line or '\r' in current_line:
+            if process_line is not None:
+                yield from process_line(current_line)
+            else:
+                yield current_line
+            current_line = ''

sky/utils/timeline.py

@@ -9,6 +9,7 @@ import json
 import os
 import threading
 import time
+import traceback
 from typing import Callable, Optional, Union
 
 import filelock
@@ -48,8 +49,9 @@ class Event:
             'ph': 'B',
             'ts': f'{time.time() * 10 ** 6: .3f}',
         })
+        event_begin['args'] = {'stack': '\n'.join(traceback.format_stack())}
         if self._message is not None:
-            event_begin['args'] = {'message': self._message}
+            event_begin['args']['message'] = self._message
         _events.append(event_begin)
 
     def end(self):

{skypilot_nightly-1.0.0.dev20241111.dist-info → skypilot_nightly-1.0.0.dev20241113.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: skypilot-nightly
-Version: 1.0.0.dev20241111
+Version: 1.0.0.dev20241113
 Summary: SkyPilot: An intercloud broker for the clouds
 Author: SkyPilot Team
 License: Apache 2.0
@@ -309,7 +309,7 @@ Runnable examples:
   - [LocalGPT](./llm/localgpt)
   - [Falcon](./llm/falcon)
   - Add yours here & see more in [`llm/`](./llm)!
-- Framework examples: [PyTorch DDP](https://github.com/skypilot-org/skypilot/blob/master/examples/resnet_distributed_torch.yaml), [DeepSpeed](./examples/deepspeed-multinode/sky.yaml), [JAX/Flax on TPU](https://github.com/skypilot-org/skypilot/blob/master/examples/tpu/tpuvm_mnist.yaml), [Stable Diffusion](https://github.com/skypilot-org/skypilot/tree/master/examples/stable_diffusion), [Detectron2](https://github.com/skypilot-org/skypilot/blob/master/examples/detectron2_docker.yaml), [Distributed](https://github.com/skypilot-org/skypilot/blob/master/examples/resnet_distributed_tf_app.py) [TensorFlow](https://github.com/skypilot-org/skypilot/blob/master/examples/resnet_app_storage.yaml), [Ray Train](examples/distributed_ray_train/ray_train.yaml), [NeMo](https://github.com/skypilot-org/skypilot/blob/master/examples/nemo/nemo.yaml), [programmatic grid search](https://github.com/skypilot-org/skypilot/blob/master/examples/huggingface_glue_imdb_grid_search_app.py), [Docker](https://github.com/skypilot-org/skypilot/blob/master/examples/docker/echo_app.yaml), [Cog](https://github.com/skypilot-org/skypilot/blob/master/examples/cog/), [Unsloth](https://github.com/skypilot-org/skypilot/blob/master/examples/unsloth/unsloth.yaml), [Ollama](https://github.com/skypilot-org/skypilot/blob/master/llm/ollama), [llm.c](https://github.com/skypilot-org/skypilot/tree/master/llm/gpt-2), [Airflow](./examples/airflow/training_workflow) and [many more (`examples/`)](./examples).
+- Framework examples: [PyTorch DDP](https://github.com/skypilot-org/skypilot/blob/master/examples/resnet_distributed_torch.yaml), [DeepSpeed](./examples/deepspeed-multinode/sky.yaml), [JAX/Flax on TPU](https://github.com/skypilot-org/skypilot/blob/master/examples/tpu/tpuvm_mnist.yaml), [Stable Diffusion](https://github.com/skypilot-org/skypilot/tree/master/examples/stable_diffusion), [Detectron2](https://github.com/skypilot-org/skypilot/blob/master/examples/detectron2_docker.yaml), [Distributed](https://github.com/skypilot-org/skypilot/blob/master/examples/resnet_distributed_tf_app.py) [TensorFlow](https://github.com/skypilot-org/skypilot/blob/master/examples/resnet_app_storage.yaml), [Ray Train](examples/distributed_ray_train/ray_train.yaml), [NeMo](https://github.com/skypilot-org/skypilot/blob/master/examples/nemo/), [programmatic grid search](https://github.com/skypilot-org/skypilot/blob/master/examples/huggingface_glue_imdb_grid_search_app.py), [Docker](https://github.com/skypilot-org/skypilot/blob/master/examples/docker/echo_app.yaml), [Cog](https://github.com/skypilot-org/skypilot/blob/master/examples/cog/), [Unsloth](https://github.com/skypilot-org/skypilot/blob/master/examples/unsloth/unsloth.yaml), [Ollama](https://github.com/skypilot-org/skypilot/blob/master/llm/ollama), [llm.c](https://github.com/skypilot-org/skypilot/tree/master/llm/gpt-2), [Airflow](./examples/airflow/training_workflow) and [many more (`examples/`)](./examples).
 
 Case Studies and Integrations: [Community Spotlights](https://blog.skypilot.co/community/)