diff --git a/examples/image_classifier/README.md b/examples/image_classifier/README.md index 5ee2f85..dd1b9a8 100644 --- a/examples/image_classifier/README.md +++ b/examples/image_classifier/README.md @@ -27,7 +27,7 @@ 2. Set the right GCP project and bucket in the [flow](./image_classifier_flow.py) where you want to upload the trained model to. 3. Run this command on the terminal (replace `your_wandb_api_key` and `your_wandb_project` with your details in this command): ```sh - WANDB_API_KEY=your_wandb_api_key WANDB_PROJECT=your_wandb_project python image_classifier_flow.py --environment=pypi run --offline False + WANDB_API_KEY=your_wandb_api_key WANDB_PROJECT=your_wandb_project python image_classifier_flow.py --environment=conda run --offline False ``` You can track the progress of the flow run on Outerbounds UI (see the url in your terminal logs). diff --git a/examples/image_classifier/image_classifier_flow.py b/examples/image_classifier/image_classifier_flow.py index 96b4fbb..fab4153 100644 --- a/examples/image_classifier/image_classifier_flow.py +++ b/examples/image_classifier/image_classifier_flow.py @@ -12,6 +12,7 @@ environment, kubernetes, pypi, + conda, nvct, ) from metaflow.cards import Markdown @@ -33,9 +34,12 @@ class ImageClassifierFlow(FlowSpec): default=True, ) - @pypi(python="3.11.9", packages={"torchvision": "0.19.1"}) + #@pypi(python="3.11.9", packages={"torchvision": "0.24.0"}) + #@conda(python="3.10", packages={"torchvision": "0.20.1"}) + @conda(python="3.10", packages={"torchvision": ""}) @card(type="default") - @kubernetes + #@kubernetes + @nvct @step def start(self): import torchvision @@ -60,9 +64,14 @@ def start(self): # Train the network # Keep @nvct decorator before @step decorator else the flow fails - @pypi( - python="3.11.9", - packages={"torch": "2.4.1", "torchvision": "0.19.1", "mozmlops": "0.1.4"}, + #@pypi( + # python="3.11.9", + # packages={"torch": "2.9.0", "torchvision": "0.24.0", "wandb": "0.22.2"}, + #) + @conda( + python="3.10", + #packages={"pytorch": "2.8.0", "torchvision": "0.20.1", "wandb": "0.22.2"}, + packages={"pytorch": "", "torchvision": "", "wandb": ""}, ) @nvct # @kubernetes @@ -147,15 +156,21 @@ def train(self): self.next(self.evaluate) # Test the model on the test data - @pypi( - python="3.11.9", - packages={ - "torch": "2.4.1", - "torchvision": "0.19.1", - }, + #@pypi( + # python="3.11.9", + # packages={ + # "torch": "2.9.0", + # "torchvision": "0.24.0", + # }, + #) + @conda( + python="3.10", + #packages={"pytorch": "2.8.0", "torchvision": "0.20.1"}, + packages={"pytorch": "", "torchvision": ""}, ) # Check https://docs.metaflow.org/api/step-decorators/kubernetes for details on @kubernetes decorator - @kubernetes(cpu=1, memory=4096) + #@kubernetes(cpu=1, memory=4096) + @nvct @step def evaluate(self): import torch @@ -195,23 +210,26 @@ def evaluate(self): ) self.next(self.upload_model_to_gcs) - @pypi(python="3.11.9", packages={"mozmlops": "0.1.4"}) - @kubernetes + #@pypi(python="3.11.9") + @conda(python="3.10") + #@kubernetes + @nvct @step def upload_model_to_gcs(self): - from mozmlops.cloud_storage_api_client import CloudStorageAPIClient + #from mozmlops.cloud_storage_api_client import CloudStorageAPIClient - print("Uploading model to gcs") + print("Skipping Uploading model to gcs") # init client - storage_client = CloudStorageAPIClient( - project_name=GCS_PROJECT_NAME, bucket_name=GCS_BUCKET_NAME - ) - storage_client.store( - data=self.model_state_dict_bytes, storage_path=MODEL_STORAGE_PATH - ) + #storage_client = CloudStorageAPIClient( + # project_name=GCS_PROJECT_NAME, bucket_name=GCS_BUCKET_NAME + #) + #storage_client.store( + # data=self.model_state_dict_bytes, storage_path=MODEL_STORAGE_PATH + #) self.next(self.end) - @kubernetes + #@kubernetes + @nvct @step def end(self): print( diff --git a/examples/image_classifier/requirements.example_image_classifier.txt b/examples/image_classifier/requirements.example_image_classifier.txt index 3a0da7e..49fdd63 100644 --- a/examples/image_classifier/requirements.example_image_classifier.txt +++ b/examples/image_classifier/requirements.example_image_classifier.txt @@ -1 +1 @@ -mozmlops==0.1.4 +outerbounds[gcp]==0.10.19