from kfp import dsl
from kfp import compiler

@dsl.component
def preprocess_data(input_path: str, output_path: str):
    """Preprocess raw data for training."""
    import pandas as pd
    df = pd.read_csv(input_path)
    # Preprocessing logic
    df.to_csv(output_path, index=False)

@dsl.component
def train_model(data_path: str, model_path: str, epochs: int = 10):
    """Train ML model on preprocessed data."""
    import tensorflow as tf
    # Training logic
    model.save(model_path)

@dsl.pipeline(name="ml-training-pipeline")
def training_pipeline(input_data: str, output_model: str):
    preprocess_task = preprocess_data(input_path=input_data, output_path="/tmp/processed.csv")
    train_task = train_model(data_path=preprocess_task.output, model_path=output_model)

compiler.Compiler().compile(training_pipeline, "pipeline.yaml")

apiVersion: serving.kserve.io/v1beta1
kind: InferenceService
metadata:
  name: sklearn-iris
spec:
  predictor:
    sklearn:
      storageUri: "gs://kfserving-examples/models/sklearn/iris"
      resources:
        requests:
          cpu: 100m
          memory: 256Mi
        limits:
          cpu: 500m
          memory: 512Mi

apiVersion: kubeflow.org/v1beta1
kind: Experiment
metadata:
  name: hyperparameter-tuning
spec:
  objective:
    type: maximize
    goal: 0.99
    objectiveMetricName: accuracy
  algorithm:
    algorithmName: bayesianoptimization
  parallelTrialCount: 3
  maxTrialCount: 12
  parameters:
    - name: learning_rate
      parameterType: double
      feasibleSpace:
        min: "0.001"
        max: "0.1"
    - name: batch_size
      parameterType: int
      feasibleSpace:
        min: "16"
        max: "128"

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: pytorch-distributed
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      template:
        spec:
          containers:
            - name: pytorch
              image: pytorch/pytorch:latest
              command: ["python", "train.py"]
    Worker:
      replicas: 4
      template:
        spec:
          containers:
            - name: pytorch
              image: pytorch/pytorch:latest
              resources:
                limits:
                  nvidia.com/gpu: 1

from kfp import dsl, compiler
from kfp.client import Client

# Define components
@dsl.component(base_image="python:3.10")
def load_data(dataset_name: str) -> str:
    import urllib.request
    url = f"https://data.example.com/{dataset_name}.csv"
    path = f"/tmp/{dataset_name}.csv"
    urllib.request.urlretrieve(url, path)
    return path

@dsl.component(base_image="tensorflow/tensorflow:2.13.0")
def train_tensorflow_model(
    data_path: str,
    epochs: int,
    learning_rate: float
) -> str:
    import tensorflow as tf
    # Load and prepare data
    # Build and train model
    model_path = "/tmp/model"
    # model.save(model_path)
    return model_path

@dsl.component(base_image="python:3.10")
def evaluate_model(model_path: str, test_data: str) -> float:
    # Load model and evaluate
    accuracy = 0.95  # placeholder
    return accuracy

# Define pipeline
@dsl.pipeline(
    name="TensorFlow Training Pipeline",
    description="End-to-end ML pipeline with TensorFlow"
)
def tf_training_pipeline(
    dataset: str = "mnist",
    epochs: int = 10,
    learning_rate: float = 0.001
):
    load_task = load_data(dataset_name=dataset)
    train_task = train_tensorflow_model(
        data_path=load_task.output,
        epochs=epochs,
        learning_rate=learning_rate
    )
    eval_task = evaluate_model(
        model_path=train_task.output,
        test_data=load_task.output
    )

# Compile and submit
compiler.Compiler().compile(tf_training_pipeline, "pipeline.yaml")

# Submit to Kubeflow
client = Client(host="http://localhost:8080")
run = client.create_run_from_pipeline_package(
    "pipeline.yaml",
    arguments={"dataset": "cifar10", "epochs": 20}
)

@dsl.component(
    base_image="python:3.10",
    packages_to_install=["mlflow", "scikit-learn"]
)
def train_with_mlflow(
    data_path: str,
    mlflow_tracking_uri: str,
    experiment_name: str
) -> str:
    import mlflow
    from sklearn.ensemble import RandomForestClassifier

    mlflow.set_tracking_uri(mlflow_tracking_uri)
    mlflow.set_experiment(experiment_name)

    with mlflow.start_run():
        # Train model
        model = RandomForestClassifier(n_estimators=100)
        # model.fit(X_train, y_train)

        # Log parameters and metrics
        mlflow.log_param("n_estimators", 100)
        mlflow.log_metric("accuracy", 0.95)

        # Log model
        mlflow.sklearn.log_model(model, "model")

        return mlflow.active_run().info.run_id

@dsl.component(base_image="pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime")
def distributed_pytorch_training(
    data_path: str,
    model_output: str,
    world_size: int = 4
):
    import torch
    import torch.distributed as dist
    from torch.nn.parallel import DistributedDataParallel as DDP

    # Initialize distributed training
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    # Create model and wrap with DDP
    model = MyModel().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])

    # Training loop
    # ...

    # Save model on rank 0
    if dist.get_rank() == 0:
        torch.save(model.module.state_dict(), model_output)

apiVersion: serving.kserve.io/v1beta1
kind: InferenceService
metadata:
  name: production-model
  annotations:
    sidecar.istio.io/inject: "true"
spec:
  predictor:
    minReplicas: 2
    maxReplicas: 10
    scaleTarget: 80
    scaleMetric: concurrency
    tensorflow:
      storageUri: "s3://models/production/v1"
      runtimeVersion: "2.13.0"
      resources:
        requests:
          cpu: "1"
          memory: "2Gi"
          nvidia.com/gpu: "1"
        limits:
          cpu: "2"
          memory: "4Gi"
          nvidia.com/gpu: "1"
  transformer:
    containers:
      - name: feature-transformer
        image: myregistry/feature-transformer:v1
        resources:
          requests:
            cpu: "500m"
            memory: "512Mi"

apiVersion: serving.kserve.io/v1beta1
kind: InferenceService
metadata:
  name: model-canary
spec:
  predictor:
    canaryTrafficPercent: 20
    tensorflow:
      storageUri: "s3://models/production/v2"  # New version
  # Previous version gets remaining 80% traffic