Skip to content

Commit

Permalink
Allow set XLA_FLAGS with --debug_dump_gcs
Browse files Browse the repository at this point in the history
  • Loading branch information
PatriosTheGreat committed Oct 23, 2024
1 parent 5c6124a commit 9d30334
Showing 1 changed file with 8 additions and 5 deletions.
13 changes: 8 additions & 5 deletions src/xpk/core/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -233,13 +233,16 @@ def parse_env_config(args, tensorboard_config, system: SystemCharacteristics):

if not args.use_pathways:
if args.debug_dump_gcs:
if 'XLA_FLAGS' in env:
if 'XLA_FLAGS' in env and '--xla_dump_to' in env['XLA_FLAGS']:
raise ValueError(
'Conflict: XLA_FLAGS defined in both --debug_dump_gcs '
'and environment file. Please choose one way to define '
'XLA_FLAGS.'
'Conflict: --xla_dump_to defined in both --debug_dump_gcs '
'and environment file via XLA_FLAGS. Please choose one way to define it.'
)
env['XLA_FLAGS'] = '--xla_dump_to=/tmp/xla_dump/'
other_flags = ''
if 'XLA_FLAGS' in env:
other_flags = env['XLA_FLAGS']

env['XLA_FLAGS'] = '--xla_dump_to=/tmp/xla_dump/ ' + other_flags

if tensorboard_config:
env['UPLOAD_DATA_TO_TENSORBOARD'] = True
Expand Down

0 comments on commit 9d30334

Please sign in to comment.