Skip to content

Minor cleanups #2581

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jysohn23 merged 2 commits into from
Oct 30, 2020
Merged

Minor cleanups #2581

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3ea1b8d
Minor cleanups
jysohn23 Oct 24, 2020
6c25e7f
Revert changes to computation_client
jysohn23 Oct 30, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,7 @@ torch_xla/csrc/aten_xla_type_default.cpp
# Below files are not deleted by "setup.py clean".

third_party/tensorflow/

# Visual Studio Code files
.vscode
.vs
4 changes: 2 additions & 2 deletions torch_xla/csrc/ir.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,13 @@ namespace {
using ShapeCache =
xla::util::Cache<xla::hash_t, xla::Shape, xla::util::HashReducer>;

struct ScapeEntry {
struct ScopeEntry {
std::string name;
size_t saved_next_id = 1;
};

struct ScopeContext {
std::vector<ScapeEntry> scopes;
std::vector<ScopeEntry> scopes;
size_t next_id = 1;
};

Expand Down
13 changes: 6 additions & 7 deletions torch_xla/distributed/xla_multiprocessing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,7 @@ def _pre_fork_setup(num_devices):
return PreForkConfig(dev_kind=dev_kind, num_devices=num_devices)


def _setup_gpu_worker(index, gindex, pf_cfg):
Copy link
Collaborator

@davidel davidel Oct 29, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IMHO when there are symmetrical APIs (gpu, tpu, cpu) it is better to have a consistent interface, even though one instance of such APIs does not actually use an argument.

Copy link
Collaborator Author

@jysohn23 jysohn23 Oct 29, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure what consistency you're referring to here. All gpu, tpu, cpu consistently don't use that argument.

Copy link
Collaborator

@davidel davidel Oct 30, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right. That argument contains the configuration info parsed before the fork.
Currently it is not used by the _setup_*_worker APIs.
IMHO it is better to pass that down, but if you want to remove the argument that is not a big deal for me.

Copy link
Collaborator Author

@jysohn23 jysohn23 Oct 30, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll keep it removed if that's fine by you. Thanks for the review Davide!

def _setup_gpu_worker(index, gindex):
os.environ[xenv.MP_DEVICE] = 'GPU:{}'.format(
_get_mp_device_ordinal(index, gindex))
os.environ[xenv.LOCAL_WORKER] = '{}:{}'.format(_LOCAL_WORKER, gindex)
Expand All @@ -228,7 +228,7 @@ def _setup_gpu_worker(index, gindex, pf_cfg):
os.environ.pop(xenv.GPU_NUM_DEVICES, None)


def _setup_cpu_worker(index, gindex, pf_cfg):
def _setup_cpu_worker(index, gindex):
task_no = 0
dev_index = _get_mp_device_ordinal(index, gindex)
os.environ[xenv.MP_DEVICE] = 'CPU:{}'.format(dev_index)
Expand All @@ -254,7 +254,7 @@ def _wants_tpu_env_config(index, gindex):
return gindex == 0


def _setup_tpu_worker(index, gindex, pf_cfg, tpu_env_config):
def _setup_tpu_worker(index, gindex, tpu_env_config):
os.environ[xenv.MP_DEVICE] = 'TPU:{}'.format(
_get_mp_device_ordinal(index, gindex))
if xenv.LOCAL_WORKER not in os.environ:
Expand Down Expand Up @@ -282,8 +282,7 @@ def _prepare_env_for_index(index, pf_cfg):
os.environ[xenv.LOCAL_ORDINAL] = str(index)

if pf_cfg.dev_kind == 'TPU':
_setup_tpu_worker(index, gindex, pf_cfg,
os.environ.get(xenv.TPU_CONFIG, None))
_setup_tpu_worker(index, gindex, os.environ.get(xenv.TPU_CONFIG, None))
if xenv.HOST_ORDINAL in os.environ:
# If xenv.HOST_ORDINAL is set, we are in a sea-of-devices TPU setup, where
# each host has local TPU devices, but not interconnected with the fast TPU
Expand All @@ -301,9 +300,9 @@ def _prepare_env_for_index(index, pf_cfg):
# to 0 in all hosts.
_setup_torch_distributed()
elif pf_cfg.dev_kind == 'GPU':
_setup_gpu_worker(index, gindex, pf_cfg)
_setup_gpu_worker(index, gindex)
elif pf_cfg.dev_kind == 'CPU':
_setup_cpu_worker(index, gindex, pf_cfg)
_setup_cpu_worker(index, gindex)
_setup_torch_distributed()
return gindex

Expand Down