From bf00090ff601185ecdbc9ce52cdc8dac7440b871 Mon Sep 17 00:00:00 2001
From: JJ Lim <jj.lim418@gmail.com>
Date: Wed, 20 Mar 2024 17:30:35 +0000
Subject: [PATCH] lint

---
 src/sagemaker/jumpstart/types.py             |  1 +
 tests/unit/sagemaker/jumpstart/constants.py  | 28 +++++++++++++++++---
 tests/unit/sagemaker/jumpstart/test_types.py | 18 +++++--------
 3 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
index ae85abd1c0..a3277bb76a 100644
--- a/src/sagemaker/jumpstart/types.py
+++ b/src/sagemaker/jumpstart/types.py
@@ -2086,6 +2086,7 @@ def get_schema_version(self) -> str:
     def get_region(self) -> str:
         return self._region
 
+
 HubContentDocument = Union[HubModelDocument, HubNotebookDocument]
 
 
diff --git a/tests/unit/sagemaker/jumpstart/constants.py b/tests/unit/sagemaker/jumpstart/constants.py
index 7e721faaf2..ce8399929d 100644
--- a/tests/unit/sagemaker/jumpstart/constants.py
+++ b/tests/unit/sagemaker/jumpstart/constants.py
@@ -7942,8 +7942,19 @@
                 "Options": ["no", "epoch", "steps"],
                 "Scope": "algorithm",
             },
-            {"Name": "save_steps", "Type": "int", "Default": 500, "Min": 1, "Scope": "algorithm"},  # noqa: E501
-            {"Name": "save_total_limit", "Type": "int", "Default": 1, "Scope": "algorithm"},  # noqa: E501
+            {
+                "Name": "save_steps",
+                "Type": "int",
+                "Default": 500,
+                "Min": 1,
+                "Scope": "algorithm",
+            },  # noqa: E501
+            {
+                "Name": "save_total_limit",
+                "Type": "int",
+                "Default": 1,
+                "Scope": "algorithm",
+            },  # noqa: E501
             {
                 "Name": "dataloader_drop_last",
                 "Type": "text",
@@ -7978,7 +7989,13 @@
                 "Options": ["constant_with_warmup", "linear"],
                 "Scope": "algorithm",
             },
-            {"Name": "warmup_steps", "Type": "int", "Default": 0, "Min": 0, "Scope": "algorithm"},  # noqa: E501
+            {
+                "Name": "warmup_steps",
+                "Type": "int",
+                "Default": 0,
+                "Min": 0,
+                "Scope": "algorithm",
+            },  # noqa: E501
             {
                 "Name": "deepspeed",
                 "Type": "text",
@@ -8103,7 +8120,10 @@
                 "Name": "huggingface-textgeneration:eval-loss",
                 "Regex": "'eval_loss': ([0-9]+\\.[0-9]+)",
             },
-            {"Name": "huggingface-textgeneration:train-loss", "Regex": "'loss': ([0-9]+\\.[0-9]+)"},  # noqa: E501
+            {
+                "Name": "huggingface-textgeneration:train-loss",
+                "Regex": "'loss': ([0-9]+\\.[0-9]+)",
+            },  # noqa: E501
         ],
         "InferenceDependencies": [],
         "TrainingDependencies": [
diff --git a/tests/unit/sagemaker/jumpstart/test_types.py b/tests/unit/sagemaker/jumpstart/test_types.py
index b9f536c4ee..0fdf179e69 100644
--- a/tests/unit/sagemaker/jumpstart/test_types.py
+++ b/tests/unit/sagemaker/jumpstart/test_types.py
@@ -942,12 +942,9 @@ def test_hub_instance_varaints():
         "SM_NUM_GPUS": "8"
     }
 
-    assert (
-        instance_variant.get_image_uri("ml.g5.2xlarge")
-        == (
-            "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1"
-            "-tgi1.4.0-gpu-py310-cu121-ubuntu20.04"
-        )
+    assert instance_variant.get_image_uri("ml.g5.2xlarge") == (
+        "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1"
+        "-tgi1.4.0-gpu-py310-cu121-ubuntu20.04"
     )
 
     instance_variant = JumpStartInstanceTypeVariants(
@@ -963,10 +960,7 @@ def test_hub_instance_varaints():
         == "meta-training/g5/v1.0.0/train-meta-textgeneration-llama-2-70b.tar.gz"
     )
 
-    assert (
-        instance_variant.get_image_uri("ml.p3.xlarge")
-        == (
-            "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training"
-            ":2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04"
-        )
+    assert instance_variant.get_image_uri("ml.p3.xlarge") == (
+        "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training"
+        ":2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04"
     )