# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# type: ignore
from pathlib import Path

from read_version import read_version
from setuptools import find_namespace_packages, setup

setup(
    name="hydra-ray-launcher",
    version=read_version("hydra_plugins/hydra_ray_launcher", "__init__.py"),
    author="Jieru Hu",
    author_email="jieru@fb.com",
    description="Hydra Ray launcher plugin",
    long_description=(Path(__file__).parent / "README.md").read_text(),
    long_description_content_type="text/markdown",
    url="https://github.com/facebookresearch/hydra/",
    packages=find_namespace_packages(include=["hydra_plugins.*"]),
    classifiers=[
        "License :: OSI Approved :: MIT License",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        # "Programming Language :: Python :: 3.9",
        "Operating System :: MacOS",
        "Operating System :: POSIX :: Linux",
    ],
    install_requires=[
        "boto3==1.17.17",
        "hydra-core>=1.1.0.dev7",
        "ray[default]==1.6.0",
        # https://github.com/aio-libs/aiohttp/issues/6203
        "aiohttp!=3.8.0",
        "cloudpickle==1.6.0",
        "pickle5==0.0.11",
    ],
    include_package_data=True,
)

plugins/hydra_ray_launcher/setup.py
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# type: ignore
from pathlib import Path

from read_version import read_version
from setuptools import find_namespace_packages, setup

setup(
    name="hydra-ray-launcher",
    version=read_version("hydra_plugins/hydra_ray_launcher", "__init__.py"),
    author="Jieru Hu",
    author_email="jieru@fb.com",
    description="Hydra Ray launcher plugin",
    long_description=(Path(__file__).parent / "README.md").read_text(),
    long_description_content_type="text/markdown",
    url="https://github.com/facebookresearch/hydra/",
    packages=find_namespace_packages(include=["hydra_plugins.*"]),
    classifiers=[
        "License :: OSI Approved :: MIT License",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Operating System :: MacOS",
        "Operating System :: POSIX :: Linux",
    ],
    install_requires=[
        "boto3==1.17.17",
        "hydra-core>=1.1.0.dev7",
        "ray[default]==1.6.0",
        # https://github.com/aio-libs/aiohttp/issues/6203
        "aiohttp!=3.8.0",
        "cloudpickle==1.6.0",
        "pickle5==0.0.11",
    ],
    include_package_data=True,
)

plugins/hydra_ray_launcher/setup.py
diff --git a/plugins/hydra_ray_launcher/integration_test_tools/README.md b/plugins/hydra_ray_launcher/integration_test_tools/README.md
new file mode 100644
index 00000000000..eb3a4a12b4a
--- /dev/null
+++ b/plugins/hydra_ray_launcher/integration_test_tools/README.md
@@ -0,0 +1,21 @@
+## Setting up a new testing AMI for ray launcher.
+
+To run the tool:
+
+- Make sure the dependencies in `setup_integration_test_ami.py` matches exactly ray launcher's `setup.py`.
+
+- Before running the tool, set up your aws profile with admin access to the Hydra test AWS account.
+```
+AWS_PROFILE=jieru python create_integration_test_ami.py
+```
+You will see a new AMI created in the output
+```commandline
+ec2.Image(id='ami-0d65d5647e065a180') current state pending
+...
+```
+Sometimes it could take hours for a new AMI to be created. Proceed to the next step once the 
+AMI becomes available.
+
+- Update the `AWS_RAY_AMI` env variable in `tests/test_ray_aws_launcher.py`
+- Run the test locally and debug if needed.
+- Create a PR and make sure all CI pass!
diff --git a/plugins/hydra_ray_launcher/integration_test_tools/create_integration_test_ami.py b/plugins/hydra_ray_launcher/integration_test_tools/create_integration_test_ami.py
index e7d670d8747..78277eab753 100644
--- a/plugins/hydra_ray_launcher/integration_test_tools/create_integration_test_ami.py
+++ b/plugins/hydra_ray_launcher/integration_test_tools/create_integration_test_ami.py
@@ -29,7 +29,7 @@ def _run_command(command: str) -> str:
     return output
 
 
-@hydra.main(config_name="create_integration_test_ami_config")
+@hydra.main(config_path=".", config_name="create_integration_test_ami_config")
 def set_up_machine(cfg: DictConfig) -> None:
     security_group_id = cfg.security_group_id
     assert security_group_id != "", "Security group cannot be empty!"
