chendl/compositional_test

1# we define a fixture function below and it will be "used" by2# referencing its name from tests3 4import os5 6import pytest7from attr import dataclass8 9 10os.environ["AWS_DEFAULT_REGION"] = "us-east-1"  # defaults region11 12 13@dataclass14class SageMakerTestEnvironment:15    framework: str16    role = "arn:aws:iam::558105141721:role/sagemaker_execution_role"17    hyperparameters = {18        "task_name": "mnli",19        "per_device_train_batch_size": 16,20        "per_device_eval_batch_size": 16,21        "do_train": True,22        "do_eval": True,23        "do_predict": True,24        "output_dir": "/opt/ml/model",25        "overwrite_output_dir": True,26        "max_steps": 500,27        "save_steps": 5500,28    }29    distributed_hyperparameters = {**hyperparameters, "max_steps": 1000}30 31    @property32    def metric_definitions(self) -> str:33        if self.framework == "pytorch":34            return [35                {"Name": "train_runtime", "Regex": "train_runtime.*=\D*(.*?)$"},36                {"Name": "eval_accuracy", "Regex": "eval_accuracy.*=\D*(.*?)$"},37                {"Name": "eval_loss", "Regex": "eval_loss.*=\D*(.*?)$"},38            ]39        else:40            return [41                {"Name": "train_runtime", "Regex": "train_runtime.*=\D*(.*?)$"},42                {"Name": "eval_accuracy", "Regex": "loss.*=\D*(.*?)]?$"},43                {"Name": "eval_loss", "Regex": "sparse_categorical_accuracy.*=\D*(.*?)]?$"},44            ]45 46    @property47    def base_job_name(self) -> str:48        return f"{self.framework}-transfromers-test"49 50    @property51    def test_path(self) -> str:52        return f"./tests/sagemaker/scripts/{self.framework}"53 54    @property55    def image_uri(self) -> str:56        if self.framework == "pytorch":57            return "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:1.7.1-transformers4.6.1-gpu-py36-cu110-ubuntu18.04"58        else:59            return "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-tensorflow-training:2.4.1-transformers4.6.1-gpu-py37-cu110-ubuntu18.04"60 61 62@pytest.fixture(scope="class")63def sm_env(request):64    request.cls.env = SageMakerTestEnvironment(framework=request.cls.framework)65