def execute_local_tests(image, pytest_cache_params):
    """
    Run the sagemaker local tests in ec2 instance for the image
    :param image: ECR url
    :param pytest_cache_params: parameters required for :param pytest_cache_util
    :return: None
    """
    account_id = os.getenv(
        "ACCOUNT_ID",
        boto3.client("sts").get_caller_identity()["Account"])
    pytest_cache_util = PytestCache(boto3.client("s3"), account_id)
    ec2_client = boto3.client("ec2",
                              config=Config(retries={"max_attempts": 10}),
                              region_name=DEFAULT_REGION)
    pytest_command, path, tag, job_type = generate_sagemaker_pytest_cmd(
        image, SAGEMAKER_LOCAL_TEST_TYPE)
    pytest_command += " --last-failed --last-failed-no-failures all "
    print(pytest_command)
    framework, _ = get_framework_and_version_from_tag(image)
    random.seed(f"{datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')}")
    ec2_key_name = f"{job_type}_{tag}_sagemaker_{random.randint(1, 1000)}"
    region = os.getenv("AWS_REGION", DEFAULT_REGION)
    ec2_ami_id = UBUNTU_18_BASE_DLAMI_US_EAST_1 if region == "us-east-1" else UBUNTU_18_BASE_DLAMI_US_WEST_2
    sm_tests_tar_name = "sagemaker_tests.tar.gz"
    ec2_test_report_path = os.path.join(UBUNTU_HOME_DIR, "test",
                                        f"{job_type}_{tag}_sm_local.xml")
    instance_id = ""
    ec2_conn = None
    try:
        key_file = generate_ssh_keypair(ec2_client, ec2_key_name)
        print(f"Launching new Instance for image: {image}")
        instance_id, ip_address = launch_sagemaker_local_ec2_instance(
            image, ec2_ami_id, ec2_key_name, region)
        ec2_conn = ec2_utils.get_ec2_fabric_connection(instance_id, key_file,
                                                       region)
        ec2_conn.put(sm_tests_tar_name, f"{UBUNTU_HOME_DIR}")
        ec2_conn.run(
            f"$(aws ecr get-login --no-include-email --region {region})")
        try:
            ec2_conn.run(f"docker pull {image}", timeout=600)
        except invoke.exceptions.CommandTimedOut as e:
            output = ec2_conn.run(
                f"docker images {image} --format '{{.Repository}}:{{.Tag}}'"
            ).stdout.strip("\n")
            if output != image:
                raise DLCSageMakerLocalTestFailure(
                    f"Image pull for {image} failed.\ndocker images output = {output}"
                ) from e
        ec2_conn.run(f"tar -xzf {sm_tests_tar_name}")
        kill_background_processes_and_run_apt_get_update(ec2_conn)
        with ec2_conn.cd(path):
            install_sm_local_dependencies(framework, job_type, image, ec2_conn,
                                          ec2_ami_id)
            pytest_cache_util.download_pytest_cache_from_s3_to_ec2(
                ec2_conn, path, **pytest_cache_params)
            # Workaround for mxnet cpu training images as test distributed
            # causes an issue with fabric ec2_connection
            if framework == "mxnet" and job_type == "training" and "cpu" in image:
                try:
                    ec2_conn.run(pytest_command, timeout=1000, warn=True)
                except exceptions.CommandTimedOut as exc:
                    print(f"Ec2 connection timed out for {image}, {exc}")
                finally:
                    print(f"Downloading Test reports for image: {image}")
                    ec2_conn.close()
                    ec2_conn_new = ec2_utils.get_ec2_fabric_connection(
                        instance_id, key_file, region)
                    ec2_conn_new.get(
                        ec2_test_report_path,
                        os.path.join("test", f"{job_type}_{tag}_sm_local.xml"))
                    output = subprocess.check_output(
                        f"cat test/{job_type}_{tag}_sm_local.xml",
                        shell=True,
                        executable="/bin/bash")
                    pytest_cache_util.upload_pytest_cache_from_ec2_to_s3(
                        ec2_conn_new, path, **pytest_cache_params)
                    if 'failures="0"' not in str(output):
                        raise ValueError(
                            f"Sagemaker Local tests failed for {image}")
            else:
                ec2_conn.run(pytest_command)
                print(f"Downloading Test reports for image: {image}")
                ec2_conn.get(
                    ec2_test_report_path,
                    os.path.join("test", f"{job_type}_{tag}_sm_local.xml"))
    finally:
        with ec2_conn.cd(path):
            pytest_cache_util.upload_pytest_cache_from_ec2_to_s3(
                ec2_conn, path, **pytest_cache_params)
        print(f"Terminating Instances for image: {image}")
        ec2_utils.terminate_instance(instance_id, region)
        print(f"Destroying ssh Key_pair for image: {image}")
        destroy_ssh_keypair(ec2_client, ec2_key_name)
        # return None here to prevent errors from multiprocessing.map(). Without this it returns some object by default
        # which is causing "cannot pickle '_thread.lock' object" error
        return None
