detect user's country and using corresponding docker mirror

Author: Zhao Wang

brainframe/cli/commands/install.py

@@ -1,5 +1,7 @@
 from argparse import ArgumentParser
 from pathlib import Path
+import json
+from urllib.request import urlopen
 
 import i18n
 
@@ -103,6 +105,13 @@ def install():
     docker_compose.download(
         install_path / "docker-compose.yml", version=args.version
     )
+    location = check_location()
+    try:
+        docker_config = get_docker_config()
+    except Exception:
+        print_utils.warning_translate("install.fail-to-load-docker-config")
+    else:
+        update_docker_mirror(location=location, docker_config=docker_config)
 
     print_utils.translate("install.downloading-images")
     docker_compose.run(install_path, ["pull"])
@@ -120,8 +129,8 @@ def install():
     # Recommend to the user to add their custom paths to environment variables
     # so that future invocations of the program will know where to look.
     if (
-        install_path != env_vars.install_path.default
-        or data_path != env_vars.data_path.default
+            install_path != env_vars.install_path.default
+            or data_path != env_vars.data_path.default
     ):
         print()
         print_utils.translate("install.set-custom-directory-env-vars")
@@ -188,3 +197,42 @@ def _parse_args():
     )
 
     return subcommand_parse_args(parser)
+
+
+def check_location() -> str:
+    url = 'http://ipinfo.io/json'
+    try:
+        response = urlopen(url, timeout=5)
+    except Exception:
+        return "US"
+    data = json.load(response)
+    return data["country"]
+
+
+def get_docker_config() -> dict:
+    with open("/etc/docker/daemon.json") as docker_config_file:
+        docker_config = json.load(docker_config_file)
+        return docker_config
+
+
+def update_docker_mirror(docker_config: dict, location: str):
+    docker_mirror_host = {"CN": "USTC", "US": "Docker Hub"}
+    if location == "CN" and \
+            ("registry-mirrors" not in docker_config.keys() or
+             docker_config["registry-mirrors"] !=
+             ["https://docker.mirrors.ustc.edu.cn/"]):
+        docker_config["registry-mirrors"] = \
+            ["https://docker.mirrors.ustc.edu.cn/"]
+    elif location != "CN" and "registry-mirrors" in docker_config.keys():
+        del docker_config["registry-mirrors"]
+    else:
+        return
+    if print_utils.ask_yes_no(
+            "install.change-docker-mirror",
+            location=location,
+            mirror_host=docker_mirror_host[location]
+    ):
+        with open("/etc/docker/daemon.json", "w") as docker_config_file:
+            json.dump(docker_config, docker_config_file, indent=4,
+                      sort_keys=True)
+        os_utils.restart_docker()

brainframe/cli/os_utils.py

@@ -77,6 +77,11 @@ def give_brainframe_group_rw_access(paths: List[Path]):
     run(["chmod", "-R", "g+rw"] + paths_str)
 
 
+def restart_docker():
+    print_utils.translate("general.restarting_docker")
+    run(["systemctl", "restart", "docker"])
+
+
 def _current_user():
     # If the SUDO_USER environment variable allows us to get the username of
     # the user running sudo instead of root. If they're not using sudo, we can

brainframe/cli/translations/general.en.yml

@@ -16,3 +16,4 @@ en:
   staging-missing-credentials: "The %{username_env_var} and %{password_env_var}
   variables must be set to access staging."
   user-not-root: "This command must be run as root!"
+  restarting_docker: "Restarting Docker..."

brainframe/cli/translations/install.en.yml

@@ -81,3 +81,10 @@ en:
   group? Doing so will allow the BrainFrame CLI to interact with Docker without
   root permissions."
   install-version: "BrainFrame:%{version} will be installed."
+  change-docker-mirror: "We have detected you are located in %{location}, do
+  you want to change docker mirror hosted by %{mirror_host} to speed up the
+  download process? Please be aware that this will overwrite you docker
+  daemon configuration and restart Docker service."
+  fail-to-load-docker-config: "Fail to load the Docker configuration file,
+  it's may caused by improper installation of Docker. Will keep using the
+  original docker mirror."

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "brainframe-cli"
-version = "0.1.5"
+version = "0.1.6.dev1"
 license = "BSD-3-Clause"
 authors = ["Aotu <info@aotu.ai>"]