Add aws named resources (#225)

Summary:
Pull Request resolved: #225

Add  aws named resources

Reviewed By: kiukchung

Differential Revision: D31483254

fbshipit-source-id: 60ab95d6ae83cacfac42f0dc61026a544d704bb7

Author: aivanou

docs/source/specs.rst

@@ -26,8 +26,15 @@ Resource
 
 .. autofunction:: get_named_resources
 
+AWS Named Resources
+^^^^^^^^^^^^^^^^^^^^^
+.. automodule:: torchx.specs.named_resources_aws
+
+.. currentmodule:: torchx.specs.named_resources_aws
+
 Macros
 ------------
+.. currentmodule:: torchx.specs
 .. autoclass:: macros
    :members:
 

torchx/specs/__init__.py

@@ -7,6 +7,7 @@
 
 from typing import Dict
 
+import torchx.specs.named_resources_aws as aws_resources
 from torchx.util.entrypoints import load_group
 
 from .api import (  # noqa: F401 F403
@@ -48,6 +49,14 @@
 def _load_named_resources() -> Dict[str, Resource]:
     resource_methods = load_group("torchx.named_resources", default={})
     materialized_resources = {}
+    default = {
+        "aws_t3.medium": aws_resources.aws_t3_medium(),
+        "aws_m5.2xlarge": aws_resources.aws_m5_2xlarge(),
+        "aws_p3.2xlarge": aws_resources.aws_p3_2xlarge(),
+        "aws_p3.8xlarge": aws_resources.aws_p3_8xlarge(),
+    }
+    for name, resource in default.items():
+        materialized_resources[name] = resource
     for resource_name, resource_method in resource_methods.items():
         materialized_resources[resource_name] = resource_method()
     materialized_resources["NULL"] = NULL_RESOURCE

torchx/specs/named_resources_aws.py

@@ -0,0 +1,70 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+r"""
+`torchx.specs.named_resources_aws` contains resource definitions that represent corresponding AWS instance types
+taken from https://aws.amazon.com/ec2/instance-types/. The resources are exposed
+via entrypoints after installing torchx lib. The mapping is stored in the `setup.py` file.
+
+The named resources currently do not specify AWS instance type capabilities but merely represent
+the equvalent resource in mem, cpu and gpu numbers.
+
+.. note::
+    These resource definitions may change in future. It is expected for each user to
+    manage their own resources. Follow https://pytorch.org/torchx/latest/specs.html#torchx.specs.get_named_resources
+    to set up named resources.
+
+Usage:
+
+    ::
+
+     from torchx.specs import named_resources
+     print(named_resources["aws_t3.medium"])
+     print(named_resources["aws_m5.2xlarge"])
+     print(named_resources["aws_p3.2xlarge"])
+     print(named_resources["aws_p3.8xlarge"])
+
+"""
+
+from torchx.specs.api import Resource
+
+GiB: int = 1024
+
+
+def aws_p3_2xlarge() -> Resource:
+    return Resource(
+        cpu=8,
+        gpu=1,
+        memMB=61 * GiB,
+        capabilities={},
+    )
+
+
+def aws_p3_8xlarge() -> Resource:
+    return Resource(
+        cpu=32,
+        gpu=4,
+        memMB=244 * GiB,
+        capabilities={},
+    )
+
+
+def aws_t3_medium() -> Resource:
+    return Resource(
+        cpu=2,
+        gpu=0,
+        memMB=4 * GiB,
+        capabilities={},
+    )
+
+
+def aws_m5_2xlarge() -> Resource:
+    return Resource(
+        cpu=8,
+        gpu=0,
+        memMB=32 * GiB,
+        capabilities={},
+    )

torchx/specs/test/api_test.py

@@ -11,6 +11,8 @@
 from typing import Dict, List, Optional, Tuple, Union
 from unittest.mock import MagicMock, patch
 
+import torchx.specs.named_resources_aws as named_resources_aws
+from torchx.specs import named_resources
 from torchx.specs.api import (
     _TERMINAL_STATES,
     MISSING,
@@ -105,6 +107,20 @@ 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(self) -> None:
+        self.assertEqual(
+            named_resources_aws.aws_m5_2xlarge(), named_resources["aws_m5.2xlarge"]
+        )
+        self.assertEqual(
+            named_resources_aws.aws_t3_medium(), named_resources["aws_t3.medium"]
+        )
+        self.assertEqual(
+            named_resources_aws.aws_p3_2xlarge(), named_resources["aws_p3.2xlarge"]
+        )
+        self.assertEqual(
+            named_resources_aws.aws_p3_8xlarge(), named_resources["aws_p3.8xlarge"]
+        )
+
 
 class RoleBuilderTest(unittest.TestCase):
     def test_defaults(self) -> None:

torchx/specs/test/named_resources_test_aws.py

@@ -0,0 +1,42 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import unittest
+
+from torchx.specs.named_resources_aws import (
+    aws_p3_2xlarge,
+    aws_p3_8xlarge,
+    aws_m5_2xlarge,
+    aws_t3_medium,
+    GiB,
+)
+
+
+class NamedResourcesTest(unittest.TestCase):
+    def test_aws_p3_2xlarge(self) -> None:
+        resource = aws_p3_2xlarge()
+        self.assertEqual(8, resource.cpu)
+        self.assertEqual(1, resource.gpu)
+        self.assertEqual(61 * GiB, resource.memMB)
+
+    def test_aws_p3_8xlarge(self) -> None:
+        resource = aws_p3_8xlarge()
+        self.assertEqual(32, resource.cpu)
+        self.assertEqual(4, resource.gpu)
+        self.assertEqual(244 * GiB, resource.memMB)
+
+    def test_aws_m5_2xlarge(self) -> None:
+        resource = aws_m5_2xlarge()
+        self.assertEqual(8, resource.cpu)
+        self.assertEqual(0, resource.gpu)
+        self.assertEqual(32 * GiB, resource.memMB)
+
+    def test_aws_t3_medium(self) -> None:
+        resource = aws_t3_medium()
+        self.assertEqual(2, resource.cpu)
+        self.assertEqual(0, resource.gpu)
+        self.assertEqual(4 * GiB, resource.memMB)