Linting

Author: bobbywlindsey

sagemaker-mlflow/sagemaker_mlflow_setup.ipynb

@@ -76,12 +76,12 @@
     "role = sagemaker.get_execution_role()\n",
     "region = sess.boto_region_name\n",
     "\n",
-    "iam_client = boto3.client('iam')\n",
+    "iam_client = boto3.client(\"iam\")\n",
     "sts_client = boto3.client(\"sts\")\n",
-    "sm_client = boto3.client('sagemaker')\n",
+    "sm_client = boto3.client(\"sagemaker\")\n",
     "account_id = sts_client.get_caller_identity()[\"Account\"]\n",
-    "tracking_server_name = 'my-setup-test3'\n",
-    "mlflow_role_name = 'mlflow-test3'"
+    "tracking_server_name = \"my-setup-test3\"\n",
+    "mlflow_role_name = \"mlflow-test3\""
    ]
   },
   {
@@ -142,26 +142,21 @@
    "outputs": [],
    "source": [
     "mlflow_trust_policy = {\n",
-    "     \"Version\": \"2012-10-17\",\n",
-    "     \"Statement\": [\n",
-    "         {\n",
-    "             \"Effect\": \"Allow\",\n",
-    "             \"Principal\": {\n",
-    "                 \"Service\": [\n",
-    "                      \"sagemaker.amazonaws.com\"\n",
-    "                 ]\n",
-    "             },\n",
-    "             \"Action\": \"sts:AssumeRole\"\n",
-    "         }\n",
-    "     ]\n",
+    "    \"Version\": \"2012-10-17\",\n",
+    "    \"Statement\": [\n",
+    "        {\n",
+    "            \"Effect\": \"Allow\",\n",
+    "            \"Principal\": {\"Service\": [\"sagemaker.amazonaws.com\"]},\n",
+    "            \"Action\": \"sts:AssumeRole\",\n",
+    "        }\n",
+    "    ],\n",
     "}\n",
     "\n",
     "# Create role for MLflow\n",
     "mlflow_role = iam_client.create_role(\n",
-    "    RoleName=mlflow_role_name,\n",
-    "    AssumeRolePolicyDocument=json.dumps(mlflow_trust_policy)\n",
+    "    RoleName=mlflow_role_name, AssumeRolePolicyDocument=json.dumps(mlflow_trust_policy)\n",
     ")\n",
-    "mlflow_role_arn = mlflow_role['Role']['Arn']\n",
+    "mlflow_role_arn = mlflow_role[\"Role\"][\"Arn\"]\n",
     "\n",
     "# Create policy for S3 and SageMaker Model Registry\n",
     "sm_s3_model_registry_policy = {\n",
@@ -177,21 +172,21 @@
     "                \"sagemaker:CreateModelPackageGroup\",\n",
     "                \"sagemaker:CreateModelPackage\",\n",
     "                \"sagemaker:UpdateModelPackage\",\n",
-    "                \"sagemaker:DescribeModelPackageGroup\"\n",
+    "                \"sagemaker:DescribeModelPackageGroup\",\n",
     "            ],\n",
-    "            \"Resource\": \"*\"\n",
+    "            \"Resource\": \"*\",\n",
     "        }\n",
-    "    ]\n",
+    "    ],\n",
     "}\n",
     "\n",
-    "mlflow_s3_sm_model_registry_iam_policy = iam_client.create_policy(PolicyName='mlflow-s3-sm-model-registry',\n",
-    "                                                                  PolicyDocument=json.dumps(sm_s3_model_registry_policy))\n",
-    "mlflow_s3_sm_model_registry_iam_policy_arn = mlflow_s3_sm_model_registry_iam_policy['Policy']['Arn']\n",
+    "mlflow_s3_sm_model_registry_iam_policy = iam_client.create_policy(\n",
+    "    PolicyName=\"mlflow-s3-sm-model-registry\", PolicyDocument=json.dumps(sm_s3_model_registry_policy)\n",
+    ")\n",
+    "mlflow_s3_sm_model_registry_iam_policy_arn = mlflow_s3_sm_model_registry_iam_policy[\"Policy\"][\"Arn\"]\n",
     "\n",
     "# Attach the policy to the MLflow role\n",
     "iam_client.attach_role_policy(\n",
-    "    RoleName=mlflow_role_name,\n",
-    "    PolicyArn=mlflow_s3_sm_model_registry_iam_policy_arn\n",
+    "    RoleName=mlflow_role_name, PolicyArn=mlflow_s3_sm_model_registry_iam_policy_arn\n",
     ")"
    ]
   },
@@ -241,11 +236,11 @@
    "source": [
     "sm_client.create_mlflow_tracking_server(\n",
     "    TrackingServerName=tracking_server_name,\n",
-    "    ArtifactStoreUri=f's3://{bucket_name}/{tracking_server_name}',\n",
-    "    TrackingServerSize='Small',\n",
-    "    MlflowVersion='2.13.2',\n",
+    "    ArtifactStoreUri=f\"s3://{bucket_name}/{tracking_server_name}\",\n",
+    "    TrackingServerSize=\"Small\",\n",
+    "    MlflowVersion=\"2.13.2\",\n",
     "    RoleArn=mlflow_role_arn,\n",
-    "    AutomaticModelRegistration=False\n",
+    "    AutomaticModelRegistration=False,\n",
     ")"
    ]
   },
@@ -256,7 +251,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tracking_server_arn = f\"arn:aws:sagemaker:{region}:{account_id}:mlflow-tracking-server/{tracking_server_name}\""
+    "tracking_server_arn = (\n",
+    "    f\"arn:aws:sagemaker:{region}:{account_id}:mlflow-tracking-server/{tracking_server_name}\"\n",
+    ")"
    ]
   },
   {
@@ -311,6 +308,7 @@
    "outputs": [],
    "source": [
     "import mlflow\n",
+    "\n",
     "mlflow.set_tracking_uri(tracking_server_arn)"
    ]
   },
@@ -348,9 +346,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sm_client.create_presigned_mlflow_tracking_server_url(\n",
-    "    TrackingServerName=tracking_server_name\n",
-    ")"
+    "sm_client.create_presigned_mlflow_tracking_server_url(TrackingServerName=tracking_server_name)"
    ]
   },
   {