|
import asyncio |
|
import contextlib |
|
import logging |
|
import os |
|
import time |
|
from typing import List |
|
|
|
import torch |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
DEBUG_COMPLETED_TIME = bool(os.environ.get('DEBUG_COMPLETED_TIME', False)) |
|
|
|
|
|
@contextlib.asynccontextmanager |
|
async def completed(trace_name='', |
|
name='', |
|
sleep_interval=0.05, |
|
streams: List[torch.cuda.Stream] = None): |
|
"""Async context manager that waits for work to complete on given CUDA |
|
streams.""" |
|
if not torch.cuda.is_available(): |
|
yield |
|
return |
|
|
|
stream_before_context_switch = torch.cuda.current_stream() |
|
if not streams: |
|
streams = [stream_before_context_switch] |
|
else: |
|
streams = [s if s else stream_before_context_switch for s in streams] |
|
|
|
end_events = [ |
|
torch.cuda.Event(enable_timing=DEBUG_COMPLETED_TIME) for _ in streams |
|
] |
|
|
|
if DEBUG_COMPLETED_TIME: |
|
start = torch.cuda.Event(enable_timing=True) |
|
stream_before_context_switch.record_event(start) |
|
|
|
cpu_start = time.monotonic() |
|
logger.debug('%s %s starting, streams: %s', trace_name, name, streams) |
|
grad_enabled_before = torch.is_grad_enabled() |
|
try: |
|
yield |
|
finally: |
|
current_stream = torch.cuda.current_stream() |
|
assert current_stream == stream_before_context_switch |
|
|
|
if DEBUG_COMPLETED_TIME: |
|
cpu_end = time.monotonic() |
|
for i, stream in enumerate(streams): |
|
event = end_events[i] |
|
stream.record_event(event) |
|
|
|
grad_enabled_after = torch.is_grad_enabled() |
|
|
|
|
|
|
|
assert (grad_enabled_before == grad_enabled_after |
|
), 'Unexpected is_grad_enabled() value change' |
|
|
|
are_done = [e.query() for e in end_events] |
|
logger.debug('%s %s completed: %s streams: %s', trace_name, name, |
|
are_done, streams) |
|
with torch.cuda.stream(stream_before_context_switch): |
|
while not all(are_done): |
|
await asyncio.sleep(sleep_interval) |
|
are_done = [e.query() for e in end_events] |
|
logger.debug( |
|
'%s %s completed: %s streams: %s', |
|
trace_name, |
|
name, |
|
are_done, |
|
streams, |
|
) |
|
|
|
current_stream = torch.cuda.current_stream() |
|
assert current_stream == stream_before_context_switch |
|
|
|
if DEBUG_COMPLETED_TIME: |
|
cpu_time = (cpu_end - cpu_start) * 1000 |
|
stream_times_ms = '' |
|
for i, stream in enumerate(streams): |
|
elapsed_time = start.elapsed_time(end_events[i]) |
|
stream_times_ms += f' {stream} {elapsed_time:.2f} ms' |
|
logger.info('%s %s %.2f ms %s', trace_name, name, cpu_time, |
|
stream_times_ms) |
|
|
|
|
|
@contextlib.asynccontextmanager |
|
async def concurrent(streamqueue: asyncio.Queue, |
|
trace_name='concurrent', |
|
name='stream'): |
|
"""Run code concurrently in different streams. |
|
|
|
:param streamqueue: asyncio.Queue instance. |
|
|
|
Queue tasks define the pool of streams used for concurrent execution. |
|
""" |
|
if not torch.cuda.is_available(): |
|
yield |
|
return |
|
|
|
initial_stream = torch.cuda.current_stream() |
|
|
|
with torch.cuda.stream(initial_stream): |
|
stream = await streamqueue.get() |
|
assert isinstance(stream, torch.cuda.Stream) |
|
|
|
try: |
|
with torch.cuda.stream(stream): |
|
logger.debug('%s %s is starting, stream: %s', trace_name, name, |
|
stream) |
|
yield |
|
current = torch.cuda.current_stream() |
|
assert current == stream |
|
logger.debug('%s %s has finished, stream: %s', trace_name, |
|
name, stream) |
|
finally: |
|
streamqueue.task_done() |
|
streamqueue.put_nowait(stream) |
|
|
