Consolidate periodic loops into one function updating Den and updatin…

…g autostop. (#873)

autostop.

runhouse/constants.py

@@ -67,6 +67,6 @@
 BULLET_UNICODE = "\u2022"
 MINUTE = 60
 HOUR = 3600
-DEFAULT_STATUS_CHECK_INTERVAL = 2 * MINUTE
+DEFAULT_STATUS_CHECK_INTERVAL = 1 * MINUTE
 INCREASED_STATUS_CHECK_INTERVAL = 1 * HOUR
-STATUS_CHECK_DELAY = 2 * MINUTE
+STATUS_CHECK_DELAY = 1 * MINUTE

runhouse/rns/rns_client.py

@@ -5,11 +5,13 @@
 import os
 import shutil
 from pathlib import Path
-from typing import Dict, Optional, Set, Union
+from typing import Any, Dict, List, Optional, Set, Union
 
 import dotenv
+import httpx
 
 import requests
+from pydantic import BaseModel
 
 from runhouse.rns.utils.api import (
  generate_uuid,
@@ -21,6 +23,17 @@
 
 logger = logging.getLogger(__name__)
 
+# This is a copy of the Pydantic model that we use to validate in Den
+class ResourceStatusData(BaseModel):
+ cluster_config: dict
+ env_resource_mapping: Dict[str, List[Dict[str, Any]]]
+ system_cpu_usage: float
+ system_memory_usage: Dict[str, Any]
+ system_disk_usage: Dict[str, Any]
+ env_servlet_processes: Dict[str, Dict[str, Any]]
+ server_pid: int
+ runhouse_version: str
+
 
 class RNSClient:
  """Manage a particular resource with the runhouse database"""
@@ -659,3 +672,26 @@ def contents(self, name_or_path, full_paths):
  return folder(name=name_or_path, path=folder_url).resources(
  full_paths=full_paths
  )
+
+ async def send_status(self, status: ResourceStatusData, cluster_rns_address: str):
+ status_data = {
+ "status": "running",
+ "resource_type": status.cluster_config.get("resource_type"),
+ "data": dict(status),
+ }
+ cluster_uri = self.format_rns_address(cluster_rns_address)
+ api_server_url = status.cluster_config.get(
+ "api_server_url", self.api_server_url
+ )
+ client = httpx.AsyncClient()
+ resp = await client.post(
+ f"{api_server_url}/resource/{cluster_uri}/cluster/status",
+ data=json.dumps(status_data),
+ headers=self.request_headers(),
+ )
+ if resp.status_code != 200:
+ logger.error(
+ f"Received [{resp.status_code}]: Failed to send cluster status info to Den: {resp.text}."
+ )
+ else:
+ logger.info("Successfully sent cluster status info to Den.")

runhouse/servers/autostop_servlet.py

@@ -0,0 +1,32 @@
+import time
+
+
+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
+
+ 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)
+
+ def update_autostop_in_sky_config(self):
+ import pickle
+
+ from sky.skylet import configs as sky_configs
+
+ autostop_mins = pickle.loads(
+ sky_configs.get_config("autostop_config")
+ ).autostop_idle_minutes
+ if autostop_mins > 0 and (
+ self._last_register is None or self._last_register < self._last_activity
+ ):
+ sky_configs.set_config("autostop_last_active_time", self._last_activity)
+ self._last_register = self._last_activity

runhouse/servers/cluster_servlet.py

@@ -3,13 +3,9 @@
 import json
 import logging
 import threading
-import time
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set, Union
 
-import requests
-from pydantic import BaseModel
-
 import runhouse
 
 from runhouse.constants import (
@@ -21,7 +17,9 @@
 
 from runhouse.globals import configs, obj_store, rns_client
 from runhouse.resources.hardware import load_cluster_config_from_file
+from runhouse.rns.rns_client import ResourceStatusData
 from runhouse.rns.utils.api import ResourceAccess
+from runhouse.servers.autostop_servlet import AutostopServlet
 from runhouse.servers.http.auth import AuthCache
 
 from runhouse.utils import sync_function
@@ -33,18 +31,6 @@ class ClusterServletError(Exception):
  pass
 
 
-# This is a copy of the Pydantic model that we use to validate in Den
-class ResourceStatusData(BaseModel):
- cluster_config: dict
- env_resource_mapping: Dict[str, List[Dict[str, Any]]]
- system_cpu_usage: float
- system_memory_usage: Dict[str, Any]
- system_disk_usage: Dict[str, Any]
- env_servlet_processes: Dict[str, Dict[str, Any]]
- server_pid: int
- runhouse_version: str
-
-
 class ClusterServlet:
  async def __init__(
  self, cluster_config: Optional[Dict[str, Any]] = None, *args, **kwargs
@@ -85,10 +71,10 @@ async def __init__(
 
  # 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) and configs.token:
- logger.debug("Creating send_status_info_to_den thread.")
+ if self.cluster_config.get("den_auth", False):
+ logger.info("Creating periodic_status_check thread.")
  post_status_thread = threading.Thread(
- target=self.send_status_info_to_den, daemon=True
+ target=self.periodic_status_check, daemon=True
  )
  post_status_thread.start()
 
@@ -260,63 +246,43 @@ async def aclear_all_references_to_env_servlet_name(self, env_servlet_name: str)
  # Cluster status functions
  ##############################################
 
- async def asend_status_info_to_den(self):
+ async def aperiodic_status_check(self):
  # Delay the start of post_status_thread, so we'll finish the cluster startup properly
  await asyncio.sleep(STATUS_CHECK_DELAY)
  while True:
- logger.info("Trying to send cluster status to Den.")
  try:
- is_config_updated = (
- await self.aupdate_status_check_interval_in_cluster_config()
- )
- interval_size = (await self.aget_cluster_config()).get(
+ await self.aupdate_status_check_interval_in_cluster_config()
+
+ cluster_config = await self.aget_cluster_config()
+ interval_size = cluster_config.get(
  "status_check_interval", DEFAULT_STATUS_CHECK_INTERVAL
  )
- den_auth = (await self.aget_cluster_config()).get(
- "den_auth", DEFAULT_STATUS_CHECK_INTERVAL
- )
- if interval_size == -1:
- if is_config_updated:
- logger.info(
- f"Disabled periodic cluster status check. For enabling it, please run "
- f"cluster.restart_server(). If you want to set up an interval size that is not the "
- f"default value {round(DEFAULT_STATUS_CHECK_INTERVAL/60,2)} please run "
- f"cluster._enable_or_update_status_check(interval_size) after restarting the server."
- )
- break
- if not den_auth:
- logger.info(
- f"Disabled periodic cluster status check because den_auth is disabled. For enabling it, please run "
- f"cluster.restart_server() and make sure that den_auth is enabled. If you want to set up an interval size that is not the "
- f"default value {round(DEFAULT_STATUS_CHECK_INTERVAL / 60, 2)} please run "
- f"cluster._enable_or_update_status_check(interval_size) after restarting the server."
- )
- break
- status: ResourceStatusData = await self.astatus()
- status_data = {
- "status": "running",
- "resource_type": status.cluster_config.get("resource_type"),
- "data": dict(status),
- }
- cluster_uri = rns_client.format_rns_address(
- (await self.aget_cluster_config()).get("name")
- )
- api_server_url = status.cluster_config.get(
- "api_server_url", rns_client.api_server_url
- )
- post_status_data_resp = requests.post(
- f"{api_server_url}/resource/{cluster_uri}/cluster/status",
- data=json.dumps(status_data),
- headers=rns_client.request_headers(),
- )
- if post_status_data_resp.status_code != 200:
- logger.error(
- f"({post_status_data_resp.status_code}) Failed to send cluster status check to Den: {post_status_data_resp.text}"
- )
- else:
+ den_auth = cluster_config.get("den_auth", False)
+
+ # Only if one of these is true, do we actually need to get the status from each EnvServlet
+ should_send_status_to_den = den_auth and interval_size != -1
+ should_update_autostop = self.autostop_servlet is not None
+ if should_send_status_to_den or should_update_autostop:
  logger.info(
- f"Successfully updated cluster status in Den. Next status check will be in {round(interval_size / 60, 2)} minutes."
+ "Performing cluster status check: potentially sending to Den or updating autostop."
  )
+ status: ResourceStatusData = await self.astatus()
+
+ if should_update_autostop:
+ function_running = any(
+ any(
+ len(resource_info["active_function_calls"]) > 0
+ for resource_info in resources
+ )
+ for resources in status.env_resource_mapping.values()
+ )
+ if function_running:
+ await self.autostop_servlet.set_last_active_time_to_now.remote()
+ await self.autostop_servlet.update_autostop_in_sky_config.remote()
+
+ if should_send_status_to_den:
+ cluster_rns_address = cluster_config.get("name")
+ await rns_client.send_status(status, cluster_rns_address)
  except Exception as e:
  logger.error(
  f"Cluster status check has failed: {e}. Please check cluster logs for more info."
@@ -328,11 +294,13 @@ async def asend_status_info_to_den(self):
  f"If a value is not provided, interval size will be set to {DEFAULT_STATUS_CHECK_INTERVAL}"
  )
  await asyncio.sleep(INCREASED_STATUS_CHECK_INTERVAL)
- finally:
+ else:
  await asyncio.sleep(interval_size)
 
- def send_status_info_to_den(self):
- asyncio.run(self.asend_status_info_to_den())
+ def periodic_status_check(self):
+ # This is only ever called once in its own thread, so we can do asyncio.run here instead of
+ # sync_function.
+ asyncio.run(self.aperiodic_status_check())
 
  async def _status_for_env_servlet(self, env_servlet_name):
  try:
@@ -420,48 +388,3 @@ async def astatus(self):
 
  def status(self):
  return sync_function(self.astatus)()
-
-
-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
- autostop_thread = threading.Thread(target=self.update_autostop, daemon=True)
- autostop_thread.start()
-
- 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)
-
- def update_autostop(self):
- import pickle
-
- from sky.skylet import configs as sky_configs
-
- while True:
-
- autostop_mins = pickle.loads(
- sky_configs.get_config("autostop_config")
- ).autostop_idle_minutes
- self._last_register = float(
- sky_configs.get_config("autostop_last_active_time")
- )
- if autostop_mins > 0 and (
- not self._last_register
- or (
- # within 2 min of autostop and there's more recent activity
- 60 * autostop_mins - (time.time() - self._last_register) < 120
- and self._last_activity > self._last_register
- )
- ):
- sky_configs.set_config("autostop_last_active_time", self._last_activity)
- self._last_register = self._last_activity
-
- time.sleep(30)