|
77 | 77 | ] |
78 | 78 |
|
79 | 79 |
|
80 | | -class RayClusterConfig(Config): |
81 | | - image: str | None = None |
82 | | - namespace: str = "ray" |
83 | | - enable_in_tree_autoscaling: bool = False |
| 80 | +class MetadataWithoutName(Config): |
| 81 | + namespace: str | None = None |
| 82 | + labels: dict[str, str] | None = None |
| 83 | + annotations: dict[str, str] | None = None |
| 84 | + |
| 85 | + |
| 86 | +class MetadataWithOptionalName(MetadataWithoutName): |
| 87 | + """This config almost matches the original Kubernetes metadata, except name which can be ommitted to be set automatically by `dagster-ray`.""" |
| 88 | + |
| 89 | + name: str | None = None |
| 90 | + |
| 91 | + |
| 92 | +class Metadata(MetadataWithoutName): |
| 93 | + name: str |
| 94 | + |
| 95 | + |
| 96 | +class RayClusterSpec(Config): |
| 97 | + """[RayCluster spec](https://ray-project.github.io/kuberay/reference/api/#rayclusterspec) configuration options. A few sensible defaults are provided for convenience.""" |
| 98 | + |
| 99 | + suspend: bool = False |
| 100 | + managed_by: str | None = None |
84 | 101 | autoscaler_options: dict[str, Any] = DEFAULT_AUTOSCALER_OPTIONS # TODO: add a dedicated Config type |
| 102 | + head_service_annotations: dict[str, str] | None = None |
| 103 | + enable_in_tree_autoscaling: bool = False |
| 104 | + gcs_fault_tolerance_options: dict[str, Any] | None = None |
85 | 105 | head_group_spec: dict[str, Any] = DEFAULT_HEAD_GROUP_SPEC # TODO: add a dedicated Config type |
| 106 | + ray_version: str | None = None |
86 | 107 | worker_group_specs: list[dict[str, Any]] = DEFAULT_WORKER_GROUP_SPECS # TODO: add a dedicated Config type |
87 | 108 |
|
88 | 109 |
|
89 | | -class RayJobConfig(Config): |
| 110 | +class RayClusterConfig(Config): |
| 111 | + kind: str = "RayCluster" |
| 112 | + api_version: str = "ray.io/v1" |
| 113 | + metadata: MetadataWithOptionalName = Field(default_factory=MetadataWithOptionalName) |
| 114 | + spec: RayClusterSpec = Field(default_factory=RayClusterSpec) |
| 115 | + |
| 116 | + |
| 117 | +class RayJobSpec(Config): |
| 118 | + """[RayJob spec](https://ray-project.github.io/kuberay/reference/api/#rayjobspec) configuration options. A few sensible defaults are provided for convenience.""" |
| 119 | + |
| 120 | + active_deadline_seconds: int = 60 * 60 * 24 # 24 hours |
| 121 | + backoff_limit: int = 0 |
| 122 | + ray_cluster_spec: RayClusterSpec = Field(default_factory=RayClusterSpec) |
| 123 | + submitter_pod_template: dict[str, Any] | None = None |
| 124 | + cluster_selector: dict[str, str] | None = None |
| 125 | + managed_by: str | None = None |
| 126 | + deletion_strategy: dict[str, Any] | None = None |
| 127 | + runtime_env_yaml: str | None = None |
| 128 | + job_id: str | None = None |
| 129 | + submission_mode: Literal["K8sJobMode", "HTTPMode", "InteractiveMode"] = "K8sJobMode" |
| 130 | + entrypoint_resources: str | None = None |
90 | 131 | entrypoint_num_cpus: float |
91 | 132 | entrypoint_memory: float |
92 | | - entrypoint_num_gpus: int |
93 | | - suspend: bool = False |
94 | | - annotations: dict[str, str] | None = None |
95 | | - labels: dict[str, str] | None = None |
| 133 | + entrypoint_num_gpus: float |
| 134 | + ttl_seconds_after_finished: int = 60 * 60 # 1 hour |
96 | 135 | shutdown_after_job_finishes: bool = True |
97 | | - ttl_seconds_after_finished: int = 60 * 10 # 10 minutes |
98 | | - active_deadline_seconds: int = 60 * 60 * 24 # 24 hours |
99 | | - submission_mode: Literal["K8sJobMode", "HTTPMode"] = "K8sJobMode" |
100 | | - runtime_env_yaml: str | None = None |
101 | | - cluster: RayClusterConfig = Field(default_factory=RayClusterConfig) |
| 136 | + suspend: bool = False |
| 137 | + |
| 138 | + |
| 139 | +class RayJobConfig(Config): |
| 140 | + kind: str = "RayJob" |
| 141 | + api_version: str = "ray.io/v1" |
| 142 | + metadata: MetadataWithOptionalName = Field(default_factory=MetadataWithOptionalName) |
| 143 | + spec: RayJobSpec = Field(default_factory=RayJobSpec) |
0 commit comments