ray/dashboard/modules/job/job_description.py

49 lines
1.7 KiB
Python

from enum import Enum
from typing import Dict, Any
from pydantic import BaseModel as PydanticBaseModel, Extra
from ray.core.generated import common_pb2
class BaseModel(PydanticBaseModel):
class Config:
arbitrary_types_allowed = True
extra = Extra.forbid
class Language(str, Enum):
PYTHON = common_pb2.Language.Name(common_pb2.PYTHON)
JAVA = common_pb2.Language.Name(common_pb2.JAVA)
CPP = common_pb2.Language.Name(common_pb2.CPP)
class RuntimeEnv(BaseModel):
# The url to download the job package archive. The archive format is
# one of “zip”, “tar”, “gztar”, “bztar”, or “xztar”. Please refer to
# https://docs.python.org/3/library/shutil.html#shutil.unpack_archive
working_dir: str
class JobDescription(BaseModel):
# The job driver language, this field determines how to start the
# driver. The value is one of the names of enum Language defined in
# common.proto, e.g. PYTHON
language: Language
# The runtime_env (RuntimeEnvDict) for the job config.
runtime_env: RuntimeEnv
# The entry to start the driver.
# PYTHON:
# - The basename of driver filename without extension in the job
# package archive.
# JAVA:
# - The driver class full name in the job package archive.
driver_entry: str
# The driver arguments in list.
# PYTHON:
# - The arguments to pass to the main() function in driver entry.
# e.g. [1, False, 3.14, "abc"]
# JAVA:
# - The arguments to pass to the driver command line.
# e.g. ["-custom-arg", "abc"]
driver_args: list = []
# The environment vars to pass to job config, type of keys should be str.
env: Dict[str, Any] = {}