Repository: skypilot-org/skypilot
Pull request: [Spot] Event callback for adding notification 
<!-- Describe the changes in this PR -->
Issue #2000: A user would like to have a customizable event callback for the spot job, so that they can implement some notification when the spot job status changes. Example:
```
event_callback:
  send_slack_message $JOB_ID $JOB_STATUS
```
<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Allows the controller to execute a customizable bash script when the spot status changes. 

Example:
Register a Slack message web hook from [here](https://api.slack.com/apps?new_app=1).
```
event_callback: |
  echo 'this is an event callback script'
  curl -X POST -H 'Content-type: application/json' --data "{\"text\":\"JOB_ID: $SKYPILOT_JOB_ID, JOB_STATUS: $JOB_STATUS\"}" https://hooks.slack.com/services/...
```

The user will receive a Slack notification when the spot status changes.
<img width="377" alt="image" src="https://github.com/skypilot-org/skypilot/assets/63817381/b4f17615-e79f-4acb-be4f-e2147c91f89d">

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[Spot] Event callback for adding notification
A user would like to have a customizable event callback for the spot jobs, so that they can implement some notification when the spot job status change.

> Contributing one idea that would enable us to implement some custom notifications when spot jobs are killed / gone into error states: can we add some event_callback configuration to a job’s yaml config? For now this could be restricted to managed spot jobs and the spot controller could be responsible for doing this.

Example:
```
event_callback:
  # some env variable will denote what the event is about, e.g. some change of status
  # custom user code can trigger notification / write logs etc
  send_slack_message $SKYPILOT_JOB_ID $JOB_STATUS
```

Hints from issue discussion:
For sending notifications (via Slack or Email), does SkyPilot have an existing mechanism? If not we need to register SkyPilot for Slack API and collect user info (slack id or email). 
> For sending notifications (via Slack or Email), does SkyPilot have an existing mechanism? If not we need to register SkyPilot for Slack API and collect user info (slack id or email).

We don't have such a mechanism in SkyPilot. As the first step, we probably don't need to support such functionality out-of-the-box. Instead, we can just support the callback field, so that the user can provide their own tool/script to do the message sending.
Thanks. Maybe something like:
```
event_callback:
  callback_script: /path/to/script $JOB_ID $JOB_STATUS
```
The controller could launch `/path/to/script/` as a subprocess with input `JOB_ID` and `JOB_STATUS`.

Commit message: running format.sh

Patch to review:
diff --git a/sky/spot/controller.py b/sky/spot/controller.py
index 1f4df44585e..7abccd7fe3a 100644
--- a/sky/spot/controller.py
+++ b/sky/spot/controller.py
@@ -7,6 +7,7 @@
 import typing
 from typing import Tuple
 
+import subprocess
 import filelock
 
 from sky import exceptions
@@ -118,7 +119,8 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
                                    start_time=time.time())
             spot_state.set_succeeded(self._job_id,
                                      task_id,
-                                     end_time=time.time())
+                                     end_time=time.time(),
+                                     callback_func=self.event_callback_func)
             return True
         usage_lib.messages.usage.update_task_id(task_id)
         task_id_env_var = task.envs[constants.TASK_ID_ENV_VAR]
@@ -173,7 +175,8 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
                 # The job is done.
                 spot_state.set_succeeded(self._job_id,
                                          task_id,
-                                         end_time=end_time)
+                                         end_time=end_time,
+                                         callback_func=self.event_callback_func)
                 logger.info(
                     f'Spot job {self._job_id} (task: {task_id}) SUCCEEDED. '
                     f'Cleaning up the spot cluster {cluster_name}.')
@@ -239,11 +242,13 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
                         'To see the details, run: '
                         f'sky spot logs --controller {self._job_id}')
 
-                    spot_state.set_failed(self._job_id,
-                                          task_id,
-                                          failure_type=spot_status_to_set,
-                                          failure_reason=failure_reason,
-                                          end_time=end_time)
+                    spot_state.set_failed(
+                        self._job_id,
+                        task_id,
+                        failure_type=spot_status_to_set,
+                        failure_reason=failure_reason,
+                        end_time=end_time,
+                        callback_func=self.event_callback_func)
                     return False
                 # Although the cluster is healthy, we fail to access the
                 # job status. Try to recover the job (will not restart the
@@ -271,6 +276,35 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
                                      task_id,
                                      recovered_time=recovered_time)
 
