Try using the CPU PJRT plugin #7184

duncantech · 2024-06-04T16:39:17Z

📚 Documentation

Follow the instructions to install the CPU PJRT plugin located here: https://github.com/pytorch/xla/blob/master/plugins/cpu/README.md

Validate that the instructions still work and report any issues that you encounter.

loganthomas · 2024-06-04T23:35:01Z

/assigntome

loganthomas · 2024-06-05T02:52:21Z

Issue 1: Cannot install torch_xla[tpu]

When trying to install the optional tpu dependencies, I get an error:

$ pip install torch~=2.3.0 torch_xla[tpu]~=2.3.0 -f https://storage.googleapis.com/libtpu-releases/index.html
zsh: no matches found: torch_xla[tpu]~=2.3.0

Issue 2: Issue building

Both pip wheel and pip install commands yield errors

$ pip wheel plugins/cpu -v
Processing ./plugins/cpu
  Running command pip subprocess to install build dependencies
  Collecting setuptools
    Using cached setuptools-70.0.0-py3-none-any.whl.metadata (5.9 kB)
  Using cached setuptools-70.0.0-py3-none-any.whl (863 kB)
  Installing collected packages: setuptools
  Successfully installed setuptools-70.0.0
  Installing build dependencies ... done
  Running command Getting requirements to build wheel
  bazel build //plugins/cpu:pjrt_c_api_cpu_plugin.so --symlink_prefix=/Users/logan/Desktop/repos/xla/plugins/cpu/bazel- --remote_default_exec_properties=cache-silo-key=dev
  Traceback (most recent call last):
    File "/Users/logan/venvs/pytorch-dev/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 353, in <module>
      main()
    File "/Users/logan/venvs/pytorch-dev/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 335, in main
      json_out['return_val'] = hook(**hook_input['kwargs'])
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/Users/logan/venvs/pytorch-dev/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 118, in get_requires_for_build_wheel
      return hook(config_settings)
             ^^^^^^^^^^^^^^^^^^^^^
    File "/private/var/folders/87/3sfr2hp948147snb3144n16w0000gn/T/pip-build-env-qq0brd10/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 325, in get_requires_for_build_wheel
      return self._get_build_requires(config_settings, requirements=['wheel'])
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/private/var/folders/87/3sfr2hp948147snb3144n16w0000gn/T/pip-build-env-qq0brd10/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 295, in _get_build_requires
      self.run_setup()
    File "/private/var/folders/87/3sfr2hp948147snb3144n16w0000gn/T/pip-build-env-qq0brd10/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 311, in run_setup
      exec(code, locals())
    File "<string>", line 10, in <module>
    File "/Users/logan/Desktop/repos/xla/plugins/cpu/../../build_util.py", line 67, in bazel_build
      subprocess.check_call(bazel_argv, stdout=sys.stdout, stderr=sys.stderr)
    File "/usr/local/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/lib/python3.12/subprocess.py", line 408, in check_call
      retcode = call(*popenargs, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/usr/local/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/lib/python3.12/subprocess.py", line 389, in call
      with Popen(*popenargs, **kwargs) as p:
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/usr/local/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/lib/python3.12/subprocess.py", line 1026, in __init__
      self._execute_child(args, executable, preexec_fn, close_fds,
    File "/usr/local/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/lib/python3.12/subprocess.py", line 1955, in _execute_child
      raise child_exception_type(errno_num, err_msg, err_filename)
  FileNotFoundError: [Errno 2] No such file or directory: 'bazel'
  error: subprocess-exited-with-error
  
  × Getting requirements to build wheel did not run successfully.
  │ exit code: 1
  ╰─> See above for output.
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  full command: /Users/logan/venvs/pytorch-dev/bin/python3.12 /Users/logan/venvs/pytorch-dev/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py get_requires_for_build_wheel /var/folders/87/3sfr2hp948147snb3144n16w0000gn/T/tmp4ptkuoqt
  cwd: /Users/logan/Desktop/repos/xla/plugins/cpu
  Getting requirements to build wheel ... error
error: subprocess-exited-with-error

× Getting requirements to build wheel did not run successfully.
│ exit code: 1
╰─> See above for output.

note: This error originates from a subprocess, and is likely not a problem with pip

$ pip install plugins/cpu -v
Using pip 24.0 from /Users/logan/venvs/pytorch-dev/lib/python3.12/site-packages/pip (python 3.12)
Processing ./plugins/cpu
  Running command pip subprocess to install build dependencies
  Collecting setuptools
    Using cached setuptools-70.0.0-py3-none-any.whl.metadata (5.9 kB)
  Using cached setuptools-70.0.0-py3-none-any.whl (863 kB)
  Installing collected packages: setuptools
  Successfully installed setuptools-70.0.0
  Installing build dependencies ... done
  Running command Getting requirements to build wheel
  bazel build //plugins/cpu:pjrt_c_api_cpu_plugin.so --symlink_prefix=/Users/logan/Desktop/repos/xla/plugins/cpu/bazel- --remote_default_exec_properties=cache-silo-key=dev
  Traceback (most recent call last):
    File "/Users/logan/venvs/pytorch-dev/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 353, in <module>
      main()
    File "/Users/logan/venvs/pytorch-dev/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 335, in main
      json_out['return_val'] = hook(**hook_input['kwargs'])
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/Users/logan/venvs/pytorch-dev/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 118, in get_requires_for_build_wheel
      return hook(config_settings)
             ^^^^^^^^^^^^^^^^^^^^^
    File "/private/var/folders/87/3sfr2hp948147snb3144n16w0000gn/T/pip-build-env-l3twedtx/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 325, in get_requires_for_build_wheel
      return self._get_build_requires(config_settings, requirements=['wheel'])
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/private/var/folders/87/3sfr2hp948147snb3144n16w0000gn/T/pip-build-env-l3twedtx/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 295, in _get_build_requires
      self.run_setup()
    File "/private/var/folders/87/3sfr2hp948147snb3144n16w0000gn/T/pip-build-env-l3twedtx/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 311, in run_setup
      exec(code, locals())
    File "<string>", line 10, in <module>
    File "/Users/logan/Desktop/repos/xla/plugins/cpu/../../build_util.py", line 67, in bazel_build
      subprocess.check_call(bazel_argv, stdout=sys.stdout, stderr=sys.stderr)
    File "/usr/local/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/lib/python3.12/subprocess.py", line 408, in check_call
      retcode = call(*popenargs, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/usr/local/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/lib/python3.12/subprocess.py", line 389, in call
      with Popen(*popenargs, **kwargs) as p:
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/usr/local/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/lib/python3.12/subprocess.py", line 1026, in __init__
      self._execute_child(args, executable, preexec_fn, close_fds,
    File "/usr/local/Cellar/python@3.12/3.12.3/Frameworks/Python.framework/Versions/3.12/lib/python3.12/subprocess.py", line 1955, in _execute_child
      raise child_exception_type(errno_num, err_msg, err_filename)
  FileNotFoundError: [Errno 2] No such file or directory: 'bazel'
  error: subprocess-exited-with-error
  
  × Getting requirements to build wheel did not run successfully.
  │ exit code: 1
  ╰─> See above for output.
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  full command: /Users/logan/venvs/pytorch-dev/bin/python3.12 /Users/logan/venvs/pytorch-dev/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py get_requires_for_build_wheel /var/folders/87/3sfr2hp948147snb3144n16w0000gn/T/tmpkgrk3zoq
  cwd: /Users/logan/Desktop/repos/xla/plugins/cpu
  Getting requirements to build wheel ... error
error: subprocess-exited-with-error

× Getting requirements to build wheel did not run successfully.
│ exit code: 1
╰─> See above for output.

note: This error originates from a subprocess, and is likely not a problem with pip.

Machine Specs:

$ python -V
Python 3.12.3

$ pip list
Package           Version
----------------- -----------
asttokens         2.4.1
decorator         5.1.1
executing         2.0.1
filelock          3.14.0
fsspec            2024.6.0
ipython           8.25.0
jedi              0.19.1
Jinja2            3.1.4
MarkupSafe        2.1.5
matplotlib-inline 0.1.7
mpmath            1.3.0
networkx          3.3
numpy             1.26.4
pandas            2.2.2
parso             0.8.4
pexpect           4.9.0
pip               24.0
prompt_toolkit    3.0.46
ptyprocess        0.7.0
pure-eval         0.2.2
Pygments          2.18.0
python-dateutil   2.9.0.post0
pytz              2024.1
scipy             1.13.1
six               1.16.0
stack-data        0.6.3
sympy             1.12.1
torch             2.2.2
torch-xla         1.0
traitlets         5.14.3
typing_extensions 4.12.1
tzdata            2024.1
wcwidth           0.2.13

$ system_profiler SPSoftwareDataType SPHardwareDataType
Software:

    System Software Overview:

      System Version: macOS 14.4.1 (23E224)
      Kernel Version: Darwin 23.4.0
      Boot Volume: Macintosh HD
      ...

Hardware:

    Hardware Overview:

      Model Name: MacBook Pro
      Model Identifier: MacBookPro15,1
      Processor Name: 6-Core Intel Core i7
      Processor Speed: 2.6 GHz
      Number of Processors: 1
      Total Number of Cores: 6
      L2 Cache (per Core): 256 KB
      L3 Cache: 9 MB
      Hyper-Threading Technology: Enabled
      Memory: 32 GB
      ...

JackCaoG · 2024-06-05T17:15:22Z

@will-cromar can you remind me if CPU plugin works? I can't remember.

will-cromar · 2024-06-05T17:21:00Z

The CPU plugin is just for testing purposes. In practice, you can debug on CPU just by setting PJRT_DEVICE=CPU without installing anything extra.

loganthomas · 2024-06-07T02:26:10Z

>>> import os
...
... # Log device type
... os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'
... os.environ['TF_CPP_VMODULE'] = 'pjrt_registry=5'
... os.environ['PJRT_DEVICE'] = 'CPU'
...
... from torch_xla.experimental import plugins
... import torch_xla_cpu_plugin
... import torch_xla.core.xla_model as xm
... import torch_xla.runtime as xr
...
... # Use dynamic plugin instead of built-in CPU support
... plugins.use_dynamic_plugins()
... plugins.register_plugin('CPU', torch_xla_cpu_plugin.CpuPlugin())
... xr.set_device_type('CPU')
...
... print(xm.xla_device())
Traceback (most recent call last):
  Cell In[1], line 8
    from torch_xla.experimental import plugins
  File ~/venvs/pytorch-dev/lib/python3.12/site-packages/torch_xla/__init__.py:1
    raise ImportError("WRONG PACKAGE. Please install the package from Neuron Repository - pip.repos.neuron.amazonaws.com")
ImportError: WRONG PACKAGE. Please install the package from Neuron Repository - pip.repos.neuron.amazonaws.com

loganthomas · 2024-06-11T12:03:30Z

@JackCaoG @will-cromar

Do I need to install a different dependency? Am I missing something that prevents me from testing this correctly?

Or, is this what you were expecting?

JackCaoG · 2024-06-11T17:46:11Z

@will-cromar is out today, let me follow up with him when he is back

will-cromar · 2024-06-14T18:44:23Z

@loganthomas your specific issue is that you have the old placeholder package (1.0) that we used to list on PyPI before we started putting stable releases there. I assume pip defaulted to that because you have Python 3.12, and we have not released stable Python 3.12 packages on PyPI yet. Let me see if we can pull the old placeholder package to reduce confusion.

I updated the example plugin docs in #7277 to reference our contributing guide for bazel setup.

will-cromar · 2024-06-15T00:31:40Z

We deleted the placeholder release. So you shouldn't get the "wrong package" message again. Closing because I tested these docs today.

duncantech added docathon-h1-2024 easy labels Jun 4, 2024

JackCaoG assigned loganthomas Jun 4, 2024

tejadhith mentioned this issue Jun 8, 2024

Try running inference on an ARM CPU #7185

Open

will-cromar closed this as completed Jun 15, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Try using the CPU PJRT plugin #7184

Try using the CPU PJRT plugin #7184

duncantech commented Jun 4, 2024

loganthomas commented Jun 4, 2024

loganthomas commented Jun 5, 2024

JackCaoG commented Jun 5, 2024

will-cromar commented Jun 5, 2024 •

edited

Loading

loganthomas commented Jun 7, 2024

loganthomas commented Jun 11, 2024

JackCaoG commented Jun 11, 2024

will-cromar commented Jun 14, 2024

will-cromar commented Jun 15, 2024

Try using the CPU PJRT plugin #7184

Try using the CPU PJRT plugin #7184

Comments

duncantech commented Jun 4, 2024

📚 Documentation

loganthomas commented Jun 4, 2024

loganthomas commented Jun 5, 2024

JackCaoG commented Jun 5, 2024

will-cromar commented Jun 5, 2024 • edited Loading

loganthomas commented Jun 7, 2024

loganthomas commented Jun 11, 2024

JackCaoG commented Jun 11, 2024

will-cromar commented Jun 14, 2024

will-cromar commented Jun 15, 2024

will-cromar commented Jun 5, 2024 •

edited

Loading