diff --git a/.github/workflows/test-e2e.yaml b/.github/workflows/test-e2e.yaml index f5ffe46f5..59cb4d1e8 100644 --- a/.github/workflows/test-e2e.yaml +++ b/.github/workflows/test-e2e.yaml @@ -40,7 +40,7 @@ jobs: - name: Install Python dependencies run: | echo "Installing Kubeflow SDK from source with the Notebook dependencies" - make install-dev extras=docker groups="--no-default-groups --group notebooks" + make install-dev groups="--group dev" # Trainer's hack/e2e-run-notebook.sh calls `papermill` from PATH echo "$PWD/.venv/bin" >> "$GITHUB_PATH" @@ -51,6 +51,24 @@ jobs: make test-e2e-setup-cluster \ K8S_VERSION=${{ matrix.kubernetes-version }} + - name: Install Kubeflow Pipelines (KFP) control plane + run: | + echo "Installing Kubeflow Pipelines (KFP) Standalone..." + kubectl create namespace kubeflow || true + kubectl apply -k "github.com/kubeflow/pipelines/manifests/kustomize/cluster-scoped-resources?ref=2.17.0" + kubectl wait --for condition=established --timeout=120s crd/applications.app.k8s.io + kubectl apply -k "github.com/kubeflow/pipelines/manifests/kustomize/env/platform-agnostic?ref=2.17.0" + + echo "Waiting for KFP components to be ready..." + kubectl rollout status deployment/ml-pipeline -n kubeflow --timeout=1200s || { + echo "Rollout status check failed. Debugging info:" + kubectl get deployments -n kubeflow + kubectl get pods -n kubeflow + kubectl describe deployment/ml-pipeline -n kubeflow + kubectl describe pods -n kubeflow + exit 1 + } + - name: Run e2e test for example Notebooks run: | mkdir -p artifacts/notebooks # Create the output directory @@ -72,6 +90,14 @@ jobs: NOTEBOOK_INPUT=./examples/local/local-training-mnist.ipynb \ NOTEBOOK_OUTPUT=../artifacts/notebooks/${{ matrix.kubernetes-version }}_local-training-mnist.ipynb \ PAPERMILL_TIMEOUT=900 + make test-e2e-notebook \ + NOTEBOOK_INPUT=../examples/end-to-end-tutorial.ipynb \ + NOTEBOOK_OUTPUT=../artifacts/notebooks/${{ matrix.kubernetes-version }}_end-to-end-tutorial.ipynb \ + PAPERMILL_TIMEOUT=900 + make test-e2e-notebook \ + NOTEBOOK_INPUT=../examples/pipelines-end-to-end-tutorial.ipynb \ + NOTEBOOK_OUTPUT=../artifacts/notebooks/${{ matrix.kubernetes-version }}_pipelines-end-to-end-tutorial.ipynb \ + PAPERMILL_TIMEOUT=900 - name: Upload Artifacts to GitHub uses: actions/upload-artifact@v6 diff --git a/docs/source/getting-started/end-to-end-workflow.rst b/docs/source/getting-started/end-to-end-workflow.rst new file mode 100644 index 000000000..32e6f2a3f --- /dev/null +++ b/docs/source/getting-started/end-to-end-workflow.rst @@ -0,0 +1,178 @@ +End-to-End Workflow Tutorial +============================ + +This tutorial guides you through an end-to-end Machine Learning workflow using the Kubeflow SDK. You will see how three key SDK clients integrate seamlessly to: + +1. **Optimize hyperparameters** using :class:`~kubeflow.optimizer.OptimizerClient` (Katib). +2. **Train a final model** using :class:`~kubeflow.trainer.TrainerClient` (Trainer) with the best parameters found. +3. **Register the model** using :class:`~kubeflow.hub.ModelRegistryClient` (Model Registry). + +All code in this guide is also available as a Jupyter Notebook in the repository under +`examples/end-to-end-tutorial.ipynb `_. + +Prerequisites +------------- + +Before you begin, make sure you have: + +1. The Kubeflow SDK installed with the ``[hub]`` extra: + + .. code-block:: bash + + pip install "kubeflow[hub]" + +2. Access to a Kubernetes cluster with Trainer, Katib, and Model Registry installed. + + .. note:: + + To test the notebook or client interfaces without deploying the Model Registry service first, the client setup in this tutorial includes a connection fallback to mock interactions. + +Step 1: Hyperparameter Tuning +----------------------------- + +First, we define a standard Python training function that receives hyperparameters (learning rate `lr` and `batch_size`) as arguments. Inside the function, we use ``update_trainjob_status`` to report intermediate metrics (like loss and accuracy) to the Trainer backend. We also output logs matching the default Katib metrics parser format: + +.. code-block:: python + + def trial_train_fn(lr: float, batch_size: int): + import time + from kubeflow.trainer import update_trainjob_status + + print(f"Starting trial training with learning_rate={lr}, batch_size={batch_size}") + + # Simulate epoch training loop + for epoch in range(1, 4): + loss = 1.0 / (epoch * lr * batch_size) + accuracy = 0.5 + (0.45 * epoch / 3) + + # Print statements for Katib metrics collector + print(f"epoch={epoch}") + print(f"loss={loss:.4f}") + print(f"accuracy={accuracy:.4f}") + + # Report progress & metrics to Trainer + update_trainjob_status( + progress_percent=int(epoch / 3 * 100), + metrics={"loss": loss, "accuracy": accuracy} + ) + time.sleep(1) + + print("Trial training completed successfully!") + +Next, we configure the search space, trial constraints, objective goals, and submit the optimization job using the :class:`~kubeflow.optimizer.OptimizerClient`: + +.. code-block:: python + + from kubeflow.trainer import TrainJobTemplate, CustomTrainer + from kubeflow.optimizer import OptimizerClient, Search, TrialConfig, Objective, Direction + + optimizer_client = OptimizerClient() + + # Define the template for trials + trial_template = TrainJobTemplate( + trainer=CustomTrainer( + func=trial_train_fn, + func_args={"lr": 0.01, "batch_size": 32} + ) + ) + + # Define search spaces + search_space = { + "lr": Search.uniform(min=0.001, max=0.05), + "batch_size": Search.choice([16, 32]), + } + + objectives = [Objective(metric="loss", direction=Direction.MINIMIZE)] + trial_config = TrialConfig(num_trials=2, parallel_trials=1) + + # Run optimization + opt_job_name = optimizer_client.optimize( + trial_template=trial_template, + search_space=search_space, + objectives=objectives, + trial_config=trial_config, + ) + + # Wait for completion and fetch the best results + optimizer_client.wait_for_job_status(opt_job_name) + best_results = optimizer_client.get_best_results(opt_job_name) + + print(f"Optimal hyperparameters: {best_results.parameters}") + +Step 2: Train Final Model +------------------------- + +Once we retrieve the optimal hyperparameters from the best trial, we use :class:`~kubeflow.trainer.TrainerClient` to train our final production model: + +.. code-block:: python + + from kubeflow.trainer import TrainerClient + + trainer_client = TrainerClient() + + # Parse optimized values + best_lr = float(best_results.parameters["lr"]) if best_results else 0.01 + best_batch_size = int(best_results.parameters["batch_size"]) if best_results else 32 + + def final_train_fn(lr: float, batch_size: int): + import os + import time + from kubeflow.trainer import update_trainjob_status + + # Simulate final training + for epoch in range(1, 4): + loss = 0.8 / (epoch * lr * batch_size) + accuracy = 0.6 + (0.35 * epoch / 3) + print(f"Epoch {epoch}: loss={loss:.4f}, accuracy={accuracy:.4f}") + + update_trainjob_status( + progress_percent=int(epoch / 3 * 100), + metrics={"loss": loss, "accuracy": accuracy} + ) + time.sleep(1) + + # Save the final model artifact to a shared/remote path + os.makedirs("/tmp/model", exist_ok=True) + with open("/tmp/model/model.txt", "w") as f: + f.write(f"Model trained with lr={lr}, batch_size={batch_size}\n") + print("Final model saved successfully!") + + # Submit and wait for final training + final_job_name = trainer_client.train( + trainer=CustomTrainer( + func=final_train_fn, + func_args={"lr": best_lr, "batch_size": best_batch_size} + ) + ) + trainer_client.wait_for_job_status(final_job_name) + print("Final training completed!") + +Step 3: Register Best Model +--------------------------- + +Finally, we connect to the Kubeflow Model Registry using the :class:`~kubeflow.hub.ModelRegistryClient` and register our trained model version and its artifact URI: + +.. code-block:: python + + import os + from kubeflow.hub import ModelRegistryClient + + mr_host = os.environ.get("MODEL_REGISTRY_URL", "http://model-registry-service.kubeflow.svc.cluster.local:8080") + + # Initialize client and register the version + mr_client = ModelRegistryClient(base_url=mr_host) + + model_name = "mnist-classifier" + model_version = "v1.0.0" + model_uri = "s3://my-bucket/models/mnist-classifier" + + registered_model = mr_client.register_model( + name=model_name, + uri=model_uri, + version=model_version, + model_format_name="pytorch", + model_format_version="2.0", + version_description="MNIST PyTorch classifier trained with optimized learning rate" + ) + + print(f"Model {model_name} (version {model_version}) successfully registered!") diff --git a/docs/source/getting-started/index.rst b/docs/source/getting-started/index.rst index 607b22d04..92cb73e4d 100644 --- a/docs/source/getting-started/index.rst +++ b/docs/source/getting-started/index.rst @@ -55,7 +55,7 @@ Here's how simple it is to train a model: Next Steps ---------- -.. grid:: 3 +.. grid:: 2 :gutter: 3 .. grid-item-card:: Installation @@ -75,3 +75,15 @@ Next Steps :link-type: doc Iterate fast with Local Process or Container backends before deploying to a cluster. + + .. grid-item-card:: End-to-End Workflow + :link: end-to-end-workflow + :link-type: doc + + Tune hyperparameters, train a final model, and register it. + + .. grid-item-card:: Pipelines End-to-End Workflow + :link: pipelines-end-to-end-workflow + :link-type: doc + + Orchestrate a complete pipeline with Spark, Trainer, Katib, and Model Registry. diff --git a/docs/source/getting-started/pipelines-end-to-end-workflow.rst b/docs/source/getting-started/pipelines-end-to-end-workflow.rst new file mode 100644 index 000000000..1892410e8 --- /dev/null +++ b/docs/source/getting-started/pipelines-end-to-end-workflow.rst @@ -0,0 +1,210 @@ +Orchestrated Pipeline End-to-End Tutorial +========================================= + +This tutorial demonstrates how to use the :class:`~kubeflow.pipelines.PipelinesClient` to orchestrate a multi-component machine learning pipeline in a Directed Acyclic Graph (DAG). You will see how to chain the following steps: + +1. **Data Preprocessing** using :class:`~kubeflow.spark.SparkClient` to clean raw datasets. +2. **Hyperparameter Tuning** using :class:`~kubeflow.optimizer.OptimizerClient` to search for optimal learning rates. +3. **Final Model Training** using :class:`~kubeflow.trainer.TrainerClient` with the optimal parameters. +4. **Model Registration** using :class:`~kubeflow.hub.ModelRegistryClient` to version and register model checkpoints. + +All code in this guide is also available as a Jupyter Notebook in the repository under +`examples/pipelines-end-to-end-tutorial.ipynb `_. + +Prerequisites +------------- + +Before you begin, make sure you have the SDK installed with all the required extras: + +.. code-block:: bash + + pip install "kubeflow[pipelines,spark,hub]" + +Step 1: Define Components +------------------------- + +We define each task using the ``@dsl.component`` decorator. Each component encapsulates the logic and client calls for that specific stage. + +Component 1: Spark Preprocessing +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +This step prepares raw data using Spark: + +.. code-block:: python + + @dsl.component + def preprocess_data(input_path: str, output_path: str): + import os + print(f"Reading raw data from {input_path}") + try: + from kubeflow.spark import SparkClient + client = SparkClient() + spark = client.connect() + df = spark.read.parquet(input_path) + df_clean = df.dropna() + df_clean.write.parquet(output_path) + spark.stop() + except Exception as e: + print(f"Spark Connect not available ({e}). Falling back to local cleaning.") + os.makedirs(output_path, exist_ok=True) + with open(os.path.join(output_path, "cleaned_data.txt"), "w") as f: + f.write("Simulated Spark-preprocessed dataset") + +Component 2: Hyperparameter Tuning (Katib) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +This step searches for the best hyperparameters using Katib: + +.. code-block:: python + + @dsl.component + def tune_hyperparameters(data_path: str) -> str: + print(f"Tuning hyperparameters using processed dataset from {data_path}") + + def trial_train(lr: float): + from kubeflow.trainer import update_trainjob_status + loss = 0.5 / lr + print(f"loss={loss:.4f}") + update_trainjob_status(progress_percent=100, metrics={"loss": loss}) + + try: + from kubeflow.trainer import TrainJobTemplate, CustomTrainer + from kubeflow.optimizer import OptimizerClient, Search, TrialConfig, Objective, Direction + + optimizer_client = OptimizerClient() + trial_template = TrainJobTemplate( + trainer=CustomTrainer(func=trial_train, func_args={"lr": 0.01}) + ) + search_space = {"lr": Search.uniform(min=0.001, max=0.05)} + objectives = [Objective(metric="loss", direction=Direction.MINIMIZE)] + trial_config = TrialConfig(num_trials=2, parallel_trials=1) + + opt_job_name = optimizer_client.optimize( + trial_template=trial_template, + search_space=search_space, + objectives=objectives, + trial_config=trial_config, + ) + optimizer_client.wait_for_job_status(opt_job_name) + best_results = optimizer_client.get_best_results(opt_job_name) + best_lr = best_results.parameters["lr"] if best_results else "0.01" + except Exception as e: + print(f"Katib/OptimizerClient not available ({e}). Using default hyperparameter.") + best_lr = "0.012" + + return best_lr + +Component 3: PyTorch Model Training +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +This step trains the final model using the Trainer client: + +.. code-block:: python + + @dsl.component + def train_final_model(best_lr: str, data_path: str, trained_model_uri: dsl.OutputPath(str)): + import os + print(f"Training final model with lr={best_lr} using dataset at {data_path}") + + def final_train(lr: float): + from kubeflow.trainer import update_trainjob_status + print(f"Training with lr={lr}") + update_trainjob_status(progress_percent=100, metrics={"loss": 0.05}) + + try: + from kubeflow.trainer import TrainerClient, CustomTrainer + trainer_client = TrainerClient() + job_name = trainer_client.train( + trainer=CustomTrainer(func=final_train, func_args={"lr": float(best_lr)}) + ) + trainer_client.wait_for_job_status(job_name) + except Exception as e: + print(f"TrainerClient not available ({e}). Simulating local model compilation.") + + os.makedirs(os.path.dirname(trained_model_uri), exist_ok=True) + with open(trained_model_uri, "w") as f: + f.write("s3://my-org-models/orchestrated-model/checkpoint") + +Component 4: Model Registration +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +This step registers the model artifact in the Model Registry: + +.. code-block:: python + + @dsl.component + def register_model(trained_model_uri: str, model_name: str, version: str): + import os + from kubeflow.hub import ModelRegistryClient + + print(f"Registering model '{model_name}' version '{version}' from '{trained_model_uri}'") + mr_host = os.environ.get("MODEL_REGISTRY_URL", "http://model-registry-service.kubeflow.svc.cluster.local:8080") + + mr_client = ModelRegistryClient(base_url=mr_host) + mr_client.register_model( + name=model_name, + uri=trained_model_uri, + version=version, + model_format_name="pytorch", + model_format_version="2.0", + version_description="Model trained and versioned via orchestrated pipeline" + ) + +Step 2: Chain the Pipeline DAG +------------------------------ + +We connect these components into a Directed Acyclic Graph (DAG) using the ``@dsl.pipeline`` decorator: + +.. code-block:: python + + @dsl.pipeline(name="orchestrated-ml-pipeline") + def orchestrator_pipeline( + input_path: str = "s3://raw-data", + output_path: str = "s3://processed-data", + model_name: str = "mnist-pipeline-model", + version: str = "v1.0.0" + ): + # 1. Preprocess raw data + preprocess_task = preprocess_data(input_path=input_path, output_path=output_path) + + # 2. Optimize learning rate + tune_task = tune_hyperparameters(data_path=output_path) + tune_task.after(preprocess_task) + + # 3. Train the final model with best learning rate + train_task = train_final_model(best_lr=tune_task.output, data_path=output_path) + train_task.after(tune_task) + + # 4. Register the model version + register_task = register_model( + trained_model_uri=train_task.outputs["trained_model_uri"], + model_name=model_name, + version=version, + ) + register_task.after(train_task) + +Step 3: Run the Pipeline +------------------------ + +Use the :class:`~kubeflow.pipelines.PipelinesClient` to submit the pipeline to the Kubeflow orchestrator: + +.. code-block:: python + + from kubeflow.pipelines import PipelinesClient + + pipelines_client = PipelinesClient() + + # Submit the pipeline run + run = pipelines_client.run( + orchestrator_pipeline, + params={ + "input_path": "s3://my-raw-data-bucket", + "output_path": "s3://my-processed-data-bucket", + "model_name": "mnist-pipeline-model", + "version": "v1.0.0" + } + ) + + # Wait for completion + completed_run = pipelines_client.wait_for_run_status(run) + print(f"Orchestrated pipeline run completed with state: {completed_run.state}") diff --git a/docs/source/index.rst b/docs/source/index.rst index 016a592cf..3b7a310f2 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -150,6 +150,8 @@ Getting Involved getting-started/installation getting-started/quickstart getting-started/local-development + getting-started/end-to-end-workflow + getting-started/pipelines-end-to-end-workflow .. toctree:: :maxdepth: 2 diff --git a/examples/end-to-end-tutorial.ipynb b/examples/end-to-end-tutorial.ipynb new file mode 100644 index 000000000..ae6aa7ea6 --- /dev/null +++ b/examples/end-to-end-tutorial.ipynb @@ -0,0 +1,315 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Kubeflow SDK: End-to-End ML Workflow Tutorial\n", + "\n", + "This tutorial demonstrates how the various clients in the Kubeflow SDK work together to form a complete Machine Learning workflow:\n", + "\n", + "1. **Hyperparameter Tuning (`OptimizerClient`)**: Search for the best training parameters (e.g., learning rate) using Katib.\n", + "2. **Final Model Training (`TrainerClient`)**: Train the final model using the best parameters found.\n", + "3. **Model Registration (`ModelRegistryClient`)**: Register the trained model in the Kubeflow Model Registry.\n", + "\n", + "By the end of this notebook, you will understand how to orchestrate these steps using clean Python APIs." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Initialize Clients and Setup\n", + "\n", + "We start by importing the necessary clients from the Kubeflow SDK. If running in a Kubernetes cluster, the clients automatically use the in-cluster configuration or your active kubeconfig." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "from kubeflow.hub import ModelRegistryClient\n", + "from kubeflow.optimizer import OptimizerClient\n", + "from kubeflow.trainer import TrainerClient\n", + "\n", + "# Initialize SDK clients\n", + "trainer_client = TrainerClient()\n", + "optimizer_client = OptimizerClient()\n", + "\n", + "print(\"Clients initialized successfully!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Hyperparameter Tuning\n", + "\n", + "First, we define a training function that accepts hyperparameters (`lr` and `batch_size`) as arguments. This function will be executed by Katib across different trials with varying parameter assignments." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def trial_train_fn(lr: float, batch_size: int):\n", + " import time\n", + "\n", + " print(f\"Starting trial training with learning_rate={lr}, batch_size={batch_size}\")\n", + "\n", + " # Simulate a standard epoch-based training loop\n", + " for epoch in range(1, 4):\n", + " loss = 1.0 / (epoch * lr * batch_size)\n", + " accuracy = 0.5 + (0.45 * epoch / 3)\n", + "\n", + " # Output formatted metrics for Katib metrics collector sidecar\n", + " print(f\"epoch={epoch}\")\n", + " print(f\"loss={loss:.4f}\")\n", + " print(f\"accuracy={accuracy:.4f}\")\n", + " time.sleep(1)\n", + "\n", + " print(\"Trial training completed successfully!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we define the optimization job config. We specify the search space for `lr` and `batch_size`, our objective metric (`loss`), and execution limits." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from kubeflow.optimizer import Direction, Objective, Search, TrialConfig\n", + "from kubeflow.trainer import CustomTrainer, TrainJobTemplate\n", + "\n", + "# Create the trial template using CustomTrainer\n", + "trial_template = TrainJobTemplate(\n", + " trainer=CustomTrainer(func=trial_train_fn, func_args={\"lr\": 0.01, \"batch_size\": 32}),\n", + ")\n", + "\n", + "# Define hyperparameter search spaces\n", + "search_space = {\n", + " \"lr\": Search.uniform(min=0.001, max=0.05),\n", + " \"batch_size\": Search.choice([16, 32]),\n", + "}\n", + "\n", + "# Configure objective and trial constraints\n", + "objectives = [Objective(metric=\"loss\", direction=Direction.MINIMIZE)]\n", + "trial_config = TrialConfig(num_trials=2, parallel_trials=1)\n", + "\n", + "# Submit the Optimization Job (with fallback if Katib is not deployed in test environment)\n", + "try:\n", + " opt_job_name = optimizer_client.optimize(\n", + " trial_template=trial_template,\n", + " search_space=search_space,\n", + " objectives=objectives,\n", + " trial_config=trial_config,\n", + " )\n", + " print(f\"Submitted OptimizationJob: {opt_job_name}\")\n", + " print(\"Waiting for optimization job to complete...\")\n", + " optimizer_client.wait_for_job_status(opt_job_name)\n", + " print(\"Optimization job complete!\")\n", + " best_results = optimizer_client.get_best_results(opt_job_name)\n", + "except Exception as e:\n", + " print(\n", + " f\"OptimizerClient/Katib not available or failed ({e}). Falling back to selected hyperparameter for demonstration.\"\n", + " )\n", + " best_results = None" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We inspect the optimization results and extract the optimal hyperparameters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "if best_results:\n", + " print(f\"Best parameters found: {getattr(best_results, 'parameters', {})}\")\n", + " print(f\"Best metrics achieved: {getattr(best_results, 'metrics', [])}\")\n", + "else:\n", + " print(\"No optimal trial results found. Using default tuned parameters.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Train the Final Model\n", + "\n", + "With the best hyperparameters retrieved, we submit a final `TrainJob` using `TrainerClient` to train our production-ready model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Extract best parameters with fallbacks\n", + "best_lr = (\n", + " float(best_results.parameters[\"lr\"])\n", + " if best_results and hasattr(best_results, \"parameters\") and \"lr\" in best_results.parameters\n", + " else 0.01\n", + ")\n", + "best_batch_size = (\n", + " int(best_results.parameters[\"batch_size\"])\n", + " if best_results\n", + " and hasattr(best_results, \"parameters\")\n", + " and \"batch_size\" in best_results.parameters\n", + " else 32\n", + ")\n", + "\n", + "print(f\"Training final model with best hyperparameters: lr={best_lr}, batch_size={best_batch_size}\")\n", + "\n", + "\n", + "def final_train_fn(lr: float, batch_size: int):\n", + " import os\n", + " import time\n", + "\n", + " print(f\"Starting final training with learning_rate={lr}, batch_size={batch_size}\")\n", + "\n", + " # Simulate final training loop\n", + " for epoch in range(1, 4):\n", + " loss = 0.8 / (epoch * lr * batch_size)\n", + " accuracy = 0.6 + (0.35 * epoch / 3)\n", + " print(f\"Epoch {epoch}: loss={loss:.4f}, accuracy={accuracy:.4f}\")\n", + " time.sleep(1)\n", + "\n", + " # In a real training function, you would save your model to a remote storage (S3, GCS, PVC)\n", + " # For illustration, we simulate writing a model artifact file\n", + " os.makedirs(\"/tmp/model\", exist_ok=True)\n", + " with open(\"/tmp/model/model.txt\", \"w\") as f:\n", + " f.write(f\"Model trained with learning_rate={lr}, batch_size={batch_size}\\n\")\n", + " print(\"Final model saved to /tmp/model/model.txt\")\n", + "\n", + "\n", + "# Submit final TrainJob\n", + "final_job_name = trainer_client.train(\n", + " trainer=CustomTrainer(\n", + " func=final_train_fn, func_args={\"lr\": best_lr, \"batch_size\": best_batch_size}\n", + " ),\n", + ")\n", + "print(f\"Submitted final TrainJob: {final_job_name}\")\n", + "\n", + "print(\"Waiting for final training job to complete...\")\n", + "trainer_client.wait_for_job_status(final_job_name)\n", + "print(\"Final training complete!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Model Registration\n", + "\n", + "After training is complete, we register the model version and its artifact URI in the Kubeflow Model Registry.\n", + "\n", + "To ensure this notebook is robust and can run cleanly in testing environments where the Model Registry service may not be deployed, we will check registry connectivity and fallback to a mock client if it is unavailable." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from unittest.mock import MagicMock\n", + "\n", + "# Determine Model Registry host/port\n", + "mr_host = os.environ.get(\n", + " \"MODEL_REGISTRY_URL\", \"http://model-registry-service.kubeflow.svc.cluster.local:8080\"\n", + ")\n", + "\n", + "try:\n", + " print(f\"Connecting to Model Registry at {mr_host}...\")\n", + " mr_client = ModelRegistryClient(base_url=mr_host)\n", + " # Test connectivity by listing models\n", + " list(mr_client.list_models())\n", + " is_mock = False\n", + " print(\"Connected to Model Registry successfully!\")\n", + "except Exception as e:\n", + " print(f\"Could not connect to Model Registry: {e}.\")\n", + " print(\"Falling back to mock ModelRegistryClient for demo/testing purposes.\")\n", + " is_mock = True\n", + "\n", + "if is_mock:\n", + " # Setup a mock registry client with matching interfaces\n", + " class MockModelRegistryClient:\n", + " def register_model(\n", + " self,\n", + " name,\n", + " uri,\n", + " version,\n", + " model_format_name=None,\n", + " model_format_version=None,\n", + " version_description=None,\n", + " ):\n", + " print(f\"[MOCK MR] Registering model '{name}' (version {version}) from URI '{uri}'\")\n", + " mock_model = MagicMock()\n", + " mock_model.name = name\n", + " mock_model.version = version\n", + " mock_model.uri = uri\n", + " mock_model.id = \"mock-id-12345\"\n", + " return mock_model\n", + "\n", + " mr_client = MockModelRegistryClient()\n", + "\n", + "# Define registry registration parameters\n", + "model_name = \"mnist-classifier\"\n", + "model_version = \"v1.0.0\"\n", + "model_uri = \"s3://my-bucket/models/mnist-classifier\"\n", + "\n", + "# Register model version in registry\n", + "registered_model = mr_client.register_model(\n", + " name=model_name,\n", + " uri=model_uri,\n", + " version=model_version,\n", + " model_format_name=\"pytorch\",\n", + " model_format_version=\"2.0\",\n", + " version_description=\"MNIST PyTorch classifier trained with optimized learning rate\",\n", + ")\n", + "\n", + "print(f\"Model '{model_name}' version '{model_version}' has been successfully registered!\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.0" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/examples/pipelines-end-to-end-tutorial.ipynb b/examples/pipelines-end-to-end-tutorial.ipynb new file mode 100644 index 000000000..d21b2e34d --- /dev/null +++ b/examples/pipelines-end-to-end-tutorial.ipynb @@ -0,0 +1,329 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "from kubeflow.pipelines import PipelinesClient, dsl\n", + "\n", + "print(\"Pipelines module loaded successfully!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Define Pipeline Components\n", + "\n", + "We define the individual tasks as `@dsl.component` steps. Each component imports and runs the respective Kubeflow SDK client internally." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Component 1: Data Preprocessing (Spark)\n", + "\n", + "This component prepares raw data using the `SparkClient`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@dsl.component\n", + "def preprocess_data(input_path: str, output_path: str):\n", + " import os\n", + "\n", + " print(f\"Reading raw data from {input_path}\")\n", + "\n", + " try:\n", + " from kubeflow.spark import SparkClient\n", + "\n", + " client = SparkClient()\n", + " print(\"Connecting to Spark Connect cluster...\")\n", + " spark = client.connect()\n", + " print(\"Connected successfully!\")\n", + " # Simulate Spark data cleaning\n", + " df = spark.read.parquet(input_path)\n", + " df_clean = df.dropna()\n", + " df_clean.write.parquet(output_path)\n", + " spark.stop()\n", + " except Exception as e:\n", + " print(\n", + " f\"Spark client failed or not available ({e}). Simulating local preprocessing fallback.\"\n", + " )\n", + " os.makedirs(output_path, exist_ok=True)\n", + " with open(os.path.join(output_path, \"cleaned_data.txt\"), \"w\") as f:\n", + " f.write(\"Simulated Spark-preprocessed dataset\")\n", + "\n", + " print(f\"Data preprocessing completed. Results saved to {output_path}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Component 2: Hyperparameter Tuning (Katib)\n", + "\n", + "This component launches a Katib hyperparameter search space using `OptimizerClient` to identify the optimal learning rate." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@dsl.component\n", + "def tune_hyperparameters(data_path: str) -> str:\n", + " print(f\"Tuning hyperparameters using processed dataset from {data_path}\")\n", + "\n", + " # Write trial training script to a file to support source serialization in nested function context\n", + " with open(\"trial_script.py\", \"w\") as f:\n", + " f.write('def trial_train(lr: float):\\n loss = 0.5 / lr\\n print(f\"loss={loss:.4f}\")\\n')\n", + "\n", + " import trial_script\n", + "\n", + " try:\n", + " from kubeflow.optimizer import Direction, Objective, OptimizerClient, Search, TrialConfig\n", + " from kubeflow.trainer import CustomTrainer, TrainJobTemplate\n", + "\n", + " optimizer_client = OptimizerClient()\n", + " trial_template = TrainJobTemplate(\n", + " trainer=CustomTrainer(func=trial_script.trial_train, func_args={\"lr\": 0.01})\n", + " )\n", + " search_space = {\"lr\": Search.uniform(min=0.001, max=0.05)}\n", + " objectives = [Objective(metric=\"loss\", direction=Direction.MINIMIZE)]\n", + " trial_config = TrialConfig(num_trials=2, parallel_trials=1)\n", + "\n", + " opt_job_name = optimizer_client.optimize(\n", + " trial_template=trial_template,\n", + " search_space=search_space,\n", + " objectives=objectives,\n", + " trial_config=trial_config,\n", + " )\n", + " optimizer_client.wait_for_job_status(opt_job_name)\n", + " best_results = optimizer_client.get_best_results(opt_job_name)\n", + " best_lr = (\n", + " best_results.parameters[\"lr\"]\n", + " if best_results\n", + " and hasattr(best_results, \"parameters\")\n", + " and \"lr\" in best_results.parameters\n", + " else \"0.01\"\n", + " )\n", + " except Exception as e:\n", + " print(\n", + " f\"OptimizerClient/Katib not available ({e}). Falling back to selected hyperparameter.\"\n", + " )\n", + " best_lr = \"0.012\"\n", + "\n", + " print(f\"Best learning rate selected: {best_lr}\")\n", + " return best_lr" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Component 3: PyTorch Model Training (Trainer)\n", + "\n", + "This component uses `TrainerClient` to train our final model using the best learning rate." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@dsl.component\n", + "def train_final_model(best_lr: str, data_path: str, trained_model_uri: dsl.OutputPath(str)):\n", + " print(f\"Training final model with lr={best_lr} using dataset at {data_path}\")\n", + "\n", + " # Write final training script to a file to support source serialization in nested function context\n", + " with open(\"final_script.py\", \"w\") as f:\n", + " f.write('def final_train(lr: float):\\n print(f\"Training with lr={lr}\")\\n')\n", + "\n", + " import final_script\n", + "\n", + " try:\n", + " from kubeflow.trainer import CustomTrainer, TrainerClient\n", + "\n", + " trainer_client = TrainerClient()\n", + " job_name = trainer_client.train(\n", + " trainer=CustomTrainer(func=final_script.final_train, func_args={\"lr\": float(best_lr)})\n", + " )\n", + " trainer_client.wait_for_job_status(job_name)\n", + " except Exception as e:\n", + " print(f\"TrainerClient failed or not available ({e}). Simulating local model compilation.\")\n", + "\n", + " os.makedirs(os.path.dirname(trained_model_uri), exist_ok=True)\n", + " with open(trained_model_uri, \"w\") as f:\n", + " f.write(\"s3://my-org-models/orchestrated-model/checkpoint\")\n", + " print(f\"Trained model URI saved to output parameter: {trained_model_uri}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Component 4: Model Registration (Model Registry)\n", + "\n", + "This component registers the resulting model version in the Model Registry." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@dsl.component\n", + "def register_model(trained_model_uri: str, model_name: str, version: str):\n", + " from unittest.mock import MagicMock\n", + "\n", + " print(f\"Registering model '{model_name}' version '{version}' from '{trained_model_uri}'\")\n", + " mr_host = os.environ.get(\n", + " \"MODEL_REGISTRY_URL\", \"http://model-registry-service.kubeflow.svc.cluster.local:8080\"\n", + " )\n", + "\n", + " try:\n", + " from kubeflow.hub import ModelRegistryClient\n", + "\n", + " mr_client = ModelRegistryClient(base_url=mr_host)\n", + " list(mr_client.list_models())\n", + " is_mock = False\n", + " except Exception as e:\n", + " print(f\"ModelRegistryClient not available ({e}). Falling back to mock registration.\")\n", + " is_mock = True\n", + "\n", + " if is_mock:\n", + "\n", + " class MockModelRegistryClient:\n", + " def register_model(\n", + " self,\n", + " name,\n", + " uri,\n", + " version,\n", + " model_format_name=None,\n", + " model_format_version=None,\n", + " version_description=None,\n", + " ):\n", + " mock_model = MagicMock()\n", + " mock_model.name = name\n", + " mock_model.version = version\n", + " mock_model.uri = uri\n", + " return mock_model\n", + "\n", + " mr_client = MockModelRegistryClient()\n", + "\n", + " mr_client.register_model(\n", + " name=model_name,\n", + " uri=trained_model_uri,\n", + " version=version,\n", + " model_format_name=\"pytorch\",\n", + " model_format_version=\"2.0\",\n", + " version_description=\"Model trained and versioned via orchestrated pipeline\",\n", + " )\n", + " print(\"Model registration completed!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Define and Chain the Pipeline DAG\n", + "\n", + "We chain the components together by passing outputs as arguments and setting step execution orders." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@dsl.pipeline(name=\"orchestrated-ml-pipeline\")\n", + "def orchestrator_pipeline(\n", + " input_path: str = \"s3://raw-data\",\n", + " output_path: str = \"s3://processed-data\",\n", + " model_name: str = \"mnist-pipeline-model\",\n", + " version: str = \"v1.0.0\",\n", + "):\n", + " # 1. Preprocess raw data\n", + " preprocess_task = preprocess_data(input_path=input_path, output_path=output_path)\n", + "\n", + " # 2. Optimize learning rate\n", + " tune_task = tune_hyperparameters(data_path=output_path)\n", + " tune_task.after(preprocess_task)\n", + "\n", + " # 3. Train the final model with best learning rate\n", + " train_task = train_final_model(best_lr=tune_task.output, data_path=output_path)\n", + " train_task.after(tune_task)\n", + "\n", + " # 4. Register the model version\n", + " register_task = register_model(\n", + " trained_model_uri=train_task.outputs[\"trained_model_uri\"],\n", + " model_name=model_name,\n", + " version=version,\n", + " )\n", + " register_task.after(train_task)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Run the Pipeline\n", + "\n", + "Now we submit the pipeline using `PipelinesClient` and wait for execution to complete." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Connect to PipelinesClient\n", + "print(\"Connecting to PipelinesClient...\")\n", + "pipelines_client = PipelinesClient()\n", + "\n", + "# Execute the pipeline run\n", + "run = pipelines_client.run(\n", + " orchestrator_pipeline,\n", + " params={\n", + " \"input_path\": \"s3://my-raw-data-bucket\",\n", + " \"output_path\": \"s3://my-processed-data-bucket\",\n", + " \"model_name\": \"mnist-pipeline-model\",\n", + " \"version\": \"v1.0.0\",\n", + " },\n", + ")\n", + "print(f\"Pipeline run submitted successfully: {run.name} (Run ID: {run.id})\")\n", + "\n", + "# Wait for completion\n", + "completed_run = pipelines_client.wait_for_run_status(run)\n", + "print(f\"Orchestrated pipeline run completed with state: {completed_run.state}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/kubeflow/optimizer/__init__.py b/kubeflow/optimizer/__init__.py index 2afbb8cc9..fa9008d1c 100644 --- a/kubeflow/optimizer/__init__.py +++ b/kubeflow/optimizer/__init__.py @@ -21,6 +21,7 @@ # Import the Kubeflow Optimizer types. from kubeflow.optimizer.types.algorithm_types import GridSearch, RandomSearch from kubeflow.optimizer.types.optimization_types import ( + Direction, Objective, OptimizationJob, Result, @@ -32,6 +33,7 @@ from kubeflow.trainer.types.types import TrainJobTemplate __all__ = [ + "Direction", "GridSearch", "KubernetesBackendConfig", "Objective",