[k8s] Logging SSH Setup on Kubernetes Provision (#3266)

* initial commit

* newline

* comments

* run linter

* reminder for down

* tentatively done with example

* formatting

* yapf

* [Storage] Storage mounting tool permissions fix (#3215)

* fix permissions

* fix permissions

* [LLM] Example for Serving Gemma (#3207)

* Add serve for gemma and fix mixtral dependency

* Add hf token

* fix model len

* Add comment

* Serve your private gemma

* fix serve yaml

* readme

* Remove chat completion due to the wrong template

* add readme

* Update llm/gemma/README.md

Co-authored-by: Zongheng Yang <[email protected]>

* address comments

* Update README.md

Co-authored-by: Zongheng Yang <[email protected]>

* Update llm/gemma/README.md

Co-authored-by: Zongheng Yang <[email protected]>

* Update llm/gemma/README.md

Co-authored-by: Zongheng Yang <[email protected]>

* Update llm/gemma/README.md

Co-authored-by: Zongheng Yang <[email protected]>

* Change to it

* Add chat API

* use HF_TOKEN env

* typo

---------

Co-authored-by: Zongheng Yang <[email protected]>

* [LLM] Add logo for Gemma (#3220)

* Minor fixes for release 0.5.0 (#3212)

* when removing cudo credential, sky check fails

* remove tips

* minor hint fix

* fix cluster version for k8s

* fix typo

* [Docker] Add retry for docker pull due to daemon not ready (#3218)

* Add retry for docker pull due to daemon not ready

* longer wait time

* longer wait time

* retry earlier

* add retry for retries as well

* longer wait time

* change wait time

* format

* Add comment

* Fix

* Fix indent for azure docker config

* Fix docker login config

* Fix comments

* More robust docker login config

* Add retry for docker check

* minor fix

* Add additional test for stop and start with docker

* Fix cancelled

* added comments

* quick fix

* finished pip issues

* fix

* fix storage error message, add example link to docs

* logging for SSH when doing kubernetes provision

* romil edits

* took out todo commnt

* removed extra file

* renamed file

* restored right version of file

* simplify things

* newline

* more formatting

* formatting

* minor fixes

* set x and logging

* fixes

* docstr

---------

Co-authored-by: Sheth <[email protected]>
Co-authored-by: Romil Bhardwaj <[email protected]>
Co-authored-by: Zhanghao Wu <[email protected]>
Co-authored-by: Zongheng Yang <[email protected]>
Co-authored-by: Romil Bhardwaj <[email protected]>

sky/provision/kubernetes/instance.py

@@ -240,7 +240,17 @@ def _wait_for_pods_to_run(namespace, new_nodes):
         time.sleep(1)
 
 
-def _run_command_on_pods(node_name, node_namespace, command):
+def _run_command_on_pods(node_name: str,
+                         node_namespace: str,
+                         command: List[str],
+                         stream_logs: bool = False):
+    """Run command on Kubernetes pods.
+
+    If `stream_logs` is True, we poll for output and error messages while the
+    command is executing, and the stdout and stderr is written to logger.info.
+    When called from the provisioner, this logger.info is written to the
+    provision.log file (see setup_provision_logging()).
+    """
     cmd_output = kubernetes.stream()(
         kubernetes.core_api().connect_get_namespaced_pod_exec,
         node_name,
@@ -250,7 +260,16 @@ def _run_command_on_pods(node_name, node_namespace, command):
         stdin=False,
         stdout=True,
         tty=False,
+        _preload_content=(not stream_logs),
         _request_timeout=kubernetes.API_TIMEOUT)
+    if stream_logs:
+        while cmd_output.is_open():
+            cmd_output.update(timeout=1)
+            if cmd_output.peek_stdout():
+                logger.info(f'{cmd_output.read_stdout().strip()}')
+            if cmd_output.peek_stderr():
+                logger.info(f'{cmd_output.read_stderr().strip()}')
+        cmd_output.close()
     return cmd_output
 
 
@@ -326,6 +345,7 @@ def _setup_ssh_in_pods(namespace: str, new_nodes: List) -> None:
         '/bin/sh',
         '-c',
         (
+            'set -x; '
             'prefix_cmd() '
             '{ if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; }; '
             'export DEBIAN_FRONTEND=noninteractive;'
@@ -351,9 +371,15 @@ def _setup_ssh_in_pods(namespace: str, new_nodes: List) -> None:
             # See https://www.educative.io/answers/error-mesg-ttyname-failed-inappropriate-ioctl-for-device  # pylint: disable=line-too-long
             '$(prefix_cmd) sed -i "s/mesg n/tty -s \\&\\& mesg n/" ~/.profile;')
     ]
-    # TODO(romilb): We need logging and surface errors here.
+    # TODO(romilb): Parallelize the setup of SSH in pods for multi-node clusters
     for new_node in new_nodes:
-        _run_command_on_pods(new_node.metadata.name, namespace, set_k8s_ssh_cmd)
+        pod_name = new_node.metadata.name
+        logger.info(f'{"-"*20}Start: Set up SSH in pod {pod_name!r} {"-"*20}')
+        _run_command_on_pods(new_node.metadata.name,
+                             namespace,
+                             set_k8s_ssh_cmd,
+                             stream_logs=True)
+        logger.info(f'{"-"*20}End: Set up SSH in pod {pod_name!r} {"-"*20}')
 
 
 def _label_pod(namespace: str, pod_name: str, label: Dict[str, str]) -> None: