Allow custom component tagging in DeviceOptions.node_name

Summary:
Modify detect_components to take a list of valid node_name prefixes instead of values.  Users can set node_name to e.g. `'sparse_component:0'`, `'sparse_component:1'`, etc.
and pass `'sparse_component:'` as a valid prefix.  Also add `Tags.SPARSE_COMPONENT` in addition to `Tags.SPARSE_SHARDED` and `Tags.SPARSE_DONT_SHARD` and update all calls to
`detect_device_components`.

Reviewed By: azzolini

Differential Revision: D6952599

fbshipit-source-id: e1b1e6b146a6bd053b295690016044fd5990c893

1 files changed, 20 insertions, 6 deletions

diff --git a/caffe2/python/layers/tags.py b/caffe2/python/layers/tags.py
index 0195563c0d..9f7a20315d 100644
--- a/caffe2/python/layers/tags.py
+++ b/caffe2/python/layers/tags.py
@@ -56,14 +56,28 @@ class Tags(object):
     PREFER_GPU = 'prefer_gpu'
     CPU_ONLY = 'cpu_only'
 
-    # The following two tags are hints to **distributed training framework**.
-    # The first tag is to determine a layer contains a sparse parameter and the
-    # parameter should be sharded and operators on those parameters should be
-    # done on distributed parameter servers. The second tag is to determine a
-    # layer contains a sparse parameters among others, and that the parameters
-    # should not be sharded (i.e. should be placed together on a node)
+    # The following three tags are hints to **distributed training framework**.
+    """
+    Indicates a layer contains a sparse shardable parameter.  The parameter
+    should be sharded nd operators on those parameters should be done on
+    distributed parameter servers.
+    """
     SPARSE_SHARDED = 'sparse_sharded'
+    """
+    Indicates a layer contains a sparse parameters among others, and that the
+    parameters should not be sharded (i.e. should be placed together on a node).
+    """
     SPARSE_DONT_SHARD = 'sparse_dont_shard'
+    """
+    Used to manually indicate a component for an operator.  Parameters for
+    all operators with the same component should be colocated on the same
+    parameter server.
+    """
+    COMPONENT = 'component:'
+    """
+    Valid tag prefixes for distributed training framework.
+    """
+    DT_TAGS = (SPARSE_SHARDED, SPARSE_DONT_SHARD, COMPONENT)
 
     # In certain cases we want to have different schema for training and
     # prediction, as an example in prediction we might need to have only