fixed formatting

Author: Narrohag

src/sagemaker/jumpstart/factory/estimator.py

@@ -315,9 +315,14 @@ def _add_hub_access_config_to_kwargs_inputs(
     dataset_uri = kwargs.specs.default_training_dataset_uri
     if isinstance(kwargs.inputs, str):
         if dataset_uri is not None and dataset_uri == kwargs.inputs:
-            kwargs.inputs = TrainingInput(s3_data=kwargs.inputs, hub_access_config=hub_access_config)
+            kwargs.inputs = TrainingInput(
+                s3_data=kwargs.inputs, hub_access_config=hub_access_config
+            )
     elif isinstance(kwargs.inputs, TrainingInput):
-        if dataset_uri is not None and dataset_uri == kwargs.inputs.config["DataSource"]["S3DataSource"]["S3Uri"]:
+        if (
+            dataset_uri is not None
+            and dataset_uri == kwargs.inputs.config["DataSource"]["S3DataSource"]["S3Uri"]
+        ):
             kwargs.inputs.add_hub_access_config(hub_access_config=hub_access_config)
     elif isinstance(kwargs.inputs, dict):
         for k, v in kwargs.inputs.items():
@@ -327,7 +332,10 @@ def _add_hub_access_config_to_kwargs_inputs(
                     training_input.add_hub_access_config(hub_access_config=hub_access_config)
                 kwargs.inputs[k] = training_input
             elif isinstance(kwargs.inputs, TrainingInput):
-                if dataset_uri is not None and dataset_uri == kwargs.inputs.config["DataSource"]["S3DataSource"]["S3Uri"]:
+                if (
+                    dataset_uri is not None
+                    and dataset_uri == kwargs.inputs.config["DataSource"]["S3DataSource"]["S3Uri"]
+                ):
                     kwargs.inputs[k].add_hub_access_config(hub_access_config=hub_access_config)
 
     return kwargs

src/sagemaker/jumpstart/types.py

@@ -1464,8 +1464,12 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
                 else None
             )
         self.model_subscription_link = json_obj.get("model_subscription_link")
-        self.default_training_dataset_key: Optional[str] = json_obj.get("default_training_dataset_key")
-        self.default_training_dataset_uri: Optional[str] = json_obj.get("default_training_dataset_uri")
+        self.default_training_dataset_key: Optional[str] = json_obj.get(
+            "default_training_dataset_key"
+        )
+        self.default_training_dataset_uri: Optional[str] = json_obj.get(
+            "default_training_dataset_uri"
+        )
 
     def to_json(self) -> Dict[str, Any]:
         """Returns json representation of JumpStartMetadataBaseFields object."""

tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py

@@ -36,7 +36,7 @@
 )
 from tests.integ.sagemaker.jumpstart.utils import (
     get_sm_session,
-    get_training_dataset_for_model_and_version
+    get_training_dataset_for_model_and_version,
 )
 
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
@@ -81,13 +81,13 @@ def test_jumpstart_hub_estimator(setup, add_model_references):
     )
 
     estimator.fit(
-        inputs = {
+        inputs={
             "training": f"s3://{get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME)}/"
             f"{get_training_dataset_for_model_and_version(model_id, model_version)}",
         }
     )
 
-     # test that we can create a JumpStartEstimator from existing job with `attach`
+    # test that we can create a JumpStartEstimator from existing job with `attach`
     estimator = JumpStartEstimator.attach(
         training_job_name=estimator.latest_training_job.name,
         model_id=model_id,
@@ -121,14 +121,13 @@ def test_jumpstart_hub_estimator_with_default_session(setup, add_model_reference
     )
 
     estimator.fit(
-        inputs = {
+        inputs={
             "training": f"s3://{get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME)}/"
             f"{get_training_dataset_for_model_and_version(model_id, model_version)}",
         }
     )
 
-
-     # test that we can create a JumpStartEstimator from existing job with `attach`
+    # test that we can create a JumpStartEstimator from existing job with `attach`
     estimator = JumpStartEstimator.attach(
         training_job_name=estimator.latest_training_job.name,
         model_id=model_id,
@@ -138,7 +137,7 @@ def test_jumpstart_hub_estimator_with_default_session(setup, add_model_reference
     # uses ml.p3.2xlarge instance
     predictor = estimator.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
-        role=get_sm_session().get_caller_identity_arn()
+        role=get_sm_session().get_caller_identity_arn(),
     )
 
     response = predictor.predict(["hello", "world"])
@@ -159,10 +158,10 @@ def test_jumpstart_hub_gated_estimator_with_eula(setup, add_model_references):
 
     estimator.fit(
         accept_eula=True,
-         inputs = {
+        inputs={
             "training": f"s3://{get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME)}/"
             f"{get_training_dataset_for_model_and_version(model_id, model_version)}",
-        }
+        },
     )
 
     estimator = JumpStartEstimator.attach(
@@ -196,14 +195,13 @@ def test_jumpstart_hub_gated_estimator_without_eula(setup, add_model_references)
     )
     with pytest.raises(Exception):
         estimator.fit(
-            inputs = {
+            inputs={
                 "training": f"s3://{get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME)}/"
                 f"{get_training_dataset_for_model_and_version(model_id, model_version)}",
             }
         )
 
 
-
 def test_instantiating_estimator(setup, add_model_references):
 
     model_id = "catboost-regression-model"