# Jupyter поддерживает работу с кластером через пакет ipyparallel
# https://ipyparallel.readthedocs.io/en/latest/
# Его можно установить через PIP
# ! pip3 install ipyparallel
# После установки в интерфейсе Jupyter должна появиться вкладка IPython Clusters.
# Если этого не произошло, то нужно сделать:
# ipcluster nbextension enable

# ipyparallel кластер можно запусить из Jupyter из вкладки IPython Clusters.
# Альтернативно можно создать кластер из консоли:
# ! ipcluster start --profile=mpi -n 16
# По умолчанию процессы для счета создаются на локальной машине.
# Здесь мы попросили создать 16 процессов и указали, что будем использовать MPI (см. ниже).

# Для работы с MPI потребуется какая-либо реализация интерфейса 
# ! sudo apt install openmpi-bin
# и вспомогательная библиотека
# ! pip3 install mpi4py

# Чтобы получить доступ к узлам кластера, нам потребуется импортировать библиотеку
import ipyparallel as ipp

# Теперь мы можем создать интерфейс для работы с этими процессами.
rc = ipp.Client(profile='mpi')
# Смотрим, какие процессы были созданы:
print(f"{rc.ids}")
# Создадим "вид", для просмотра данных процессов
view = rc[:]
print(view)
# Следующая строка нужна для использования магии Jupyter
view.activate()
# Теперь мы можем выполнить содержимое ячейки на всех с помощью заклинания %%px. 

"""
Explanation: MPI через ipyparallel
End of explanation
"""


%%px
from mpi4py import MPI
import numpy as np

def psum(a):
    locsum = np.sum(a)
    rcvBuf = np.array(0.0,'d')
    MPI.COMM_WORLD.Allreduce([locsum, MPI.DOUBLE],
        [rcvBuf, MPI.DOUBLE],
        op=MPI.SUM)
    return rcvBuf

%pxresult

# Запускаем содержимое файла (идентично предыдущей ячейки) на каждой узле.
# view.run('psum.py')

# Импорт numpy нужен, так как ячейкой выше мы сделали это на удаленных машинах.
import numpy as np
# Рассылаем массив на все узлы кластера равными порциями.
view.scatter('a',np.arange(63,dtype='float'))
# Выводим содержимое массива `a` на каждой узле.
view['a']

# Вызываем написанную нами функцию суммирования:
%px totalsum = psum(a)

# Смотрим результат
%pxresult

# Аналогично заклинанию выше.
# view.execute('totalsum = psum(a)')

# Выводим результат, получившийся на каждой машине:
view['totalsum']

"""
Explanation: Используем MPI
End of explanation
"""


! pip install mpi4py
! pip3 install ipyparallel

! ipcluster start --profile=mpi -n 2 --daemonize

"""
Explanation: MPI на Google Colab
End of explanation
"""


# Для работы с NVidia GPU проще всего использовать numba.cuda.
# Устанавливается она как обычная numba
! pip3 install numba
# но для доступа к CUDA должно быть установлено соответствующее окружение, например для Ubuntu
# ! sudo apt install nvidia-cuda-toolkit
# Проверить корректность установки можно командой
! numba -s | grep CUDA
# Документация доступна здесь:
# https://numba.pydata.org/numba-doc/latest/cuda/index.html
# Развернутые сведения об устройстве можно получить командой
# ! clinfo

# Импортируем необходимую библиотеку.
import numba.cuda as cuda
# Проверяем доступность numba.cuda.
print(f"{cuda.is_available()}")
# Перечисляем доступные устройства.
cuda.detect() 

# Названия функций для математический операций нужно импортировать.
import math
import numpy as np
# Пишем простую функцию для работы на GPU.
@cuda.jit
def cudasqrt(x, y):
    i = cuda.grid(1) # Оси в CUDA нумеруются с 1-ой.
    if i>=cuda.gridsize(1): return
    y[i] = math.sqrt(x[i])
# Считаем корни
x = np.arange(10, dtype=np.float32)**2
y = np.empty_like(x)
cudasqrt[1, 100](x, y) # Обязательно указываем [число блоков, число потоков на блок].
print(y)

"""
Explanation: Далее действуем как указанно выше для запуска MPI через Jupyter.
CUDA через Numba
Если вы запускаете ноутбук через Google Colab, то для доступа к GPU вам нужно изменить настройки по-умолчанию: Меню > Среда выполнения > Сменить среду выполнения > Аппаратный ускоритель > GPU.
End of explanation
"""


