Skip to content

Commit 0706acf

Browse files
ndacostafacebook-github-bot
authored andcommitted
Allow field path env variables
Summary: # Context Currently the kubernetes scheduler translate all the env of the role into something like this. ``` - name: FOO value: bar ``` In some cases we may also need env variables of type: ``` - name: FOO valueFrom: fieldRef: fieldPath: bar.path ``` This diff adds the ability to specify env in the role that will be translated to the second option. # Changes - Create constant placeholder to insert in the value of env in roles if we want them to be translated to the second option - If placeholder is present in the value, convert to code below stripping the placeholder ``` V1EnvVar( name=key, value_from=V1EnvVarSource( field_ref=V1ObjectFieldSelector(field_path=value) ), ) ``` if not, translate to ``` V1EnvVar(name=key, value=value) ``` - update tests Differential Revision: D59865500
1 parent 7df8ea5 commit 0706acf

File tree

2 files changed

+35
-5
lines changed

2 files changed

+35
-5
lines changed

torchx/schedulers/kubernetes_scheduler.py

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,8 @@
169169

170170
LABEL_INSTANCE_TYPE = "node.kubernetes.io/instance-type"
171171

172+
PLACEHOLDER_FIELD_PATH = "[FIELD_PATH]"
173+
172174

173175
def sanitize_for_serialization(obj: object) -> object:
174176
from kubernetes import client
@@ -183,7 +185,9 @@ def role_to_pod(name: str, role: Role, service_account: Optional[str]) -> "V1Pod
183185
V1ContainerPort,
184186
V1EmptyDirVolumeSource,
185187
V1EnvVar,
188+
V1EnvVarSource,
186189
V1HostPathVolumeSource,
190+
V1ObjectFieldSelector,
187191
V1ObjectMeta,
188192
V1PersistentVolumeClaimVolumeSource,
189193
V1Pod,
@@ -303,9 +307,20 @@ def role_to_pod(name: str, role: Role, service_account: Optional[str]) -> "V1Pod
303307
image=role.image,
304308
name=name,
305309
env=[
306-
V1EnvVar(
307-
name=name,
308-
value=value,
310+
(
311+
V1EnvVar(
312+
name=name,
313+
value_from=V1EnvVarSource(
314+
field_ref=V1ObjectFieldSelector(
315+
field_path=value.strip(PLACEHOLDER_FIELD_PATH)
316+
)
317+
),
318+
)
319+
if value.startswith(PLACEHOLDER_FIELD_PATH)
320+
else V1EnvVar(
321+
name=name,
322+
value=value,
323+
)
309324
)
310325
for name, value in role.env.items()
311326
],

torchx/schedulers/test/kubernetes_scheduler_test.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
KubernetesOpts,
2929
KubernetesScheduler,
3030
LABEL_INSTANCE_TYPE,
31+
PLACEHOLDER_FIELD_PATH,
3132
role_to_pod,
3233
)
3334
from torchx.specs import AppState
@@ -75,7 +76,7 @@ def _test_app(num_replicas: int = 1) -> specs.AppDef:
7576
"--rank0-env",
7677
specs.macros.rank0_env,
7778
],
78-
env={"FOO": "bar"},
79+
env={"FOO": "bar", "FOO_FIELD_PATH": f"{PLACEHOLDER_FIELD_PATH}bar"},
7980
resource=specs.Resource(
8081
cpu=2,
8182
memMB=3000,
@@ -149,7 +150,9 @@ def test_role_to_pod(self) -> None:
149150
V1ContainerPort,
150151
V1EmptyDirVolumeSource,
151152
V1EnvVar,
153+
V1EnvVarSource,
152154
V1HostPathVolumeSource,
155+
V1ObjectFieldSelector,
153156
V1ObjectMeta,
154157
V1Pod,
155158
V1PodSpec,
@@ -188,7 +191,15 @@ def test_role_to_pod(self) -> None:
188191
],
189192
image="pytorch/torchx:latest",
190193
name="name",
191-
env=[V1EnvVar(name="FOO", value="bar")],
194+
env=[
195+
V1EnvVar(name="FOO", value="bar"),
196+
V1EnvVar(
197+
name="FOO_FIELD_PATH",
198+
value_from=V1EnvVarSource(
199+
field_ref=V1ObjectFieldSelector(field_path="bar")
200+
),
201+
),
202+
],
192203
resources=resources,
193204
ports=[V1ContainerPort(name="foo", container_port=1234)],
194205
security_context=V1SecurityContext(),
@@ -303,6 +314,10 @@ def test_submit_dryrun(self) -> None:
303314
env:
304315
- name: FOO
305316
value: bar
317+
- name: FOO_FIELD_PATH
318+
valueFrom:
319+
fieldRef:
320+
fieldPath: bar
306321
- name: TORCHX_RANK0_HOST
307322
value: localhost
308323
image: pytorch/torchx:latest

0 commit comments

Comments
 (0)