transformers

1
# we define a fixture function below and it will be "used" by
2
# referencing its name from tests
3

4
import os
5

6
import pytest
7
from attr import dataclass
8

9

10
os.environ["AWS_DEFAULT_REGION"] = "us-east-1"  # defaults region
11

12

13
@dataclass
14
class SageMakerTestEnvironment:
15
    framework: str
16
    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
    @property
32
    def metric_definitions(self) -> str:
33
        if self.framework == "pytorch":
34
            return [
35
                {"Name": "train_runtime", "Regex": r"train_runtime.*=\D*(.*?)$"},
36
                {"Name": "eval_accuracy", "Regex": r"eval_accuracy.*=\D*(.*?)$"},
37
                {"Name": "eval_loss", "Regex": r"eval_loss.*=\D*(.*?)$"},
38
            ]
39
        else:
40
            return [
41
                {"Name": "train_runtime", "Regex": r"train_runtime.*=\D*(.*?)$"},
42
                {"Name": "eval_accuracy", "Regex": r"loss.*=\D*(.*?)]?$"},
43
                {"Name": "eval_loss", "Regex": r"sparse_categorical_accuracy.*=\D*(.*?)]?$"},
44
            ]
45

46
    @property
47
    def base_job_name(self) -> str:
48
        return f"{self.framework}-transfromers-test"
49

50
    @property
51
    def test_path(self) -> str:
52
        return f"./tests/sagemaker/scripts/{self.framework}"
53

54
    @property
55
    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")
63
def sm_env(request):
64
    request.cls.env = SageMakerTestEnvironment(framework=request.cls.framework)
65