diff --git a/plugins/hydra_ray_launcher/integration_test_tools/create_integration_test_ami_config.yaml b/plugins/hydra_ray_launcher/integration_test_tools/create_integration_test_ami_config.yaml
index 47085e7c2be..bea49b03b15 100644
--- a/plugins/hydra_ray_launcher/integration_test_tools/create_integration_test_ami_config.yaml
+++ b/plugins/hydra_ray_launcher/integration_test_tools/create_integration_test_ami_config.yaml
@@ -2,6 +2,7 @@ security_group_id: sg-095ac7c26aa0d33bb
 python_versions:
   - 3.7
   - 3.8
+  - 3.9
 ray_yaml:
   cluster_name: ray_test_base_AMI
   min_workers: 0
diff --git a/plugins/hydra_ray_launcher/integration_test_tools/setup_integration_test_ami.py b/plugins/hydra_ray_launcher/integration_test_tools/setup_integration_test_ami.py
index 317296992a7..72d955a5b26 100644
--- a/plugins/hydra_ray_launcher/integration_test_tools/setup_integration_test_ami.py
+++ b/plugins/hydra_ray_launcher/integration_test_tools/setup_integration_test_ami.py
@@ -9,6 +9,7 @@
     "ray[default]==1.6.0",
     "cloudpickle==1.6.0",
     "pickle5==0.0.11",
+    "aiohttp!=3.8",
 ]
 
 
diff --git a/plugins/hydra_ray_launcher/news/1205.feature b/plugins/hydra_ray_launcher/news/1205.feature
new file mode 100644
index 00000000000..20440f60a4f
--- /dev/null
+++ b/plugins/hydra_ray_launcher/news/1205.feature
@@ -0,0 +1 @@
+Add support for python 3.9
diff --git a/plugins/hydra_ray_launcher/setup.py b/plugins/hydra_ray_launcher/setup.py
index b0a71f58e97..4a64a977977 100644
--- a/plugins/hydra_ray_launcher/setup.py
+++ b/plugins/hydra_ray_launcher/setup.py
@@ -19,7 +19,7 @@
         "License :: OSI Approved :: MIT License",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
-        # "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.9",
         "Operating System :: MacOS",
         "Operating System :: POSIX :: Linux",
     ],
diff --git a/plugins/hydra_ray_launcher/tests/test_ray_aws_launcher.py b/plugins/hydra_ray_launcher/tests/test_ray_aws_launcher.py
index e1c89def69a..2b9a2f557ca 100644
--- a/plugins/hydra_ray_launcher/tests/test_ray_aws_launcher.py
+++ b/plugins/hydra_ray_launcher/tests/test_ray_aws_launcher.py
@@ -19,6 +19,7 @@
     LauncherTestSuite,
 )
 from hydra.test_utils.test_utils import chdir_hydra_root, chdir_plugin_root
