Skip to content

Commit

Permalink
schedulers: add macro support to metadata variables (#921)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: #921

This adds torchx.spec.macros support to the metadata variable values. This allows use cases like setting smcBridge to the app_id eg:
```
job_spec.roles[0].metadata["mast"] = {
        "HpcTaskGroupSpec": {
            "smcBridge": {
                "portName": "thrift",
                "smcTier": specs.macros.app_id,
            }
        }
    }
```

Reviewed By: manav-a

Differential Revision: D57519702
  • Loading branch information
jason-b-akers committed Jun 17, 2024
1 parent 2ec3673 commit f414e8c
Showing 1 changed file with 18 additions and 0 deletions.
18 changes: 18 additions & 0 deletions torchx/specs/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import copy
import json
import re
import typing
from dataclasses import asdict, dataclass, field
from datetime import datetime
from enum import Enum
Expand Down Expand Up @@ -189,8 +190,25 @@ def apply(self, role: "Role") -> "Role":
role = copy.deepcopy(role)
role.args = [self.substitute(arg) for arg in role.args]
role.env = {key: self.substitute(arg) for key, arg in role.env.items()}
role.metadata = self._apply_nested(role.metadata)

return role

def _apply_nested(self, d: typing.Dict[str, Any]) -> typing.Dict[str, Any]:
stack = [d]
while stack:
current_dict = stack.pop()
for k, v in current_dict.items():
if isinstance(v, dict):
stack.append(v)
elif isinstance(v, str):
current_dict[k] = self.substitute(v)
elif isinstance(v, list):
for i in range(len(v)):
if isinstance(v[i], str):
v[i] = self.substitute(v[i])
return d

def substitute(self, arg: str) -> str:
"""
substitute applies the values to the template arg.
Expand Down

0 comments on commit f414e8c

Please sign in to comment.