{"action":"created","comment":{"url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/comments/228656589","pull_request_review_id":168955066,"id":228656589,"node_id":"MDI0OlB1bGxSZXF1ZXN0UmV2aWV3Q29tbWVudDIyODY1NjU4OQ==","diff_hunk":"@@ -13,16 +13,127 @@\n \n from __future__ import absolute_import\n \n+import json\n import logging\n \n import sagemaker_containers.beta.framework as framework\n \n+\n logger = logging.getLogger(__name__)\n \n \n+SAGEMAKER_PARAMETER_SERVER_NUM = 'sagemaker_parameter_server_num'\n+\n+\n+def _is_host_master(hosts, current_host):\n+    return current_host == hosts[0]\n+\n+\n+def _build_tf_config(hosts, current_host, ps_num=0, ps_task=False):\n+    \"\"\"\n+    Builds a dictionary containing cluster information based on number of hosts and number of\n+    parameter servers.\n+\n+    Args:\n+        hosts (list[str]): List of host name in the cluster\n+        current_host (str): Current host name\n+        ps_num (int): Number of parameter servers\n+        ps_task (bool): Set to True if this config is built for a ps server process\n+\n+    Returns:\n+        dict[str: dict]: A dictionary describing the cluster setup for distributed training.\n+        For more information regarding TF_CONFIG:\n+        https://cloud.google.com/ml-engine/docs/tensorflow/distributed-training-details\n+\n+    \"\"\"\n+\n+    # Assign the first host as the master. Rest of the hosts if any will be worker hosts.\n+    # The first ps_num hosts will also have a parameter task assign to them.\n+    masters = hosts[:1]\n+    workers = hosts[1:]\n+    ps = hosts[:ps_num] if len(hosts) > 1 and ps_num > 0 else None\n+\n+    def host_addresses(hosts, port=2222):\n+        return [\"{}:{}\".format(host, port) for host in hosts]\n+    tf_config = {\n+        \"cluster\": {\n+            \"master\": host_addresses(masters)\n+        },\n+        \"environment\": \"cloud\"\n+    }\n+\n+    if ps is not None:\n+        tf_config[\"cluster\"][\"ps\"] = host_addresses(ps, port='2223')\n+\n+    if len(workers) > 0:\n+        tf_config[\"cluster\"][\"worker\"] = host_addresses(workers)\n+\n+    if ps_task:\n+        if ps is None:\n+            raise ValueError(\n+                \"Can not have a ps task if there are no parameter servers in the cluster\")\n+        task_type = \"ps\"\n+        task_index = ps.index(current_host)\n+    elif _is_host_master(hosts, current_host):\n+        task_type = \"master\"\n+        task_index = 0\n+    else:\n+        task_type = \"worker\"\n+        task_index = workers.index(current_host)\n+\n+    tf_config[\"task\"] = {\"index\": task_index, \"type\": task_type}\n+    return tf_config\n+\n+\n+def _get_env_vars_with_tf_config(env, ps_task):","path":"src/sagemaker_tensorflow_container/training.py","position":76,"original_position":76,"commit_id":"d0812a1a0acaf744e3673dc7881d683b386fcadc","original_commit_id":"d0812a1a0acaf744e3673dc7881d683b386fcadc","user":{"login":"laurenyu","id":6631887,"node_id":"MDQ6VXNlcjY2MzE4ODc=","avatar_url":"https://avatars2.githubusercontent.com/u/6631887?v=4","gravatar_id":"","url":"https://api.github.com/users/laurenyu","html_url":"https://github.com/laurenyu","followers_url":"https://api.github.com/users/laurenyu/followers","following_url":"https://api.github.com/users/laurenyu/following{/other_user}","gists_url":"https://api.github.com/users/laurenyu/gists{/gist_id}","starred_url":"https://api.github.com/users/laurenyu/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/laurenyu/subscriptions","organizations_url":"https://api.github.com/users/laurenyu/orgs","repos_url":"https://api.github.com/users/laurenyu/repos","events_url":"https://api.github.com/users/laurenyu/events{/privacy}","received_events_url":"https://api.github.com/users/laurenyu/received_events","type":"User","site_admin":false},"body":"I think prefixing with \"get\" is considered a little unpythonic or whatever - anyway, you can make it just `_env_vars_with_tf_config`","created_at":"2018-10-26T20:28:00Z","updated_at":"2018-10-26T20:29:04Z","html_url":"https://github.com/aws/sagemaker-tensorflow-container/pull/98#discussion_r228656589","pull_request_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/98","author_association":"MEMBER","_links":{"self":{"href":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/comments/228656589"},"html":{"href":"https://github.com/aws/sagemaker-tensorflow-container/pull/98#discussion_r228656589"},"pull_request":{"href":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/98"}}},"pull_request":{"url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/98","id":226173891,"node_id":"MDExOlB1bGxSZXF1ZXN0MjI2MTczODkx","html_url":"https://github.com/aws/sagemaker-tensorflow-container/pull/98","diff_url":"https://github.com/aws/sagemaker-tensorflow-container/pull/98.diff","patch_url":"https://github.com/aws/sagemaker-tensorflow-container/pull/98.patch","issue_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/issues/98","number":98,"state":"open","locked":false,"title":"Add distributed training support","user":{"login":"icywang86rui","id":24926915,"node_id":"MDQ6VXNlcjI0OTI2OTE1","avatar_url":"https://avatars0.githubusercontent.com/u/24926915?v=4","gravatar_id":"","url":"https://api.github.com/users/icywang86rui","html_url":"https://github.com/icywang86rui","followers_url":"https://api.github.com/users/icywang86rui/followers","following_url":"https://api.github.com/users/icywang86rui/following{/other_user}","gists_url":"https://api.github.com/users/icywang86rui/gists{/gist_id}","starred_url":"https://api.github.com/users/icywang86rui/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/icywang86rui/subscriptions","organizations_url":"https://api.github.com/users/icywang86rui/orgs","repos_url":"https://api.github.com/users/icywang86rui/repos","events_url":"https://api.github.com/users/icywang86rui/events{/privacy}","received_events_url":"https://api.github.com/users/icywang86rui/received_events","type":"User","site_admin":false},"body":"*Issue #, if available:*\r\n\r\n*Description of changes:*\r\nThis change sets up parameter servers for users in distributed training context. We read the desired number of parameter servers from hyperparameter 'sagemaker_parameter_server_num'. Use it to construct TF_CONFIG and launch ps and worker processes.\r\n\r\nBy submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.\r\n","created_at":"2018-10-26T17:37:43Z","updated_at":"2018-10-26T20:29:04Z","closed_at":null,"merged_at":null,"merge_commit_sha":"33aefaef8589867183292614b233ea96c3ee35b3","assignee":null,"assignees":[],"requested_reviewers":[{"login":"mvsusp","id":142074,"node_id":"MDQ6VXNlcjE0MjA3NA==","avatar_url":"https://avatars0.githubusercontent.com/u/142074?v=4","gravatar_id":"","url":"https://api.github.com/users/mvsusp","html_url":"https://github.com/mvsusp","followers_url":"https://api.github.com/users/mvsusp/followers","following_url":"https://api.github.com/users/mvsusp/following{/other_user}","gists_url":"https://api.github.com/users/mvsusp/gists{/gist_id}","starred_url":"https://api.github.com/users/mvsusp/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/mvsusp/subscriptions","organizations_url":"https://api.github.com/users/mvsusp/orgs","repos_url":"https://api.github.com/users/mvsusp/repos","events_url":"https://api.github.com/users/mvsusp/events{/privacy}","received_events_url":"https://api.github.com/users/mvsusp/received_events","type":"User","site_admin":false}],"requested_teams":[],"labels":[],"milestone":null,"commits_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/98/commits","review_comments_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/98/comments","review_comment_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/comments{/number}","comments_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/issues/98/comments","statuses_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/statuses/d0812a1a0acaf744e3673dc7881d683b386fcadc","head":{"label":"icywang86rui:distributed","ref":"distributed","sha":"d0812a1a0acaf744e3673dc7881d683b386fcadc","user":{"login":"icywang86rui","id":24926915,"node_id":"MDQ6VXNlcjI0OTI2OTE1","avatar_url":"https://avatars0.githubusercontent.com/u/24926915?v=4","gravatar_id":"","url":"https://api.github.com/users/icywang86rui","html_url":"https://github.com/icywang86rui","followers_url":"https://api.github.com/users/icywang86rui/followers","following_url":"https://api.github.com/users/icywang86rui/following{/other_user}","gists_url":"https://api.github.com/users/icywang86rui/gists{/gist_id}","starred_url":"https://api.github.com/users/icywang86rui/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/icywang86rui/subscriptions","organizations_url":"https://api.github.com/users/icywang86rui/orgs","repos_url":"https://api.github.com/users/icywang86rui/repos","events_url":"https://api.github.com/users/icywang86rui/events{/privacy}","received_events_url":"https://api.github.com/users/icywang86rui/received_events","type":"User","site_admin":false},"repo":{"id":147874768,"node_id":"MDEwOlJlcG9zaXRvcnkxNDc4NzQ3Njg=","name":"sagemaker-tensorflow-container","full_name":"icywang86rui/sagemaker-tensorflow-container","private":false,"owner":{"login":"icywang86rui","id":24926915,"node_id":"MDQ6VXNlcjI0OTI2OTE1","avatar_url":"https://avatars0.githubusercontent.com/u/24926915?v=4","gravatar_id":"","url":"https://api.github.com/users/icywang86rui","html_url":"https://github.com/icywang86rui","followers_url":"https://api.github.com/users/icywang86rui/followers","following_url":"https://api.github.com/users/icywang86rui/following{/other_user}","gists_url":"https://api.github.com/users/icywang86rui/gists{/gist_id}","starred_url":"https://api.github.com/users/icywang86rui/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/icywang86rui/subscriptions","organizations_url":"https://api.github.com/users/icywang86rui/orgs","repos_url":"https://api.github.com/users/icywang86rui/repos","events_url":"https://api.github.com/users/icywang86rui/events{/privacy}","received_events_url":"https://api.github.com/users/icywang86rui/received_events","type":"User","site_admin":false},"html_url":"https://github.com/icywang86rui/sagemaker-tensorflow-container","description":"This support code is used for making the TensorFlow framework run on Amazon SageMaker.","fork":true,"url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container","forks_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/forks","keys_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/keys{/key_id}","collaborators_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/teams","hooks_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/hooks","issue_events_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/issues/events{/number}","events_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/events","assignees_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/assignees{/user}","branches_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/branches{/branch}","tags_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/tags","blobs_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/git/refs{/sha}","trees_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/git/trees{/sha}","statuses_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/statuses/{sha}","languages_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/languages","stargazers_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/stargazers","contributors_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/contributors","subscribers_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/subscribers","subscription_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/subscription","commits_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/commits{/sha}","git_commits_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/git/commits{/sha}","comments_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/comments{/number}","issue_comment_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/issues/comments{/number}","contents_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/contents/{+path}","compare_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/compare/{base}...{head}","merges_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/merges","archive_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/downloads","issues_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/issues{/number}","pulls_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/pulls{/number}","milestones_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/milestones{/number}","notifications_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/labels{/name}","releases_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/releases{/id}","deployments_url":"https://api.github.com/repos/icywang86rui/sagemaker-tensorflow-container/deployments","created_at":"2018-09-07T21:06:01Z","updated_at":"2018-09-07T21:06:04Z","pushed_at":"2018-10-26T17:59:39Z","git_url":"git://github.com/icywang86rui/sagemaker-tensorflow-container.git","ssh_url":"git@github.com:icywang86rui/sagemaker-tensorflow-container.git","clone_url":"https://github.com/icywang86rui/sagemaker-tensorflow-container.git","svn_url":"https://github.com/icywang86rui/sagemaker-tensorflow-container","homepage":null,"size":13472,"stargazers_count":0,"watchers_count":0,"language":"Python","has_issues":false,"has_projects":true,"has_downloads":true,"has_wiki":false,"has_pages":false,"forks_count":0,"mirror_url":null,"archived":false,"open_issues_count":0,"license":{"key":"apache-2.0","name":"Apache License 2.0","spdx_id":"Apache-2.0","url":"https://api.github.com/licenses/apache-2.0","node_id":"MDc6TGljZW5zZTI="},"forks":0,"open_issues":0,"watchers":0,"default_branch":"master"}},"base":{"label":"aws:script-mode","ref":"script-mode","sha":"a1916a82616fcee3096b466b1bd8edbcbe0e0bf6","user":{"login":"aws","id":2232217,"node_id":"MDEyOk9yZ2FuaXphdGlvbjIyMzIyMTc=","avatar_url":"https://avatars3.githubusercontent.com/u/2232217?v=4","gravatar_id":"","url":"https://api.github.com/users/aws","html_url":"https://github.com/aws","followers_url":"https://api.github.com/users/aws/followers","following_url":"https://api.github.com/users/aws/following{/other_user}","gists_url":"https://api.github.com/users/aws/gists{/gist_id}","starred_url":"https://api.github.com/users/aws/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/aws/subscriptions","organizations_url":"https://api.github.com/users/aws/orgs","repos_url":"https://api.github.com/users/aws/repos","events_url":"https://api.github.com/users/aws/events{/privacy}","received_events_url":"https://api.github.com/users/aws/received_events","type":"Organization","site_admin":false},"repo":{"id":118646184,"node_id":"MDEwOlJlcG9zaXRvcnkxMTg2NDYxODQ=","name":"sagemaker-tensorflow-container","full_name":"aws/sagemaker-tensorflow-container","private":false,"owner":{"login":"aws","id":2232217,"node_id":"MDEyOk9yZ2FuaXphdGlvbjIyMzIyMTc=","avatar_url":"https://avatars3.githubusercontent.com/u/2232217?v=4","gravatar_id":"","url":"https://api.github.com/users/aws","html_url":"https://github.com/aws","followers_url":"https://api.github.com/users/aws/followers","following_url":"https://api.github.com/users/aws/following{/other_user}","gists_url":"https://api.github.com/users/aws/gists{/gist_id}","starred_url":"https://api.github.com/users/aws/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/aws/subscriptions","organizations_url":"https://api.github.com/users/aws/orgs","repos_url":"https://api.github.com/users/aws/repos","events_url":"https://api.github.com/users/aws/events{/privacy}","received_events_url":"https://api.github.com/users/aws/received_events","type":"Organization","site_admin":false},"html_url":"https://github.com/aws/sagemaker-tensorflow-container","description":"This support code is used for making the TensorFlow framework run on Amazon SageMaker.","fork":false,"url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container","forks_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/forks","keys_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/keys{/key_id}","collaborators_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/teams","hooks_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/hooks","issue_events_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/issues/events{/number}","events_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/events","assignees_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/assignees{/user}","branches_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/branches{/branch}","tags_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/tags","blobs_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/git/refs{/sha}","trees_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/git/trees{/sha}","statuses_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/statuses/{sha}","languages_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/languages","stargazers_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/stargazers","contributors_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/contributors","subscribers_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/subscribers","subscription_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/subscription","commits_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/commits{/sha}","git_commits_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/git/commits{/sha}","comments_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/comments{/number}","issue_comment_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/issues/comments{/number}","contents_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/contents/{+path}","compare_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/compare/{base}...{head}","merges_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/merges","archive_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/downloads","issues_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/issues{/number}","pulls_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls{/number}","milestones_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/milestones{/number}","notifications_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/labels{/name}","releases_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/releases{/id}","deployments_url":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/deployments","created_at":"2018-01-23T17:41:21Z","updated_at":"2018-10-25T21:18:06Z","pushed_at":"2018-10-26T17:59:41Z","git_url":"git://github.com/aws/sagemaker-tensorflow-container.git","ssh_url":"git@github.com:aws/sagemaker-tensorflow-container.git","clone_url":"https://github.com/aws/sagemaker-tensorflow-container.git","svn_url":"https://github.com/aws/sagemaker-tensorflow-container","homepage":null,"size":12920,"stargazers_count":81,"watchers_count":81,"language":"Python","has_issues":true,"has_projects":false,"has_downloads":true,"has_wiki":false,"has_pages":false,"forks_count":41,"mirror_url":null,"archived":false,"open_issues_count":12,"license":{"key":"apache-2.0","name":"Apache License 2.0","spdx_id":"Apache-2.0","url":"https://api.github.com/licenses/apache-2.0","node_id":"MDc6TGljZW5zZTI="},"forks":41,"open_issues":12,"watchers":81,"default_branch":"master"}},"_links":{"self":{"href":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/98"},"html":{"href":"https://github.com/aws/sagemaker-tensorflow-container/pull/98"},"issue":{"href":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/issues/98"},"comments":{"href":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/issues/98/comments"},"review_comments":{"href":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/98/comments"},"review_comment":{"href":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/comments{/number}"},"commits":{"href":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/pulls/98/commits"},"statuses":{"href":"https://api.github.com/repos/aws/sagemaker-tensorflow-container/statuses/d0812a1a0acaf744e3673dc7881d683b386fcadc"}},"author_association":"NONE"}}
{"actor":{"display_login":"laurenyu"}}