Format all projects with latest ruff

classifier-e2e/pipelines/__init__.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 #
 
+from .deploy import deploy
 from .feature_engineering import feature_engineering
 from .inference import inference
 from .training import training
-from .deploy import deploy

classifier-e2e/pipelines/deploy.py

@@ -1,7 +1,13 @@
-from zenml import pipeline, get_pipeline_context
-from steps import data_loader, inference_preprocessor
 import random
-from steps import deploy_endpoint, predict_on_endpoint, shutdown_endpoint
+
+from steps import (
+    data_loader,
+    deploy_endpoint,
+    inference_preprocessor,
+    predict_on_endpoint,
+    shutdown_endpoint,
+)
+from zenml import get_pipeline_context, pipeline
 
 
 @pipeline

classifier-e2e/pipelines/feature_engineering.py

@@ -15,15 +15,14 @@
 # limitations under the License.
 #
 
-from typing import List, Optional
 import random
+from typing import List, Optional
 
 from steps import (
     data_loader,
     data_preprocessor,
     data_splitter,
 )
-
 from zenml import pipeline
 from zenml.logger import get_logger
 
@@ -59,7 +58,7 @@ def feature_engineering(
     # Link all the steps together by calling them and passing the output
     # of one step as the input of the next step.
     if random_state is None:
-        random_state = random.randint(0,1000)
+        random_state = random.randint(0, 1000)
     raw_data = data_loader(random_state=random_state, target=target)
     dataset_trn, dataset_tst = data_splitter(
         dataset=raw_data,

classifier-e2e/pipelines/inference.py

@@ -20,7 +20,6 @@
     inference_predict,
     inference_preprocessor,
 )
-
 from zenml import get_pipeline_context, pipeline
 from zenml.logger import get_logger
 
@@ -41,7 +40,9 @@ def inference(random_state: str, target: str):
         target: Name of target column in dataset.
     """
     # Get the production model artifact
-    model = get_pipeline_context().model.get_artifact("breast_cancer_classifier")
+    model = get_pipeline_context().model.get_artifact(
+        "breast_cancer_classifier"
+    )
 
     # Get the preprocess pipeline artifact associated with this version
     preprocess_pipeline = get_pipeline_context().model.get_artifact(

classifier-e2e/pipelines/training.py

@@ -19,13 +19,13 @@
 from uuid import UUID
 
 from steps import model_evaluator, model_promoter, model_trainer
+from zenml import pipeline
+from zenml.client import Client
+from zenml.logger import get_logger
 
 from pipelines import (
     feature_engineering,
 )
-from zenml import pipeline
-from zenml.client import Client
-from zenml.logger import get_logger
 
 logger = get_logger(__name__)
 
@@ -58,7 +58,9 @@ def training(
 
     # Execute Feature Engineering Pipeline
     if train_dataset_id is None or test_dataset_id is None:
-        dataset_trn, dataset_tst = feature_engineering(random_state=random_state)
+        dataset_trn, dataset_tst = feature_engineering(
+            random_state=random_state
+        )
     else:
         client = Client()
         dataset_trn = client.get_artifact_version(

classifier-e2e/run.py

@@ -26,7 +26,6 @@
     inference,
     training,
 )
-
 from zenml.client import Client
 from zenml.logger import get_logger
 
@@ -204,7 +203,7 @@ def main(
                 "test_dataset_id"
             ] = test_dataset_artifact_version.id
 
-        run_args_train["random_state"] = random.randint(0,1000)
+        run_args_train["random_state"] = random.randint(0, 1000)
 
         # Run the SGD pipeline
         pipeline_args = {}
@@ -224,7 +223,9 @@ def main(
             config_folder, f"training_xgboost{custom_training_suffix}.yaml"
         )
         training.with_options(**pipeline_args)(**run_args_train)
-        logger.info("Training pipeline with XGBoost finished successfully!\n\n")
+        logger.info(
+            "Training pipeline with XGBoost finished successfully!\n\n"
+        )
 
     if inference_pipeline:
         run_args_inference = {}

classifier-e2e/steps/__init__.py

@@ -9,20 +9,20 @@
 from .data_splitter import (
     data_splitter,
 )
+from .deploy_endpoint import deploy_endpoint
 from .inference_predict import (
     inference_predict,
 )
 from .inference_preprocessor import (
     inference_preprocessor,
 )
+from .misc_endpoint import predict_on_endpoint, shutdown_endpoint
 from .model_evaluator import (
     model_evaluator,
 )
-from .model_trainer import (
-    model_trainer,
-)
 from .model_promoter import (
     model_promoter,
 )
-from .deploy_endpoint import deploy_endpoint
-from .misc_endpoint import predict_on_endpoint, shutdown_endpoint
+from .model_trainer import (
+    model_trainer,
+)

classifier-e2e/steps/data_loader.py

@@ -18,7 +18,6 @@
 import pandas as pd
 from sklearn.datasets import load_breast_cancer
 from typing_extensions import Annotated
-
 from zenml import step
 from zenml.logger import get_logger
 

classifier-e2e/steps/data_preprocessor.py

@@ -22,7 +22,6 @@
 from sklearn.preprocessing import MinMaxScaler
 from typing_extensions import Annotated
 from utils.preprocess import ColumnsDropper, DataFrameCaster, NADropper
-
 from zenml import log_artifact_metadata, step
 
 

classifier-e2e/steps/data_splitter.py

@@ -20,7 +20,6 @@
 import pandas as pd
 from sklearn.model_selection import train_test_split
 from typing_extensions import Annotated
-
 from zenml import step
 
 

classifier-e2e/steps/deploy_endpoint.py

@@ -1,24 +1,24 @@
-from typing_extensions import Annotated
+from datetime import datetime
 
 import sagemaker
-from sagemaker.image_uris import retrieve
 from sagemaker import Predictor
-
-from zenml import step, get_step_context, ArtifactConfig, log_artifact_metadata
-from datetime import datetime
-
+from sagemaker.image_uris import retrieve
+from typing_extensions import Annotated
 from utils.aws import get_aws_config
 from utils.sagemaker_materializer import SagemakerPredictorMaterializer
+from zenml import ArtifactConfig, get_step_context, log_artifact_metadata, step
 
 
 @step(
     enable_cache=False,
     output_materializers=[SagemakerPredictorMaterializer],
 )
-def deploy_endpoint() -> Annotated[
-    Predictor,
-    ArtifactConfig(name="sagemaker_endpoint", is_deployment_artifact=True),
-]:
+def deploy_endpoint() -> (
+    Annotated[
+        Predictor,
+        ArtifactConfig(name="sagemaker_endpoint", is_deployment_artifact=True),
+    ]
+):
     role, session, region = get_aws_config()
 
     model = get_step_context().model._get_model_version()

classifier-e2e/steps/inference_predict.py

@@ -19,7 +19,6 @@
 
 import pandas as pd
 from typing_extensions import Annotated
-
 from zenml import step
 from zenml.logger import get_logger
 

classifier-e2e/steps/inference_preprocessor.py

@@ -18,7 +18,6 @@
 import pandas as pd
 from sklearn.pipeline import Pipeline
 from typing_extensions import Annotated
-
 from zenml import step
 
 

classifier-e2e/steps/misc_endpoint.py

@@ -1,9 +1,7 @@
+import pandas as pd
+from sagemaker.predictor import Predictor
 from typing_extensions import Annotated
-
-
 from zenml import step
-from sagemaker.predictor import Predictor
-import pandas as pd
 
 
 @step

classifier-e2e/steps/model_evaluator.py

@@ -18,20 +18,18 @@
 from typing import Optional
 
 import pandas as pd
+import wandb
 from sklearn.base import ClassifierMixin
 from sklearn.metrics import confusion_matrix
-
 from zenml import (
+    get_step_context,
     log_artifact_metadata,
-    step,
     log_model_metadata,
-    get_step_context,
+    step,
 )
-from zenml.logger import get_logger
-import wandb
 from zenml.client import Client
 from zenml.exceptions import StepContextError
-
+from zenml.logger import get_logger
 
 logger = get_logger(__name__)
 

classifier-e2e/steps/model_promoter.py

@@ -15,10 +15,10 @@
 # limitations under the License.
 #
 
+from sklearn.metrics import accuracy_score
 from zenml import get_step_context, step
 from zenml.client import Client
 from zenml.logger import get_logger
-from sklearn.metrics import accuracy_score
 
 logger = get_logger(__name__)
 
@@ -58,7 +58,9 @@ def model_promoter(accuracy: float, stage: str = "production") -> bool:
         try:
             stage_model = client.get_model_version(current_model.name, stage)
             # We compare their metrics
-            prod_classifier = stage_model.get_artifact("breast_cancer_classifier")
+            prod_classifier = stage_model.get_artifact(
+                "breast_cancer_classifier"
+            )
             if prod_classifier:
                 # and recompute metrics for current prod model using current test set
                 prod_classifier = prod_classifier.load()

classifier-e2e/steps/model_trainer.py

@@ -21,15 +21,18 @@
 from sklearn.base import ClassifierMixin
 from sklearn.linear_model import SGDClassifier
 from typing_extensions import Annotated
-
+from utils.sagemaker_materializer import SagemakerMaterializer
 from zenml import ArtifactConfig, step
 from zenml.logger import get_logger
 
-from utils.sagemaker_materializer import SagemakerMaterializer
-
 logger = get_logger(__name__)
 
-@step(output_materializers=[SagemakerMaterializer,])
+
+@step(
+    output_materializers=[
+        SagemakerMaterializer,
+    ]
+)
 def model_trainer(
     dataset_trn: pd.DataFrame,
     model_type: str = "sgd",

classifier-e2e/utils/__init__.py

@@ -1,16 +1,16 @@
 # Apache Software License 2.0
-# 
+#
 # Copyright (c) ZenML GmbH 2024. All rights reserved.
-# 
+#
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
-# 
+#
 # http://www.apache.org/licenses/LICENSE-2.0
-# 
+#
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-# 
+#

classifier-e2e/utils/aws.py

@@ -1,5 +1,5 @@
 import os
-from typing import Tuple, Any
+from typing import Any, Tuple
 
 import boto3
 import sagemaker

classifier-e2e/utils/sagemaker_materializer.py

@@ -1,17 +1,18 @@
 import os
+import tarfile
+import tempfile
 from typing import Type, Union
 
+import joblib
+from sagemaker import Predictor
+from sklearn.base import ClassifierMixin
+from sklearn.linear_model import SGDClassifier
+from xgboost import XGBClassifier
 from zenml.enums import ArtifactType
 from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.built_in_materializer import BuiltInMaterializer
-from sklearn.linear_model import SGDClassifier
-from xgboost import XGBClassifier
-import tarfile
-import tempfile
-import joblib
-from sklearn.base import ClassifierMixin
-from sagemaker import Predictor
+
 
 class SagemakerMaterializer(BaseMaterializer):
     ASSOCIATED_TYPES = (ClassifierMixin,)
@@ -77,9 +78,7 @@ class SagemakerPredictorMaterializer(BaseMaterializer):
     ASSOCIATED_TYPES = (Predictor,)
     ASSOCIATED_ARTIFACT_TYPE = ArtifactType.SERVICE
 
-    def load(
-        self, data_type: Type[Predictor]
-    ) -> Predictor:
+    def load(self, data_type: Type[Predictor]) -> Predictor:
         """Read from artifact store."""
         return Predictor(endpoint_name=BuiltInMaterializer(self.uri).load(str))