feat: add decoupled workflows config

[jwemmons]

Description

Refactors workflows startup script to implement new decoupled setup flow. Logic branching determined by enableWorkflowsConfigV2 feature flag.

If feature flag enabled:

• Pulls new config files from project s3 bucket into .workflows_setup
• Writes user's requirements.txt and startup.sh into plugins/user-requirements.txt and plugins/user-startup.sh.

Otherwise, existing logic is used.

Type of Change

• Image update - Bug fix
• Image update - New feature
• Image update - Breaking change
• SMD image build tool update
• Documentation update

Release Information

Does this change need to be included in patch version releases? By default, any pull requests will only be added to the next SMD image minor version release once they are merged in template folder. Only critical bug fix or security update should be applied to new patch versions of existed image minor versions.

• Yes (Critical bug fix or security update)
• No (New feature or non-critical change)
• N/A (Not an image update)

If yes, please explain why:
While not a bug fix/security update, feature flag will be enabled world-wide, so this change should be present in all active image versions.

How Has This Been Tested?

1. Modified following files:

etc/sagemaker-ui/workflows/workflow_client.py
etc/sagemaker-ui/workflows/start-workflows-container.sh

2. Stopped workflows local runner and re-ran post-startup script:

python /etc/sagemaker-ui/workflows/workflow_client.py stop-local-runner
bash /etc/sagemaker-ui/sagemaker_ui_post_startup.sh

3. Observed existing startup behavior in .workflows_setup (as feature flag not yet enabled) and containers started:

sagemaker-user@default:~/.workflows_setup$ ls plugins
amzn_SagemakerWorkflowsOperator-1.0-py3-none-any.whl  amzn_sagemaker_studio-1.0.19-py3-none-any.whl

sagemaker-user@default:~/.workflows_setup$ cat requirements/requirements.txt 
--find-links /usr/local/airflow/plugins
--no-index

amzn_sagemaker_studio
amzn_SagemakerWorkflowsOperator
# Place any requirements you'd like included in your workflows environment here
# Note that you will need to restart your space for changes to take effect.
# For example:
# numpy==1.26.4

sagemaker-user@default:~/.workflows_setup$ cat startup/startup.sh 
#!/bin/sh
# Place any special instructions you'd like run during your workflows environment startup here
# Note that you will need to restart your space for changes to take effect.
# For example:
# pip install dbt-core

sagemaker-user@default:~/.workflows_setup$ docker ps
CONTAINER ID   IMAGE                                                            COMMAND                  CREATED         STATUS                   PORTS     NAMES
5a8ffc80156a   058264401727.dkr.ecr.us-east-2.amazonaws.com/mwaa_image:latest   "python3 -m mwaa.ent…"   7 minutes ago   Up 7 minutes                       mwaa-292-webserver
e8e6d0d6c0b3   058264401727.dkr.ecr.us-east-2.amazonaws.com/mwaa_image:latest   "python3 -m mwaa.ent…"   7 minutes ago   Up 7 minutes                       mwaa-292-scheduler
176625f24f34   058264401727.dkr.ecr.us-east-2.amazonaws.com/postgres:13         "docker-entrypoint.s…"   7 minutes ago   Up 7 minutes (healthy)             mwaa-292-db

4. Manually uploaded new config files to project s3 bucket. Then, modified workflow_client.py again, hardcoding check_config_v2_enabled to print "True" once feature flag value found:

for resource in environment_resources:
    if resource['name'] == "enableWorkflowsConfigV2":
        # print(resource['value'] == "true")
        print("True")
        return

