reformat

Author: SifeiLi

src/sagemaker/amazon/amazon_estimator.py

@@ -133,14 +133,18 @@ def data_location(self, data_location: str):
 
         if not data_location.startswith("s3://"):
             raise ValueError(
-                'Expecting an S3 URL beginning with "s3://". Got "{}"'.format(data_location)
+                'Expecting an S3 URL beginning with "s3://". Got "{}"'.format(
+                    data_location
+                )
             )
         if data_location[-1] != "/":
             data_location = data_location + "/"
         self._data_location = data_location
 
     @classmethod
-    def _prepare_init_params_from_job_description(cls, job_details, model_channel_name=None):
+    def _prepare_init_params_from_job_description(
+        cls, job_details, model_channel_name=None
+    ):
         """Convert the job description to init params that can be handled by the class constructor.
 
         Args:
@@ -168,7 +172,9 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
         del init_params["image_uri"]
         return init_params
 
-    def prepare_workflow_for_training(self, records=None, mini_batch_size=None, job_name=None):
+    def prepare_workflow_for_training(
+        self, records=None, mini_batch_size=None, job_name=None
+    ):
         """Calls _prepare_for_training. Used when setting up a workflow.
 
         Args:
@@ -194,7 +200,9 @@ def _prepare_for_training(self, records, mini_batch_size=None, job_name=None):
                 specified, one is generated, using the base name given to the
                 constructor if applicable.
         """
-        super(AmazonAlgorithmEstimatorBase, self)._prepare_for_training(job_name=job_name)
+        super(AmazonAlgorithmEstimatorBase, self)._prepare_for_training(
+            job_name=job_name
+        )
 
         feature_dim = None
 
@@ -260,7 +268,9 @@ def fit(
                 will be unassociated.
                 * `TrialComponentDisplayName` is used for display in Studio.
         """
-        self._prepare_for_training(records, job_name=job_name, mini_batch_size=mini_batch_size)
+        self._prepare_for_training(
+            records, job_name=job_name, mini_batch_size=mini_batch_size
+        )
 
         experiment_config = check_and_get_run_experiment_config(experiment_config)
         self.latest_training_job = _TrainingJob.start_new(
@@ -269,12 +279,14 @@ def fit(
         if wait:
             self.latest_training_job.wait(logs=logs)
 
-    def record_set(self,
-                   train,
-                   labels=None,
-                   channel="train",
-                   encrypt=False,
-                   distribution="ShardedByS3Key"):
+    def record_set(
+        self,
+        train,
+        labels=None,
+        channel="train",
+        encrypt=False,
+        distribution="ShardedByS3Key",
+    ):
         """Build a :class:`~RecordSet` from a numpy :class:`~ndarray` matrix and label vector.
 
         For the 2D ``ndarray`` ``train``, each row is converted to a
@@ -311,7 +323,9 @@ def record_set(self,
         )
         parsed_s3_url = urlparse(self.data_location)
         bucket, key_prefix = parsed_s3_url.netloc, parsed_s3_url.path
-        key_prefix = key_prefix + "{}-{}/".format(type(self).__name__, sagemaker_timestamp())
+        key_prefix = key_prefix + "{}-{}/".format(
+            type(self).__name__, sagemaker_timestamp()
+        )
         key_prefix = key_prefix.lstrip("/")
         logger.debug("Uploading to bucket %s and key_prefix %s", bucket, key_prefix)
         manifest_s3_file = upload_numpy_to_s3_shards(
@@ -338,7 +352,9 @@ def _get_default_mini_batch_size(self, num_records: int):
             )
             return 1
 
-        return min(self.DEFAULT_MINI_BATCH_SIZE, max(1, int(num_records / self.instance_count)))
+        return min(
+            self.DEFAULT_MINI_BATCH_SIZE, max(1, int(num_records / self.instance_count))
+        )
 
 
 class RecordSet(object):
@@ -447,7 +463,10 @@ def _build_shards(num_shards, array):
     shard_size = int(array.shape[0] / num_shards)
     if shard_size == 0:
         raise ValueError("Array length is less than num shards")
-    shards = [array[i * shard_size : i * shard_size + shard_size] for i in range(num_shards - 1)]
+    shards = [
+        array[i * shard_size : i * shard_size + shard_size]
+        for i in range(num_shards - 1)
+    ]
     shards.append(array[(num_shards - 1) * shard_size :])
     return shards
 
@@ -494,7 +513,9 @@ def upload_numpy_to_s3_shards(
         manifest_str = json.dumps(
             [{"prefix": "s3://{}/{}".format(bucket, key_prefix)}] + uploaded_files
         )
-        s3.Object(bucket, manifest_key).put(Body=manifest_str.encode("utf-8"), **extra_put_kwargs)
+        s3.Object(bucket, manifest_key).put(
+            Body=manifest_str.encode("utf-8"), **extra_put_kwargs
+        )
         return "s3://{}/{}".format(bucket, manifest_key)
     except Exception as ex:  # pylint: disable=broad-except
         try: