Convert CodeArtifact integration to simply generate an AWS CLI command to log into CodeArtifact

Author: akuma12

src/sagemaker/processing.py

@@ -22,7 +22,7 @@
 import pathlib
 import logging
 from textwrap import dedent
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 from copy import copy
 import re
 
@@ -1845,19 +1845,18 @@ def _pack_and_upload_code(
 
         return s3_runproc_sh, inputs, job_name
 
-    def _get_codeartifact_index(self, codeartifact_repo_arn: str, codeartifact_client: Any = None):
-        """Build an authenticated codeartifact index url based on the arn provided.
+    def _get_codeartifact_command(self, codeartifact_repo_arn: str) -> str:
+        """Build an AWS CLI CodeArtifact command to configure pip.
 
         The codeartifact_repo_arn property must follow the form
         # `arn:${Partition}:codeartifact:${Region}:${Account}:repository/${Domain}/${Repository}`
         https://docs.aws.amazon.com/codeartifact/latest/ug/python-configure-pip.html
         https://docs.aws.amazon.com/service-authorization/latest/reference/list_awscodeartifact.html#awscodeartifact-resources-for-iam-policies
-        
+
         Args:
             codeartifact_repo_arn: arn of the codeartifact repository
-            codeartifact_client: boto3 client for codeartifact (used for testing)
         Returns:
-            authenticated codeartifact index url
+            codeartifact command string
         """
 
         arn_regex = (
@@ -1866,7 +1865,9 @@ def _get_codeartifact_index(self, codeartifact_repo_arn: str, codeartifact_clien
         )
         m = re.match(arn_regex, codeartifact_repo_arn)
         if not m:
-            raise ValueError("invalid CodeArtifact repository arn {}".format(codeartifact_repo_arn))
+            raise ValueError(
+                "invalid CodeArtifact repository arn {}".format(codeartifact_repo_arn)
+            )
         domain = m.group("domain")
         owner = m.group("account")
         repository = m.group("repository")
@@ -1880,28 +1881,8 @@ def _get_codeartifact_index(self, codeartifact_repo_arn: str, codeartifact_clien
             repository,
             region,
         )
-        try:
-            if not codeartifact_client:
-                codeartifact_client = self.sagemaker_session.boto_session.client("codeartifact", region_name=region)
-            
-            auth_token_response = codeartifact_client.get_authorization_token(domain=domain, domainOwner=owner)
-            token = auth_token_response["authorizationToken"]
-            endpoint_response = codeartifact_client.get_repository_endpoint(
-                domain=domain, domainOwner=owner, repository=repository, format="pypi"
-            )
-            unauthenticated_index = endpoint_response["repositoryEndpoint"]
-            return re.sub(
-                "https://",
-                "https://aws:{}@".format(token),
-                re.sub(
-                    "{}/?$".format(repository),
-                    "{}/simple/".format(repository),
-                    unauthenticated_index,
-                ),
-            )
-        except Exception as e:
-            logger.error("failed to configure pip to use codeartifact: %s", e, exc_info=True)
-            raise RuntimeError("failed to configure pip to use codeartifact")
+
+        return f"aws codeartifact login --tool pip --domain {domain} --domain-owner {owner} --repository {repository} --region {region}" # pylint: disable=line-too-long
 
     def _generate_framework_script(
         self, user_script: str, codeartifact_repo_arn: str = None
@@ -1920,10 +1901,12 @@ def _generate_framework_script(
                 logged into before installing dependencies (default: None).
         """
         if codeartifact_repo_arn:
-            index = self._get_codeartifact_index(codeartifact_repo_arn)
-            index_option = "-i {}".format(index)
+            codeartifact_login_command = self._get_codeartifact_command(
+                codeartifact_repo_arn
+            )
         else:
-            index_option = ""
+            codeartifact_login_command = \
+                "echo 'CodeArtifact repository not specified. Skipping login.'"
 
         return dedent(
             """\
@@ -1936,16 +1919,23 @@ def _generate_framework_script(
             set -e
 
             if [[ -f 'requirements.txt' ]]; then
+                # Optionally log into CodeArtifact
+                if ! hash aws 2>/dev/null; then
+                    echo "AWS CLI is not installed. Skipping CodeArtifact login."
+                else
+                    {codeartifact_login_command}
+                fi
+
                 # Some py3 containers has typing, which may breaks pip install
                 pip uninstall --yes typing
 
-                pip install -r requirements.txt {index_option}
+                pip install -r requirements.txt
             fi
 
             {entry_point_command} {entry_point} "$@"
         """
         ).format(
-            index_option=index_option,
+            codeartifact_login_command=codeartifact_login_command,
             entry_point_command=" ".join(self.command),
             entry_point=user_script,
         )
@@ -2039,7 +2029,9 @@ def _create_and_upload_runproc(
         from sagemaker.workflow.utilities import _pipeline_config, hash_object
 
         if _pipeline_config and _pipeline_config.pipeline_name:
-            runproc_file_str = self._generate_framework_script(user_script, codeartifact_repo_arn)
+            runproc_file_str = self._generate_framework_script(
+                user_script, codeartifact_repo_arn
+            )
             runproc_file_hash = hash_object(runproc_file_str)
             s3_uri = s3.s3_path_join(
                 "s3://",

tests/unit/test_processing.py

@@ -13,13 +13,11 @@
 from __future__ import absolute_import
 
 import copy
-import datetime
 
-import boto3
-from botocore.stub import Stubber
 import pytest
 from mock import Mock, patch, MagicMock
 from packaging import version
+from textwrap import dedent
 
 from sagemaker import LocalSession
 from sagemaker.dataset_definition.inputs import (
@@ -1106,28 +1104,8 @@ def test_pyspark_processor_configuration_path_pipeline_config(
 
 
 @patch("sagemaker.workflow.utilities._pipeline_config", MOCKED_PIPELINE_CONFIG)
-def test_get_codeartifact_index(pipeline_session):
+def test_get_codeartifact_command(pipeline_session):
     codeartifact_repo_arn = "arn:aws:codeartifact:us-west-2:012345678901:repository/test-domain/test-repository"
-    codeartifact_url = "test-domain-012345678901.d.codeartifact.us-west-2.amazonaws.com/pypi/test-repository/simple/"
-
-    client = boto3.client('codeartifact', region_name=REGION)
-    stubber = Stubber(client)
-        
-    get_auth_token_response = {
-        "authorizationToken": "mocked_token",
-        "expiration": datetime.datetime(2045, 1, 1, 0, 0, 0)
-    }
-    auth_token_expected_params = {"domain": "test-domain", "domainOwner": "012345678901"}
-    stubber.add_response("get_authorization_token", get_auth_token_response, auth_token_expected_params)
-
-    get_repo_endpoint_response = {"repositoryEndpoint": f"https://{codeartifact_url}"}
-    repo_endpoint_expected_params = {
-        "domain": "test-domain",
-        "domainOwner": "012345678901",
-        "repository": "test-repository",
-        "format": "pypi"
-    }
-    stubber.add_response("get_repository_endpoint", get_repo_endpoint_response, repo_endpoint_expected_params)
 
     processor = PyTorchProcessor(
         role=ROLE,
@@ -1138,35 +1116,14 @@ def test_get_codeartifact_index(pipeline_session):
         sagemaker_session=pipeline_session,
     )
 
-    with stubber:
-        codeartifact_index = processor._get_codeartifact_index(codeartifact_repo_arn=codeartifact_repo_arn, codeartifact_client=client)
+    codeartifact_command = processor._get_codeartifact_command(codeartifact_repo_arn=codeartifact_repo_arn)
 
-    assert codeartifact_index == f"https://aws:mocked_token@{codeartifact_url}"
+    assert codeartifact_command == "aws codeartifact login --tool pip --domain test-domain --domain-owner 012345678901 --repository test-repository --region us-west-2"
 
 
 @patch("sagemaker.workflow.utilities._pipeline_config", MOCKED_PIPELINE_CONFIG)
-def test_get_codeartifact_index_bad_repo_arn(pipeline_session):
+def test_get_codeartifact_command_bad_repo_arn(pipeline_session):
     codeartifact_repo_arn = "arn:aws:codeartifact:us-west-2:012345678901:repository/test-domain"
-    codeartifact_url = "test-domain-012345678901.d.codeartifact.us-west-2.amazonaws.com/pypi/test-repository/simple/"
-
-    client = boto3.client('codeartifact', region_name=REGION)
-    stubber = Stubber(client)
-        
-    get_auth_token_response = {
-        "authorizationToken": "mocked_token",
-        "expiration": datetime.datetime(2045, 1, 1, 0, 0, 0)
-    }
-    auth_token_expected_params = {"domain": "test-domain", "domainOwner": "012345678901"}
-    stubber.add_response("get_authorization_token", get_auth_token_response, auth_token_expected_params)
-
-    get_repo_endpoint_response = {"repositoryEndpoint": f"https://{codeartifact_url}"}
-    repo_endpoint_expected_params = {
-        "domain": "test-domain",
-        "domainOwner": "012345678901",
-        "repository": "test-repository",
-        "format": "pypi"
-    }
-    stubber.add_response("get_repository_endpoint", get_repo_endpoint_response, repo_endpoint_expected_params)
 
     processor = PyTorchProcessor(
         role=ROLE,
@@ -1177,35 +1134,52 @@ def test_get_codeartifact_index_bad_repo_arn(pipeline_session):
         sagemaker_session=pipeline_session,
     )
 
-    with stubber:
-        with pytest.raises(ValueError):
-            processor._get_codeartifact_index(codeartifact_repo_arn=codeartifact_repo_arn, codeartifact_client=client)
-
+    with pytest.raises(ValueError):
+        processor._get_codeartifact_command(codeartifact_repo_arn=codeartifact_repo_arn)
 
 @patch("sagemaker.workflow.utilities._pipeline_config", MOCKED_PIPELINE_CONFIG)
-def test_get_codeartifact_index_client_error(pipeline_session):
-    codeartifact_repo_arn = "arn:aws:codeartifact:us-west-2:012345678901:repository/test-domain/test-repository"
-    codeartifact_url = "test-domain-012345678901.d.codeartifact.us-west-2.amazonaws.com/pypi/test-repository/simple/"
-
-    client = boto3.client('codeartifact', region_name=REGION)
-    stubber = Stubber(client)
-        
-    get_auth_token_response = {
-        "authorizationToken": "mocked_token",
-        "expiration": datetime.datetime(2045, 1, 1, 0, 0, 0)
-    }
-    auth_token_expected_params = {"domain": "test-domain", "domainOwner": "012345678901"}
-    stubber.add_client_error("get_authorization_token", service_error_code="404", expected_params=auth_token_expected_params)
-
-    get_repo_endpoint_response = {"repositoryEndpoint": f"https://{codeartifact_url}"}
-    repo_endpoint_expected_params = {
-        "domain": "test-domain",
-        "domainOwner": "012345678901",
-        "repository": "test-repository",
-        "format": "pypi"
-    }
-    stubber.add_response("get_repository_endpoint", get_repo_endpoint_response, repo_endpoint_expected_params)
+def test_generate_framework_script(pipeline_session):
+    processor = PyTorchProcessor(
+        role=ROLE,
+        instance_type="ml.m4.xlarge",
+        framework_version="2.0.1",
+        py_version="py310",
+        instance_count=1,
+        sagemaker_session=pipeline_session,
+    )
+
+    framework_script = processor._generate_framework_script(user_script="process.py")
 
+    assert framework_script == dedent(
+        """\
+        #!/bin/bash
+
+        cd /opt/ml/processing/input/code/
+        tar -xzf sourcedir.tar.gz
+
+        # Exit on any error. SageMaker uses error code to mark failed job.
+        set -e
+
+        if [[ -f 'requirements.txt' ]]; then
+            # Optionally log into CodeArtifact
+            if ! hash aws 2>/dev/null; then
+                echo "AWS CLI is not installed. Skipping CodeArtifact login."
+            else
+                echo 'CodeArtifact repository not specified. Skipping login.'
+            fi
+
+            # Some py3 containers has typing, which may breaks pip install
+            pip uninstall --yes typing
+
+            pip install -r requirements.txt
+        fi
+
+        python process.py "$@"
+    """
+    )
+    
+@patch("sagemaker.workflow.utilities._pipeline_config", MOCKED_PIPELINE_CONFIG)
+def test_generate_framework_script_with_codeartifact(pipeline_session):
     processor = PyTorchProcessor(
         role=ROLE,
         instance_type="ml.m4.xlarge",
@@ -1215,10 +1189,38 @@ def test_get_codeartifact_index_client_error(pipeline_session):
         sagemaker_session=pipeline_session,
     )
 
-    with stubber:
-        with pytest.raises(RuntimeError):
-            processor._get_codeartifact_index(codeartifact_repo_arn=codeartifact_repo_arn, codeartifact_client=client)
+    framework_script = processor._generate_framework_script(
+        user_script="process.py",
+        codeartifact_repo_arn="arn:aws:codeartifact:us-west-2:012345678901:repository/test-domain/test-repository"
+    )
+
+    assert framework_script == dedent(
+        """\
+        #!/bin/bash
+
+        cd /opt/ml/processing/input/code/
+        tar -xzf sourcedir.tar.gz
+
+        # Exit on any error. SageMaker uses error code to mark failed job.
+        set -e
 
+        if [[ -f 'requirements.txt' ]]; then
+            # Optionally log into CodeArtifact
+            if ! hash aws 2>/dev/null; then
+                echo "AWS CLI is not installed. Skipping CodeArtifact login."
+            else
+                "aws codeartifact login --tool pip --domain test-domain --domain-owner 012345678901 --repository test-repository --region us-west-2"
+            fi
+
+            # Some py3 containers has typing, which may breaks pip install
+            pip uninstall --yes typing
+
+            pip install -r requirements.txt
+        fi
+
+        python process.py "$@"
+    """
+    )
 
 def _get_script_processor(sagemaker_session):
     return ScriptProcessor(