Rename named_resource to named_resources, make named_resources statically initialize resources (#63)

Summary:
Pull Request resolved: #63

Rename ``named_resource`` to ``named_resources``, make ``named_resources`` statically initialize resources

Reviewed By: kiukchung

Differential Revision: D29153123

fbshipit-source-id: ed8718838c71571f91ba233abfe492c1f8741394

Author: aivanou

docs/source/components/base.rst

@@ -4,8 +4,6 @@ Base
 .. automodule:: torchx.components.base
 .. currentmodule:: torchx.components.base
 
-.. autofunction:: named_resource
-
 .. automodule:: torchx.components.base.binary_component
 .. currentmodule:: torchx.components.base.binary_component
 .. autofunction:: binary_component

docs/source/configure.rst

@@ -70,12 +70,22 @@ representation:
 
 ::
 
- [torchx.schedulers]
+ [torchx.named_resources]
  gpu_x_2 = my_module.resources:gpu_x_2
 
 
 The named resource after that can be used in the following manner:
 
+.. code-block:: python
+
+  # my_module.component
+  from torchx.specs import AppDef, Role, named_resources
+  from torchx.components.base import torch_dist_role
+
+  app = AppDef(name="test_app", roles=[Role(.., resource=named_resources("gpu_x_2"))])
+
+or using ``torch.distributed.run``:
+
 .. code-block:: python
 
   # my_module.component

docs/source/specs.rst

@@ -24,6 +24,8 @@ Resource
 .. autoclass:: Resource
    :members:
 
+.. autofunction:: get_named_resources
+
 Macros
 ------------
 .. autoclass:: macros

examples/apps/lightning_classy_vision/component.py

@@ -14,7 +14,7 @@
 from typing import Optional
 
 import torchx.specs.api as torchx
-from torchx.components.base import named_resource
+from torchx.components.base import named_resources
 from torchx.components.base.binary_component import binary_component
 
 
@@ -50,7 +50,7 @@ def trainer(
             data_path,
         ],
         image=image,
-        resource=named_resource(resource)
+        resource=named_resources[resource]
         if resource
         else torchx.Resource(cpu=1, gpu=0, memMB=1024),
     )
@@ -85,7 +85,7 @@ def interpret(
             output_path,
         ],
         image=image,
-        resource=named_resource(resource)
+        resource=named_resources[resource]
         if resource
         else torchx.Resource(cpu=1, gpu=0, memMB=1024),
     )

torchx/components/base/__init__.py

@@ -12,7 +12,7 @@
 """
 from typing import Any, Dict, List, Optional, Union
 
-from torchx.specs import named_resource
+from torchx.specs import named_resources
 from torchx.specs.api import NULL_RESOURCE, Resource, RetryPolicy, Role
 from torchx.util.entrypoints import load
 
@@ -23,7 +23,7 @@ def _resolve_resource(resource: Union[str, Resource]) -> Resource:
     if isinstance(resource, Resource):
         return resource
     else:
-        return named_resource(resource.lower())
+        return named_resources[resource.upper()]
 
 
 def torch_dist_role(

torchx/components/base/test/lib_test.py

@@ -18,7 +18,7 @@ def test_torch_dist_role_str_resources(self) -> None:
         expected_resources = {}
         with patch.object(metadata, "entry_points") as entry_points_mock:
             entry_points_mock.return_value = {}
-            with self.assertRaises(ValueError):
+            with self.assertRaises(KeyError):
                 torch_dist_role(
                     name="dist_role",
                     image="test_image",

torchx/specs/__init__.py

@@ -6,16 +6,28 @@
 # LICENSE file in the root directory of this source tree.
 
 from .api import *  # noqa: F401 F403
-from torchx.util.entrypoints import load
+from torchx.util.entrypoints import load_group
 
 GiB: int = 1024
 
 
-def named_resource(resource: str) -> Resource:
+def _load_named_resources() -> Dict[str, Resource]:
+    resource_methods = load_group("torchx.named_resources", default={})
+    materialized_resources = {}
+    for resource_name, resource_method in resource_methods.items():
+        materialized_resources[resource_name] = resource_method()
+    materialized_resources["NULL"] = NULL_RESOURCE
+    return materialized_resources
+
+
+named_resources: Dict[str, Resource] = _load_named_resources()
+
+
+def get_named_resources(res: str) -> Resource:
     """
-    Gets resource object based on the string definition registered via entrypoints.txt.
+    Get resource object based on the string definition registered via entrypoints.txt.
 
-    Torchx implements named resource registration mechanism, which consists of
+    Torchx implements ``named_resource`` registration mechanism, which consists of
     the following steps:
 
     1. Create a module and define your resource retrieval function:
@@ -33,20 +45,15 @@ def gpu_x_1() -> Dict[str, Resource]:
 
     ::
 
-     [torchx.schedulers]
+     [torchx.named_resources]
      gpu_x_1 = my_module.resources:gpu_x_1
 
     The ``gpu_x_1`` can be used as string argument to this function:
 
     ::
 
-     resource = named_resource("gpu_x_1")
+     from torchx.specs import named_resources
+     resource = named_resources["gpu_x_1"]
 
     """
-    try:
-        resource_fn = load("torchx.resources", name=resource.lower())
-    except KeyError:
-        raise ValueError(
-            f"Resource: {resource} is not registered, check entrypoints.txt file."
-        )
-    return resource_fn()
+    return named_resources[res]