r"""Run command in subprocess.
(command)
| 64 | |
| 65 | |
| 66 | def launch(command): |
| 67 | r"""Run command in subprocess.""" |
| 68 | visible_devices = _query_visible_devices() |
| 69 | local_world_size_str = str(len(visible_devices)) |
| 70 | strategy = os.getenv("COLLECTIVE_STRATEGY", "hb") |
| 71 | |
| 72 | signal.signal(signal.SIGINT, sigintHandler) |
| 73 | signal.signal(signal.SIGHUP, sigintHandler) |
| 74 | signal.signal(signal.SIGTERM, sigintHandler) |
| 75 | |
| 76 | if strategy == "hb": |
| 77 | port = int(os.getenv("HB_RUN_BASE_PORT", "20001")) |
| 78 | device_to_ports = [] |
| 79 | for d in visible_devices: |
| 80 | device_to_ports.append([d, port]) |
| 81 | port += 1 |
| 82 | |
| 83 | if len(device_to_ports) < 1: |
| 84 | os.environ["CUDA_VISIBLE_DEVICES"] = "" |
| 85 | os.environ["HB_OP_OPTIMIZATION_DISABLED"] = "1" |
| 86 | if callable(command): |
| 87 | command() |
| 88 | return |
| 89 | subprocess.check_call(command) |
| 90 | return |
| 91 | |
| 92 | if len(device_to_ports) == 1: |
| 93 | os.environ["CUDA_VISIBLE_DEVICES"] = "0" |
| 94 | if callable(command): |
| 95 | command() |
| 96 | return |
| 97 | subprocess.check_call(command) |
| 98 | return |
| 99 | |
| 100 | tf_config = json.loads(os.getenv("TF_CONFIG", "{}")) |
| 101 | if tf_config: |
| 102 | task = tf_config["task"] |
| 103 | task_type = task["type"] |
| 104 | task_id = int(task["index"]) |
| 105 | cluster = tf_config["cluster"] |
| 106 | else: |
| 107 | task_type = "chief" |
| 108 | task_id = 0 |
| 109 | cluster = {"chief": ["127.0.0.1:20000"]} |
| 110 | |
| 111 | workers = [] |
| 112 | if "chief" in cluster: |
| 113 | workers.extend(cluster["chief"]) |
| 114 | if "worker" in cluster: |
| 115 | workers.extend(cluster["worker"]) |
| 116 | worker_hosts = [w.split(":")[0] for w in workers] |
| 117 | new_workers = [ |
| 118 | f"{h}:{p}" for h in worker_hosts for _, p in device_to_ports] |
| 119 | new_cluster = cluster.copy() |
| 120 | if "chief" in cluster: |
| 121 | new_cluster["chief"] = [new_workers[0]] |
| 122 | if len(new_workers) > 1: |
| 123 | new_cluster["worker"] = new_workers[1:] |
no test coverage detected