def get(self, category, key): if _internal_kv_initialized(): value = _internal_kv_get(_make_key(category, key)) if value is None: raise ValueError( "Registry value for {}/{} doesn't exist.".format( category, key)) return pickle.loads(value) else: return pickle.loads(self._to_flush[(category, key)])
def _wait_for_reader(self): """Checks for backpressure by the downstream reader.""" if self.max_size <= 0: # Unlimited queue return if self.write_item_offset - self.cached_remote_offset <= self.max_size: return # Hasn't reached max size remote_offset = internal_kv._internal_kv_get(self.read_ack_key) if remote_offset is None: # logger.debug("[writer] Waiting for reader to start...") while remote_offset is None: time.sleep(0.01) remote_offset = internal_kv._internal_kv_get(self.read_ack_key) remote_offset = int(remote_offset) if self.write_item_offset - remote_offset > self.max_size: logger.debug( "[writer] Waiting for reader to catch up {} to {} - {}".format( remote_offset, self.write_item_offset, self.max_size)) while self.write_item_offset - remote_offset > self.max_size: time.sleep(0.01) remote_offset = int( internal_kv._internal_kv_get(self.read_ack_key)) self.cached_remote_offset = remote_offset
def get_actor(name): """Get a named actor which was previously created. If the actor doesn't exist, an exception will be raised. Args: name: The name of the named actor. Returns: The ActorHandle object corresponding to the name. """ actor_name = _calculate_key(name) pickled_state = _internal_kv_get(actor_name) if pickled_state is None: raise ValueError("The actor with name={} doesn't exist".format(name)) handle = pickle.loads(pickled_state) return handle
def fetch_package(pkg_uri: str, pkg_file: Path) -> int: """Fetch a package from a given uri. This function is used to fetch a pacakge from the given uri to local filesystem. Args: pkg_uri (str): The uri of the package to download. pkg_file (pathlib.Path): The path in local filesystem to download the package. Returns: The number of bytes downloaded. """ (protocol, pkg_name) = _parse_uri(pkg_uri) if protocol in (Protocol.GCS, Protocol.PIN_GCS): code = _internal_kv_get(pkg_uri) code = code or b"" pkg_file.write_bytes(code) return len(code) else: raise NotImplementedError(f"Protocol {protocol} is not supported")
async def _perform_iteration(self, publisher): """Get any changes to the log files and push updates to kv.""" while True: try: formatted_status_string = internal_kv._internal_kv_get( DEBUG_AUTOSCALING_STATUS) stats = self._get_all_stats() # Report stats only when metrics collection is enabled. if not self._metrics_collection_disabled: cluster_stats = (json.loads( formatted_status_string.decode()) if formatted_status_string else {}) records_reported = self._record_stats(stats, cluster_stats) self._metrics_agent.record_reporter_stats(records_reported) await publisher.publish_resource_usage(self._key, jsonify_asdict(stats)) except Exception: logger.exception("Error publishing node physical stats.") await asyncio.sleep(reporter_consts.REPORTER_UPDATE_INTERVAL_MS / 1000)
async def get_serve_info(self) -> Dict[str, Any]: # Conditionally import serve to prevent ModuleNotFoundError from serve # dependencies when only ray[default] is installed (#17712) try: from ray.serve.controller import SNAPSHOT_KEY as SERVE_SNAPSHOT_KEY from ray.serve.constants import SERVE_CONTROLLER_NAME except Exception: return {} # Serve wraps Ray's internal KV store and specially formats the keys. # These are the keys we are interested in: # SERVE_CONTROLLER_NAME(+ optional random letters):SERVE_SNAPSHOT_KEY serve_keys = _internal_kv_list( SERVE_CONTROLLER_NAME, namespace=ray_constants.KV_NAMESPACE_SERVE) serve_snapshot_keys = filter(lambda k: SERVE_SNAPSHOT_KEY in str(k), serve_keys) deployments_per_controller: List[Dict[str, Any]] = [] for key in serve_snapshot_keys: val_bytes = _internal_kv_get( key, namespace=ray_constants.KV_NAMESPACE_SERVE ) or "{}".encode("utf-8") deployments_per_controller.append( json.loads(val_bytes.decode("utf-8"))) # Merge the deployments dicts of all controllers. deployments: Dict[str, Any] = { k: v for d in deployments_per_controller for k, v in d.items() } # Replace the keys (deployment names) with their hashes to prevent # collisions caused by the automatic conversion to camelcase by the # dashboard agent. deployments = { hashlib.sha1(name.encode()).hexdigest(): info for name, info in deployments.items() } return deployments
async def _connect_to_dashboard(self): """ Connect to the dashboard. If the dashboard is not started, then this method will never returns. Returns: The ReportEventServiceStub object. """ while True: try: # TODO: Use async version if performance is an issue dashboard_rpc_address = internal_kv._internal_kv_get( dashboard_consts.DASHBOARD_RPC_ADDRESS, namespace=ray_constants.KV_NAMESPACE_DASHBOARD) if dashboard_rpc_address: logger.info("Report events to %s", dashboard_rpc_address) options = (("grpc.enable_http_proxy", 0), ) channel = utils.init_grpc_channel(dashboard_rpc_address, options=options, asynchronous=True) return event_pb2_grpc.ReportEventServiceStub(channel) except Exception: logger.exception("Connect to dashboard failed.") await asyncio.sleep( event_consts.RETRY_CONNECT_TO_DASHBOARD_INTERVAL_SECONDS)
def _get(self, key): return self._deserialize( ray_kv._internal_kv_get(self._format_key(self._serialize(key))))
def contains(self, category, key): if _internal_kv_initialized(): value = _internal_kv_get(_make_key(category, key)) return value is not None else: return (category, key) in self._to_flush
def get(self, key: str) -> bytes: ret = internal_kv._internal_kv_get(key) return ret
def download_and_unpack_package( pkg_uri: str, base_directory: str, logger: Optional[logging.Logger] = default_logger, ) -> str: """Download the package corresponding to this URI and unpack it. Will be written to a directory named {base_directory}/{uri}. """ pkg_file = Path(_get_local_path(base_directory, pkg_uri)) with FileLock(str(pkg_file) + ".lock"): if logger is None: logger = default_logger logger.debug(f"Fetching package for URI: {pkg_uri}") local_dir = pkg_file.with_suffix("") assert local_dir != pkg_file, "Invalid pkg_file!" if local_dir.exists(): assert local_dir.is_dir(), f"{local_dir} is not a directory" else: protocol, pkg_name = parse_uri(pkg_uri) if protocol == Protocol.GCS: # Download package from the GCS. code = _internal_kv_get(pkg_uri) if code is None: raise IOError(f"Failed to fetch URI {pkg_uri} from GCS.") code = code or b"" pkg_file.write_bytes(code) unzip_package( package_path=pkg_file, target_dir=local_dir, remove_top_level_directory=False, unlink_zip=True, logger=logger) elif protocol in Protocol.remote_protocols(): # Download package from remote URI tp = None if protocol == Protocol.S3: try: from smart_open import open import boto3 except ImportError: raise ImportError( "You must `pip install smart_open` and " "`pip install boto3` to fetch URIs in s3 " "bucket.") tp = {"client": boto3.client("s3")} elif protocol == Protocol.GS: try: from smart_open import open from google.cloud import storage # noqa: F401 except ImportError: raise ImportError( "You must `pip install smart_open` and " "`pip install google-cloud-storage` " "to fetch URIs in Google Cloud Storage bucket.") else: try: from smart_open import open except ImportError: raise ImportError( "You must `pip install smart_open` " f"to fetch {protocol.value.upper()} URIs.") with open(pkg_uri, "rb", transport_params=tp) as package_zip: with open(pkg_file, "wb") as fin: fin.write(package_zip.read()) unzip_package( package_path=pkg_file, target_dir=local_dir, remove_top_level_directory=True, unlink_zip=True, logger=logger) else: raise NotImplementedError( f"Protocol {protocol} is not supported") return str(local_dir)
def get(self, key: str) -> bytes: key = self.__concat_key_with_prefixes(key) ret = internal_kv._internal_kv_get(key) return ret
def get_status(self, job_id: str) -> JobStatus: pickled_status = _internal_kv_get( self.JOB_STATUS_KEY.format(job_id=job_id)) assert pickled_status is not None, f"Status not found for {job_id}" return pickle.loads(pickled_status)