|
8 | 8 | DEFAULT_REPLICA_GROUP_NAME, |
9 | 9 | ApplyRunPlanInput, |
10 | 10 | JobSpec, |
| 11 | + JobSubmission, |
11 | 12 | RunSpec, |
12 | 13 | ) |
13 | 14 | from dstack._internal.server.schemas.runs import GetRunPlanRequest, ListRunsRequest |
@@ -36,8 +37,21 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[IncludeExcludeD |
36 | 37 | current_resource_excludes["jobs"] = { |
37 | 38 | "__all__": { |
38 | 39 | "job_spec": get_job_spec_excludes([job.job_spec for job in current_resource.jobs]), |
| 40 | + "job_submissions": { |
| 41 | + "__all__": get_job_submission_excludes( |
| 42 | + [ |
| 43 | + submission |
| 44 | + for job in current_resource.jobs |
| 45 | + for submission in job.job_submissions |
| 46 | + ] |
| 47 | + ), |
| 48 | + }, |
39 | 49 | } |
40 | 50 | } |
| 51 | + if current_resource.latest_job_submission is not None: |
| 52 | + current_resource_excludes["latest_job_submission"] = get_job_submission_excludes( |
| 53 | + [current_resource.latest_job_submission] |
| 54 | + ) |
41 | 55 | return {"plan": apply_plan_excludes} |
42 | 56 |
|
43 | 57 |
|
@@ -104,3 +118,23 @@ def get_job_spec_excludes(job_specs: list[JobSpec]) -> IncludeExcludeDictType: |
104 | 118 | probe_excludes["until_ready"] = True |
105 | 119 |
|
106 | 120 | return spec_excludes |
| 121 | + |
| 122 | + |
| 123 | +def get_job_submission_excludes(job_submissions: list[JobSubmission]) -> IncludeExcludeDictType: |
| 124 | + submission_excludes: IncludeExcludeDictType = {} |
| 125 | + |
| 126 | + if any(s.job_runtime_data is not None for s in job_submissions): |
| 127 | + jrd_excludes = {} |
| 128 | + if all( |
| 129 | + s.job_runtime_data is None or s.job_runtime_data.username is None |
| 130 | + for s in job_submissions |
| 131 | + ): |
| 132 | + jrd_excludes["username"] = True |
| 133 | + if all( |
| 134 | + s.job_runtime_data is None or s.job_runtime_data.working_dir is None |
| 135 | + for s in job_submissions |
| 136 | + ): |
| 137 | + jrd_excludes["working_dir"] = True |
| 138 | + submission_excludes["job_runtime_data"] = jrd_excludes |
| 139 | + |
| 140 | + return submission_excludes |
0 commit comments