5. Repeated step 2
6. Observed:
   a) config files pulled from s3 bucket during post-startup:

   + bash /etc/sagemaker-ui/workflows/start-workflows-container.sh
   ...
   Reading state information... Done
   docker-ce-cli is already the newest version (5:29.1.3-1~ubuntu.22.04~jammy).
   docker-compose-plugin is already the newest version (2.29.2-1~ubuntu.22.04~jammy).
   0 upgraded, 0 newly installed, 0 to remove and 7 not upgraded.
   delete: ../../../home/sagemaker-user/.workflows_setup/plugins/amzn_SagemakerWorkflowsOperator-1.0-py3-none-any.whl
   delete: ../../../home/sagemaker-user/.workflows_setup/plugins/amzn_sagemaker_studio-1.0.19-py3-none-any.whl
   download: s3://<BUCKET NAME>/<KEY>/sys/workflows/plugins/apache_airflow_providers_amazon-9.9.0-py3-none-any.whl to ../../../home/sagemaker-user/.workflows_setup/plugins/apache_airflow_providers_amazon-9.9.0-py3-none-any.whl
   download: s3://<BUCKET NAME>/<KEY>/sys/workflows/plugins/sagemaker_studio-1.0.23-py3-none-any.whl to ../../../home/sagemaker-user/.workflows_setup/plugins/sagemaker_studio-1.0.23-py3-none-any.whl
   download: s3://<BUCKET NAME>/<KEY>/sys/workflows/requirements.txt to ../../../home/sagemaker-user/.workflows_setup/requirements/requirements.txt
   download: s3://<BUCKET NAME>/<KEY>/sys/workflows/startup.sh to ../../../home/sagemaker-user/.workflows_setup/startup/startup.sh
   

   b) New config files present in .workflows_setup and containers started:

   sagemaker-user@default:~/.workflows_setup$ ls plugins/
   apache_airflow_providers_amazon-9.9.0-py3-none-any.whl  sagemaker_studio-1.0.23-py3-none-any.whl  user-requirements.txt  user-startup.sh
   
   sagemaker-user@default:~/.workflows_setup$ cat requirements/requirements.txt 
   # Place any requirements you'd like included in your workflows environment here
   # Note that you will need to restart your space for changes to take effect.
   # For example:
   # numpy==1.26.4
   
   --constraint /usr/local/airflow/plugins/constraints.txt
   --find-links /usr/local/airflow/plugins
   --no-index
   
   sagemaker-studio
   apache-airflow-providers-amazon
   
   sagemaker-user@default:~/.workflows_setup$ cat startup/startup.sh 
   #!/bin/bash
   
   # Install the user's requirements for them
   pip install /usr/local/airflow/plugins/user-requirements.txt \
       --constraint /usr/local/airflow/plugins/constraints.txt
   
   # Run the user's startup script for them
   bash /usr/local/airflow/plugins/user-startup.sh
   
   sagemaker-user@default:~/.workflows_setup$ docker ps
   CONTAINER ID   IMAGE                                                            COMMAND                  CREATED         STATUS                   PORTS     NAMES
   faae9476956b   058264401727.dkr.ecr.us-east-2.amazonaws.com/mwaa_image:latest   "python3 -m mwaa.ent…"   4 minutes ago   Up 4 minutes                       mwaa-292-webserver
   3c84f9a07cc6   058264401727.dkr.ecr.us-east-2.amazonaws.com/mwaa_image:latest   "python3 -m mwaa.ent…"   4 minutes ago   Up 4 minutes                       mwaa-292-scheduler
   1a1e536a6feb   058264401727.dkr.ecr.us-east-2.amazonaws.com/postgres:13         "docker-entrypoint.s…"   4 minutes ago   Up 4 minutes (healthy)             mwaa-292-db
   

Checklist:

• My code follows the style guidelines of this project
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works

Test Screenshots (if applicable):

Related Issues

N/A

Additional Notes

• NOTE: pull from s3 also deletes any existing plugins (note use of --delete flag). This was done to remove the old system plugins in existing spaces.

[reganbaum]

we should probably wrap these calls in a try/catch, and if it fails, then assume the FF isn't enabled as a fallback.

It's also probably useful for debugging purposes later on to include some echo statements throughout this logic so if we ever need to get the logs from the customer, we can see what happened more easily. Could we add in the try/catch and echo statements (e.g., print if the v2config is enabled, when it starts/finishes something, after receiving the blueprint id, etc)

[jwemmons]

There is overall exception handling for this block here: https://github.com/aws/sagemaker-distribution/blob/main/template/v2/dirs/etc/sagemaker-ui/workflows/start-workflows-container.sh#L173

I don't know that falling back to the pre-FF behavior is the correct approach. Say this new logic has already successfully ran previously and now the s3 calls fail. Now you'll end up copying/using the old config files instead of the new ones and your local and shared environments will differ. I think the current logic of catching the failure and exiting immediately is more appropriate.

Will add additional logging though.