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

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,22 @@ 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[str]:
+        """Iterates through the names of the registered named_resources.
+
+        Usage:
+
+        .. doctest::
+
+            from torchx import specs
+
+            for resource_name in specs.named_resources:
+                resource = specs.resource(h=resource_name)
+                assert isinstance(resource, specs.Resource)
+
+        """
+        for key in _named_resource_factories:
+            yield (key)
 
 
 named_resources: _NamedResourcesLibrary = _NamedResourcesLibrary()

torchx/specs/test/api_test.py

@@ -361,6 +361,16 @@ 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 = set()
+        for resource_name in named_resources:
+            # just make sure we can create the resource using the name
+            self.assertIsInstance(resource(h=resource_name), Resource)
+            registered_named_resources.add(resource_name)
+
+        # 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"]