Exemplo n.º 2
0
def execute_local_tests(image, ec2_client):
    """
    Run the sagemaker local tests in ec2 instance for the image
    :param image: ECR url
    :param ec2_client: boto3_obj
    :return: None
    """
    pytest_command, path, tag, job_type = generate_sagemaker_pytest_cmd(
        image, SAGEMAKER_LOCAL_TEST_TYPE)
    print(pytest_command)
    framework, _ = get_framework_and_version_from_tag(image)
    random.seed(f"{datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')}")
    ec2_key_name = f"{job_type}_{tag}_sagemaker_{random.randint(1, 1000)}"
    region = os.getenv("AWS_REGION", DEFAULT_REGION)
    sm_tests_tar_name = "sagemaker_tests.tar.gz"
    ec2_test_report_path = os.path.join(UBUNTU_HOME_DIR, "test",
                                        f"{job_type}_{tag}_sm_local.xml")
    try:
        key_file = generate_ssh_keypair(ec2_client, ec2_key_name)
        print(f"Launching new Instance for image: {image}")
        instance_id, ip_address = launch_sagemaker_local_ec2_instance(
            image, UBUNTU_16_BASE_DLAMI_US_EAST_1 if region == "us-east-1" else
            UBUNTU_16_BASE_DLAMI_US_WEST_2, ec2_key_name, region)
        ec2_conn = ec2_utils.get_ec2_fabric_connection(instance_id, key_file,
                                                       region)
        ec2_conn.put(sm_tests_tar_name, f"{UBUNTU_HOME_DIR}")
        ec2_conn.run(
            f"$(aws ecr get-login --no-include-email --region {region})")
        ec2_conn.run(f"docker pull {image}")
        ec2_conn.run(f"tar -xzf {sm_tests_tar_name}")
        with ec2_conn.cd(path):
            install_sm_local_dependencies(framework, job_type, image, ec2_conn)
            # Workaround for mxnet cpu training images as test distributed
            # causes an issue with fabric ec2_connection
            if framework == "mxnet" and job_type == "training" and "cpu" in image:
                try:
                    ec2_conn.run(pytest_command, timeout=1000, warn=True)
                except exceptions.CommandTimedOut as exc:
                    print(f"Ec2 connection timed out for {image}, {exc}")
                finally:
                    print(f"Downloading Test reports for image: {image}")
                    ec2_conn.close()
                    ec2_conn_new = ec2_utils.get_ec2_fabric_connection(
                        instance_id, key_file, region)
                    ec2_conn_new.get(
                        ec2_test_report_path,
                        os.path.join("test", f"{job_type}_{tag}_sm_local.xml"))
                    output = subprocess.check_output(
                        f"cat test/{job_type}_{tag}_sm_local.xml",
                        shell=True,
                        executable="/bin/bash")
                    if 'failures="0"' not in str(output):
                        raise ValueError(
                            f"Sagemaker Local tests failed for {image}")
            else:
                ec2_conn.run(pytest_command)
                print(f"Downloading Test reports for image: {image}")
                ec2_conn.get(
                    ec2_test_report_path,
                    os.path.join("test", f"{job_type}_{tag}_sm_local.xml"))
    finally:
        print(f"Terminating Instances for image: {image}")
        ec2_utils.terminate_instance(instance_id, region)
        print(f"Destroying ssh Key_pair for image: {image}")
        destroy_ssh_keypair(ec2_client, ec2_key_name)