(torchx/specs) Make torchx.specs.named_resources iterable (#1163)

Summary:

There's a bunch of use-cases where it is useful to be able to iterate over all the registered named resources. Implementing the `__iter__()` method makes this possible.

NOTE: To keep the API consistent, the `_NamedResourcesLibrary.__iter__()` method materializes the resource objec by calling the factory function. This means that if you iterate over `named_resources` early in the program, you end up losing out on any lazy materialization.

Reviewed By: AbishekS

Differential Revision: D86984810

Author: kiukchung

torchx/specs/__init__.py

@@ -14,7 +14,7 @@
 import difflib
 
 import os
-from typing import Callable, Dict, Mapping, Optional
+from typing import Callable, Dict, Iterator, Mapping, Optional
 
 from torchx.specs.api import (
     ALL,
@@ -113,8 +113,9 @@ def __getitem__(self, key: str) -> Resource:
     def __contains__(self, key: str) -> bool:
         return key in _named_resource_factories
 
-    def __iter__(self) -> None:
-        raise NotImplementedError("named resources doesn't support iterating")
+    def __iter__(self) -> Iterator[tuple[str, Resource]]:
+        for key in _named_resource_factories:
+            yield (key, self[key])
 
 
 named_resources: _NamedResourcesLibrary = _NamedResourcesLibrary()

torchx/specs/test/api_test.py

@@ -361,6 +361,15 @@ def test_copy_resource(self) -> None:
         self.assertEqual(new_resource.capabilities["new_key"], "new_value")
         self.assertEqual(resource.capabilities["test_key"], "test_value")
 
+    def test_named_resources_iterator(self) -> None:
+        registered_named_resources = {}
+        for name, resource_obj in named_resources:
+            self.assertEqual(resource(h=name), resource_obj)
+            registered_named_resources[name] = resource_obj
+
+        # validate that the for-loop was not vacuous
+        self.assertTrue(registered_named_resources)
+
     def test_named_resources(self) -> None:
         self.assertEqual(
             named_resources_aws.aws_m5_2xlarge(), named_resources["aws_m5.2xlarge"]