+    def event_callback_func(self, task_id: int, state: str, comment: str = ''):
+
+        callback_str = self._dag.tasks[task_id].event_callback
+        if callback_str is None:
+            return
+
+        if '$JOB_ID' in callback_str:
+            callback_str = callback_str.replace('$JOB_ID', str(self._job_id))
+        if '$TASK_ID' in callback_str:
+            callback_str = callback_str.replace('$TASK_ID', str(task_id))
+        if '$JOB_STATUS' in callback_str:
+            callback_str = callback_str.replace('$JOB_STATUS', state)
+        if '$COMMENT' in callback_str:
+            callback_str = callback_str.replace('$COMMENT', comment)
+
+        if callback_str.find('controller_log:') == 0:
+            print_str = callback_str[len('controller_log:'):].strip()
+            logger.info(print_str)
+        elif callback_str.find('callback_script:') == 0:
+            run_str = callback_str[len('callback_script:'):].strip()
+            result = subprocess.run(run_str,
+                                    capture_output=True,
+                                    shell=True,
+                                    check=True,
+                                    text=True)
+            logger.info(f'stdout: {result.stdout.strip()}')
+        else:
+            logger.info(f'Unrecognized callback_str: {callback_str}')
+
     def run(self):
         """Run controller logic and handle exceptions."""
         task_id = 0
@@ -292,7 +326,8 @@ def run(self):
                 self._job_id,
                 task_id=task_id,
                 failure_type=spot_state.SpotStatus.FAILED_PRECHECKS,
-                failure_reason=failure_reason)
+                failure_reason=failure_reason,
+                callback_func=self.event_callback_func)
         except exceptions.SpotJobReachedMaxRetriesError as e:
             # Please refer to the docstring of self._run for the cases when
             # this exception can occur.
@@ -303,7 +338,8 @@ def run(self):
                 self._job_id,
                 task_id=task_id,
                 failure_type=spot_state.SpotStatus.FAILED_NO_RESOURCE,
-                failure_reason=common_utils.format_exception(e))
+                failure_reason=common_utils.format_exception(e),
+                callback_func=self.event_callback_func)
         except (Exception, SystemExit) as e:  # pylint: disable=broad-except
             logger.error(traceback.format_exc())
             msg = ('Unexpected error occurred: '
@@ -313,7 +349,8 @@ def run(self):
                 self._job_id,
                 task_id=task_id,
                 failure_type=spot_state.SpotStatus.FAILED_CONTROLLER,
-                failure_reason=msg)
+                failure_reason=msg,
+                callback_func=self.event_callback_func)
         finally:
             # This will set all unfinished tasks to CANCELLING, and will not
             # affect the jobs in terminal states.
diff --git a/sky/spot/spot_state.py b/sky/spot/spot_state.py
index fe357b43d82..ddc8cd30a9f 100644
--- a/sky/spot/spot_state.py
+++ b/sky/spot/spot_state.py
@@ -5,7 +5,7 @@
 import pathlib
 import sqlite3
 import time
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, Callable
 
 import colorama
 
@@ -345,7 +345,11 @@ def set_recovered(job_id: int, task_id: int, recovered_time: float):
     logger.info('==== Recovered. ====')
 
 
