Skip to content

Commit 5058b6b

Browse files
schedulers: add macro support to metadata variables
Differential Revision: D57519702 Pull Request resolved: #921
1 parent 2ec3673 commit 5058b6b

File tree

1 file changed

+18
-0
lines changed

1 file changed

+18
-0
lines changed

torchx/specs/api.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import copy
1111
import json
1212
import re
13+
import typing
1314
from dataclasses import asdict, dataclass, field
1415
from datetime import datetime
1516
from enum import Enum
@@ -189,8 +190,25 @@ def apply(self, role: "Role") -> "Role":
189190
role = copy.deepcopy(role)
190191
role.args = [self.substitute(arg) for arg in role.args]
191192
role.env = {key: self.substitute(arg) for key, arg in role.env.items()}
193+
role.metadata = self._apply_nested(role.metadata)
194+
192195
return role
193196

197+
def _apply_nested(self, d: typing.Dict[str, Any]) -> typing.Dict[str, Any]:
198+
stack = [d]
199+
while stack:
200+
current_dict = stack.pop()
201+
for k, v in current_dict.items():
202+
if isinstance(v, dict):
203+
stack.append(v)
204+
elif isinstance(v, str):
205+
current_dict[k] = self.substitute(v)
206+
elif isinstance(v, list):
207+
for i in range(len(v)):
208+
if isinstance(v[i], str):
209+
v[i] = self.substitute(v[i])
210+
return d
211+
194212
def substitute(self, arg: str) -> str:
195213
"""
196214
substitute applies the values to the template arg.

0 commit comments

Comments
 (0)