From 8256865b06de013b1ed60d207bd5bd184b99acd2 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Tue, 12 Nov 2024 14:53:49 -0500 Subject: [PATCH] Adds comet-ml plugin example Signed-off-by: Thomas J. Fan --- docs/integrations/index.md | 3 + examples/comet_ml_plugin/Dockerfile | 27 +++ examples/comet_ml_plugin/README.md | 36 ++++ .../comet_ml_plugin/__init__.py | 0 .../comet_ml_plugin/comet_ml_example.py | 157 ++++++++++++++++++ examples/comet_ml_plugin/requirements.in | 1 + 6 files changed, 224 insertions(+) create mode 100644 examples/comet_ml_plugin/Dockerfile create mode 100644 examples/comet_ml_plugin/README.md create mode 100644 examples/comet_ml_plugin/comet_ml_plugin/__init__.py create mode 100644 examples/comet_ml_plugin/comet_ml_plugin/comet_ml_example.py create mode 100644 examples/comet_ml_plugin/requirements.in diff --git a/docs/integrations/index.md b/docs/integrations/index.md index 7d48d1531..c89d0e29d 100644 --- a/docs/integrations/index.md +++ b/docs/integrations/index.md @@ -18,6 +18,8 @@ Flytekit functionality. For comparison, these plugins can be thought of like :header-rows: 0 :widths: 20 30 +* - {doc}`Comet ` + - `comet-ml`: Comet’s machine learning platform. * - {doc}`DBT ` - Run and test your `dbt` pipelines in Flyte. * - {doc}`Dolt ` @@ -195,6 +197,7 @@ constructs natively within other orchestration tools. :hidden: :caption: Flytekit plugins +Comet DBT Dolt DuckDB diff --git a/examples/comet_ml_plugin/Dockerfile b/examples/comet_ml_plugin/Dockerfile new file mode 100644 index 000000000..4e04f77bd --- /dev/null +++ b/examples/comet_ml_plugin/Dockerfile @@ -0,0 +1,27 @@ +FROM python:3.11-slim-bookworm +LABEL org.opencontainers.image.source https://github.com/flyteorg/flytesnacks + +WORKDIR /root +ENV VENV /opt/venv +ENV LANG C.UTF-8 +ENV LC_ALL C.UTF-8 +ENV PYTHONPATH /root + +WORKDIR /root + +ENV VENV /opt/venv +# Virtual environment +RUN python3 -m venv ${VENV} +ENV PATH="${VENV}/bin:$PATH" + +# Install Python dependencies +COPY requirements.in /root +RUN pip install -r /root/requirements.in + +# Copy the actual code +COPY . /root + +# This tag is supplied by the build script and will be used to determine the version +# when registering tasks, workflows, and launch plans +ARG tag +ENV FLYTE_INTERNAL_IMAGE $tag diff --git a/examples/comet_ml_plugin/README.md b/examples/comet_ml_plugin/README.md new file mode 100644 index 000000000..40a962a52 --- /dev/null +++ b/examples/comet_ml_plugin/README.md @@ -0,0 +1,36 @@ +(comet_ml)= + +# Comet ML + +```{eval-rst} +.. tags:: Integration, Data, Metrics, Intermediate +``` + +Comet’s machine learning platform integrates with your existing infrastructure and tools so you can manage, visualize, and optimize models from training runs to production monitoring. This plugin integrates Flyte with Comet by configuring links between the two platforms. + +To install the plugin, run: + +```bash +pip install flytekitplugins-comet-ml +``` + +Comet requires an API key to authenticate with their platform. In the above example, a secret is created using +[Flyte's Secrets manager](https://docs.flyte.org/en/latest/user_guide/productionizing/secrets.html). + +To enable linking from the Flyte side panel to Comet.ml, add the following to Flyte's configuration: + +```yaml +plugins: + logs: + dynamic-log-links: + - comet-ml-execution-id: + displayName: Comet + templateUris: "{{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .executionName }}{{ .nodeId }}{{ .taskRetryAttempt }}{{ .taskConfig.link_suffix }}" + - comet-ml-custom-id: + displayName: Comet + templateUris: "{{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .taskConfig.experiment_key }}" +``` + +```{auto-examples-toc} +comet_ml_example +``` diff --git a/examples/comet_ml_plugin/comet_ml_plugin/__init__.py b/examples/comet_ml_plugin/comet_ml_plugin/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/examples/comet_ml_plugin/comet_ml_plugin/comet_ml_example.py b/examples/comet_ml_plugin/comet_ml_plugin/comet_ml_example.py new file mode 100644 index 000000000..af5c92e71 --- /dev/null +++ b/examples/comet_ml_plugin/comet_ml_plugin/comet_ml_example.py @@ -0,0 +1,157 @@ +# %% [markdown] +# (comet_ml_example)= +# +# # Comet Example +# Comet’s machine learning platform integrates with your existing infrastructure and +# tools so you can manage, visualize, and optimize models from training runs to +# production monitoring. This plugin integrates Flyte with Comet by configuring +# links between the two platforms. +import os +import os.path + +from flytekit import ( + ImageSpec, + Secret, + current_context, + task, + workflow, +) +from flytekit.types.directory import FlyteDirectory +from flytekitplugins.comet_ml import comet_ml_login + +# %% [markdown] +# First, we specify the project and workspace that we will use with Comet's platform +# Please update `PROJECT_NAME` and `WORKSPACE` to the values associated with your account. +# %% +PROJECT_NAME = "flytekit-comet-ml-v1" +WORKSPACE = "thomas-unionai" + +# %% [markdown] +# W&B requires an API key to authenticate with Comet. In the above example, +# the secret is created using +# [Flyte's Secrets manager](https://docs.flyte.org/en/latest/user_guide/productionizing/secrets.html). +# %% +secret = Secret(key="comet-ml-key", group="comet-ml-group") + +# %% [markdown] +# Next, we use `ImageSpec` to construct a container that contains the dependencies for this +# task: +# %% + +REGISTRY = os.getenv("REGISTRY", "localhost:30000") +image = ImageSpec( + name="comet-ml", + packages=[ + "torch==2.3.1", + "comet-ml==3.43.2", + "lightning==2.3.0", + "flytekitplugins-comet-ml", + "torchvision", + ], + builder="default", + registry=REGISTRY, +) + + +# %% [markdown] +# Here, we use a Flyte task to download the dataset and cache it: +# %% +@task(cache=True, cache_version="2", container_image=image) +def get_dataset() -> FlyteDirectory: + from torchvision.datasets import MNIST + + ctx = current_context() + dataset_dir = os.path.join(ctx.working_directory, "datasetset") + os.makedirs(dataset_dir, exist_ok=True) + + # Download training and evaluation dataset + MNIST(dataset_dir, train=True, download=True) + MNIST(dataset_dir, train=False, download=True) + + return dataset_dir + + +# %% +# The `comet_ml_login` decorator calls `comet_ml.init` and configures it to use Flyte's +# execution id as the Comet's experiment key. The body of the task is PyTorch Lightning +# training code, where we pass `CometLogger` into the `Trainer`'s `logger`. +@task( + secret_requests=[secret], + container_image=image, +) +@comet_ml_login( + project_name=PROJECT_NAME, + workspace=WORKSPACE, + secret=secret, +) +def train_lightning(dataset: FlyteDirectory, hidden_layer_size: int): + import pytorch_lightning as pl + import torch + import torch.nn.functional as F + from pytorch_lightning import Trainer + from pytorch_lightning.loggers import CometLogger + from torch.utils.data import DataLoader + from torchvision import transforms + from torchvision.datasets import MNIST + + class Model(pl.LightningModule): + def __init__(self, layer_size=784, hidden_layer_size=256): + super().__init__() + self.save_hyperparameters() + self.layers = torch.nn.Sequential( + torch.nn.Linear(layer_size, hidden_layer_size), + torch.nn.Linear(hidden_layer_size, 10), + ) + + def forward(self, x): + return torch.relu(self.layers(x.view(x.size(0), -1))) + + def training_step(self, batch, batch_nb): + x, y = batch + loss = F.cross_entropy(self(x), y) + self.logger.log_metrics({"train_loss": loss}, step=batch_nb) + return loss + + def validation_step(self, batch, batch_nb): + x, y = batch + y_hat = self.forward(x) + loss = F.cross_entropy(y_hat, y) + self.logger.log_metrics({"val_loss": loss}, step=batch_nb) + return loss + + def configure_optimizers(self): + return torch.optim.Adam(self.parameters(), lr=0.02) + + dataset.download() + train_ds = MNIST(dataset, train=True, download=False, transform=transforms.ToTensor()) + eval_ds = MNIST(dataset, train=False, download=False, transform=transforms.ToTensor()) + train_loader = DataLoader(train_ds, batch_size=32) + eval_loader = DataLoader(eval_ds, batch_size=32) + + comet_logger = CometLogger() + comet_logger.log_hyperparams({"batch_size": 32}) + + model = Model(hidden_layer_size=hidden_layer_size) + trainer = Trainer(max_epochs=1, fast_dev_run=True, logger=comet_logger) + trainer.fit(model, train_loader, eval_loader) + + +@workflow +def main(hidden_layer_size: int = 32): + dataset = get_dataset() + train_lightning(dataset=dataset, hidden_layer_size=hidden_layer_size) + + +# %% [markdown] +# To enable dynamic log links, add plugin to Flyte's configuration file: +# ```yaml +# plugins: +# logs: +# dynamic-log-links: +# - comet-ml-execution-id: +# displayName: Comet +# templateUris: "{{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .executionName }}{{ .nodeId }}{{ .taskRetryAttempt }}{{ .taskConfig.link_suffix }}" +# - comet-ml-custom-id: +# displayName: Comet +# templateUris: "{{ .taskConfig.host }}/{{ .taskConfig.workspace }}/{{ .taskConfig.project_name }}/{{ .taskConfig.experiment_key }}" +# ``` diff --git a/examples/comet_ml_plugin/requirements.in b/examples/comet_ml_plugin/requirements.in new file mode 100644 index 000000000..bfadcb009 --- /dev/null +++ b/examples/comet_ml_plugin/requirements.in @@ -0,0 +1 @@ +flytekitplugins-comet-ml