anyscale 0.24.86__py3-none-any.whl

anyscale/utils/workspace_utils.py

@@ -0,0 +1,65 @@
+from functools import partial
+import os
+
+from anyscale.anyscale_pydantic import BaseModel
+from anyscale.client.openapi_client.api.default_api import DefaultApi
+from anyscale.util import get_working_dir
+
+
+class SSHKeyInfo(BaseModel):
+    user: str
+    key_path: str
+
+
+class WorkspaceCommandContext(BaseModel):
+    ssh: SSHKeyInfo
+    working_dir: str
+    head_ip: str
+    project_id: str
+
+
+def _get_ssh_key_info(
+    session_id: str, api_client: DefaultApi, ssh_dir: str = "~/.ssh",
+) -> SSHKeyInfo:
+    # TODO (yiran): cleanup SSH keys if session no longer exists.
+    def _write_ssh_key(name: str, ssh_key: str) -> str:
+        key_path = os.path.join(os.path.expanduser(ssh_dir), f"{name}.pem")
+        os.makedirs(os.path.dirname(key_path), exist_ok=True)
+
+        with open(key_path, "w", opener=partial(os.open, mode=0o600)) as f:
+            f.write(ssh_key)
+
+        return key_path
+
+    ssh_key = api_client.get_session_ssh_key_api_v2_sessions_session_id_ssh_key_get(
+        session_id
+    ).result
+
+    key_path = _write_ssh_key(ssh_key.key_name, ssh_key.private_key)
+
+    # since v2 stack the host user is ubuntu.
+    return SSHKeyInfo(user="ubuntu", key_path=key_path)
+
+
+def extract_workspace_parameters(
+    cluster_id: str, project_id: str, api_client: DefaultApi, ssh_dir: str,
+) -> WorkspaceCommandContext:
+    """
+    This function pulls relevant fields from the legacy cluster config
+    to populate a `WorkspaceCommandContext`.
+
+    NOTE: Any additional fields that need to go into the WorkspaceCommandContext
+    should be passed via actual API calls, and not pulled from cluster configs.
+    """
+
+    head_ip = api_client.get_session_head_ip_api_v2_sessions_session_id_head_ip_get(
+        cluster_id
+    ).result.head_ip
+
+    ssh_info = _get_ssh_key_info(cluster_id, api_client, ssh_dir)
+    return WorkspaceCommandContext(
+        ssh=ssh_info,
+        project_id=project_id,
+        head_ip=head_ip,
+        working_dir=get_working_dir(project_id, api_client) + "/",
+    )

anyscale/version.py

@@ -0,0 +1 @@
+__version__ = "0.24.86"

anyscale/webterminal/bash-preexec.sh

