Метод query

Метод query класса Stream проверяет статус выполнения всех операций, поставленных в очередь в данный поток CUDA. Метод не блокирует выполнение программы и возвращает булево значение: True, если все операции завершены, и False, если в потоке ещё есть выполняющиеся операции. Это полезно для неблокирующей проверки состояния потока без ожидания завершения всех операций.

Синтаксис

stream.query()

Метод не принимает параметров и возвращает значение типа bool.

Пример

Давайте создадим поток и проверим его состояние сразу после создания:

import torch

stream = torch.cuda.Stream()
res = stream.query()
print(res)

Результат выполнения кода:

True

Поток только что создан и не содержит операций, поэтому метод возвращает True.

Пример

Давайте запустим операцию в потоке и сразу проверим его состояние:

import torch

torch.manual_seed(0)
stream = torch.cuda.Stream()

t = torch.randn(1000, 1000, device='cuda')
with torch.cuda.stream(stream):
	t = t * 2

res = stream.query()
print(res)

Результат выполнения кода может быть таким:

False

Операция умножения была поставлена в очередь, но ещё не завершилась, поэтому метод возвращает False.

Пример

Давайте дождёмся завершения операций в потоке с помощью метода synchronize и затем проверим состояние:

import torch

torch.manual_seed(0)
stream = torch.cuda.Stream()

t = torch.randn(1000, 1000, device='cuda')
with torch.cuda.stream(stream):
	t = t * 2

stream.synchronize()
res = stream.query()
print(res)

Результат выполнения кода:

True

После синхронизации все операции в потоке завершены, и метод возвращает True.

Пример

Давайте используем метод query в цикле для ожидания завершения операций без блокировки:

import torch
import time

torch.manual_seed(0)
stream = torch.cuda.Stream()

t = torch.randn(1000, 1000, device='cuda')
with torch.cuda.stream(stream):
	t = t.matmul(t)

while not stream.query():
	time.sleep(0.01)

print("All operations completed")

Результат выполнения кода:

"All operations completed"

Цикл проверяет состояние потока каждые 10 миллисекунд и продолжает работу, как только все операции завершены.

Смотрите также

метод synchronize,
который блокирует выполнение до завершения всех операций в потоке
метод wait_stream,
который ожидает завершения операций другого потока
метод record_event,
который записывает событие в поток для последующей синхронизации
класс Stream,
который представляет поток CUDA для асинхронного выполнения операций