-
Notifications
You must be signed in to change notification settings - Fork 116
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Adds comet-ml plugin example #1702
Open
thomasjpfan
wants to merge
16
commits into
flyteorg:master
Choose a base branch
from
thomasjpfan:thomasjpfan/comet-ml-init
base: master
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
16 commits
Select commit
Hold shift + click to select a range
d6deff0
Adds comet-ml plugin example
thomasjpfan 8d22818
Remove unneeded arguments
thomasjpfan bcb5cff
Use new comet-ml name
thomasjpfan 0c90ad7
Use environment variable
thomasjpfan 8adb29d
Use double quotes
thomasjpfan e54796e
Use flyte name
thomasjpfan 066bf9e
Use quotes
thomasjpfan 8dbc6bf
Add to toc
thomasjpfan acbbe8a
Address comments
thomasjpfan 82639e5
Adds comet-ml dockerfile
thomasjpfan bdb39f7
Fixes docker image
thomasjpfan c16048b
Merge remote-tracking branch 'upstream/master' into thomasjpfan/comet…
thomasjpfan 87390bc
Merge remote-tracking branch 'upstream/master' into thomasjpfan/comet…
thomasjpfan b61f903
Fix merge
thomasjpfan 708cefc
Merge remote-tracking branch 'upstream/master' into thomasjpfan/comet…
thomasjpfan 2486869
Use default image builder
thomasjpfan File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
``` |
Empty file.
157 changes: 157 additions & 0 deletions
157
examples/comet_ml_plugin/comet_ml_plugin/comet_ml_example.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 }} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. same thing |
||
# ``` |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
flytekitplugins-comet-ml |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We should add double quote here.