Skip to content

Commit a9dcb60

Browse files
colinjcr4victor
colinjc
and
r4victor
authored
Exclude backward incompatible fields from rest plugin calls (#2767)
* Exclude backward incompatible fields from plugin calls * Refactor compatibility functions * Handle gateways and volumes compatibility --------- Co-authored-by: Victor Skvortsov <vds003@gmail.com>
1 parent b1815ba commit a9dcb60

File tree

10 files changed

+306
-231
lines changed

10 files changed

+306
-231
lines changed

‎src/dstack/_internal/core/compatibility/__init__.py‎

Whitespace-only changes.

‎src/dstack/_internal/core/compatibility/fleets.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
from typing import Any, Dict, Optional
2+
3+
from dstack._internal.core.models.fleets import ApplyFleetPlanInput, FleetSpec
4+
from dstack._internal.core.models.instances import Instance
5+
6+
7+
def get_get_plan_excludes(fleet_spec: FleetSpec) -> Dict:
8+
get_plan_excludes = {}
9+
spec_excludes = get_fleet_spec_excludes(fleet_spec)
10+
if spec_excludes:
11+
get_plan_excludes["spec"] = spec_excludes
12+
return get_plan_excludes
13+
14+
15+
def get_apply_plan_excludes(plan_input: ApplyFleetPlanInput) -> Dict:
16+
apply_plan_excludes = {}
17+
spec_excludes = get_fleet_spec_excludes(plan_input.spec)
18+
if spec_excludes:
19+
apply_plan_excludes["spec"] = spec_excludes
20+
current_resource = plan_input.current_resource
21+
if current_resource is not None:
22+
current_resource_excludes = {}
23+
apply_plan_excludes["current_resource"] = current_resource_excludes
24+
if all(map(_should_exclude_instance_cpu_arch, current_resource.instances)):
25+
current_resource_excludes["instances"] = {
26+
"__all__": {"instance_type": {"resources": {"cpu_arch"}}}
27+
}
28+
return {"plan": apply_plan_excludes}
29+
30+
31+
def get_create_fleet_excludes(fleet_spec: FleetSpec) -> Dict:
32+
create_fleet_excludes = {}
33+
spec_excludes = get_fleet_spec_excludes(fleet_spec)
34+
if spec_excludes:
35+
create_fleet_excludes["spec"] = spec_excludes
36+
return create_fleet_excludes
37+
38+
39+
def get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[Dict]:
40+
"""
41+
Returns `fleet_spec` exclude mapping to exclude certain fields from the request.
42+
Use this method to exclude new fields when they are not set to keep
43+
clients backward-compatibility with older servers.
44+
"""
45+
spec_excludes: Dict[str, Any] = {}
46+
configuration_excludes: Dict[str, Any] = {}
47+
profile_excludes: set[str] = set()
48+
profile = fleet_spec.profile
49+
if profile.fleets is None:
50+
profile_excludes.add("fleets")
51+
if fleet_spec.configuration.tags is None:
52+
configuration_excludes["tags"] = True
53+
if profile.tags is None:
54+
profile_excludes.add("tags")
55+
if profile.startup_order is None:
56+
profile_excludes.add("startup_order")
57+
if profile.stop_criteria is None:
58+
profile_excludes.add("stop_criteria")
59+
if configuration_excludes:
60+
spec_excludes["configuration"] = configuration_excludes
61+
if profile_excludes:
62+
spec_excludes["profile"] = profile_excludes
63+
if spec_excludes:
64+
return spec_excludes
65+
return None
66+
67+
68+
def _should_exclude_instance_cpu_arch(instance: Instance) -> bool:
69+
try:
70+
return instance.instance_type.resources.cpu_arch is None
71+
except AttributeError:
72+
return True

‎src/dstack/_internal/core/compatibility/gateways.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
from typing import Dict
2+
3+
from dstack._internal.core.models.gateways import GatewayConfiguration, GatewaySpec
4+
5+
6+
def get_gateway_spec_excludes(gateway_spec: GatewaySpec) -> Dict:
7+
"""
8+
Returns `gateway_spec` exclude mapping to exclude certain fields from the request.
9+
Use this method to exclude new fields when they are not set to keep
10+
clients backward-compatibility with older servers.
11+
"""
12+
spec_excludes = {}
13+
spec_excludes["configuration"] = _get_gateway_configuration_excludes(
14+
gateway_spec.configuration
15+
)
16+
return spec_excludes
17+
18+
19+
def get_create_gateway_excludes(configuration: GatewayConfiguration) -> Dict:
20+
"""
21+
Returns an exclude mapping to exclude certain fields from the create gateway request.
22+
Use this method to exclude new fields when they are not set to keep
23+
clients backward-compatibility with older servers.
24+
"""
25+
create_gateway_excludes = {}
26+
create_gateway_excludes["configuration"] = _get_gateway_configuration_excludes(configuration)
27+
return create_gateway_excludes
28+
29+
30+
def _get_gateway_configuration_excludes(configuration: GatewayConfiguration) -> Dict:
31+
configuration_excludes = {}
32+
if configuration.tags is None:
33+
configuration_excludes["tags"] = True
34+
return configuration_excludes

‎src/dstack/_internal/core/compatibility/runs.py‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
from typing import Any, Dict, Optional
2+
3+
from dstack._internal.core.models.configurations import ServiceConfiguration
4+
from dstack._internal.core.models.runs import ApplyRunPlanInput, JobSubmission, RunSpec
5+
from dstack._internal.server.schemas.runs import GetRunPlanRequest
6+
7+
8+
def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
9+
"""
10+
Returns `plan` exclude mapping to exclude certain fields from the request.
11+
Use this method to exclude new fields when they are not set to keep
12+
clients backward-compatibility with older servers.
13+
"""
14+
apply_plan_excludes = {}
15+
run_spec_excludes = get_run_spec_excludes(plan.run_spec)
16+
if run_spec_excludes is not None:
17+
apply_plan_excludes["run_spec"] = run_spec_excludes
18+
current_resource = plan.current_resource
19+
if current_resource is not None:
20+
current_resource_excludes = {}
21+
current_resource_excludes["status_message"] = True
22+
apply_plan_excludes["current_resource"] = current_resource_excludes
23+
current_resource_excludes["run_spec"] = get_run_spec_excludes(current_resource.run_spec)
24+
job_submissions_excludes = {}
25+
current_resource_excludes["jobs"] = {
26+
"__all__": {"job_submissions": {"__all__": job_submissions_excludes}}
27+
}
28+
job_submissions = [js for j in current_resource.jobs for js in j.job_submissions]
29+
if all(map(_should_exclude_job_submission_jpd_cpu_arch, job_submissions)):
30+
job_submissions_excludes["job_provisioning_data"] = {
31+
"instance_type": {"resources": {"cpu_arch"}}
32+
}
33+
if all(map(_should_exclude_job_submission_jrd_cpu_arch, job_submissions)):
34+
job_submissions_excludes["job_runtime_data"] = {
35+
"offer": {"instance": {"resources": {"cpu_arch"}}}
36+
}
37+
if all(js.exit_status is None for js in job_submissions):
38+
job_submissions_excludes["exit_status"] = True
39+
latest_job_submission = current_resource.latest_job_submission
40+
if latest_job_submission is not None:
41+
latest_job_submission_excludes = {}
42+
current_resource_excludes["latest_job_submission"] = latest_job_submission_excludes
43+
if _should_exclude_job_submission_jpd_cpu_arch(latest_job_submission):
44+
latest_job_submission_excludes["job_provisioning_data"] = {
45+
"instance_type": {"resources": {"cpu_arch"}}
46+
}
47+
if _should_exclude_job_submission_jrd_cpu_arch(latest_job_submission):
48+
latest_job_submission_excludes["job_runtime_data"] = {
49+
"offer": {"instance": {"resources": {"cpu_arch"}}}
50+
}
51+
if latest_job_submission.exit_status is None:
52+
latest_job_submission_excludes["exit_status"] = True
53+
return {"plan": apply_plan_excludes}
54+
55+
56+
def get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[Dict]:
57+
"""
58+
Excludes new fields when they are not set to keep
59+
clients backward-compatibility with older servers.
60+
"""
61+
get_plan_excludes = {}
62+
run_spec_excludes = get_run_spec_excludes(request.run_spec)
63+
if run_spec_excludes is not None:
64+
get_plan_excludes["run_spec"] = run_spec_excludes
65+
if request.max_offers is None:
66+
get_plan_excludes["max_offers"] = True
67+
return get_plan_excludes
68+
69+
70+
def get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
71+
"""
72+
Returns `run_spec` exclude mapping to exclude certain fields from the request.
73+
Use this method to exclude new fields when they are not set to keep
74+
clients backward-compatibility with older servers.
75+
"""
76+
spec_excludes: dict[str, Any] = {}
77+
configuration_excludes: dict[str, Any] = {}
78+
profile_excludes: set[str] = set()
79+
configuration = run_spec.configuration
80+
profile = run_spec.profile
81+
82+
if configuration.fleets is None:
83+
configuration_excludes["fleets"] = True
84+
if profile is not None and profile.fleets is None:
85+
profile_excludes.add("fleets")
86+
if configuration.tags is None:
87+
configuration_excludes["tags"] = True
88+
if profile is not None and profile.tags is None:
89+
profile_excludes.add("tags")
90+
if isinstance(configuration, ServiceConfiguration) and not configuration.rate_limits:
91+
configuration_excludes["rate_limits"] = True
92+
if configuration.shell is None:
93+
configuration_excludes["shell"] = True
94+
if configuration.priority is None:
95+
configuration_excludes["priority"] = True
96+
if configuration.startup_order is None:
97+
configuration_excludes["startup_order"] = True
98+
if profile is not None and profile.startup_order is None:
99+
profile_excludes.add("startup_order")
100+
if configuration.stop_criteria is None:
101+
configuration_excludes["stop_criteria"] = True
102+
if profile is not None and profile.stop_criteria is None:
103+
profile_excludes.add("stop_criteria")
104+
105+
if configuration_excludes:
106+
spec_excludes["configuration"] = configuration_excludes
107+
if profile_excludes:
108+
spec_excludes["profile"] = profile_excludes
109+
if spec_excludes:
110+
return spec_excludes
111+
return None
112+
113+
114+
def _should_exclude_job_submission_jpd_cpu_arch(job_submission: JobSubmission) -> bool:
115+
try:
116+
return job_submission.job_provisioning_data.instance_type.resources.cpu_arch is None
117+
except AttributeError:
118+
return True
119+
120+
121+
def _should_exclude_job_submission_jrd_cpu_arch(job_submission: JobSubmission) -> bool:
122+
try:
123+
return job_submission.job_runtime_data.offer.instance.resources.cpu_arch is None
124+
except AttributeError:
125+
return True

‎src/dstack/_internal/core/compatibility/volumes.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
from typing import Dict
2+
3+
from dstack._internal.core.models.volumes import VolumeConfiguration, VolumeSpec
4+
5+
6+
def get_volume_spec_excludes(volume_spec: VolumeSpec) -> Dict:
7+
"""
8+
Returns `volume_spec` exclude mapping to exclude certain fields from the request.
9+
Use this method to exclude new fields when they are not set to keep
10+
clients backward-compatibility with older servers.
11+
"""
12+
spec_excludes = {}
13+
spec_excludes["configuration"] = _get_volume_configuration_excludes(volume_spec.configuration)
14+
return spec_excludes
15+
16+
17+
def get_create_volume_excludes(configuration: VolumeConfiguration) -> Dict:
18+
"""
19+
Returns an exclude mapping to exclude certain fields from the create volume request.
20+
Use this method to exclude new fields when they are not set to keep
21+
clients backward-compatibility with older servers.
22+
"""
23+
create_volume_excludes = {}
24+
create_volume_excludes["configuration"] = _get_volume_configuration_excludes(configuration)
25+
return create_volume_excludes
26+
27+
28+
def _get_volume_configuration_excludes(configuration: VolumeConfiguration) -> Dict:
29+
configuration_excludes = {}
30+
if configuration.tags is None:
31+
configuration_excludes["tags"] = True
32+
return configuration_excludes

‎src/dstack/api/server/_fleets.py‎

Lines changed: 9 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
1-
from typing import Any, Dict, List, Optional, Union
1+
from typing import List, Union
22

33
from pydantic import parse_obj_as
44

5+
from dstack._internal.core.compatibility.fleets import (
6+
get_apply_plan_excludes,
7+
get_create_fleet_excludes,
8+
get_get_plan_excludes,
9+
)
510
from dstack._internal.core.models.fleets import ApplyFleetPlanInput, Fleet, FleetPlan, FleetSpec
6-
from dstack._internal.core.models.instances import Instance
711
from dstack._internal.server.schemas.fleets import (
812
ApplyFleetPlanRequest,
913
CreateFleetRequest,
@@ -34,7 +38,7 @@ def get_plan(
3438
spec: FleetSpec,
3539
) -> FleetPlan:
3640
body = GetFleetPlanRequest(spec=spec)
37-
body_json = body.json(exclude=_get_get_plan_excludes(spec))
41+
body_json = body.json(exclude=get_get_plan_excludes(spec))
3842
resp = self._request(f"/api/project/{project_name}/fleets/get_plan", body=body_json)
3943
return parse_obj_as(FleetPlan.__response__, resp.json())
4044

@@ -46,7 +50,7 @@ def apply_plan(
4650
) -> Fleet:
4751
plan_input = ApplyFleetPlanInput.__response__.parse_obj(plan)
4852
body = ApplyFleetPlanRequest(plan=plan_input, force=force)
49-
body_json = body.json(exclude=_get_apply_plan_excludes(plan_input))
53+
body_json = body.json(exclude=get_apply_plan_excludes(plan_input))
5054
resp = self._request(f"/api/project/{project_name}/fleets/apply", body=body_json)
5155
return parse_obj_as(Fleet.__response__, resp.json())
5256

@@ -66,74 +70,6 @@ def create(
6670
spec: FleetSpec,
6771
) -> Fleet:
6872
body = CreateFleetRequest(spec=spec)
69-
body_json = body.json(exclude=_get_create_fleet_excludes(spec))
73+
body_json = body.json(exclude=get_create_fleet_excludes(spec))
7074
resp = self._request(f"/api/project/{project_name}/fleets/create", body=body_json)
7175
return parse_obj_as(Fleet.__response__, resp.json())
72-
73-
74-
def _get_get_plan_excludes(fleet_spec: FleetSpec) -> Dict:
75-
get_plan_excludes = {}
76-
spec_excludes = _get_fleet_spec_excludes(fleet_spec)
77-
if spec_excludes:
78-
get_plan_excludes["spec"] = spec_excludes
79-
return get_plan_excludes
80-
81-
82-
def _get_apply_plan_excludes(plan_input: ApplyFleetPlanInput) -> Dict:
83-
apply_plan_excludes = {}
84-
spec_excludes = _get_fleet_spec_excludes(plan_input.spec)
85-
if spec_excludes:
86-
apply_plan_excludes["spec"] = spec_excludes
87-
current_resource = plan_input.current_resource
88-
if current_resource is not None:
89-
current_resource_excludes = {}
90-
apply_plan_excludes["current_resource"] = current_resource_excludes
91-
if all(map(_should_exclude_instance_cpu_arch, current_resource.instances)):
92-
current_resource_excludes["instances"] = {
93-
"__all__": {"instance_type": {"resources": {"cpu_arch"}}}
94-
}
95-
return {"plan": apply_plan_excludes}
96-
97-
98-
def _should_exclude_instance_cpu_arch(instance: Instance) -> bool:
99-
try:
100-
return instance.instance_type.resources.cpu_arch is None
101-
except AttributeError:
102-
return True
103-
104-
105-
def _get_create_fleet_excludes(fleet_spec: FleetSpec) -> Dict:
106-
create_fleet_excludes = {}
107-
spec_excludes = _get_fleet_spec_excludes(fleet_spec)
108-
if spec_excludes:
109-
create_fleet_excludes["spec"] = spec_excludes
110-
return create_fleet_excludes
111-
112-
113-
def _get_fleet_spec_excludes(fleet_spec: FleetSpec) -> Optional[Dict]:
114-
"""
115-
Returns `fleet_spec` exclude mapping to exclude certain fields from the request.
116-
Use this method to exclude new fields when they are not set to keep
117-
clients backward-compatibility with older servers.
118-
"""
119-
spec_excludes: Dict[str, Any] = {}
120-
configuration_excludes: Dict[str, Any] = {}
121-
profile_excludes: set[str] = set()
122-
profile = fleet_spec.profile
123-
if profile.fleets is None:
124-
profile_excludes.add("fleets")
125-
if fleet_spec.configuration.tags is None:
126-
configuration_excludes["tags"] = True
127-
if profile.tags is None:
128-
profile_excludes.add("tags")
129-
if profile.startup_order is None:
130-
profile_excludes.add("startup_order")
131-
if profile.stop_criteria is None:
132-
profile_excludes.add("stop_criteria")
133-
if configuration_excludes:
134-
spec_excludes["configuration"] = configuration_excludes
135-
if profile_excludes:
136-
spec_excludes["profile"] = profile_excludes
137-
if spec_excludes:
138-
return spec_excludes
139-
return None

0 commit comments

Comments
 (0)