-def set_succeeded(job_id: int, task_id: int, end_time: float):
+def set_succeeded(job_id: int,
+                  task_id: int,
+                  end_time: float,
+                  callback_func: Optional[Callable[[int, str, str],
+                                                   None]] = None):
     """Set the task to succeeded, if it is in a non-terminal state."""
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         cursor.execute(
@@ -355,6 +359,8 @@ def set_succeeded(job_id: int, task_id: int, end_time: float):
             WHERE spot_job_id=(?) AND task_id=(?)
             AND end_at IS null""",
             (SpotStatus.SUCCEEDED.value, end_time, job_id, task_id))
+    if callback_func is not None:
+        callback_func(task_id, 'SUCCEEDED', '')
     logger.info('Job succeeded.')
 
 
@@ -362,7 +368,8 @@ def set_failed(job_id: int,
                task_id: Optional[int],
                failure_type: SpotStatus,
                failure_reason: str,
-               end_time: Optional[float] = None):
+               end_time: Optional[float] = None,
+               callback_func: Optional[Callable[[int, str, str], None]] = None):
     """Set an entire job or task to failed, if they are in non-terminal states.
 
     Args:
@@ -401,6 +408,8 @@ def set_failed(job_id: int,
             {set_str}
             WHERE spot_job_id=(?){task_str} AND end_at IS null""",
             (*list(fields_to_set.values()), job_id))
+    if callback_func is not None and task_id is not None:
+        callback_func(task_id, 'FAILED', failure_reason)
     logger.info(failure_reason)
 
 
diff --git a/sky/task.py b/sky/task.py
index bc2ff18b66e..c06ce5a9881 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -82,6 +82,7 @@ def __init__(
         num_nodes: Optional[int] = None,
         # Advanced:
         docker_image: Optional[str] = None,
+        event_callback: str = '',
     ):
         """Initializes a Task.
 
@@ -147,6 +148,8 @@ def __init__(
         self.workdir = workdir
         self.docker_image = (docker_image if docker_image else
                              'gpuci/miniforge-cuda:11.4-devel-ubuntu18.04')
+        self.event_callback = event_callback.strip(
+        ) if event_callback is not None else ''
         # Ignore type error due to a mypy bug.
         # https://github.com/python/mypy/issues/3004
         self.num_nodes = num_nodes  # type: ignore
@@ -242,6 +245,7 @@ def from_yaml_config(config: Dict[str, Any]) -> 'Task':
             setup=config.pop('setup', None),
             num_nodes=config.pop('num_nodes', None),
             envs=config.pop('envs', None),
+            event_callback=config.pop('event_callback', None),
         )
 
         # Create lists to store storage objects inlined in file_mounts.
@@ -829,6 +833,7 @@ def add_if_not_none(key, value, no_empty: bool = False):
 
         add_if_not_none('setup', self.setup)
         add_if_not_none('workdir', self.workdir)
+        add_if_not_none('event_callback', self.event_callback)
         add_if_not_none('run', self.run)
         add_if_not_none('envs', self.envs, no_empty=True)
 
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index f3c4b8f7137..549bfcb9403 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -150,6 +150,9 @@ def get_task_schema():
             'workdir': {
                 'type': 'string',
             },
+            'event_callback': {
+                'type': 'string',
+            },
             'num_nodes': {
                 'type': 'integer',
             },


Code review comments:
diff_hunk: @@ -271,6 +276,35 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
                                      task_id,
                                      recovered_time=recovered_time)
 
+    def event_callback_func(self, task_id: int, state: str, comment: str = ''):
+
+        callback_str = self._dag.tasks[task_id].event_callback
+        if callback_str is None:
+            return
+
+        if '$JOB_ID' in callback_str:
+            callback_str = callback_str.replace('$JOB_ID', str(self._job_id))
+        if '$TASK_ID' in callback_str:
+            callback_str = callback_str.replace('$TASK_ID', str(task_id))
+        if '$JOB_STATUS' in callback_str:
+            callback_str = callback_str.replace('$JOB_STATUS', state)
+        if '$COMMENT' in callback_str:
+            callback_str = callback_str.replace('$COMMENT', comment)
path: sky/spot/controller.py
text: @user1:
How about we just set the environment variable in the subprocess call?
```
subprocess.run(..., env=dict(os.environ, JOB_ID=self._job_id, ...), ...)
```

diff_hunk: @@ -147,6 +148,8 @@ def __init__(
         self.workdir = workdir
         self.docker_image = (docker_image if docker_image else
                              'gpuci/miniforge-cuda:11.4-devel-ubuntu18.04')
+        self.event_callback = event_callback.strip(
+        ) if event_callback is not None else ''
path: sky/task.py
text: @user1:
This may not be needed if we are only taking a script for the `event_callback`

diff_hunk: @@ -271,6 +276,35 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
                                      task_id,
                                      recovered_time=recovered_time)
 
+    def event_callback_func(self, task_id: int, state: str, comment: str = ''):
+
+        callback_str = self._dag.tasks[task_id].event_callback
+        if callback_str is None:
+            return
+
+        if '$JOB_ID' in callback_str:
+            callback_str = callback_str.replace('$JOB_ID', str(self._job_id))
+        if '$TASK_ID' in callback_str:
+            callback_str = callback_str.replace('$TASK_ID', str(task_id))
+        if '$JOB_STATUS' in callback_str:
+            callback_str = callback_str.replace('$JOB_STATUS', state)
+        if '$COMMENT' in callback_str:
+            callback_str = callback_str.replace('$COMMENT', comment)
+
+        if callback_str.find('controller_log:') == 0:
+            print_str = callback_str[len('controller_log:'):].strip()
+            logger.info(print_str)
path: sky/spot/controller.py
text: @user1:
Maybe we can get rid of this, as the controller log is normally not a user-facing log, and we probably have enough loggings in it.
If a user would like to log something to the controller log, they can also specify the bash script with `echo hello world` in it.

diff_hunk: @@ -271,6 +276,35 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
                                      task_id,
                                      recovered_time=recovered_time)
 
