mirror of
https://github.com/vale981/ray
synced 2025-03-09 12:56:46 -04:00
e54d3117a4
## Why are these changes needed? This is part of redis removal project. In this PR all direct usage of redis got removed except function table. Function table will be migrated in the next PR ## Related issue number #19443
114 lines
2.6 KiB
Python
114 lines
2.6 KiB
Python
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Any, Dict, Optional
|
|
import pickle
|
|
from ray import ray_constants
|
|
from ray.experimental.internal_kv import (
|
|
_internal_kv_initialized,
|
|
_internal_kv_get,
|
|
_internal_kv_put,
|
|
)
|
|
|
|
# NOTE(edoakes): constant should be considered a public API because it's
|
|
# exposed in the snapshot API.
|
|
JOB_ID_METADATA_KEY = "job_submission_id"
|
|
|
|
JOBS_API_PREFIX = "/api/jobs/"
|
|
JOBS_API_ROUTE_LOGS = JOBS_API_PREFIX + "logs"
|
|
JOBS_API_ROUTE_SUBMIT = JOBS_API_PREFIX + "submit"
|
|
JOBS_API_ROUTE_STOP = JOBS_API_PREFIX + "stop"
|
|
JOBS_API_ROUTE_STATUS = JOBS_API_PREFIX + "status"
|
|
JOBS_API_ROUTE_PACKAGE = JOBS_API_PREFIX + "package"
|
|
|
|
|
|
class JobStatus(str, Enum):
|
|
def __str__(self):
|
|
return f"{self.value}"
|
|
|
|
DOES_NOT_EXIST = "DOES_NOT_EXIST"
|
|
PENDING = "PENDING"
|
|
RUNNING = "RUNNING"
|
|
STOPPED = "STOPPED"
|
|
SUCCEEDED = "SUCCEEDED"
|
|
FAILED = "FAILED"
|
|
|
|
|
|
class JobStatusStorageClient:
|
|
"""
|
|
Handles formatting of status storage key given job id.
|
|
"""
|
|
JOB_STATUS_KEY = "_ray_internal_job_status_{job_id}"
|
|
|
|
def __init__(self):
|
|
assert _internal_kv_initialized()
|
|
|
|
def put_status(self, job_id: str, status: JobStatus):
|
|
assert isinstance(status, JobStatus)
|
|
_internal_kv_put(
|
|
self.JOB_STATUS_KEY.format(job_id=job_id),
|
|
pickle.dumps(status),
|
|
namespace=ray_constants.KV_NAMESPACE_JOB)
|
|
|
|
def get_status(self, job_id: str) -> JobStatus:
|
|
pickled_status = _internal_kv_get(
|
|
self.JOB_STATUS_KEY.format(job_id=job_id),
|
|
namespace=ray_constants.KV_NAMESPACE_JOB)
|
|
if pickled_status is None:
|
|
return JobStatus.DOES_NOT_EXIST
|
|
else:
|
|
return pickle.loads(pickled_status)
|
|
|
|
|
|
# ==== Get Package ====
|
|
|
|
|
|
@dataclass
|
|
class GetPackageResponse:
|
|
package_exists: bool
|
|
|
|
|
|
# ==== Job Submit ====
|
|
|
|
|
|
@dataclass
|
|
class JobSubmitRequest:
|
|
# Dict to setup execution environment.
|
|
runtime_env: Dict[str, Any]
|
|
# Command to start execution, ex: "python script.py"
|
|
entrypoint: str
|
|
# Optional job_id to specify for the job. If the job_id is not specified,
|
|
# one will be generated. If a job with the same job_id already exists, it
|
|
# will be rejected.
|
|
job_id: Optional[str]
|
|
# Metadata to pass in to the JobConfig.
|
|
metadata: Dict[str, str]
|
|
|
|
|
|
@dataclass
|
|
class JobSubmitResponse:
|
|
job_id: str
|
|
|
|
|
|
# ==== Job Stop ====
|
|
|
|
|
|
@dataclass
|
|
class JobStopResponse:
|
|
stopped: bool
|
|
|
|
|
|
# ==== Job Status ====
|
|
|
|
|
|
@dataclass
|
|
class JobStatusResponse:
|
|
job_status: JobStatus
|
|
|
|
|
|
# ==== Job Logs ====
|
|
|
|
|
|
# TODO(jiaodong): Support log streaming #19415
|
|
@dataclass
|
|
class JobLogsResponse:
|
|
logs: str
|