diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py index 5bfb7b8991b..6cf186a1349 100644 --- a/torch_xla/__init__.py +++ b/torch_xla/__init__.py @@ -77,6 +77,7 @@ def _setup_default_env(): os.environ.setdefault('TPU_ML_PLATFORM', 'PyTorch/XLA') # This is used for ML Framework Telemetry. os.environ.setdefault('TPU_ML_PLATFORM_VERSION', __version__) + os.environ.setdefault('ENABLE_RUNTIME_UPTIME_TELEMETRY', '1') if tpu.version() == 4: os.environ.setdefault('TPU_MEGACORE', 'megacore_dense') diff --git a/torch_xla/_internal/tpu.py b/torch_xla/_internal/tpu.py index 241bf469d4a..61db1a8f8c9 100644 --- a/torch_xla/_internal/tpu.py +++ b/torch_xla/_internal/tpu.py @@ -17,6 +17,7 @@ import torch_xla.core.xla_env_vars as xenv import torch_xla.core.xla_model as xm from torch_xla.experimental import plugins +from torch_xla.version import __version__ _GCE_METADATA_ROOT_URL = 'http://metadata.google.internal/computeMetadata/v1' _ACCELERATOR_TYPE_TO_HOST_BOUNDS = { @@ -347,5 +348,9 @@ def physical_chip_count(self): def client_create_options(self): return { 'max_inflight_computations': - xu.getenv_as('XLA_TPU_MAX_INFLIGHT_COMPUTATIONS', int, 4) + xu.getenv_as('XLA_TPU_MAX_INFLIGHT_COMPUTATIONS', int, 4), + 'ml_framework_name': + 'PyTorch/XLA', + 'ml_framework_version': + __version__ }