[Bug] `LookupError: <ContextVar name='shell_parent'` after TPU VM Upgrade

## 🐛 Bug

After TPU VM upgrade, transformers stuck in loading/downloading model
seems the tqdm broken in the new VM

### To Reproduce

Select TPU accelerator
```python
import torch
import torch_xla
from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
```

### Expected behavior
Transformer download the model.

### Current behavior

```python
---------------------------------------------------------------------------
KeyError Traceback (most recent call last)
File /usr/local/lib/python3.12/site-packages/traitlets/traitlets.py:632, in TraitType.get(self, obj, cls)
631 try:
--> 632 value = obj._trait_values[self.name]
633 except KeyError:
634 # Check for a dynamic initializer.

KeyError: 'layout'

During handling of the above exception, another exception occurred:

LookupError Traceback (most recent call last)
Cell In[1], line 5
2 import torch_xla
3 from transformers import AutoTokenizer, AutoModelForCausalLM
----> 5 model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-3-mini-4k-instruct")

File /usr/local/lib/python3.12/site-packages/transformers/models/auto/auto_factory.py:604, in _BaseAutoModelClass.from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
602 if model_class.config_class == config.sub_configs.get("text_config", None):
603 config = config.get_text_config()
--> 604 return model_class.from_pretrained(
605 pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
606 )
607 raise ValueError(
608 f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.\n"
609 f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping)}."
610 )

File /usr/local/lib/python3.12/site-packages/transformers/modeling_utils.py:288, in restore_default_dtype.._wrapper(*args, **kwargs)
286 old_dtype = torch.get_default_dtype()
287 try:
--> 288 return func(*args, **kwargs)
289 finally:
290 torch.set_default_dtype(old_dtype)

File /usr/local/lib/python3.12/site-packages/transformers/modeling_utils.py:5030, in PreTrainedModel.from_pretrained(cls, pretrained_model_name_or_path, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, use_safetensors, weights_only, *model_args, **kwargs)
5020 if (
5021 gguf_file
5022 and device_map is not None
5023 and ((isinstance(device_map, dict) and "disk" in device_map.values()) or "disk" in device_map)
5024 ):
5025 raise RuntimeError(
5026 "One or more modules is configured to be mapped to disk. Disk offload is not supported for models "
5027 "loaded from GGUF files."
5028 )
-> 5030 checkpoint_files, sharded_metadata = _get_resolved_checkpoint_files(
5031 pretrained_model_name_or_path=pretrained_model_name_or_path,
5032 subfolder=subfolder,
5033 variant=variant,
5034 gguf_file=gguf_file,
5035 from_tf=from_tf,
5036 from_flax=from_flax,
5037 use_safetensors=use_safetensors,
5038 cache_dir=cache_dir,
5039 force_download=force_download,
5040 proxies=proxies,
5041 local_files_only=local_files_only,
5042 token=token,
5043 user_agent=user_agent,
5044 revision=revision,
5045 commit_hash=commit_hash,
5046 is_remote_code=cls._auto_class is not None,
5047 transformers_explicit_filename=transformers_explicit_filename,
5048 )
5050 is_sharded = sharded_metadata is not None
5051 is_quantized = hf_quantizer is not None

File /usr/local/lib/python3.12/site-packages/transformers/modeling_utils.py:1308, in _get_resolved_checkpoint_files(pretrained_model_name_or_path, subfolder, variant, gguf_file, from_tf, from_flax, use_safetensors, cache_dir, force_download, proxies, local_files_only, token, user_agent, revision, commit_hash, is_remote_code, transformers_explicit_filename)
1306 sharded_metadata = None
1307 if is_sharded:
-> 1308 checkpoint_files, sharded_metadata = get_checkpoint_shard_files(
1309 pretrained_model_name_or_path,
1310 resolved_archive_file,
1311 cache_dir=cache_dir,
1312 force_download=force_download,
1313 proxies=proxies,
1314 local_files_only=local_files_only,
1315 token=token,
1316 user_agent=user_agent,
1317 revision=revision,
1318 subfolder=subfolder,
1319 _commit_hash=commit_hash,
1320 )
1321 else:
1322 checkpoint_files = [resolved_archive_file] if pretrained_model_name_or_path is not None else None

File /usr/local/lib/python3.12/site-packages/transformers/utils/hub.py:1119, in get_checkpoint_shard_files(pretrained_model_name_or_path, index_filename, cache_dir, force_download, proxies, resume_download, local_files_only, token, user_agent, revision, subfolder, _commit_hash, **deprecated_kwargs)
1115 return shard_filenames, sharded_metadata
1117 # At this stage pretrained_model_name_or_path is a model identifier on the Hub. Try to get everything from cache,
1118 # or download the files
-> 1119 cached_filenames = cached_files(
1120 pretrained_model_name_or_path,
1121 shard_filenames,
1122 cache_dir=cache_dir,
1123 force_download=force_download,
1124 proxies=proxies,
1125 resume_download=resume_download,
1126 local_files_only=local_files_only,
1127 token=token,
1128 user_agent=user_agent,
1129 revision=revision,
1130 subfolder=subfolder,
1131 _commit_hash=_commit_hash,
1132 )
1134 return cached_filenames, sharded_metadata

File /usr/local/lib/python3.12/site-packages/transformers/utils/hub.py:566, in cached_files(path_or_repo_id, filenames, cache_dir, force_download, resume_download, proxies, token, revision, local_files_only, subfolder, repo_type, user_agent, _raise_exceptions_for_gated_repo, _raise_exceptions_for_missing_entries, _raise_exceptions_for_connection_errors, _commit_hash, **deprecated_kwargs)
563 # Any other Exception type should now be re-raised, in order to provide helpful error messages and break the execution flow
564 # (EntryNotFoundError will be treated outside this block and correctly re-raised if needed)
565 elif not isinstance(e, EntryNotFoundError):
--> 566 raise e
568 resolved_files = [
569 _get_cache_file_to_return(path_or_repo_id, filename, cache_dir, revision) for filename in full_filenames
570 ]
571 # If there are any missing file and the flag is active, raise

File /usr/local/lib/python3.12/site-packages/transformers/utils/hub.py:493, in cached_files(path_or_repo_id, filenames, cache_dir, force_download, resume_download, proxies, token, revision, local_files_only, subfolder, repo_type, user_agent, _raise_exceptions_for_gated_repo, _raise_exceptions_for_missing_entries, _raise_exceptions_for_connection_errors, _commit_hash, **deprecated_kwargs)
478 hf_hub_download(
479 path_or_repo_id,
480 filenames[0],
(...) 490 local_files_only=local_files_only,
491 )
492 else:
--> 493 snapshot_download(
494 path_or_repo_id,
495 allow_patterns=full_filenames,
496 repo_type=repo_type,
497 revision=revision,
498 cache_dir=cache_dir,
499 user_agent=user_agent,
500 force_download=force_download,
501 proxies=proxies,
502 resume_download=resume_download,
503 token=token,
504 local_files_only=local_files_only,
505 )
507 except Exception as e:
508 # We cannot recover from them
509 if isinstance(e, RepositoryNotFoundError) and not isinstance(e, GatedRepoError):

File /usr/local/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py:114, in validate_hf_hub_args.._inner_fn(*args, **kwargs)
111 if check_use_auth_token:
112 kwargs = smoothly_deprecate_use_auth_token(fn_name=fn.__name__, has_token=has_token, kwargs=kwargs)
--> 114 return fn(*args, **kwargs)

File /usr/local/lib/python3.12/site-packages/huggingface_hub/_snapshot_download.py:332, in snapshot_download(repo_id, repo_type, revision, cache_dir, local_dir, library_name, library_version, user_agent, proxies, etag_timeout, force_download, token, local_files_only, allow_patterns, ignore_patterns, max_workers, tqdm_class, headers, endpoint, local_dir_use_symlinks, resume_download)
330 _inner_hf_hub_download(file)
331 else:
--> 332 thread_map(
333 _inner_hf_hub_download,
334 filtered_repo_files,
335 desc=tqdm_desc,
336 max_workers=max_workers,
337 # User can use its own tqdm class or the default one from `huggingface_hub.utils`
338 tqdm_class=tqdm_class or hf_tqdm,
339 )
341 if local_dir is not None:
342 return str(os.path.realpath(local_dir))

File /usr/local/lib/python3.12/site-packages/tqdm/contrib/concurrent.py:69, in thread_map(fn, *iterables, **tqdm_kwargs)
55 """
56 Equivalent of `list(map(fn, *iterables))`
57 driven by `concurrent.futures.ThreadPoolExecutor`.
(...) 66 [default: max(32, cpu_count() + 4)].
67 """
68 from concurrent.futures import ThreadPoolExecutor
---> 69 return _executor_map(ThreadPoolExecutor, fn, *iterables, **tqdm_kwargs)

File /usr/local/lib/python3.12/site-packages/tqdm/contrib/concurrent.py:51, in _executor_map(PoolExecutor, fn, *iterables, **tqdm_kwargs)
47 with ensure_lock(tqdm_class, lock_name=lock_name) as lk:
48 # share lock in case workers are already using `tqdm`
49 with PoolExecutor(max_workers=max_workers, initializer=tqdm_class.set_lock,
50 initargs=(lk,)) as ex:
---> 51 return list(tqdm_class(ex.map(fn, *iterables, chunksize=chunksize), **kwargs))

File /usr/local/lib/python3.12/site-packages/tqdm/notebook.py:250, in tqdm_notebook.__iter__(self)
248 try:
249 it = super().__iter__()
--> 250 for obj in it:
251 # return super(tqdm...) will not catch exception
252 yield obj
253 # NB: except ... [ as ...] breaks IPython async KeyboardInterrupt

File /usr/local/lib/python3.12/site-packages/tqdm/std.py:1181, in tqdm.__iter__(self)
1178 time = self._time
1180 try:
-> 1181 for obj in iterable:
1182 yield obj
1183 # Update and possibly print the progressbar.
1184 # Note: does not call self.update(1) for speed optimisation.

File /usr/local/lib/python3.12/concurrent/futures/_base.py:619, in Executor.map..result_iterator()
616 while fs:
617 # Careful not to keep a reference to the popped future
618 if timeout is None:
--> 619 yield _result_or_cancel(fs.pop())
620 else:
621 yield _result_or_cancel(fs.pop(), end_time - time.monotonic())

File /usr/local/lib/python3.12/concurrent/futures/_base.py:317, in _result_or_cancel(***failed resolving arguments***)
315 try:
316 try:
--> 317 return fut.result(timeout)
318 finally:
319 fut.cancel()

File /usr/local/lib/python3.12/concurrent/futures/_base.py:456, in Future.result(self, timeout)
454 raise CancelledError()
455 elif self._state == FINISHED:
--> 456 return self.__get_result()
457 else:
458 raise TimeoutError()

File /usr/local/lib/python3.12/concurrent/futures/_base.py:401, in Future.__get_result(self)
399 if self._exception:
400 try:
--> 401 raise self._exception
402 finally:
403 # Break a reference cycle with the exception in self._exception
404 self = None

File /usr/local/lib/python3.12/concurrent/futures/thread.py:59, in _WorkItem.run(self)
56 return
58 try:
---> 59 result = self.fn(*self.args, **self.kwargs)
60 except BaseException as exc:
61 self.future.set_exception(exc)

File /usr/local/lib/python3.12/site-packages/huggingface_hub/_snapshot_download.py:306, in snapshot_download.._inner_hf_hub_download(repo_file)
305 def _inner_hf_hub_download(repo_file: str):
--> 306 return hf_hub_download(
307 repo_id,
308 filename=repo_file,
309 repo_type=repo_type,
310 revision=commit_hash,
311 endpoint=endpoint,
312 cache_dir=cache_dir,
313 local_dir=local_dir,
314 local_dir_use_symlinks=local_dir_use_symlinks,
315 library_name=library_name,
316 library_version=library_version,
317 user_agent=user_agent,
318 proxies=proxies,
319 etag_timeout=etag_timeout,
320 resume_download=resume_download,
321 force_download=force_download,
322 token=token,
323 headers=headers,
324 )

File /usr/local/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py:114, in validate_hf_hub_args.._inner_fn(*args, **kwargs)
111 if check_use_auth_token:
112 kwargs = smoothly_deprecate_use_auth_token(fn_name=fn.__name__, has_token=has_token, kwargs=kwargs)
--> 114 return fn(*args, **kwargs)

File /usr/local/lib/python3.12/site-packages/huggingface_hub/file_download.py:1010, in hf_hub_download(repo_id, filename, subfolder, repo_type, revision, library_name, library_version, cache_dir, local_dir, user_agent, force_download, proxies, etag_timeout, token, local_files_only, headers, endpoint, resume_download, force_filename, local_dir_use_symlinks)
990 return _hf_hub_download_to_local_dir(
991 # Destination
992 local_dir=local_dir,
(...) 1007 local_files_only=local_files_only,
1008 )
1009 else:
-> 1010 return _hf_hub_download_to_cache_dir(
1011 # Destination
1012 cache_dir=cache_dir,
1013 # File info
1014 repo_id=repo_id,
1015 filename=filename,
1016 repo_type=repo_type,
1017 revision=revision,
1018 # HTTP info
1019 endpoint=endpoint,
1020 etag_timeout=etag_timeout,
1021 headers=hf_headers,
1022 proxies=proxies,
1023 token=token,
1024 # Additional options
1025 local_files_only=local_files_only,
1026 force_download=force_download,
1027 )

File /usr/local/lib/python3.12/site-packages/huggingface_hub/file_download.py:1171, in _hf_hub_download_to_cache_dir(cache_dir, repo_id, filename, repo_type, revision, endpoint, etag_timeout, headers, proxies, token, local_files_only, force_download)
1168 # Local file doesn't exist or etag isn't a match => retrieve file from remote (or cache)
1170 with WeakFileLock(lock_path):
-> 1171 _download_to_tmp_and_move(
1172 incomplete_path=Path(blob_path + ".incomplete"),
1173 destination_path=Path(blob_path),
1174 url_to_download=url_to_download,
1175 proxies=proxies,
1176 headers=headers,
1177 expected_size=expected_size,
1178 filename=filename,
1179 force_download=force_download,
1180 etag=etag,
1181 xet_file_data=xet_file_data,
1182 )
1183 if not os.path.exists(pointer_path):
1184 _create_symlink(blob_path, pointer_path, new_blob=True)

File /usr/local/lib/python3.12/site-packages/huggingface_hub/file_download.py:1723, in _download_to_tmp_and_move(incomplete_path, destination_path, url_to_download, proxies, headers, expected_size, filename, force_download, etag, xet_file_data)
1721 if xet_file_data is not None and is_xet_available():
1722 logger.debug("Xet Storage is enabled for this repo. Downloading file from Xet Storage..")
-> 1723 xet_get(
1724 incomplete_path=incomplete_path,
1725 xet_file_data=xet_file_data,
1726 headers=headers,
1727 expected_size=expected_size,
1728 displayed_filename=filename,
1729 )
1730 else:
1731 if xet_file_data is not None and not constants.HF_HUB_DISABLE_XET:

File /usr/local/lib/python3.12/site-packages/huggingface_hub/file_download.py:615, in xet_get(incomplete_path, xet_file_data, headers, expected_size, displayed_filename, _tqdm_bar)
612 if len(displayed_filename) > 40:
613 displayed_filename = f"{displayed_filename[:40]}(…)"
--> 615 progress_cm = _get_progress_bar_context(
616 desc=displayed_filename,
617 log_level=logger.getEffectiveLevel(),
618 total=expected_size,
619 initial=0,
620 name="huggingface_hub.xet_get",
621 _tqdm_bar=_tqdm_bar,
622 )
624 with progress_cm as progress:
626 def progress_updater(progress_bytes: float):

File /usr/local/lib/python3.12/site-packages/huggingface_hub/utils/tqdm.py:299, in _get_progress_bar_context(desc, log_level, total, initial, unit, unit_scale, name, _tqdm_bar)
294 return nullcontext(_tqdm_bar)
295 # ^ `contextlib.nullcontext` mimics a context manager that does nothing
296 # Makes it easier to use the same code path for both cases but in the later
297 # case, the progress bar is not closed when exiting the context manager.
--> 299 return tqdm(
300 unit=unit,
301 unit_scale=unit_scale,
302 total=total,
303 initial=initial,
304 desc=desc,
305 disable=is_tqdm_disabled(log_level=log_level),
306 name=name,
307 )

File /usr/local/lib/python3.12/site-packages/huggingface_hub/utils/tqdm.py:225, in tqdm.__init__(self, *args, **kwargs)
223 if are_progress_bars_disabled(name):
224 kwargs["disable"] = True
--> 225 super().__init__(*args, **kwargs)

File /usr/local/lib/python3.12/site-packages/tqdm/notebook.py:234, in tqdm_notebook.__init__(self, *args, **kwargs)
232 unit_scale = 1 if self.unit_scale is True else self.unit_scale or 1
233 total = self.total * unit_scale if self.total else self.total
--> 234 self.container = self.status_printer(self.fp, total, self.desc, self.ncols)
235 self.container.pbar = proxy(self)
236 self.displayed = False

File /usr/local/lib/python3.12/site-packages/tqdm/notebook.py:110, in tqdm_notebook.status_printer(_, total, desc, ncols)
108 raise ImportError(WARN_NOIPYW)
109 if total:
--> 110 pbar = IProgress(min=0, max=total)
111 else: # No total? Show info style bar with no progress tqdm status
112 pbar = IProgress(min=0, max=1)

File /usr/local/lib/python3.12/site-packages/ipywidgets/widgets/widget_float.py:26, in _Float.__init__(self, value, **kwargs)
24 if value is not None:
25 kwargs['value'] = value
---> 26 super().__init__(**kwargs)

File /usr/local/lib/python3.12/site-packages/ipywidgets/widgets/widget_description.py:35, in DescriptionWidget.__init__(self, *args, **kwargs)
33 kwargs.setdefault('tooltip', kwargs['description_tooltip'])
34 del kwargs['description_tooltip']
---> 35 super().__init__(*args, **kwargs)

File /usr/local/lib/python3.12/site-packages/ipywidgets/widgets/widget.py:506, in Widget.__init__(self, **kwargs)
503 super().__init__(**kwargs)
505 Widget._call_widget_constructed(self)
--> 506 self.open()

File /usr/local/lib/python3.12/site-packages/ipywidgets/widgets/widget.py:525, in Widget.open(self)
523 """Open a comm to the frontend if one isn't already open."""
524 if self.comm is None:
--> 525 state, buffer_paths, buffers = _remove_buffers(self.get_state())
527 args = dict(target_name='jupyter.widget',
528 data={'state': state, 'buffer_paths': buffer_paths},
529 buffers=buffers,
530 metadata={'version': __protocol_version__}
531 )
532 if self._model_id is not None:

File /usr/local/lib/python3.12/site-packages/ipywidgets/widgets/widget.py:615, in Widget.get_state(self, key, drop_defaults)
613 for k in keys:
614 to_json = self.trait_metadata(k, 'to_json', self._trait_to_json)
--> 615 value = to_json(getattr(self, k), self)
616 if not drop_defaults or not self._compare(value, traits[k].default_value):
617 state[k] = value

File /usr/local/lib/python3.12/site-packages/traitlets/traitlets.py:687, in TraitType.__get__(self, obj, cls)
685 return self
686 else:
--> 687 return t.cast(G, self.get(obj, cls))

File /usr/local/lib/python3.12/site-packages/traitlets/traitlets.py:635, in TraitType.get(self, obj, cls)
632 value = obj._trait_values[self.name]
633 except KeyError:
634 # Check for a dynamic initializer.
--> 635 default = obj.trait_defaults(self.name)
636 if default is Undefined:
637 warn(
638 "Explicit using of Undefined as the default value "
639 "is deprecated in traitlets 5.0, and may cause "
(...) 642 stacklevel=2,
643 )

File /usr/local/lib/python3.12/site-packages/traitlets/traitlets.py:1897, in HasTraits.trait_defaults(self, *names, **metadata)
1894 raise TraitError(f"'{n}' is not a trait of '{type(self).__name__}' instances")
1896 if len(names) == 1 and len(metadata) == 0:
-> 1897 return t.cast(Sentinel, self._get_trait_default_generator(names[0])(self))
1899 trait_names = self.trait_names(**metadata)
1900 trait_names.extend(names)

File /usr/local/lib/python3.12/site-packages/traitlets/traitlets.py:602, in TraitType.default(self, obj)
600 return t.cast(G, self.default_value)
601 elif hasattr(self, "make_dynamic_default"):
--> 602 return t.cast(G, self.make_dynamic_default())
603 else:
604 # Undefined will raise in TraitType.get
605 return t.cast(G, self.default_value)

File /usr/local/lib/python3.12/site-packages/ipywidgets/widgets/trait_types.py:409, in InstanceDict.make_dynamic_default(self)
408 def make_dynamic_default(self):
--> 409 return self.klass(*(self.default_args or ()),
410 **(self.default_kwargs or {}))

File /usr/local/lib/python3.12/site-packages/ipywidgets/widgets/widget_layout.py:86, in Layout.__init__(self, **kwargs)
83 for side in ['top', 'right', 'bottom', 'left']:
84 kwargs.setdefault(f'border_{side}', border)
---> 86 super().__init__(**kwargs)

File /usr/local/lib/python3.12/site-packages/ipywidgets/widgets/widget.py:506, in Widget.__init__(self, **kwargs)
503 super().__init__(**kwargs)
505 Widget._call_widget_constructed(self)
--> 506 self.open()

File /usr/local/lib/python3.12/site-packages/ipywidgets/widgets/widget.py:535, in Widget.open(self)
532 if self._model_id is not None:
533 args['comm_id'] = self._model_id
--> 535 self.comm = comm.create_comm(**args)

File /usr/local/lib/python3.12/site-packages/ipywidgets/comm.py:33, in create_comm(*args, **kwargs)
31 return Comm(*args, **kwargs)
32 else:
---> 33 return comm.create_comm(*args, **kwargs)

File /usr/local/lib/python3.12/site-packages/ipykernel/ipkernel.py:52, in _create_comm(*args, **kwargs)
50 def _create_comm(*args, **kwargs):
51 """Create a new Comm."""
---> 52 return BaseComm(*args, **kwargs)

File /usr/local/lib/python3.12/site-packages/comm/base_comm.py:65, in BaseComm.__init__(self, target_name, data, metadata, buffers, comm_id, primary, target_module, topic, _open_data, _close_data, **kwargs)
61 self._closed = True
63 if self.primary:
64 # I am primary, open my peer.
---> 65 self.open(data=data, metadata=metadata, buffers=buffers)
66 else:
67 self._closed = False

File /usr/local/lib/python3.12/site-packages/comm/base_comm.py:103, in BaseComm.open(self, data, metadata, buffers)
101 comm_manager.register_comm(self)
102 try:
--> 103 self.publish_msg(
104 "comm_open",
105 data=data,
106 metadata=metadata,
107 buffers=buffers,
108 target_name=self.target_name,
109 target_module=self.target_module,
110 )
111 self._closed = False
112 except Exception:

File /usr/local/lib/python3.12/site-packages/ipykernel/comm/comm.py:42, in BaseComm.publish_msg(self, msg_type, data, metadata, buffers, **keys)
34 self.kernel = Kernel.instance()
36 assert self.kernel.session is not None
37 self.kernel.session.send(
38 self.kernel.iopub_socket,
39 msg_type,
40 content,
41 metadata=json_clean(metadata),
---> 42 parent=self.kernel.get_parent(),
43 ident=self.topic,
44 buffers=buffers,
45 )

File /usr/local/lib/python3.12/site-packages/ipykernel/kernelbase.py:797, in Kernel.get_parent(self, channel)
795 if channel == "control":
796 return self._control_parent
--> 797 return self._shell_parent.get()

LookupError:
```