fix black-check

Author: Brent Millare

tests/data/xgboost_abalone/abalone.py

@@ -5,35 +5,39 @@
 
 import xgboost as xgb
 
-model_filename = 'xgboost-model'
+model_filename = "xgboost-model"
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
 
     # Sagemaker specific arguments. Defaults are set in the environment variables.
-    parser.add_argument('--model_dir', type=str, default=os.environ.get('SM_MODEL_DIR','/opt/ml/model'))
-    parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAIN','/opt/ml/input/data/abalone'))
+    parser.add_argument(
+        "--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
+    )
+    parser.add_argument(
+        "--train",
+        type=str,
+        default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/abalone"),
+    )
 
     args, _ = parser.parse_known_args()
 
-    dtrain = get_dmatrix(args.train, 'libsvm')
+    dtrain = get_dmatrix(args.train, "libsvm")
 
     params = {
-        'max_depth': 5,
-        'eta': 0.2,
-        'gamma': 4,
-        'min_child_weight': 6,
-        'subsample': 0.7,
-        'verbosity': 2,
-        'objective': 'reg:squarederror',
-        'tree_method': 'auto',
-        'predictor': 'auto',
+        "max_depth": 5,
+        "eta": 0.2,
+        "gamma": 4,
+        "min_child_weight": 6,
+        "subsample": 0.7,
+        "verbosity": 2,
+        "objective": "reg:squarederror",
+        "tree_method": "auto",
+        "predictor": "auto",
     }
 
-    booster = xgb.train(params=params,
-                        dtrain=dtrain,
-                        num_boost_round=50)
-    booster.save_model(args.model_dir + '/' + model_filename)
+    booster = xgb.train(params=params, dtrain=dtrain, num_boost_round=50)
+    booster.save_model(args.model_dir + "/" + model_filename)
 
 
 def model_fn(model_dir):

tests/integ/test_xgboost.py

@@ -75,13 +75,10 @@ def test_training_with_network_isolation(
 
         train_input = xgboost.sagemaker_session.upload_data(
             path=os.path.join(DATA_DIR, "xgboost_abalone", "abalone"),
-            key_prefix="integ-test-data/xgboost_abalone/abalone"
+            key_prefix="integ-test-data/xgboost_abalone/abalone",
         )
         job_name = unique_name_from_base(base_job_name)
-        xgboost.fit(
-            inputs={"train": train_input},
-            job_name=job_name
-        )
+        xgboost.fit(inputs={"train": train_input}, job_name=job_name)
         assert sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=job_name)[
             "EnableNetworkIsolation"
         ]