dm_control/dm_control/_render/base.py at main · google-deepmind/dm_control · GitHub

History

166 lines (132 loc) · 5.29 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ============================================================================

"""Base class for OpenGL context handlers.

`ContextBase` defines a common API that OpenGL rendering contexts should conform

to. In addition, it provides a `make_current` context manager that:

1. Makes this OpenGL context current within the appropriate rendering thread.

2. Yields an object exposing a `call` method that can be used to execute OpenGL

calls within the rendering thread.

See the docstring for `dm_control.utils.render_executor` for further details

regarding rendering threads.

"""

import abc

import atexit

import collections

import contextlib

import sys

import weakref

from absl import logging

from dm_control._render import executor

import numpy as np

_CURRENT_CONTEXT_FOR_THREAD = collections.defaultdict(lambda: None)

_CURRENT_THREAD_FOR_CONTEXT = collections.defaultdict(lambda: None)

class ContextBase(metaclass=abc.ABCMeta):

"""Base class for managing OpenGL contexts."""

def __init__(self,

max_width,

max_height,

render_executor_class=executor.RenderExecutor):

"""Initializes this context."""

logging.debug('Using render executor class: %s',

render_executor_class.__name__)

self._render_executor = render_executor_class()

self._refcount = 0

self_weakref = weakref.ref(self)

def _free_at_exit():

if self_weakref():

self_weakref()._free_unconditionally() # pylint: disable=protected-access

atexit.register(_free_at_exit)

with self._render_executor.execution_context() as ctx:

ctx.call(self._platform_init, max_width, max_height)

self._patients = []

def keep_alive(self, obj):

self._patients.append(obj)

def dont_keep_alive(self, obj):

try:

self._patients.remove(obj)

except ValueError:

pass

def increment_refcount(self):

self._refcount += 1

def decrement_refcount(self):

self._refcount -= 1

@property

def terminated(self):

return self._render_executor.terminated

@property

def thread(self):

return self._render_executor.thread

def _free_on_executor_thread(self): # pylint: disable=missing-function-docstring

current_ctx = _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread]

if current_ctx is not None:

del _CURRENT_THREAD_FOR_CONTEXT[current_ctx]

del _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread]

self._platform_make_current()

try:

dummy = []

while self._patients:

patient = self._patients.pop()

assert sys.getrefcount(patient) == sys.getrefcount(dummy)

if hasattr(patient, 'free'):

patient.free()

del patient

finally:

self._platform_free()

def free(self):

"""Frees resources associated with this context if its refcount is zero."""

if self._refcount == 0:

self._free_unconditionally()

def _free_unconditionally(self):

self._render_executor.terminate(self._free_on_executor_thread)

def __del__(self):

self._free_unconditionally()

@contextlib.contextmanager

def make_current(self):

"""Context manager that makes this Renderer's OpenGL context current.

Yields:

An object that exposes a `call` method that can be used to call a

function on the dedicated rendering thread.

Raises:

RuntimeError: If this context is already current on another thread.

"""

with self._render_executor.execution_context() as ctx:

if _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread] != id(self):

if _CURRENT_THREAD_FOR_CONTEXT[id(self)]:

raise RuntimeError(

'Cannot make context {!r} current on thread {!r}: '

'this context is already current on another thread {!r}.'

.format(self, self._render_executor.thread,

_CURRENT_THREAD_FOR_CONTEXT[id(self)]))

else:

current_context = (

_CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread])

if current_context:

del _CURRENT_THREAD_FOR_CONTEXT[current_context]

_CURRENT_THREAD_FOR_CONTEXT[id(self)] = self._render_executor.thread

_CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread] = id(self)

ctx.call(self._platform_make_current)

yield ctx

def to_pixels(self, buffer):

"""Converts the buffer to pixels."""

return np.flipud(buffer)

@abc.abstractmethod

def _platform_init(self, max_width, max_height):

"""Performs an implementation-specific context initialization."""

@abc.abstractmethod

def _platform_make_current(self):

"""Make the OpenGL context current on the executing thread."""

@abc.abstractmethod

def _platform_free(self):

"""Performs an implementation-specific context cleanup."""

dm_control/dm_control/_render/base.py at main · google-deepmind/dm_control · GitHub

Uh oh!

Expand file tree

Breadcrumbs

base.py

Latest commit

History

Breadcrumbs

base.py

File metadata and controls