Fix diabetes example memory limit (#137)

Signed-off-by: Ketan Umare <[email protected]>

cookbook/case_studies/pima_diabetes/diabetes.py

@@ -10,7 +10,7 @@
 import joblib
 import pandas as pd
 from dataclasses_json import dataclass_json
-from flytekit import task, workflow
+from flytekit import Resources, task, workflow
 from flytekit.types.file import FlyteFile
 from flytekit.types.schema import FlyteSchema
 from sklearn.metrics import accuracy_score
@@ -63,7 +63,7 @@
 # columns and converts it to a typed schema.
 # An example CSV file is available at
 # `https://raw.githubusercontent.com/jbrownlee/Datasets/master/pima-indians-diabetes.data.csv<https://raw.githubusercontent.com/jbrownlee/Datasets/master/pima-indians-diabetes.data.csv>`
-@task(cache_version="1.0", cache=True, memory_limit="200Mi")
+@task(cache_version="1.0", cache=True, limits=Resources(mem="200Mi"))
 def split_traintest_dataset(
     dataset: FlyteFile[typing.TypeVar("csv")], seed: int, test_split_ratio: float
 ) -> (
@@ -121,7 +121,7 @@ class XGBoostModelHyperparams(object):
 )
 
 
-@task(cache_version="1.0", cache=True, memory_limit="200Mi")
+@task(cache_version="1.0", cache=True, limits=Resources(mem="200Mi"))
 def fit(
     x: FlyteSchema[FEATURE_COLUMNS],
     y: FlyteSchema[CLASSES_COLUMNS],
@@ -151,7 +151,7 @@ def fit(
     return (fname,)
 
 
-@task(cache_version="1.0", cache=True, memory_limit="200Mi")
+@task(cache_version="1.0", cache=True, limits=Resources(mem="200Mi"))
 def predict(
     x: FlyteSchema[FEATURE_COLUMNS], model_ser: FlyteFile[MODELSER_JOBLIB],
 ) -> FlyteSchema[CLASSES_COLUMNS]:
@@ -170,7 +170,7 @@ def predict(
     return y_pred_df
 
 
-@task(cache_version="1.0", cache=True, memory_limit="200Mi")
+@task(cache_version="1.0", cache=True, limits=Resources(mem="200Mi"))
 def score(
     predictions: FlyteSchema[CLASSES_COLUMNS], y: FlyteSchema[CLASSES_COLUMNS]
 ) -> float:

cookbook/core/intermediate/map_task.py

@@ -10,7 +10,9 @@
 """
 
 import typing
-from flytekit import task, workflow, map_task, TaskMetadata
+
+from flytekit import TaskMetadata, map_task, task, workflow
+
 
 # %%
 # Note that this is the single task that we'll use in our map task. It can only accept one input and produce one output

cookbook/docs/conf.py

@@ -6,14 +6,14 @@
 
 # -- Path setup --------------------------------------------------------------
 
-import re
 import logging
 
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
 import os
+import re
 import sys
 
 from sphinx_gallery.sorting import ExplicitOrder, FileNameSortKey
@@ -157,7 +157,11 @@ def __call__(self, filename):
     "nav_links": [
         {"href": "https://flyte.readthedocs.io/", "internal": False, "title": "Flyte"},
         {"href": "index", "internal": True, "title": "Tutorials"},
-        {"href": "https://flytekit.readthedocs.io/en/latest/", "internal": False, "title": "Flytekit Python Reference"},
+        {
+            "href": "https://flytekit.readthedocs.io/en/latest/",
+            "internal": False,
+            "title": "Flytekit Python Reference",
+        },
     ],
 }
 

cookbook/plugins/k8s_spark/dataframe_passing.py

@@ -47,7 +47,7 @@ def create_spark_df() -> my_schema:
     """
     sess = flytekit.current_context().spark_session
     return sess.createDataFrame(
-        [("Alice", 5), ("Bob", 10), ("Charlie", 15), ], my_schema.column_names(),
+        [("Alice", 5), ("Bob", 10), ("Charlie", 15),], my_schema.column_names(),
     )
 
 

cookbook/plugins/papermilltasks/simple.py

@@ -30,7 +30,9 @@
 #    The example task is shown below
 nb = NotebookTask(
     name="simple-nb",
-    notebook_path=os.path.join(pathlib.Path(__file__).parent.absolute(), "nb-simple.ipynb"),
+    notebook_path=os.path.join(
+        pathlib.Path(__file__).parent.absolute(), "nb-simple.ipynb"
+    ),
     inputs=kwtypes(v=float),
     outputs=kwtypes(square=float),
 )
@@ -47,13 +49,15 @@
 def square_root_task(f: float) -> float:
     return math.sqrt(f)
 
+
 #%%
 # You can now treat the notebook task as a regular task
 @workflow
 def nb_to_python_wf(f: float) -> float:
     out = nb(v=f)
     return square_root_task(f=out.square)
 
+
 #%%
 # You can execute the task locally as well
 if __name__ == "__main__":
@@ -64,4 +68,4 @@ def nb_to_python_wf(f: float) -> float:
 # ^^^^^^^^^^^^^^^^^^
 # On executing you should see 3 outputs instead of the expected one. This is because this task generates 2 implicit outputs.
 # One of them is the executed notebook (captured) and a rendered (HTML) of the executed notebook. In this case they are called
-# ``nb-simple-out.ipynb`` and ``nb-simple-out.html`` respectively
+# ``nb-simple-out.ipynb`` and ``nb-simple-out.html`` respectively

cookbook/plugins/pod/pod.py

@@ -29,11 +29,11 @@
 from flytekitplugins.pod import Pod
 from kubernetes.client.models import (
     V1Container,
+    V1EmptyDirVolumeSource,
     V1PodSpec,
-    V1VolumeMount,
     V1ResourceRequirements,
     V1Volume,
-    V1EmptyDirVolumeSource,
+    V1VolumeMount,
 )
 
 _SHARED_DATA_PATH = "/data/message.txt"