@@ -0,0 +1,370 @@
+# DONT CHANGE THE LOCATION OF THIS FILE WITHOUT UPDATING THE utils.py
+
+# bash-preexec.sh -- Bash support for ZSH-like 'preexec' and 'precmd' functions.
+# https://github.com/rcaloras/bash-preexec
+#
+#
+# 'preexec' functions are executed before each interactive command is
+# executed, with the interactive command as its argument. The 'precmd'
+# function is executed before each prompt is displayed.
+#
+# Author: Ryan Caloras (ryan@bashhub.com)
+# Forked from Original Author: Glyph Lefkowitz
+#
+# V0.5.0
+#
+
+# General Usage:
+#
+#  1. Source this file at the end of your bash profile so as not to interfere
+#     with anything else that's using PROMPT_COMMAND.
+#
+#  2. Add any precmd or preexec functions by appending them to their arrays:
+#       e.g.
+#       precmd_functions+=(my_precmd_function)
+#       precmd_functions+=(some_other_precmd_function)
+#
+#       preexec_functions+=(my_preexec_function)
+#
+#  3. Consider changing anything using the DEBUG trap or PROMPT_COMMAND
+#     to use preexec and precmd instead. Preexisting usages will be
+#     preserved, but doing so manually may be less surprising.
+#
+#  Note: This module requires two Bash features which you must not otherwise be
+#  using: the "DEBUG" trap, and the "PROMPT_COMMAND" variable. If you override
+#  either of these after bash-preexec has been installed it will most likely break.
+
+# Tell shellcheck what kind of file this is.
+# shellcheck shell=bash
+
+# Make sure this is bash that's running and return otherwise.
+# Use POSIX syntax for this line:
+if [ -z "${BASH_VERSION-}" ]; then
+	return 1
+fi
+
+# We only support Bash 3.1+.
+# Note: BASH_VERSINFO is first available in Bash-2.0.
+if [[ -z "${BASH_VERSINFO-}" || BASH_VERSINFO[0] -lt 3 || BASH_VERSINFO[0] -eq 3 && BASH_VERSINFO[1] -lt 1 ]]; then
+	return 1
+fi
+
+# Avoid duplicate inclusion
+if [[ -n "${bash_preexec_imported:-}" ]]; then
+	return 0
+fi
+bash_preexec_imported="defined"
+
+# WARNING: This variable is no longer used and should not be relied upon.
+# Use ${bash_preexec_imported} instead.
+# shellcheck disable=SC2034
+__bp_imported="${bash_preexec_imported}"
+
+# Should be available to each precmd and preexec
+# functions, should they want it. $? and $_ are available as $? and $_, but
+# $PIPESTATUS is available only in a copy, $BP_PIPESTATUS.
+# TODO: Figure out how to restore PIPESTATUS before each precmd or preexec
+# function.
+__bp_last_ret_value="$?"
+BP_PIPESTATUS=("${PIPESTATUS[@]}")
+__bp_last_argument_prev_command="$_"
+
+__bp_inside_precmd=0
+__bp_inside_preexec=0
+
+# Initial PROMPT_COMMAND string that is removed from PROMPT_COMMAND post __bp_install
+__bp_install_string=$'__bp_trap_string="$(trap -p DEBUG)"\ntrap - DEBUG\n__bp_install'
+
+# Fails if any of the given variables are readonly
+# Reference https://stackoverflow.com/a/4441178
+__bp_require_not_readonly() {
+	local var
+	for var; do
+		if ! (unset "$var" 2>/dev/null); then
+			echo "bash-preexec requires write access to ${var}" >&2
+			return 1
+		fi
+	done
+}
+
+# Remove ignorespace and or replace ignoreboth from HISTCONTROL
+# so we can accurately invoke preexec with a command from our
+# history even if it starts with a space.
+__bp_adjust_histcontrol() {
+	local histcontrol
+	histcontrol="${HISTCONTROL:-}"
+	histcontrol="${histcontrol//ignorespace/}"
+	# Replace ignoreboth with ignoredups
+	if [[ "$histcontrol" == *"ignoreboth"* ]]; then
+		histcontrol="ignoredups:${histcontrol//ignoreboth/}"
+	fi
+	export HISTCONTROL="$histcontrol"
+}
+
+# This variable describes whether we are currently in "interactive mode";
+# i.e. whether this shell has just executed a prompt and is waiting for user
+# input.  It documents whether the current command invoked by the trace hook is
+# run interactively by the user; it's set immediately after the prompt hook,
+# and unset as soon as the trace hook is run.
+__bp_preexec_interactive_mode=""
+
+# These arrays are used to add functions to be run before, or after, prompts.
+declare -a precmd_functions
+declare -a preexec_functions
+
+# Trims leading and trailing whitespace from $2 and writes it to the variable
+# name passed as $1
+__bp_trim_whitespace() {
+	local var=${1:?} text=${2:-}
+	text="${text#"${text%%[![:space:]]*}"}" # remove leading whitespace characters
+	text="${text%"${text##*[![:space:]]}"}" # remove trailing whitespace characters
+	printf -v "$var" '%s' "$text"
+}
+
+# Trims whitespace and removes any leading or trailing semicolons from $2 and
+# writes the resulting string to the variable name passed as $1. Used for
+# manipulating substrings in PROMPT_COMMAND
+__bp_sanitize_string() {
+	local var=${1:?} text=${2:-} sanitized
+	__bp_trim_whitespace sanitized "$text"
+	sanitized=${sanitized%;}
+	sanitized=${sanitized#;}
+	__bp_trim_whitespace sanitized "$sanitized"
+	printf -v "$var" '%s' "$sanitized"
+}
+
+# This function is installed as part of the PROMPT_COMMAND;
+# It sets a variable to indicate that the prompt was just displayed,
+# to allow the DEBUG trap to know that the next command is likely interactive.
+__bp_interactive_mode() {
+	__bp_preexec_interactive_mode="on"
+}
+
+# This function is installed as part of the PROMPT_COMMAND.
+# It will invoke any functions defined in the precmd_functions array.
+__bp_precmd_invoke_cmd() {
+	# Save the returned value from our last command, and from each process in
+	# its pipeline. Note: this MUST be the first thing done in this function.
+	# BP_PIPESTATUS may be unused, ignore
+	# shellcheck disable=SC2034
+
+	__bp_last_ret_value="$?" BP_PIPESTATUS=("${PIPESTATUS[@]}")
+
+	# Don't invoke precmds if we are inside an execution of an "original
+	# prompt command" by another precmd execution loop. This avoids infinite
+	# recursion.
+	if ((__bp_inside_precmd > 0)); then
+		return
+	fi
+	local __bp_inside_precmd=1
+
+	# Invoke every function defined in our function array.
+	local precmd_function
+	for precmd_function in "${precmd_functions[@]}"; do
+
+		# Only execute this function if it actually exists.
+		# Test existence of functions with: declare -[Ff]
+		if type -t "$precmd_function" 1>/dev/null; then
+			__bp_set_ret_value "$__bp_last_ret_value" "$__bp_last_argument_prev_command"
+			# Quote our function invocation to prevent issues with IFS
+			"$precmd_function"
+		fi
+	done
+
+	__bp_set_ret_value "$__bp_last_ret_value"
+}
+
+# Sets a return value in $?. We may want to get access to the $? variable in our
+# precmd functions. This is available for instance in zsh. We can simulate it in bash
+# by setting the value here.
+__bp_set_ret_value() {
+	return ${1:+"$1"}
+}
+
+__bp_in_prompt_command() {
+
+	local prompt_command_array
+	IFS=$'\n;' read -rd '' -a prompt_command_array <<<"${PROMPT_COMMAND:-}"
+
+	local trimmed_arg
+	__bp_trim_whitespace trimmed_arg "${1:-}"
+
+	local command trimmed_command
+	for command in "${prompt_command_array[@]:-}"; do
+		__bp_trim_whitespace trimmed_command "$command"
+		if [[ "$trimmed_command" == "$trimmed_arg" ]]; then
+			return 0
+		fi
+	done
+
+	return 1
+}
+
+# This function is installed as the DEBUG trap.  It is invoked before each
+# interactive prompt display.  Its purpose is to inspect the current
+# environment to attempt to detect if the current command is being invoked
+# interactively, and invoke 'preexec' if so.
+__bp_preexec_invoke_exec() {
+
+	# Save the contents of $_ so that it can be restored later on.
+	# https://stackoverflow.com/questions/40944532/bash-preserve-in-a-debug-trap#40944702
+	__bp_last_argument_prev_command="${1:-}"
+	# Don't invoke preexecs if we are inside of another preexec.
+	if ((__bp_inside_preexec > 0)); then
+		return
+	fi
+	local __bp_inside_preexec=1
+
+	# Checks if the file descriptor is not standard out (i.e. '1')
+	# __bp_delay_install checks if we're in test. Needed for bats to run.
+	# Prevents preexec from being invoked for functions in PS1
+	if [[ ! -t 1 && -z "${__bp_delay_install:-}" ]]; then
+		return
+	fi
+
+	if [[ -n "${COMP_LINE:-}" ]]; then
+		# We're in the middle of a completer. This obviously can't be
+		# an interactively issued command.
+		return
+	fi
+	if [[ -z "${__bp_preexec_interactive_mode:-}" ]]; then
+		# We're doing something related to displaying the prompt.  Let the
+		# prompt set the title instead of me.
+		return
+	else
+		# If we're in a subshell, then the prompt won't be re-displayed to put
+		# us back into interactive mode, so let's not set the variable back.
+		# In other words, if you have a subshell like
+		#   (sleep 1; sleep 2)
+		# You want to see the 'sleep 2' as a set_command_title as well.
+		if [[ 0 -eq "${BASH_SUBSHELL:-}" ]]; then
+			__bp_preexec_interactive_mode=""
+		fi
+	fi
+
+	if __bp_in_prompt_command "${BASH_COMMAND:-}"; then
+		# If we're executing something inside our prompt_command then we don't
+		# want to call preexec. Bash prior to 3.1 can't detect this at all :/
+		__bp_preexec_interactive_mode=""
+		return
+	fi
+
+	local this_command
+	this_command=$(
+		export LC_ALL=C
+		HISTTIMEFORMAT='' builtin history 1 | sed '1 s/^ *[0-9][0-9]*[* ] //'
+	)
+
+	# Sanity check to make sure we have something to invoke our function with.
+	if [[ -z "$this_command" ]]; then
+		return
+	fi
+
+	# Invoke every function defined in our function array.
+	local preexec_function
+	local preexec_function_ret_value
+	local preexec_ret_value=0
+	for preexec_function in "${preexec_functions[@]:-}"; do
+
+		# Only execute each function if it actually exists.
+		# Test existence of function with: declare -[fF]
+		if type -t "$preexec_function" 1>/dev/null; then
+			__bp_set_ret_value "${__bp_last_ret_value:-}"
+			# Quote our function invocation to prevent issues with IFS
+			"$preexec_function" "$this_command"
+			preexec_function_ret_value="$?"
+			if [[ "$preexec_function_ret_value" != 0 ]]; then
+				preexec_ret_value="$preexec_function_ret_value"
+			fi
+		fi
+	done
+
+	# Restore the last argument of the last executed command, and set the return
+	# value of the DEBUG trap to be the return code of the last preexec function
+	# to return an error.
+	# If `extdebug` is enabled a non-zero return value from any preexec function
+	# will cause the user's command not to execute.
+	# Run `shopt -s extdebug` to enable
+	__bp_set_ret_value "$preexec_ret_value" "$__bp_last_argument_prev_command"
+}
+
+__bp_install() {
+	# Exit if we already have this installed.
+	if [[ "${PROMPT_COMMAND:-}" == *"__bp_precmd_invoke_cmd"* ]]; then
+		return 1
+	fi
+
+	trap '__bp_preexec_invoke_exec "$_"' DEBUG
+
+	# Preserve any prior DEBUG trap as a preexec function
+	local prior_trap
+	# we can't easily do this with variable expansion. Leaving as sed command.
+	# shellcheck disable=SC2001
+	prior_trap=$(sed "s/[^']*'\(.*\)'[^']*/\1/" <<<"${__bp_trap_string:-}")
+	unset __bp_trap_string
+	if [[ -n "$prior_trap" ]]; then
+		eval '__bp_original_debug_trap() {
+          '"$prior_trap"'
+        }'
+		preexec_functions+=(__bp_original_debug_trap)
+	fi
+
+	# Adjust our HISTCONTROL Variable if needed.
+	__bp_adjust_histcontrol
+
+	# Issue #25. Setting debug trap for subshells causes sessions to exit for
+	# backgrounded subshell commands (e.g. (pwd)& ). Believe this is a bug in Bash.
+	#
+	# Disabling this by default. It can be enabled by setting this variable.
+	if [[ -n "${__bp_enable_subshells:-}" ]]; then
+
+		# Set so debug trap will work be invoked in subshells.
+		set -o functrace >/dev/null 2>&1
+		shopt -s extdebug >/dev/null 2>&1
+	fi
+
+	local existing_prompt_command
+	# Remove setting our trap install string and sanitize the existing prompt command string
+	existing_prompt_command="${PROMPT_COMMAND:-}"
+	existing_prompt_command="${existing_prompt_command//${__bp_install_string}[;$'\n']/}" # Edge case of appending to PROMPT_COMMAND
+	existing_prompt_command="${existing_prompt_command//$__bp_install_string/}"
+	__bp_sanitize_string existing_prompt_command "$existing_prompt_command"
+
+	# Install our hooks in PROMPT_COMMAND to allow our trap to know when we've
+	# actually entered something.
+	PROMPT_COMMAND=$'__bp_precmd_invoke_cmd\n'
+	if [[ -n "$existing_prompt_command" ]]; then
+		PROMPT_COMMAND+=${existing_prompt_command}$'\n'
+	fi
+	PROMPT_COMMAND+='__bp_interactive_mode'
+
+	# Add two functions to our arrays for convenience
+	# of definition.
+	precmd_functions+=(precmd)
+	preexec_functions+=(preexec)
+
+	# Invoke our two functions manually that were added to $PROMPT_COMMAND
+	__bp_precmd_invoke_cmd
+	__bp_interactive_mode
+}
+
+# Sets an installation string as part of our PROMPT_COMMAND to install
+# after our session has started. This allows bash-preexec to be included
+# at any point in our bash profile.
+__bp_install_after_session_init() {
+	# bash-preexec needs to modify these variables in order to work correctly
+	# if it can't, just stop the installation
+	__bp_require_not_readonly PROMPT_COMMAND HISTCONTROL HISTTIMEFORMAT || return
+
+	local sanitized_prompt_command
+	__bp_sanitize_string sanitized_prompt_command "${PROMPT_COMMAND:-}"
+	if [[ -n "$sanitized_prompt_command" ]]; then
+		PROMPT_COMMAND=${sanitized_prompt_command}$'\n'
+	fi
+	PROMPT_COMMAND+=${__bp_install_string}
+}
+
+# Run our install so long as we're not delaying it.
+if [[ -z "${__bp_delay_install:-}" ]]; then
+	__bp_install_after_session_init
+fi

anyscale/webterminal/command_persister.py

@@ -0,0 +1,164 @@
+import asyncio
+from collections import defaultdict
+import json
+import os
+import time
+from typing import Dict, List
+
+from anyscale.authenticate import get_auth_api_client
+from anyscale.client.openapi_client.api.default_api import DefaultApi
+from anyscale.client.openapi_client.models.session_command_types import (
+    SessionCommandTypes,
+)
+from anyscale.commands.session_commands_hidden import session_upload_command_logs_impl
+from anyscale.webterminal.utils import Command
+
+
+KILLED_STATUS_CODE = 130
+MINUTE_IN_SECONDS = 60
+IDLE_TERMINATION_DIR = os.environ.get(
+    "IDLE_TERMINATION_DIR", "/tmp/anyscale/idle_termination_reports"
+)
+
+
+class CommandPersister:
+    """
+        CommandPersister manages persisting the
+        state of all commands to our backend.
+        It functions as a queue.
+    """
+
+    def __init__(self, cli_token: str, host: str, session_id: str):
+        self.api_client: DefaultApi = get_auth_api_client(
+            cli_token=cli_token, host=host
+        ).api_client
+        self.cli_token = cli_token
+        self.host = host
+        self.session_id = session_id
+        self.queue: List[Command] = []
+        self.running_commands: Dict[str, Command] = {}
+        self.last_uploaded_time: Dict[str, float] = defaultdict(float)
+        self.is_active = True
+        # on first load, report an activity to signal healthy start
+        update_idle_termination_status(time.time())
+
+    def shutdown(self) -> None:
+        self.is_active = False
+
+    async def persist_commands(self) -> None:  # noqa: PLR0912
+        while self.is_active:
+            for session_command in self.running_commands.values():
+                self.upload_command_log(session_command)
+            if len(self.queue) == 0:
+                await asyncio.sleep(1)
+                continue
+            cmd = self.queue.pop(0)
+            try:
+                if cmd.exec_command:
+                    # Mark exec commands as running by default so that we do not try to create this command on the
+                    # backend ( it is already created).
+                    # When the exec command is finished, command_persister will treat it like other WT commands.
+                    self.running_commands[cmd.scid] = cmd
+                if cmd.scid not in self.running_commands:
+                    if cmd.finished_at:
+                        update_idle_termination_status(cmd.finished_at.timestamp())
+                    self.api_client.create_session_command_api_v2_sessions_session_id_create_session_command_post(
+                        session_id=self.session_id,
+                        external_terminal_command={
+                            "scid": cmd.scid,
+                            "command": cmd.command,
+                            "created_at": cmd.created_at,
+                            "finished_at": cmd.finished_at,
+                            "status_code": cmd.status_code,
+                            "web_terminal_tab_id": cmd.term_name,
+                            "type": SessionCommandTypes.WEBTERMINAL,
+                        },
+                    )
+                    if not cmd.finished:
+                        self.running_commands[cmd.scid] = cmd
+                elif cmd.finished:
+                    if cmd.finished_at:
+                        update_idle_termination_status(cmd.finished_at.timestamp())
+                    self.api_client.finish_session_command_api_v2_session_commands_session_command_id_finish_post(
+                        session_command_id=cmd.scid,
+                        session_command_finish_options={
+                            "status_code": cmd.status_code,
+                            "stop": False,
+                            "terminate": False,
+                            "finished_at": cmd.finished_at,
+                            "killed_at": cmd.finished_at
+                            if cmd.status_code == KILLED_STATUS_CODE
+                            else None,
+                        },
+                    )
+                    del self.running_commands[cmd.scid]
+                self.upload_command_log(cmd)
+                if cmd.finished:
+                    del self.last_uploaded_time[cmd.scid]
+            except Exception as e:  # noqa: BLE001
+                with open(f"/tmp/{cmd.scid}.err", "a") as err_file:
+                    err_file.write(f"{e!s}\n")
+            if len(self.queue) > 0:
+                # If the queue is not empty continue faster through it.
+                await asyncio.sleep(0.5)
+            else:
+                await asyncio.sleep(1)
+
+    def upload_command_log(self, session_command: Command) -> None:
+        """
+        We upload logs to S3 if any of the below conditions is true:
+        1. Session command terminates
+        2. Session command runs the first time
+        3. The duration from the last uploaded time is larger than the specified period
+        """
+        current_time_in_s = time.time()
+        if (
+            session_command.finished
+            or current_time_in_s - self.last_uploaded_time[session_command.scid]
+            > MINUTE_IN_SECONDS * 2
+        ):
+            session_upload_command_logs_impl(
+                command_id=session_command.scid,
+                cli_token=self.cli_token,
+                host=self.host,
+            )
+            # We store the timestamp in which logs have been successfully
+            # uploaded to avoid re-uploading large logs that take some time
+            # to upload.
+            finish_upload_time_in_s = time.time()
+            self.last_uploaded_time[session_command.scid] = finish_upload_time_in_s
+
+    def enqueue_command(self, cmd: Command) -> None:
+        """
+            enqueue_command adds a Command object to be persisted to our backend.
+            Only enqueue the command to be peristed if it is
+                - just created (therefore not in self.running_commands)
+                  we need to create the entry on the backend.
+                - finished
+                  we need to update the entry on the backend.
+        """
+        if cmd.scid not in self.running_commands or cmd.finished:
+            self.queue.append(cmd)
+
+    def kill_running_commands_for_terminal(self, term_name: str) -> None:
+        """
+            When we exit a terminal tab, the zsh process is killed and
+            all of its children process are also killed. This function then
+            enqueues them to be marked as finished.
+        """
+        commands = list(self.running_commands.values())
+        for c in commands:
+            if c.term_name == term_name:
+                c.finish(KILLED_STATUS_CODE)
+                self.enqueue_command(c)
+
+
+def update_idle_termination_status(timestamp):
+    """
+        Updates the idle termination status for the web-terminal by storing the last
+        activity date as the last finished time of a successful command.
+    """
+    if not os.path.exists(IDLE_TERMINATION_DIR) or not timestamp:
+        return
+    with open(f"{IDLE_TERMINATION_DIR}/web_terminal.json", "w") as f:
+        f.write(json.dumps({"last_activity_timestamp": timestamp}))