torchx
torchx copied to clipboard
pytorch
support dataclasses for component args
Summary:
This change adds the ability to provide component arguments via a single dataclass (which can still be combined with positional varargs).
One motivation for supporting dataclasses is to facilitate composition of components, eg consider a scenario where component2 is an extension of component1 – it takes the same args as component1 + some extra ones, calls component1 and does some extra work on top with the additional arguments.
Before
def component1(arg1: str, arg2: int, ...argN: str): ...
# need to repeat all arguments, unclear which ones are specific to component2
def component2(arg1: str, arg2: int, ... argN: str, otherArg: str):
# need to pass all component1-specific args to component1 explicitly
app_def = component1(arg1, arg2, ... argN)
return do_something(app_def, otherArg)
After
dataclass
class Comp1Args:
arg1: str
arg2: int
...
argN: str
# separation of args for the 2 components is explicit via the dataclasses
# no need to spell out all the args thanks to inheritance
dataclass
class Comp2Args(Comp1Args):
other_arg: str
def component1(args: Comp1Args): ...
def component2(args: Comp2Args):
# no need to spell out all the args when calling component1
app_def = component1(args)
return do_something(app_def, other_arg)
Differential Revision: D75320316
