Generate requirements.txt file for ACME when running with Launchpad.

PiperOrigin-RevId: 427233844

setup.py

@@ -14,12 +14,31 @@
 # limitations under the License.
 
 """Setup file for installing the BLE."""
+import os
 import pathlib
 import setuptools
 
 current_directory = pathlib.Path(__file__).parent
 description = (current_directory / 'README.md').read_text()
 
+core_requirements = [
+    'absl-py',
+    'flax',
+    'gcsfs',
+    'gin-config',
+    'gym',
+    'jax >= 0.2.28',
+    'jaxlib >= 0.1.76',
+    'opensimplex <= 0.3.0',
+    's2sphere',
+    'scikit-learn',
+    'tensorflow',
+    'tensorflow-datasets >= 4.4.0',
+    'tensorflow-probability',
+    'transitions',
+    'zarr',
+]
+
 dopamine_requirements = [
     'dopamine-rl >= 4.0.0',
 ]
@@ -32,29 +51,52 @@
     'rlax',
 ]
 
+
+def generate_requirements_file(path=None):
+  """Generates requirements.txt file needed for running Acme.
+
+  It is used by Launchpad GCP runtime to generate Acme requirements to be
+  installed inside the docker image. Acme itself is not installed from pypi,
+  but instead sources are copied over to reflect any local changes made to
+  the codebase.
+  Args:
+    path: path to the requirements.txt file to generate.
+  """
+  if not path:
+    path = os.path.join(os.path.dirname(__file__), 'acme/requirements.txt')
+  with open(path, 'w') as f:
+    for package in set(core_requirements + dopamine_requirements +
+                       acme_requirements):
+      f.write(f'{package}\n')
+
+
+class BuildPy(setuptools.command.build_py.build_py):
+
+  def run(self):
+    generate_requirements_file()
+    setuptools.command.build_py.build_py.run(self)
+
+
+class Develop(setuptools.command.develop.develop):
+
+  def run(self):
+    generate_requirements_file()
+    setuptools.command.develop.develop.run(self)
+
+cmdclass = {
+    'build_py': BuildPy,
+    'develop': Develop,
+}
+
+
 setuptools.setup(
     name='balloon_learning_environment',
     long_description=description,
     long_description_content_type='text/markdown',
     version='0.1.1',
+    cmdclass=cmdclass,
     packages=setuptools.find_packages(),
-    install_requires=[
-        'absl-py',
-        'flax',
-        'gcsfs',
-        'gin-config',
-        'gym',
-        'jax >= 0.2.28',
-        'jaxlib >= 0.1.76',
-        'opensimplex <= 0.3.0',
-        's2sphere',
-        'scikit-learn',
-        'tensorflow',
-        'tensorflow-datasets >= 4.4.0',
-        'tensorflow-probability',
-        'transitions',
-        'zarr',
-    ],
+    install_requires=core_requirements,
     extras_require={
         'dopamine': dopamine_requirements,
         'acme': acme_requirements,