+    def event_callback_func(self, task_id: int, state: str, comment: str = ''):
+
+        callback_str = self._dag.tasks[task_id].event_callback
+        if callback_str is None:
+            return
+
+        if '$JOB_ID' in callback_str:
+            callback_str = callback_str.replace('$JOB_ID', str(self._job_id))
+        if '$TASK_ID' in callback_str:
+            callback_str = callback_str.replace('$TASK_ID', str(task_id))
+        if '$JOB_STATUS' in callback_str:
+            callback_str = callback_str.replace('$JOB_STATUS', state)
+        if '$COMMENT' in callback_str:
+            callback_str = callback_str.replace('$COMMENT', comment)
+
+        if callback_str.find('controller_log:') == 0:
+            print_str = callback_str[len('controller_log:'):].strip()
+            logger.info(print_str)
+        elif callback_str.find('callback_script:') == 0:
+            run_str = callback_str[len('callback_script:'):].strip()
+            result = subprocess.run(run_str,
+                                    capture_output=True,
+                                    shell=True,
+                                    check=True,
+                                    text=True)
+            logger.info(f'stdout: {result.stdout.strip()}')
path: sky/spot/controller.py
text: @user1:
An idea for this (better UX):
```
logger.info(f'=== START: event callback for {state!r} ===')
log_lib.run_bash_command_with_log(...)
logger.info(f'=== END: event callback for {state!r} ===')
```

diff_hunk: @@ -271,6 +276,35 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
                                      task_id,
                                      recovered_time=recovered_time)
 
+    def event_callback_func(self, task_id: int, state: str, comment: str = ''):
+
+        callback_str = self._dag.tasks[task_id].event_callback
+        if callback_str is None:
+            return
+
+        if '$JOB_ID' in callback_str:
+            callback_str = callback_str.replace('$JOB_ID', str(self._job_id))
+        if '$TASK_ID' in callback_str:
+            callback_str = callback_str.replace('$TASK_ID', str(task_id))
+        if '$JOB_STATUS' in callback_str:
+            callback_str = callback_str.replace('$JOB_STATUS', state)
+        if '$COMMENT' in callback_str:
+            callback_str = callback_str.replace('$COMMENT', comment)
+
+        if callback_str.find('controller_log:') == 0:
+            print_str = callback_str[len('controller_log:'):].strip()
+            logger.info(print_str)
+        elif callback_str.find('callback_script:') == 0:
+            run_str = callback_str[len('callback_script:'):].strip()
path: sky/spot/controller.py
text: @user1:
Do you think we can just let the `event_callback` to just take a bash script instead of having this prefix?

diff_hunk: @@ -271,6 +276,35 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
                                      task_id,
                                      recovered_time=recovered_time)
 
+    def event_callback_func(self, task_id: int, state: str, comment: str = ''):
+
+        callback_str = self._dag.tasks[task_id].event_callback
+        if callback_str is None:
+            return
+
+        if '$JOB_ID' in callback_str:
+            callback_str = callback_str.replace('$JOB_ID', str(self._job_id))
+        if '$TASK_ID' in callback_str:
+            callback_str = callback_str.replace('$TASK_ID', str(task_id))
+        if '$JOB_STATUS' in callback_str:
+            callback_str = callback_str.replace('$JOB_STATUS', state)
+        if '$COMMENT' in callback_str:
+            callback_str = callback_str.replace('$COMMENT', comment)
+
+        if callback_str.find('controller_log:') == 0:
+            print_str = callback_str[len('controller_log:'):].strip()
+            logger.info(print_str)
+        elif callback_str.find('callback_script:') == 0:
+            run_str = callback_str[len('callback_script:'):].strip()
+            result = subprocess.run(run_str,
path: sky/spot/controller.py
text: @user1:
`subprocess.run` does not support running a bash script. Maybe we should call the `log_lib.run_bash_command_with_log` function (may need to slightly modify the function's argument to make the `job_id` and `job_owner` optional.