+from omegaconf import OmegaConf
 from pytest import fixture, mark
 
 from hydra_plugins.hydra_ray_launcher.ray_aws_launcher import (  # type: ignore
@@ -54,7 +55,7 @@
     aws_not_configured = True
 
 
-ami = os.environ.get("AWS_RAY_AMI", "ami-0d03f5ce1006a7ed5")
+ami = os.environ.get("AWS_RAY_AMI", "ami-0436072b623a028fa")
 security_group_id = os.environ.get("AWS_RAY_SECURITY_GROUP", "sg-0a12b09a5ff961aee")
 subnet_id = os.environ.get("AWS_RAY_SUBNET", "subnet-acd2cfe7")
 instance_role = os.environ.get(
@@ -116,32 +117,33 @@
 chdir_plugin_root()
 
 
+def run_command(commands: str) -> str:
+    log.info(f"running: {commands}")
+    output = subprocess.getoutput(commands)
+    log.info(f"outputs: {output}")
+    return output
+
+
 def build_ray_launcher_wheel(tmp_wheel_dir: str) -> str:
     chdir_hydra_root()
     plugin = "hydra_ray_launcher"
     os.chdir(Path("plugins") / plugin)
     log.info(f"Build wheel for {plugin}, save wheel to {tmp_wheel_dir}.")
-    subprocess.getoutput(
-        f"python setup.py sdist bdist_wheel && cp dist/*.whl {tmp_wheel_dir}"
-    )
+    run_command(f"python setup.py sdist bdist_wheel && cp dist/*.whl {tmp_wheel_dir}")
     log.info("Download all plugin dependency wheels.")
-    subprocess.getoutput(f"pip download . -d {tmp_wheel_dir}")
-    plugin_wheel = subprocess.getoutput("ls dist/*.whl").split("/")[-1]
+    run_command(f"pip download . -d {tmp_wheel_dir}")
+    plugin_wheel = run_command("ls dist/*.whl").split("/")[-1]
     chdir_hydra_root()
     return plugin_wheel
 
 
 def build_core_wheel(tmp_wheel_dir: str) -> str:
     chdir_hydra_root()
-    subprocess.getoutput(
-        f"python setup.py sdist bdist_wheel && cp dist/*.whl {tmp_wheel_dir}"
-    )
+    run_command(f"python setup.py sdist bdist_wheel && cp dist/*.whl {tmp_wheel_dir}")
 
     # download dependency wheel for hydra-core
-    subprocess.getoutput(
-        f"pip download -r requirements/requirements.txt -d {tmp_wheel_dir}"
-    )
-    wheel = subprocess.getoutput("ls dist/*.whl").split("/")[-1]
+    run_command(f"pip download -r requirements/requirements.txt -d {tmp_wheel_dir}")
+    wheel = run_command("ls dist/*.whl").split("/")[-1]
     return wheel
 
 
@@ -161,12 +163,16 @@ def upload_and_install_wheels(
     sdk.run_on_cluster(
         connect_config,
         cmd=f"pip install --no-index --find-links={temp_remote_wheel_dir} {temp_remote_wheel_dir}{core_wheel}",
+        with_output=True,
+    )
+    log.info(f"Install plugin wheel {plugin_wheel} ")
+    log.info(
+        f"pip install --no-index --find-links={temp_remote_wheel_dir} {temp_remote_wheel_dir}{plugin_wheel}"
     )
-
-    log.info(f"Install plugin wheel {plugin_wheel}")
     sdk.run_on_cluster(
         connect_config,
         cmd=f"pip install --no-index --find-links={temp_remote_wheel_dir} {temp_remote_wheel_dir}{plugin_wheel}",
+        with_output=True,
     )
 
 
@@ -214,8 +220,8 @@ def manage_cluster() -> Generator[None, None, None]:
 
     # build all the wheels
     tmpdir = tempfile.mkdtemp()
-    plugin_wheel = build_ray_launcher_wheel(tmpdir)
     core_wheel = build_core_wheel(tmpdir)
+    plugin_wheel = build_ray_launcher_wheel(tmpdir)
     connect_config = {
         "cluster_name": cluster_name,
         "provider": {
@@ -233,14 +239,28 @@ def manage_cluster() -> Generator[None, None, None]:
         "head_node": ray_nodes_conf,
         "worker_nodes": ray_nodes_conf,
     }
+
+    # save connect_config as yaml, this could be useful for debugging
+    # you can run `ray attach <connect_config>.yaml` and log on to the AWS cluster for debugging.
+    conf = OmegaConf.create(connect_config)
+    with tempfile.NamedTemporaryFile(suffix=".yaml", delete=False) as fp:
+        OmegaConf.save(config=conf, f=fp.name, resolve=True)
+        log.info(f"Saving config to {fp.name}")
+
     sdk.create_or_update_cluster(
         connect_config,
     )
     sdk.run_on_cluster(
-        connect_config, run_env="auto", cmd=f"mkdir -p {temp_remote_dir}"
+        connect_config,
+        run_env="auto",
+        cmd=f"mkdir -p {temp_remote_dir}",
+        with_output=True,
     )
     sdk.run_on_cluster(
-        connect_config, run_env="auto", cmd=f"mkdir -p {temp_remote_wheel_dir}"
+        connect_config,
+        run_env="auto",
+        cmd=f"mkdir -p {temp_remote_wheel_dir}",
+        with_output=True,
     )
     upload_and_install_wheels(tmpdir, connect_config, core_wheel, plugin_wheel)
     validate_lib_version(connect_config)

[Ray-Plugin] Add support for Python 3.9
Python 3.9 support depends on https://github.com/ray-project/ray/issues/12788

Related to #1062 
