Skip to content

Conversation

ManfeiBai
Copy link
Collaborator

@ManfeiBai ManfeiBai commented Mar 7, 2023

Update code to catch info from metadata


accelerator_type = tpu_env[xenv.ACCELERATOR_TYPE]
if version() == 4:
if tpu_env['ACCELERATOR_TYPE'].startswith('v4') or tpu_env['ACCELERATOR_TYPE'].startswith('v5'):
Copy link
Collaborator

Choose a reason for hiding this comment

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

@will-cromar what's the right place to add a test for this one? I felt like we had similar tests somwehre.

Copy link
Collaborator

Choose a reason for hiding this comment

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

The relevant unit test is here:

def test_tpu_env_from_gce_metadata(self, tpu_env_yaml, expected_env,

Can we just do tpu.version() >= 4 here?

Copy link
Collaborator Author

@ManfeiBai ManfeiBai Mar 8, 2023

Choose a reason for hiding this comment

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

Thanks, updated

@ManfeiBai ManfeiBai marked this pull request as ready for review March 8, 2023 05:08
Copy link
Collaborator

@suexu1025 suexu1025 left a comment

Choose a reason for hiding this comment

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

one question regarding #ACCELERATOR_TYPE match

Copy link
Contributor

@yeounoh yeounoh left a comment

Choose a reason for hiding this comment

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

LGTM

Copy link
Collaborator

@will-cromar will-cromar left a comment

Choose a reason for hiding this comment

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

Thanks!

@ManfeiBai ManfeiBai merged commit 4d67a04 into pytorch:master Mar 28, 2023
ManfeiBai added a commit that referenced this pull request Mar 28, 2023
@ManfeiBai ManfeiBai deleted the v5tpufile branch March 28, 2023 22:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants