def write(self, *args, **kwargs):

return sys.stderr.write(*args, **kwargs)

def writeln(self, *args, **kwargs):

if args or kwargs:

sys.stderr.write(*args, **kwargs)

def __init__(self, exclude_list):

super(TestLoader, self).__init__()

self._exclude_list = exclude_list

def _match_path(self, path, full_path, pattern):

if not fnmatch.fnmatch(path, pattern):

return False

module_name = full_path.replace('/', '.').rstrip('.py')

if any(module_name.endswith(x) for x in self._exclude_list):

return False

testcases = [x.rstrip() for x in open(filename, 'r').readlines() if x]

# Remove comments and blanks after comments are removed.

testcases = [x.partition('#')[0].strip() for x in testcases]

def run_tests(self):

# Import absl inside run, where dependencies have been loaded already.

from absl import app # pylint: disable=g-import-not-at-top

def main(_):

# pybullet imports multiprocessing in their setup.py, which causes an

# issue when we import multiprocessing.pool.dummy down the line because

# the PYTHONPATH has changed.

for module in [

'multiprocessing', 'multiprocessing.pool', 'multiprocessing.dummy',

'multiprocessing.pool.dummy'

]:

if module in sys.modules:

del sys.modules[module]

# Reimport multiprocessing to avoid spurious error printouts. See

# https://bugs.python.org/issue15881.

import multiprocessing as _ # pylint: disable=g-import-not-at-top

import tensorflow as tf # pylint: disable=g-import-not-at-top

# Sets all GPUs to 1GB of memory. The process running the bulk of the unit

# tests allocates all GPU memory because by default TensorFlow allocates

# all GPU memory during initialization. This causes tests in

# run_seperately to fail with out of memory errors because they are run as

# a subprocess of the process holding the GPU memory.

gpus = tf.config.experimental.list_physical_devices('GPU')

for gpu in gpus:

tf.config.set_logical_device_configuration(

gpu, [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])

run_separately = load_test_list('test_individually.txt')

broken_tests = load_test_list(FLAGS.broken_tests)

test_loader = TestLoader(exclude_list=run_separately + broken_tests)

test_suite = test_loader.discover('tf_agents', pattern='*_test.py')

stderr = StderrWrapper()

result = unittest.TextTestResult(stderr, descriptions=True, verbosity=2)

test_suite.run(result)

external_test_failures = []

for test in run_separately:

filename = 'tf_agents/%s.py' % test.replace('.', '/')

try:

subprocess.check_call([sys.executable, filename])

except subprocess.CalledProcessError as e:

external_test_failures.append(e)

result.printErrors()

for failure in external_test_failures:

stderr.writeln(str(failure))

final_output = (

'Tests run: {} grouped and {} external. '.format(

result.testsRun, len(run_separately)) +

'Errors: {} Failures: {} External failures: {}.'.format(

len(result.errors),

len(result.failures),

len(external_test_failures)))

header = '=' * len(final_output)

stderr.writeln(header)

stderr.writeln(final_output)

stderr.writeln(header)

if result.wasSuccessful() and not external_test_failures:

return 0

else:

return 1

# Run inside absl.app.run to ensure flags parsing is done.

from tf_agents.system import system_multiprocessing as multiprocessing # pylint: disable=g-import-not-at-top

"""This class is needed in order to create OS specific wheels."""

def has_ext_modules(self):

"""Returns list of required packages."""

required_packages = [

'absl-py >= 0.6.1',

'cloudpickle >= 1.3',

'gin-config >= 0.4.0',

'gym >= 0.17.0, <=0.23.0',

'numpy >= 1.13.3',

'pillow',

'six >= 1.10.0',

'protobuf >= 3.11.3',

'wrapt >= 1.11.1',

'typing-extensions >= 3.7.4.3',

# Used by gym >= 0.22.0. Only installed as a dependency when gym[all] is

# installed or if gym[*] (where * is an environment which lists pygame as

# a dependency).

'pygame == 2.1.0',

]

add_additional_packages(required_packages)

"""Adds additional required packages."""

if FLAGS.release:

tfp_version = TFP_VERSION

else:

tfp_version = TFP_NIGHTLY

if FLAGS.tfp_version:

tfp_version = FLAGS.tfp_version

"""Returns list of packages needed when testing."""

test_packages = [

'mock >= 2.0.0',

'opencv-python >= 3.4.1.15',

'pybullet',

'scipy >= 1.1.0',

'tensorflow_datasets'

]

"""Returns list of required packages if using reverb."""

reverb_packages = []

if FLAGS.release:

tf_version = TENSORFLOW_VERSION

reverb_version = REVERB_VERSION

rlds_version = RLDS_VERSION

else:

tf_version = TENSORFLOW_NIGHTLY

reverb_version = REVERB_NIGHTLY

rlds_version = RLDS_NIGHTLY

# Overrides required versions if FLAGS are set.

if FLAGS.tf_version:

tf_version = FLAGS.tf_version

if FLAGS.reverb_version:

reverb_version = FLAGS.reverb_version

if FLAGS.rlds_version:

rlds_version = FLAGS.rlds_version

reverb_packages.append(rlds_version)

reverb_packages.append(reverb_version)

reverb_packages.append(tf_version)

"""Returns the version and project name to associate with the build."""

__dev_version__ = tf_agents_version.__dev_version__ # pylint: disable=invalid-name

__rel_version__ = tf_agents_version.__rel_version__ # pylint: disable=invalid-name

if FLAGS.release:

version = __rel_version__

project_name = 'tf-agents'

else:

version = __dev_version__

version += datetime.datetime.now().strftime('%Y%m%d')

project_name = 'tf-agents-nightly'

"""Triggers build, install, and other features of `setuptools.setup`."""

# Builds the long description from the README.

root_path = os.path.abspath(os.path.dirname(__file__))

with codecs.open(os.path.join(root_path, 'README.md'), encoding='utf-8') as f:

long_description = f.read()

version, project_name = get_version()

test_packages = get_test_packages()

setup(

name=project_name,

version=version,

description='TF-Agents: A Reinforcement Learning Library for TensorFlow',

long_description=long_description,

long_description_content_type='text/markdown',

author='Google LLC',

author_email='no-reply@google.com',

url='https://github.com/tensorflow/agents',

license='Apache 2.0',

include_package_data=True,

packages=find_packages(),

install_requires=get_required_packages(),

tests_require=test_packages,

extras_require={

'tests': test_packages,

'reverb': get_reverb_packages(),

},

# Supports Python 3 only.

python_requires='>=3',

# Add in any packaged data.

zip_safe=False,

distclass=BinaryDistribution,

cmdclass={

'test': Test,

},

classifiers=[

'Development Status :: 3 - Alpha',

'Intended Audience :: Developers',

'Intended Audience :: Education',

'Intended Audience :: Science/Research',

'License :: OSI Approved :: Apache Software License',

'Programming Language :: Python :: 3',

'Programming Language :: Python :: 3.7',

'Programming Language :: Python :: 3.8',

'Programming Language :: Python :: 3.9',

'Programming Language :: Python :: 3.10',

'Topic :: Scientific/Engineering',

'Topic :: Scientific/Engineering :: Mathematics',

'Topic :: Scientific/Engineering :: Artificial Intelligence',

'Topic :: Software Development',

'Topic :: Software Development :: Libraries',

'Topic :: Software Development :: Libraries :: Python Modules',

],

keywords='tensorflow agents reinforcement learning machine bandits',