kfp-tekton/sdk/python/tests/compiler/testdata/trusted_ai.py

95 lines
4.1 KiB
Python

# Copyright 2021 kubeflow.org
#
# 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.
from kfp import components, dsl
fairness_check_ops = components.load_component_from_url(
'https://raw.githubusercontent.com/Trusted-AI/AIF360/master/mlops/kubeflow/bias_detector_pytorch/component.yaml')
robustness_check_ops = components.load_component_from_url(
'https://raw.githubusercontent.com/Trusted-AI/adversarial-robustness-toolbox/main/utils/'
'mlops/kubeflow/robustness_evaluation_fgsm_pytorch/component.yaml')
@dsl.pipeline(
name="Launch trusted-ai pipeline",
description="An example for trusted-ai integration."
)
def trusted_ai(
namespace: str = "kubeflow",
fgsm_attack_epsilon: str = '0.2',
model_class_file: str = 'PyTorchModel.py',
model_class_name: str = 'ThreeLayerCNN',
feature_testset_path: str = 'processed_data/X_test.npy',
label_testset_path: str = 'processed_data/y_test.npy',
loss_fn: str = 'torch.nn.CrossEntropyLoss()',
optimizer: str = 'torch.optim.Adam(model.parameters(), lr=0.001)',
clip_values: str = '(0, 1)',
nb_classes: str = '2',
input_shape: str = '(1,3,64,64)'):
job_manifest = {
"apiVersion": "batch/v1",
"kind": "Job",
"metadata": {
"name": "trusted-ai-train-job-{{workflow.name}}",
"namespace": namespace
},
"spec": {
"ttlSecondsAfterFinished": 100,
"template": {
"metadata": {
"annotations": {
"sidecar.istio.io/inject": "false"
}
},
"spec": {
"restartPolicy": "Never",
"containers": [
{"name": "classification-training",
"image": "aipipeline/gender-classification:latest",
"command": [
"python", "-u", "gender_classification_training.py", "--data_bucket", "mlpipeline",
"--result_bucket", "mlpipeline"
],
"env": [{'name': 'S3_ENDPOINT', 'value': 'minio-service.kubeflow:9000'}]
}
],
}
}
}
}
train_step = dsl.ResourceOp(
name="trust-ai-train-step",
k8s_resource=job_manifest,
action='create',
success_condition='status.succeeded > 0',
failure_condition='status.failed > 0')
robustness_check = robustness_check_ops(model_id='training-example',
epsilon=fgsm_attack_epsilon,
model_class_file=model_class_file,
model_class_name=model_class_name,
feature_testset_path=feature_testset_path,
label_testset_path=label_testset_path,
loss_fn=loss_fn,
optimizer=optimizer,
clip_values=clip_values,
nb_classes=nb_classes,
input_shape=input_shape,
data_bucket_name='mlpipeline',
result_bucket_name='mlpipeline').after(train_step).set_image_pull_policy("Always")
if __name__ == '__main__':
from kfp_tekton.compiler import TektonCompiler
TektonCompiler().compile(trusted_ai, __file__.replace('.py', '.yaml'))