def generate_sagemaker_pytest_cmd(image, sagemaker_test_type):
"""
Parses the image ECR url and returns appropriate pytest command
:param image: ECR url of image
:param sagemaker_test_type: local or remote test type
:return: <tuple> pytest command to be run, path where it should be executed, image tag
"""
region = os.getenv("AWS_REGION", DEFAULT_REGION)
account_id = os.getenv("ACCOUNT_ID", image.split(".")[0])
print("image name {}".format(image))
sm_remote_docker_base_name, tag = image.split("/")[1].split(":")
sm_local_docker_repo_uri = image.split(":")[0]
# Assign instance type
instance_type = assign_sagemaker_remote_job_instance_type(image)
# Get path to test directory
find_path = sm_remote_docker_base_name.split("-")
# NOTE: We are relying on the fact that repos are defined as <context>-<framework>-<job_type> in our infrastructure
framework, framework_version = get_framework_and_version_from_tag(image)
framework_major_version = framework_version.split(".")[0]
job_type = get_job_type_from_image(image)
path = os.path.join("test", "sagemaker_tests", framework, job_type)
aws_id_arg = "--aws-id"
docker_base_arg = "--docker-base-name"
instance_type_arg = "--instance-type"
accelerator_type_arg = "--accelerator-type"
framework_version_arg = "--framework-version"
eia_arg = "ml.eia1.large"
processor = ("neuron" if "neuron" in image else "gpu"
if "gpu" in image else "eia" if "eia" in image else "cpu")
py_version = re.search(r"py\d+", tag).group()
sm_local_py_version = "37" if py_version == "py37" else "38" if py_version == "py38" else "2" if py_version == "py27" else "3"
if framework == "tensorflow" and job_type == "inference":
# Tf Inference tests have an additional sub directory with test
integration_path = os.path.join("test", "integration",
sagemaker_test_type)
else:
integration_path = os.path.join("integration", sagemaker_test_type)
# Conditions for modifying tensorflow SageMaker pytest commands
if framework == "tensorflow" and sagemaker_test_type == SAGEMAKER_REMOTE_TEST_TYPE:
if job_type == "inference":
aws_id_arg = "--registry"
docker_base_arg = "--repo"
instance_type_arg = "--instance-types"
framework_version_arg = "--versions"
integration_path = os.path.join(
integration_path,
"test_tfs.py") if processor != "eia" else os.path.join(
integration_path, "test_ei.py")
if framework == "tensorflow" and job_type == "training":
aws_id_arg = "--account-id"
test_report = os.path.join(os.getcwd(), "test", f"{job_type}_{tag}.xml")
local_test_report = os.path.join(UBUNTU_HOME_DIR, "test",
f"{job_type}_{tag}_sm_local.xml")
# Explanation of why we need the if-condition below:
# We have separate Pipeline Actions that run EFA tests, which have the env variable "EFA_DEDICATED=True" configured
# so that those Actions only run the EFA tests.
# However, there is no such dedicated CB job dedicated to EFA tests in the PR context. This means that when in the
# PR context, setting "DISABLE_EFA_TESTS" to True should skip EFA tests, but setting it to False should enable
# not just the EFA tests, but also all other tests as well.
if is_pr_context():
efa_tests_disabled = os.getenv("DISABLE_EFA_TESTS",
"False").lower() == "true"
efa_flag = "-m \"not efa\"" if efa_tests_disabled else ""
else:
efa_dedicated = os.getenv("EFA_DEDICATED", "False").lower() == "true"
efa_flag = '--efa' if efa_dedicated else '-m \"not efa\"'
region_list = ",".join(SAGEMAKER_EXECUTION_REGIONS)
sagemaker_regions_list = f"--sagemaker-regions {region_list}"
remote_pytest_cmd = (
f"pytest -rA {integration_path} --region {region} --processor {processor} {docker_base_arg} "
f"{sm_remote_docker_base_name} --tag {tag} {framework_version_arg} {framework_version} "
f"{aws_id_arg} {account_id} {instance_type_arg} {instance_type} {efa_flag} {sagemaker_regions_list} --junitxml {test_report}"
)
if processor == "eia":
remote_pytest_cmd += f"{accelerator_type_arg} {eia_arg}"
local_pytest_cmd = (
f"pytest -s -v {integration_path} {docker_base_arg} "
f"{sm_local_docker_repo_uri} --tag {tag} --framework-version {framework_version} "
f"--processor {processor} {aws_id_arg} {account_id} --junitxml {local_test_report}"
)
if framework == "tensorflow" and job_type != "inference":
local_pytest_cmd = f"{local_pytest_cmd} --py-version {sm_local_py_version} --region {region}"
if framework == "tensorflow" and job_type == "training":
path = os.path.join(os.path.dirname(path),
f"{framework}{framework_major_version}_training")
if "huggingface" in framework and job_type == "inference":
path = os.path.join("test", "sagemaker_tests", "huggingface",
"inference")
return (
remote_pytest_cmd if sagemaker_test_type == SAGEMAKER_REMOTE_TEST_TYPE
else local_pytest_cmd,
path,
tag,
job_type,
)
def generate_sagemaker_pytest_cmd(image, sagemaker_test_type):
"""
Parses the image ECR url and returns appropriate pytest command
:param image: ECR url of image
:param sagemaker_test_type: local or remote test type
:return: <tuple> pytest command to be run, path where it should be executed, image tag
"""
reruns = 4
region = os.getenv("AWS_REGION", DEFAULT_REGION)
account_id = os.getenv("ACCOUNT_ID", image.split(".")[0])
print("image name {}".format(image))
sm_remote_docker_base_name, tag = image.split("/")[1].split(":")
sm_local_docker_repo_uri = image.split(":")[0]
# Assign instance type
instance_type = assign_sagemaker_remote_job_instance_type(image)
# Get path to test directory
find_path = sm_remote_docker_base_name.split("-")
# NOTE: We are relying on the fact that repos are defined as <context>-<framework>-<job_type> in our infrastructure
framework, framework_version = get_framework_and_version_from_tag(image)
job_type = get_job_type_from_image(image)
path = os.path.join("test", "sagemaker_tests", framework, job_type)
aws_id_arg = "--aws-id"
docker_base_arg = "--docker-base-name"
instance_type_arg = "--instance-type"
accelerator_type_arg = "--accelerator-type"
eia_arg = "ml.eia1.large"
framework_version = re.search(r"\d+(\.\d+){2}", tag).group()
framework_major_version = framework_version.split(".")[0]
processor = "gpu" if "gpu" in image else "eia" if "eia" in image else "cpu"
py_version = re.search(r"py\d+", tag).group()
sm_local_py_version = "37" if py_version == "py37" else "2" if py_version == "py27" else "3"
if framework == "tensorflow" and job_type == "inference":
# Tf Inference tests have an additional sub directory with test
integration_path = os.path.join("test", "integration",
sagemaker_test_type)
else:
integration_path = os.path.join("integration", sagemaker_test_type)
# Conditions for modifying tensorflow SageMaker pytest commands
if framework == "tensorflow" and sagemaker_test_type == SAGEMAKER_REMOTE_TEST_TYPE:
if job_type == "inference":
aws_id_arg = "--registry"
docker_base_arg = "--repo"
instance_type_arg = "--instance-types"
integration_path = os.path.join(
integration_path,
"test_tfs.py") if processor != "eia" else os.path.join(
integration_path, "test_ei.py")
if framework == "tensorflow" and job_type == "training":
aws_id_arg = "--account-id"
test_report = os.path.join(os.getcwd(), "test", f"{job_type}_{tag}.xml")
local_test_report = os.path.join(UBUNTU_HOME_DIR, "test",
f"{job_type}_{tag}_sm_local.xml")
is_py3 = " python3 -m "
remote_pytest_cmd = (
f"pytest {integration_path} --region {region} {docker_base_arg} "
f"{sm_remote_docker_base_name} --tag {tag} {aws_id_arg} {account_id} "
f"{instance_type_arg} {instance_type} --junitxml {test_report}")
if processor == "eia":
remote_pytest_cmd += (f" {accelerator_type_arg} {eia_arg}")
local_pytest_cmd = (
f"{is_py3} pytest -v {integration_path} {docker_base_arg} "
f"{sm_local_docker_repo_uri} --tag {tag} --framework-version {framework_version} "
f"--processor {processor} {aws_id_arg} {account_id} --junitxml {local_test_report}"
)
if framework == "tensorflow" and job_type != "inference":
local_pytest_cmd = f"{local_pytest_cmd} --py-version {sm_local_py_version} --region {region}"
if framework == "tensorflow" and job_type == "training":
path = os.path.join(os.path.dirname(path),
f"{framework}{framework_major_version}_training")
return (
remote_pytest_cmd if sagemaker_test_type == SAGEMAKER_REMOTE_TEST_TYPE
else local_pytest_cmd,
path,
tag,
job_type,
)
