chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,774 @@
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import enum
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import AsyncIterator, Dict, Optional, Tuple
|
||||
|
||||
import aiohttp.web
|
||||
from aiohttp.client import ClientResponse
|
||||
from aiohttp.web import Request, Response, StreamResponse
|
||||
|
||||
import ray
|
||||
from ray import NodeID
|
||||
from ray._common.network_utils import build_address
|
||||
from ray._common.pydantic_compat import BaseModel, Extra, Field, validator
|
||||
from ray._common.utils import get_or_create_event_loop, load_class
|
||||
from ray._private.authentication.http_token_authentication import (
|
||||
get_auth_headers_if_auth_enabled,
|
||||
)
|
||||
from ray._private.ray_constants import KV_NAMESPACE_DASHBOARD
|
||||
from ray._private.runtime_env.packaging import (
|
||||
package_exists,
|
||||
pin_runtime_env_uri,
|
||||
upload_package_to_gcs,
|
||||
)
|
||||
from ray.dashboard.consts import (
|
||||
DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX,
|
||||
GCS_RPC_TIMEOUT_SECONDS,
|
||||
RAY_CLUSTER_ACTIVITY_HOOK,
|
||||
TRY_TO_GET_AGENT_INFO_INTERVAL_SECONDS,
|
||||
WAIT_AVAILABLE_AGENT_TIMEOUT,
|
||||
)
|
||||
from ray.dashboard.modules.job.common import (
|
||||
JobDeleteResponse,
|
||||
JobInfoStorageClient,
|
||||
JobLogsResponse,
|
||||
JobStopResponse,
|
||||
JobSubmitRequest,
|
||||
JobSubmitResponse,
|
||||
http_uri_components_to_uri,
|
||||
)
|
||||
from ray.dashboard.modules.job.pydantic_models import JobDetails, JobType
|
||||
from ray.dashboard.modules.job.utils import (
|
||||
find_job_by_ids,
|
||||
get_driver_jobs,
|
||||
parse_and_validate_request,
|
||||
)
|
||||
from ray.dashboard.modules.version import CURRENT_VERSION, VersionResponse
|
||||
from ray.dashboard.subprocesses.module import SubprocessModule
|
||||
from ray.dashboard.subprocesses.routes import SubprocessRouteTable as routes
|
||||
from ray.dashboard.subprocesses.utils import ResponseType
|
||||
from ray.dashboard.utils import get_head_node_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
|
||||
class RayActivityStatus(str, enum.Enum):
|
||||
ACTIVE = "ACTIVE"
|
||||
INACTIVE = "INACTIVE"
|
||||
ERROR = "ERROR"
|
||||
|
||||
|
||||
class RayActivityResponse(BaseModel, extra=Extra.allow):
|
||||
"""
|
||||
Pydantic model used to inform if a particular Ray component can be considered
|
||||
active, and metadata about observation.
|
||||
"""
|
||||
|
||||
is_active: RayActivityStatus = Field(
|
||||
...,
|
||||
description=(
|
||||
"Whether the corresponding Ray component is considered active or inactive, "
|
||||
"or if there was an error while collecting this observation."
|
||||
),
|
||||
)
|
||||
reason: Optional[str] = Field(
|
||||
None, description="Reason if Ray component is considered active or errored."
|
||||
)
|
||||
timestamp: float = Field(
|
||||
...,
|
||||
description=(
|
||||
"Timestamp of when this observation about the Ray component was made. "
|
||||
"This is in the format of seconds since unix epoch."
|
||||
),
|
||||
)
|
||||
last_activity_at: Optional[float] = Field(
|
||||
None,
|
||||
description=(
|
||||
"Timestamp when last actvity of this Ray component finished in format of "
|
||||
"seconds since unix epoch. This field does not need to be populated "
|
||||
"for Ray components where it is not meaningful."
|
||||
),
|
||||
)
|
||||
|
||||
@validator("reason", always=True)
|
||||
def reason_required(cls, v, values, **kwargs):
|
||||
if "is_active" in values and values["is_active"] != RayActivityStatus.INACTIVE:
|
||||
if v is None:
|
||||
raise ValueError(
|
||||
'Reason is required if is_active is "active" or "error"'
|
||||
)
|
||||
return v
|
||||
|
||||
|
||||
class JobAgentSubmissionClient:
|
||||
"""A local client for submitting and interacting with jobs on a specific node
|
||||
in the remote cluster.
|
||||
Submits requests over HTTP to the job agent on the specific node using the REST API.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dashboard_agent_address: str,
|
||||
):
|
||||
self._agent_address = dashboard_agent_address
|
||||
self._session = aiohttp.ClientSession()
|
||||
|
||||
def _get_headers(self):
|
||||
"""Get auth headers if token authentication is enabled."""
|
||||
return get_auth_headers_if_auth_enabled({})
|
||||
|
||||
async def _raise_error(self, resp: ClientResponse):
|
||||
status = resp.status
|
||||
error_text = await resp.text()
|
||||
raise RuntimeError(f"Request failed with status code {status}: {error_text}.")
|
||||
|
||||
async def submit_job_internal(self, req: JobSubmitRequest) -> JobSubmitResponse:
|
||||
logger.debug(f"Submitting job with submission_id={req.submission_id}.")
|
||||
|
||||
async with self._session.post(
|
||||
f"{self._agent_address}/api/job_agent/jobs/",
|
||||
json=dataclasses.asdict(req),
|
||||
headers=self._get_headers(),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
result_json = await resp.json()
|
||||
return JobSubmitResponse(**result_json)
|
||||
else:
|
||||
await self._raise_error(resp)
|
||||
|
||||
async def stop_job_internal(self, job_id: str) -> JobStopResponse:
|
||||
logger.debug(f"Stopping job with job_id={job_id}.")
|
||||
|
||||
async with self._session.post(
|
||||
f"{self._agent_address}/api/job_agent/jobs/{job_id}/stop",
|
||||
headers=self._get_headers(),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
result_json = await resp.json()
|
||||
return JobStopResponse(**result_json)
|
||||
else:
|
||||
await self._raise_error(resp)
|
||||
|
||||
async def delete_job_internal(self, job_id: str) -> JobDeleteResponse:
|
||||
logger.debug(f"Deleting job with job_id={job_id}.")
|
||||
|
||||
async with self._session.delete(
|
||||
f"{self._agent_address}/api/job_agent/jobs/{job_id}",
|
||||
headers=self._get_headers(),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
result_json = await resp.json()
|
||||
return JobDeleteResponse(**result_json)
|
||||
else:
|
||||
await self._raise_error(resp)
|
||||
|
||||
async def get_job_logs_internal(self, job_id: str) -> JobLogsResponse:
|
||||
async with self._session.get(
|
||||
f"{self._agent_address}/api/job_agent/jobs/{job_id}/logs",
|
||||
headers=self._get_headers(),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
result_json = await resp.json()
|
||||
return JobLogsResponse(**result_json)
|
||||
else:
|
||||
await self._raise_error(resp)
|
||||
|
||||
async def tail_job_logs(self, job_id: str) -> AsyncIterator[str]:
|
||||
"""Get an iterator that follows the logs of a job."""
|
||||
ws = await self._session.ws_connect(
|
||||
f"{self._agent_address}/api/job_agent/jobs/{job_id}/logs/tail",
|
||||
headers=self._get_headers(),
|
||||
)
|
||||
|
||||
while True:
|
||||
msg = await ws.receive()
|
||||
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
yield msg.data
|
||||
elif msg.type == aiohttp.WSMsgType.CLOSED:
|
||||
logger.info(
|
||||
f"WebSocket to job agent closed for job {job_id} "
|
||||
f"with close code {ws.close_code}"
|
||||
)
|
||||
break
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
logger.warning(
|
||||
f"WebSocket to job agent received an error message "
|
||||
f"while tailing logs for job {job_id}: {ws.exception()!r}. "
|
||||
)
|
||||
pass
|
||||
|
||||
async def close(self, ignore_error=True):
|
||||
try:
|
||||
await self._session.close()
|
||||
except Exception:
|
||||
if not ignore_error:
|
||||
raise
|
||||
|
||||
|
||||
class JobHead(SubprocessModule):
|
||||
"""Runs on the head node of a Ray cluster and handles Ray Jobs APIs.
|
||||
|
||||
NOTE(architkulkarni): Please keep this class in sync with the OpenAPI spec at
|
||||
`doc/source/cluster/running-applications/job-submission/openapi.yml`.
|
||||
We currently do not automatically check that the OpenAPI
|
||||
spec is in sync with the implementation. If any changes are made to the
|
||||
paths in the @route decorators or in the Responses returned by the
|
||||
methods (or any nested fields in the Responses), you will need to find the
|
||||
corresponding field of the OpenAPI yaml file and update it manually. Also,
|
||||
bump the version number in the yaml file and in this class's `get_version`.
|
||||
"""
|
||||
|
||||
# Time that we sleep while tailing logs while waiting for
|
||||
# the supervisor actor to start. We don't know which node
|
||||
# to read the logs from until then.
|
||||
WAIT_FOR_SUPERVISOR_ACTOR_INTERVAL_S = 1
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._job_info_client = None
|
||||
|
||||
# To make sure that the internal KV is initialized by getting the lazy property
|
||||
assert self.gcs_client is not None
|
||||
assert ray.experimental.internal_kv._internal_kv_initialized()
|
||||
|
||||
# It contains all `JobAgentSubmissionClient` that
|
||||
# `JobHead` has ever used, and will not be deleted
|
||||
# from it unless `JobAgentSubmissionClient` is no
|
||||
# longer available (the corresponding agent process is dead)
|
||||
# {node_id: JobAgentSubmissionClient}
|
||||
self._agents: Dict[NodeID, JobAgentSubmissionClient] = dict()
|
||||
|
||||
async def get_target_agent(
|
||||
self, timeout_s: float = WAIT_AVAILABLE_AGENT_TIMEOUT
|
||||
) -> JobAgentSubmissionClient:
|
||||
"""
|
||||
Get a `JobAgentSubmissionClient`, which is a client for interacting with jobs
|
||||
via an agent process.
|
||||
|
||||
Args:
|
||||
timeout_s: The timeout for the operation.
|
||||
|
||||
Returns:
|
||||
A `JobAgentSubmissionClient` for interacting with jobs via an agent process.
|
||||
|
||||
Raises:
|
||||
TimeoutError: If the operation times out.
|
||||
"""
|
||||
return await self._get_head_node_agent(timeout_s)
|
||||
|
||||
async def _get_head_node_agent_once(self) -> JobAgentSubmissionClient:
|
||||
head_node_id_hex = await get_head_node_id(self.gcs_client)
|
||||
|
||||
if not head_node_id_hex:
|
||||
raise Exception("Head node id has not yet been persisted in GCS")
|
||||
|
||||
head_node_id = NodeID.from_hex(head_node_id_hex)
|
||||
|
||||
if head_node_id not in self._agents:
|
||||
ip, http_port, _ = await self._fetch_agent_info(head_node_id)
|
||||
agent_http_address = f"http://{build_address(ip, http_port)}"
|
||||
self._agents[head_node_id] = JobAgentSubmissionClient(agent_http_address)
|
||||
|
||||
return self._agents[head_node_id]
|
||||
|
||||
async def _get_head_node_agent(self, timeout_s: float) -> JobAgentSubmissionClient:
|
||||
"""Retrieves HTTP client for `JobAgent` running on the Head node. If the head
|
||||
node does not have an agent, it will retry every
|
||||
`TRY_TO_GET_AGENT_INFO_INTERVAL_SECONDS` seconds indefinitely.
|
||||
|
||||
Args:
|
||||
timeout_s: The timeout for the operation.
|
||||
|
||||
Returns:
|
||||
A `JobAgentSubmissionClient` for interacting with jobs via the head node's agent process.
|
||||
|
||||
Raises:
|
||||
TimeoutError: If the operation times out.
|
||||
"""
|
||||
timeout_point = time.time() + timeout_s
|
||||
exception = None
|
||||
while time.time() < timeout_point:
|
||||
try:
|
||||
return await self._get_head_node_agent_once()
|
||||
except Exception as e:
|
||||
exception = e
|
||||
logger.exception(
|
||||
f"Failed to get head node agent, retrying in {TRY_TO_GET_AGENT_INFO_INTERVAL_SECONDS} seconds..."
|
||||
)
|
||||
await asyncio.sleep(TRY_TO_GET_AGENT_INFO_INTERVAL_SECONDS)
|
||||
raise TimeoutError(
|
||||
f"Failed to get head node agent within {timeout_s} seconds. The last exception is {exception}"
|
||||
)
|
||||
|
||||
async def _fetch_agent_info(self, target_node_id: NodeID) -> Tuple[str, int, int]:
|
||||
"""
|
||||
Fetches agent info by the Node ID. May raise exception if there's network error or the
|
||||
agent info is not found.
|
||||
|
||||
Returns: (ip, http_port, grpc_port)
|
||||
"""
|
||||
key = f"{DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX}{target_node_id.hex()}"
|
||||
value = await self.gcs_client.async_internal_kv_get(
|
||||
key,
|
||||
namespace=KV_NAMESPACE_DASHBOARD,
|
||||
timeout=GCS_RPC_TIMEOUT_SECONDS,
|
||||
)
|
||||
if not value:
|
||||
raise KeyError(
|
||||
f"Agent info not found in internal KV for node {target_node_id}. "
|
||||
"It's possible that the agent didn't launch successfully due to "
|
||||
"port conflicts or other issues. Please check `dashboard_agent.log` "
|
||||
"for more details."
|
||||
)
|
||||
return json.loads(value.decode())
|
||||
|
||||
@routes.get("/api/version")
|
||||
async def get_version(self, req: Request) -> Response:
|
||||
# NOTE(edoakes): CURRENT_VERSION should be bumped and checked on the
|
||||
# client when we have backwards-incompatible changes.
|
||||
resp = VersionResponse(
|
||||
version=CURRENT_VERSION,
|
||||
ray_version=ray.__version__,
|
||||
ray_commit=ray.__commit__,
|
||||
session_name=self.session_name,
|
||||
)
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)),
|
||||
content_type="application/json",
|
||||
status=aiohttp.web.HTTPOk.status_code,
|
||||
)
|
||||
|
||||
@routes.get("/api/packages/{protocol}/{package_name}")
|
||||
async def get_package(self, req: Request) -> Response:
|
||||
package_uri = http_uri_components_to_uri(
|
||||
protocol=req.match_info["protocol"],
|
||||
package_name=req.match_info["package_name"],
|
||||
)
|
||||
|
||||
logger.debug(f"Adding temporary reference to package {package_uri}.")
|
||||
try:
|
||||
pin_runtime_env_uri(package_uri)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
if not package_exists(package_uri):
|
||||
return Response(
|
||||
text=f"Package {package_uri} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
return Response()
|
||||
|
||||
@routes.put("/api/packages/{protocol}/{package_name}")
|
||||
async def upload_package(self, req: Request):
|
||||
package_uri = http_uri_components_to_uri(
|
||||
protocol=req.match_info["protocol"],
|
||||
package_name=req.match_info["package_name"],
|
||||
)
|
||||
logger.info(f"Uploading package {package_uri} to the GCS.")
|
||||
try:
|
||||
data = await req.read()
|
||||
await get_or_create_event_loop().run_in_executor(
|
||||
None,
|
||||
upload_package_to_gcs,
|
||||
package_uri,
|
||||
data,
|
||||
)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(status=aiohttp.web.HTTPOk.status_code)
|
||||
|
||||
@routes.post("/api/jobs/")
|
||||
async def submit_job(self, req: Request) -> Response:
|
||||
result = await parse_and_validate_request(req, JobSubmitRequest)
|
||||
# Request parsing failed, returned with Response object.
|
||||
if isinstance(result, Response):
|
||||
return result
|
||||
else:
|
||||
submit_request: JobSubmitRequest = result
|
||||
|
||||
try:
|
||||
job_agent_client = await self.get_target_agent()
|
||||
resp = await job_agent_client.submit_job_internal(submit_request)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"Timed out waiting for an available job agent to submit the job."
|
||||
)
|
||||
return Response(
|
||||
text="No available agent to submit job, please try again later.",
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
logger.warning("Failed to submit job due to an invalid request.")
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to submit job due to unexpected exception.")
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)),
|
||||
content_type="application/json",
|
||||
status=aiohttp.web.HTTPOk.status_code,
|
||||
)
|
||||
|
||||
@routes.post("/api/jobs/{job_or_submission_id}/stop")
|
||||
async def stop_job(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only stop submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
job_agent_client = await self.get_target_agent()
|
||||
resp = await job_agent_client.stop_job_internal(job.submission_id)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)), content_type="application/json"
|
||||
)
|
||||
|
||||
@routes.delete("/api/jobs/{job_or_submission_id}")
|
||||
async def delete_job(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only delete submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
job_agent_client = await self.get_target_agent()
|
||||
resp = await job_agent_client.delete_job_internal(job.submission_id)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(resp)), content_type="application/json"
|
||||
)
|
||||
|
||||
@routes.get("/api/jobs/{job_or_submission_id}")
|
||||
async def get_job_info(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
return Response(
|
||||
text=json.dumps(job.dict()),
|
||||
content_type="application/json",
|
||||
)
|
||||
|
||||
# TODO(rickyx): This endpoint's logic is also mirrored in state API's endpoint.
|
||||
# We should eventually unify the backend logic (and keep the logic in sync before
|
||||
# that).
|
||||
@routes.get("/api/jobs/")
|
||||
async def list_jobs(self, req: Request) -> Response:
|
||||
(driver_jobs, submission_job_drivers), submission_jobs = await asyncio.gather(
|
||||
get_driver_jobs(self.gcs_client), self._job_info_client.get_all_jobs()
|
||||
)
|
||||
|
||||
submission_jobs = [
|
||||
JobDetails(
|
||||
**dataclasses.asdict(job),
|
||||
submission_id=submission_id,
|
||||
job_id=submission_job_drivers.get(submission_id).id
|
||||
if submission_id in submission_job_drivers
|
||||
else None,
|
||||
driver_info=submission_job_drivers.get(submission_id),
|
||||
type=JobType.SUBMISSION,
|
||||
)
|
||||
for submission_id, job in submission_jobs.items()
|
||||
]
|
||||
return Response(
|
||||
text=json.dumps(
|
||||
[
|
||||
*[submission_job.dict() for submission_job in submission_jobs],
|
||||
*[job_info.dict() for job_info in driver_jobs.values()],
|
||||
]
|
||||
),
|
||||
content_type="application/json",
|
||||
)
|
||||
|
||||
@routes.get("/api/jobs/{job_or_submission_id}/logs")
|
||||
async def get_job_logs(self, req: Request) -> Response:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only get logs of submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
job_agent_client = self.get_job_driver_agent_client(job)
|
||||
payload = (
|
||||
await job_agent_client.get_job_logs_internal(job.submission_id)
|
||||
if job_agent_client
|
||||
else JobLogsResponse("")
|
||||
)
|
||||
return Response(
|
||||
text=json.dumps(dataclasses.asdict(payload)),
|
||||
content_type="application/json",
|
||||
)
|
||||
except Exception:
|
||||
return Response(
|
||||
text=traceback.format_exc(),
|
||||
status=aiohttp.web.HTTPInternalServerError.status_code,
|
||||
)
|
||||
|
||||
@routes.get(
|
||||
"/api/jobs/{job_or_submission_id}/logs/tail", resp_type=ResponseType.WEBSOCKET
|
||||
)
|
||||
async def tail_job_logs(self, req: Request) -> StreamResponse:
|
||||
job_or_submission_id = req.match_info["job_or_submission_id"]
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
if not job:
|
||||
return Response(
|
||||
text=f"Job {job_or_submission_id} does not exist",
|
||||
status=aiohttp.web.HTTPNotFound.status_code,
|
||||
)
|
||||
|
||||
if job.type is not JobType.SUBMISSION:
|
||||
return Response(
|
||||
text="Can only get logs of submission type jobs",
|
||||
status=aiohttp.web.HTTPBadRequest.status_code,
|
||||
)
|
||||
|
||||
ws = aiohttp.web.WebSocketResponse()
|
||||
await ws.prepare(req)
|
||||
|
||||
driver_agent_http_address = None
|
||||
while driver_agent_http_address is None:
|
||||
job = await find_job_by_ids(
|
||||
self.gcs_client,
|
||||
self._job_info_client,
|
||||
job_or_submission_id,
|
||||
)
|
||||
driver_agent_http_address = job.driver_agent_http_address
|
||||
status = job.status
|
||||
if status.is_terminal() and driver_agent_http_address is None:
|
||||
# Job exited before supervisor actor started.
|
||||
return ws
|
||||
|
||||
await asyncio.sleep(self.WAIT_FOR_SUPERVISOR_ACTOR_INTERVAL_S)
|
||||
|
||||
job_agent_client = self.get_job_driver_agent_client(job)
|
||||
|
||||
async for lines in job_agent_client.tail_job_logs(job.submission_id):
|
||||
await ws.send_str(lines)
|
||||
|
||||
return ws
|
||||
|
||||
def get_job_driver_agent_client(
|
||||
self, job: JobDetails
|
||||
) -> Optional[JobAgentSubmissionClient]:
|
||||
if job.driver_agent_http_address is None:
|
||||
return None
|
||||
|
||||
driver_node_id = job.driver_node_id
|
||||
if driver_node_id not in self._agents:
|
||||
self._agents[driver_node_id] = JobAgentSubmissionClient(
|
||||
job.driver_agent_http_address
|
||||
)
|
||||
|
||||
return self._agents[driver_node_id]
|
||||
|
||||
@routes.get("/api/component_activities")
|
||||
async def get_component_activities(
|
||||
self, req: aiohttp.web.Request
|
||||
) -> aiohttp.web.Response:
|
||||
timeout = req.query.get("timeout", None)
|
||||
if timeout and timeout.isdigit():
|
||||
timeout = int(timeout)
|
||||
else:
|
||||
timeout = 30
|
||||
|
||||
# Get activity information for driver
|
||||
driver_activity_info = await self._get_job_activity_info(timeout=timeout)
|
||||
resp = {"driver": dict(driver_activity_info)}
|
||||
|
||||
if RAY_CLUSTER_ACTIVITY_HOOK in os.environ:
|
||||
try:
|
||||
cluster_activity_callable = load_class(
|
||||
os.environ[RAY_CLUSTER_ACTIVITY_HOOK]
|
||||
)
|
||||
external_activity_output = cluster_activity_callable()
|
||||
assert isinstance(external_activity_output, dict), (
|
||||
f"Output of hook {os.environ[RAY_CLUSTER_ACTIVITY_HOOK]} "
|
||||
"should be Dict[str, RayActivityResponse]. Got "
|
||||
f"output: {external_activity_output}"
|
||||
)
|
||||
for component_type in external_activity_output:
|
||||
try:
|
||||
component_activity_output = external_activity_output[
|
||||
component_type
|
||||
]
|
||||
# Parse and validate output to type RayActivityResponse
|
||||
component_activity_output = RayActivityResponse(
|
||||
**dict(component_activity_output)
|
||||
)
|
||||
resp[component_type] = dict(component_activity_output)
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Failed to get activity status of {component_type} "
|
||||
f"from user hook {os.environ[RAY_CLUSTER_ACTIVITY_HOOK]}."
|
||||
)
|
||||
resp[component_type] = {
|
||||
"is_active": RayActivityStatus.ERROR,
|
||||
"reason": repr(e),
|
||||
"timestamp": datetime.now().timestamp(),
|
||||
}
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
"Failed to get activity status from user "
|
||||
f"hook {os.environ[RAY_CLUSTER_ACTIVITY_HOOK]}."
|
||||
)
|
||||
resp["external_component"] = {
|
||||
"is_active": RayActivityStatus.ERROR,
|
||||
"reason": repr(e),
|
||||
"timestamp": datetime.now().timestamp(),
|
||||
}
|
||||
|
||||
return aiohttp.web.Response(
|
||||
text=json.dumps(resp),
|
||||
content_type="application/json",
|
||||
status=aiohttp.web.HTTPOk.status_code,
|
||||
)
|
||||
|
||||
async def _get_job_activity_info(self, timeout: int) -> RayActivityResponse:
|
||||
# Returns if there is Ray activity from drivers (job).
|
||||
# Drivers in namespaces that start with _ray_internal_ are not
|
||||
# considered activity.
|
||||
# This includes the _ray_internal_dashboard job that gets automatically
|
||||
# created with every cluster
|
||||
try:
|
||||
reply = await self.gcs_client.async_get_all_job_info(
|
||||
skip_submission_job_info_field=True,
|
||||
skip_is_running_tasks_field=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
num_active_drivers = 0
|
||||
latest_job_end_time = 0
|
||||
for job_table_entry in reply.values():
|
||||
is_dead = bool(job_table_entry.is_dead)
|
||||
in_internal_namespace = job_table_entry.config.ray_namespace.startswith(
|
||||
"_ray_internal_"
|
||||
)
|
||||
latest_job_end_time = (
|
||||
max(latest_job_end_time, job_table_entry.end_time)
|
||||
if job_table_entry.end_time
|
||||
else latest_job_end_time
|
||||
)
|
||||
if not is_dead and not in_internal_namespace:
|
||||
num_active_drivers += 1
|
||||
|
||||
current_timestamp = datetime.now().timestamp()
|
||||
# Latest job end time must be before or equal to the current timestamp.
|
||||
# Job end times may be provided in epoch milliseconds. Check if this
|
||||
# is true, and convert to seconds
|
||||
if latest_job_end_time > current_timestamp:
|
||||
latest_job_end_time = latest_job_end_time / 1000
|
||||
assert current_timestamp >= latest_job_end_time, (
|
||||
f"Most recent job end time {latest_job_end_time} must be "
|
||||
f"before or equal to the current timestamp {current_timestamp}"
|
||||
)
|
||||
|
||||
is_active = (
|
||||
RayActivityStatus.ACTIVE
|
||||
if num_active_drivers > 0
|
||||
else RayActivityStatus.INACTIVE
|
||||
)
|
||||
return RayActivityResponse(
|
||||
is_active=is_active,
|
||||
reason=f"Number of active drivers: {num_active_drivers}"
|
||||
if num_active_drivers
|
||||
else None,
|
||||
timestamp=current_timestamp,
|
||||
# If latest_job_end_time == 0, no jobs have finished yet so don't
|
||||
# populate last_activity_at
|
||||
last_activity_at=latest_job_end_time if latest_job_end_time else None,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("Failed to get activity status of Ray drivers.")
|
||||
return RayActivityResponse(
|
||||
is_active=RayActivityStatus.ERROR,
|
||||
reason=repr(e),
|
||||
timestamp=datetime.now().timestamp(),
|
||||
)
|
||||
|
||||
async def run(self):
|
||||
await super().run()
|
||||
if not self._job_info_client:
|
||||
self._job_info_client = JobInfoStorageClient(self.gcs_client)
|
||||
Reference in New Issue
Block a user