Skip to content

Commit

Permalink
Simplify Autostop More
Browse files Browse the repository at this point in the history
  • Loading branch information
dongreenberg committed Jun 15, 2024
1 parent 0943f0a commit 86b3c0a
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 24 deletions.
35 changes: 20 additions & 15 deletions runhouse/servers/autostop_servlet.py
Original file line number Diff line number Diff line change
@@ -1,25 +1,30 @@
import time
import logging
import subprocess

logger = logging.getLogger(__name__)


class AutostopServlet:
"""A helper class strictly to run SkyPilot methods on OnDemandClusters inside SkyPilot's conda env."""

def __init__(self):
self._last_activity = time.time()
self._last_register = None
self._activity_registered = False

def set_last_active_time_to_now(self):
self._last_activity = time.time()

def set_autostop(self, value=None):
from sky.skylet import autostop_lib

self.set_last_active_time_to_now()
autostop_lib.set_autostop(value, None, True)
self._activity_registered = True

def update_autostop_in_sky_config(self):
from sky.skylet.autostop_lib import set_last_active_time_to_now

if self._last_register is None or self._last_register < self._last_activity:
set_last_active_time_to_now()
self._last_register = self._last_activity
SKY_VENV = "~/skypilot-runtime"
SKY_AUTOSTOP_CMD = "from sky.skylet.autostop_lib import set_last_active_time_to_now; set_last_active_time_to_now()"
SKY_CMD = f"{SKY_VENV}/bin/python -c '{SKY_AUTOSTOP_CMD}'"

if self._activity_registered:
logger.debug(
"Activity registered, updating last active time in SkyConfig with command: {SKY_CMD}"
)
subprocess.run(SKY_CMD, shell=True, check=True)
self._activity_registered = False
else:
logger.debug(
"No activity registered, not updating last active time in SkyConfig"
)
14 changes: 5 additions & 9 deletions runhouse/servers/cluster_servlet.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,19 +64,15 @@ async def __init__(
max_concurrency=1000,
resources={f"node:{current_ip}": 0.001},
num_cpus=0,
runtime_env={"env_vars": {"VIRTUAL_ENV": "skypilot-runtime"}},
)
.remote()
)

# Only send for clusters that have den_auth enabled and if we are logged in with a user's token
# to authenticate the request
if self.cluster_config.get("den_auth", False):
logger.info("Creating periodic_status_check thread.")
post_status_thread = threading.Thread(
target=self.periodic_status_check, daemon=True
)
post_status_thread.start()
logger.info("Creating periodic_status_check thread.")
post_status_thread = threading.Thread(
target=self.periodic_status_check, daemon=True
)
post_status_thread.start()

##############################################
# Cluster config state storage methods
Expand Down

0 comments on commit 86b3c0a

Please sign in to comment.