integrate-ai 9.18.0__tar.gz → 9.19.0__tar.gz

{integrate_ai-9.18.0 → integrate_ai-9.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: integrate_ai
-Version: 9.18.0
+Version: 9.19.0
 Summary: integrate.ai
 Author: integrate.ai
 License: Copyright (C) 2022 integrate.ai, Inc. All rights reserved.

{integrate_ai-9.18.0 → integrate_ai-9.19.0}/integrate_ai/onprem_node.py

@@ -11,8 +11,9 @@ from integrate_ai.utils.typer_utils import (
     TogglePromptOption,
 )
 
-
 app = typer.Typer(no_args_is_help=True)
+# Default white list for AWS addr to query EC2 metadata
+no_proxy_addr = "169.254.169.254,169.254.170.2,/var/run/docker.sock"
 
 
 @app.command()
@@ -36,15 +37,39 @@ def install(
     if os.geteuid() != 0:
         rich.print("This command must be run as root.")
         sys.exit(1)
+    try:
+        # clean up rm /tmp/ecs-anywhere-install.sh if there is any
+        if os.path.exists("/tmp/ecs-anywhere-install.sh"):
+            subprocess.run("rm /tmp/ecs-anywhere-install.sh", shell=True, check=True, text=True, capture_output=True)
+
+        # register taskrunner id
+        client = RestClient(token=token)
+        response = client.register_on_prem_taskrunner(taskrunner_name)
+
+        # install ecs anywhere
+        rich.print("Installing ECS Agent..")
+        install_ecs_agent(response)
+        rich.print("---ECS Agent Installed---")
+        taskrunners = client.get_taskrunner_info(taskrunner_name=taskrunner_name)
+        prov_info = taskrunners[0].get("provisioning_info")
+        runtime_info = prov_info.get("runtime_info")
+        if runtime_info:
+            proxy_addr = runtime_info.get("iai_forward_proxy_http_url")
+            rich.print(f"Proxy info from API : {proxy_addr}")
+            # configure proxy
+            if proxy_addr:
+                rich.print(f"Configuring proxy {proxy_addr}")
+                if proxy_addr:
+                    configure_proxy(proxy_addr=proxy_addr)
+        rich.print("OnPrem node installed successfully!")
+    except Exception as e:
+        rich.print(
+            f"On-Prem register command failed with error: {e}, Logs can be found in ecsanywhere_output.txt and proxy_config.txt"
+        )
+        raise e
 
-    # clean up rm /tmp/ecs-anywhere-install.sh if there is any
-    if os.path.exists("/tmp/ecs-anywhere-install.sh"):
-        subprocess.run("rm /tmp/ecs-anywhere-install.sh", shell=True, check=True, text=True, capture_output=True)
-
-    # register taskrunner id
-    client = RestClient(token=token)
-    response = client.register_on_prem_taskrunner(taskrunner_name)
 
+def install_ecs_agent(response):
     # install ecs anywhere
     cmd = 'sudo curl --proto "https" -o "/tmp/ecs-anywhere-install.sh" "https://amazon-ecs-agent.s3.amazonaws.com/ecs-anywhere-install-latest.sh"'
     cmd += " && sudo bash /tmp/ecs-anywhere-install.sh"
@@ -53,14 +78,75 @@ def install(
     cmd += f' --activation-id "{response["activation_id"]}"'
     cmd += f' --activation-code "{response["activation_code"]}"'
     cmd += "> ecsanywhere_output.txt"
-
     try:
         rich.print("Registering...")
-        subprocess.run(cmd, shell=True, check=True)
+        subprocess.run(cmd, shell=True, check=True, timeout=1200)
         rich.print("Agent registered successfully.")
         rich.print("Output is saved in ecsanywhere_output.txt. The file contains instance id, please do not delete.")
     except subprocess.CalledProcessError as e:
-        rich.print(f"Command failed with error: {e.stderr}, Logs can be found in ecsanywhere_output.txt.")
+        message = f"ECS Installation command failed with error: {e.stderr}, Logs can be found in ecsanywhere_output.txt"
+        rich.print(message)
+        raise Exception(message) from e
+
+
+# Add proxy addresses to the ecs config so traffic can go through the proxy by default
+# See : https://docs.aws.amazon.com/AmazonECS/latest/developerguide/http_proxy_config.html
+def configure_proxy(proxy_addr: str):
+    rich.print("Configuring Proxy config in the host")
+    # proxy ssm agent
+    ssm_config_cmd = "sudo mkdir -p /etc/systemd/system/amazon-ssm-agent.service.d"
+    ssm_config_cmd += "&& { cat > /etc/systemd/system/amazon-ssm-agent.service.d/http-proxy.conf <<-EOF \n"
+    ssm_config_cmd += "[Service] \n"
+    ssm_config_cmd += f"Environment=HTTPS_PROXY={proxy_addr} \n"
+    ssm_config_cmd += f"Environment=NO_PROXY={no_proxy_addr} \n"
+    ssm_config_cmd += "EOF\n"
+    ssm_config_cmd += "}"
+
+    # proxy docker service
+    docker_svc_cmd = "sudo mkdir -p /etc/systemd/system/docker.service.d"
+    docker_svc_cmd += "&& { cat > /etc/systemd/system/docker.service.d/http-proxy.conf <<-EOF \n"
+    docker_svc_cmd += "[Service] \n"
+    docker_svc_cmd += f"Environment=HTTPS_PROXY={proxy_addr} \n"
+    docker_svc_cmd += f"Environment=NO_PROXY={no_proxy_addr} \n"
+    docker_svc_cmd += "EOF\n"
+    docker_svc_cmd += "}"
+
+    # proxy ecs-init
+    ecs_init_config_cmd = "sudo mkdir -p /etc/systemd/system/ecs.service.d"
+    ecs_init_config_cmd += "&& { cat > /etc/systemd/system/ecs.service.d/http-proxy.conf <<-EOF \n"
+    ecs_init_config_cmd += "[Service] \n"
+    ecs_init_config_cmd += f"Environment=HTTPS_PROXY={proxy_addr} \n"
+    ecs_init_config_cmd += f"Environment=NO_PROXY={no_proxy_addr} \n"
+    ecs_init_config_cmd += "EOF\n"
+    ecs_init_config_cmd += "}"
+
+    # proxy ecs
+    ecs_config_cmd = "echo patching ecs config"
+    ecs_config_cmd += "&& { cat >> /etc/ecs/ecs.config <<-EOF \n"
+    ecs_config_cmd += "[Service] \n"
+    ecs_config_cmd += f"HTTPS_PROXY={proxy_addr} \n"
+    ecs_config_cmd += f"NO_PROXY={no_proxy_addr} \n"
+    ecs_config_cmd += "EOF\n"
+    ecs_config_cmd += "}"
+
+    # restart all affected services
+    restart_cmd = (
+        "sudo systemctl daemon-reload && sudo systemctl restart docker.service && sudo systemctl restart "
+        "amazon-ssm-agent && sudo systemctl restart ecs"
+    )
+
+    # completion_command
+    completion_command = "echo proxy config successfully updated"
+
+    command_to_execute = f"{ssm_config_cmd} && {docker_svc_cmd} && {ecs_init_config_cmd} && {ecs_config_cmd} && {restart_cmd} && {completion_command}"
+    try:
+        # rich.print(f"Configuring Proxy command : {command_to_execute}")
+        subprocess.run(command_to_execute, shell=True, check=True, timeout=1200)
+        rich.print("Proxy configured successfully")
+    except subprocess.CalledProcessError as e:
+        message = f"Proxy configuration failed with error: {e.stderr}"
+        rich.print(message)
+        raise Exception(message) from e
 
 
 @app.command()
@@ -121,6 +207,11 @@ def uninstall(
         remove_directories_cmd = (
             "sudo rm -rf /var/lib/ecs /etc/ecs /var/lib/amazon/ssm /var/log/ecs /var/log/amazon/ssm"
         )
+        remove_directories_cmd += " && sudo rm -rf /etc/systemd/system/amazon-ssm-agent.service.d/http-proxy.conf"
+        remove_directories_cmd += " && sudo rm -rf /etc/systemd/system/docker.service.d/http-proxy.conf"
+        remove_directories_cmd += " && sudo rm -rf /etc/systemd/system/ecs.service.d/http-proxy.conf"
+        remove_directories_cmd += " && sudo rm -rf /etc/ecs/ecs.config"
+
         subprocess.run(remove_directories_cmd, shell=True, check=True, text=True)
         rich.print("Leftover directories removed")
 
@@ -128,6 +219,7 @@ def uninstall(
         remove_output_cmd = "sudo rm ecsanywhere_output.txt"
         subprocess.run(remove_output_cmd, shell=True, check=True, text=True)
         rich.print("ecsanywhere_output.txt removed")
+        rich.print("proxy entries removed")
 
     except subprocess.CalledProcessError as e:
         rich.print(f"Command failed with error: {e.stderr}")

{integrate_ai-9.18.0 → integrate_ai-9.19.0}/integrate_ai/utils/rest_client.py

@@ -174,6 +174,21 @@ class RestClient:
 
         return response.json()
 
+    def get_taskrunner_info(self, taskrunner_name) -> Dict[str, Any]:
+        """Get taskrunner info.
+        Args:
+            taskrunner_name (str): Name of the taskrunner
+        Returns:
+                List of Taskrunner info
+        Raises:
+            IntegrateAIException: Customized IntegrateAI exception for the HTTP Exception.
+        """
+        url = f"{self.api_url}/taskrunners"
+        params = {"name": taskrunner_name}
+        response = requests.get(url, headers=self._headers(), params=params)
+        check_for_IntegrateAIException(response=response)
+        return response.json()
+
     def register_on_prem_taskrunner(self, taskrunner_name) -> Dict[str, Any]:
         """Get package versions.
         Args:
@@ -182,6 +197,7 @@ class RestClient:
                 A Dict of the form
                 ```
             {
+                "proxy_info" : {"iai_forward_proxy_http_url" : "http://iai-proxy-ca-central-1-prod.integrateai.net:8888"}
                 "activation_id": <activation_id>,
                 "activation_code": <activation_code>,
                 "region": <region>,

{integrate_ai-9.18.0 → integrate_ai-9.19.0}/integrate_ai.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: integrate-ai
-Version: 9.18.0
+Version: 9.19.0
 Summary: integrate.ai
 Author: integrate.ai
 License: Copyright (C) 2022 integrate.ai, Inc. All rights reserved.

{integrate_ai-9.18.0 → integrate_ai-9.19.0}/setup.py

@@ -70,5 +70,5 @@ iai client train --token <IAI_TOKEN> --session <SESSION_ID> --train_path <PATH_T
         'integrate_ai.utils',
     ),
     'python_requires': '>=3.7.5',
-    'version': '9.18.0',
+    'version': '9.19.0',
 })