Skip to content

Commit eef90a2

Browse files
kiukchungfacebook-github-bot
authored andcommitted
(torchx-fb/named_resources) semi-automate registration of fractional resources for gpu types
Summary: Context: D82556202 ``` register_fractionals def zion_4s_80g(fractional: float = 1.0) -> Resource: ... ``` Creates: 1. `def zion_4s_80g_1(): return zion_4s_80g(fractional=EIGHTH)` 1. `zion_4s_80g_2(): return zion_4s_80g(fractional=QUARTER)` 1. `zion_4s_80g_4(): return zion_4s_80g(fractional=HALF)` 1. `zion_4s_80g_8(): return zion_4s_80g(fractional=WHOLE)` so that we don't have to manually do this NOTE: each of the fractional types (including WHOLE) will have `resource.capabilities[IS_FRACTIONAL] == True` so that we can use it as a way to determine whether a resource originates from a fractional named resource. Reviewed By: ajauhri Differential Revision: D82594325
1 parent c534e45 commit eef90a2

File tree

2 files changed

+14
-4
lines changed

2 files changed

+14
-4
lines changed

torchx/specs/__init__.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -52,14 +52,19 @@
5252

5353
from torchx.util.modules import import_attr
5454

55-
AWS_NAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = import_attr(
55+
GiB: int = 1024
56+
57+
ResourceFactory = Callable[[], Resource]
58+
59+
AWS_NAMED_RESOURCES: Mapping[str, ResourceFactory] = import_attr(
5660
"torchx.specs.named_resources_aws", "NAMED_RESOURCES", default={}
5761
)
58-
GENERIC_NAMED_RESOURCES: Mapping[str, Callable[[], Resource]] = import_attr(
62+
GENERIC_NAMED_RESOURCES: Mapping[str, ResourceFactory] = import_attr(
5963
"torchx.specs.named_resources_generic", "NAMED_RESOURCES", default={}
6064
)
61-
62-
GiB: int = 1024
65+
FB_NAMED_RESOURCES: Mapping[str, ResourceFactory] = import_attr(
66+
"torchx.specs.fb.named_resources", "NAMED_RESOURCES", default={}
67+
)
6368

6469

6570
def _load_named_resources() -> Dict[str, Callable[[], Resource]]:
@@ -69,6 +74,7 @@ def _load_named_resources() -> Dict[str, Callable[[], Resource]]:
6974
for name, resource in {
7075
**GENERIC_NAMED_RESOURCES,
7176
**AWS_NAMED_RESOURCES,
77+
**FB_NAMED_RESOURCES,
7278
**resource_methods,
7379
}.items():
7480
materialized_resources[name] = resource

torchx/specs/api.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ class Resource:
8383
memMB: MB of ram
8484
capabilities: additional hardware specs (interpreted by scheduler)
8585
devices: a list of named devices with their quantities
86+
tags: metadata tags for the resource (not interpreted by schedulers)
87+
used to add non-functional information about resources (e.g. whether it is an alias of another resource)
8688
8789
Note: you should prefer to use named_resources instead of specifying the raw
8890
resource requirement directly.
@@ -93,6 +95,7 @@ class Resource:
9395
memMB: int
9496
capabilities: Dict[str, Any] = field(default_factory=dict)
9597
devices: Dict[str, int] = field(default_factory=dict)
98+
tags: Dict[str, object] = field(default_factory=dict)
9699

97100
@staticmethod
98101
def copy(original: "Resource", **capabilities: Any) -> "Resource":
@@ -101,6 +104,7 @@ def copy(original: "Resource", **capabilities: Any) -> "Resource":
101104
are present in the original resource and as parameter, the one from parameter
102105
will be used.
103106
"""
107+
104108
res_capabilities = dict(original.capabilities)
105109
res_capabilities.update(capabilities)
106110
return Resource(

0 commit comments

Comments
 (0)