flytekitplugins-optuna 1.15.0__tar.gz

flytekitplugins_optuna-1.15.0/PKG-INFO

@@ -0,0 +1,28 @@
+Metadata-Version: 2.2
+Name: flytekitplugins-optuna
+Version: 1.15.0
+Summary: Optuna plugin for flytekit
+Author: flyteorg
+Author-email: admin@flyte.org
+License: apache2
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Requires-Dist: flytekit>=1.15.0
+Requires-Dist: optuna<5.0.0,>=4.0.0
+Requires-Dist: typing-extensions<5.0,>=4.10
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

flytekitplugins_optuna-1.15.0/README.md

@@ -0,0 +1,188 @@
+# Fully Parallelized Wrapper Around Optuna Using Flyte
+
+## Overview
+
+This documentation provides a guide to a fully parallelized Flyte plugin for Optuna. This wrapper leverages Flyte's scalable and distributed workflow orchestration capabilities to parallelize Optuna's hyperparameter optimization across multiple trials efficiently.
+
+![Timeline](timeline.png)
+
+
+## Features
+
+- **Ease of Use**: This plugin requires no external data storage or experiment tracking.
+- **Parallelized Trial Execution**: Enables concurrent execution of Optuna trials, dramatically speeding up optimization tasks.
+- **Scalability**: Leverages Flyte’s ability to scale horizontally to handle large-scale hyperparameter tuning jobs.
+- **Flexible Integration**: Compatible with various machine learning frameworks and training pipelines.
+
+## Installation
+
+- Install `flytekit`
+- Install `flytekitplugins.optuna`
+
+## Getting Started
+
+### Prerequisites
+
+- A Flyte deployment configured and running.
+- Python 3.9 or later.
+- Familiarity with Flyte and asynchronous programming.
+
+### Define the Objective Function
+
+The objective function defines the problem to be optimized. It should include the hyperparameters to be tuned and return a value to minimize or maximize.
+
+```python
+import math
+
+import flytekit as fl
+
+image = fl.ImageSpec(packages=["flytekitplugins.optuna"])
+
+@fl.task(container_image=image)
+async def objective(x: float, y: int, z: int, power: int) -> float:
+    return math.log((((x - 5) ** 2) + (y + 4) ** 4 + (3 * z - 3) ** 2)) ** power
+
+```
+
+### Configure the Flyte Workflow
+
+The Flyte workflow orchestrates the parallel execution of Optuna trials. Below is an example:
+
+```python
+import flytekit as fl
+from flytekitplugins.optuna import Optimizer, suggest
+
+@fl.eager(container_image=image)
+async def train(concurrency: int, n_trials: int) -> float:
+
+    optimizer = Optimizer(objective=objective, concurrency=concurrency, n_trials=n_trials)
+
+    await optimizer(
+        x=suggest.float(low=-10, high=10),
+        y=suggest.integer(low=-10, high=10),
+        z=suggest.category([-5, 0, 3, 6, 9]),
+        power=2,
+    )
+
+    print(optimizer.study.best_value)
+
+```
+
+### Register and Execute the Workflow
+
+Submit the workflow to Flyte for execution:
+
+```bash
+pyflyte register files .
+pyflyte run --name train
+```
+
+### Monitor Progress
+
+You can monitor the progress of the trials via the Flyte Console. Each trial runs as a separate task, and the results are aggregated by the Optuna wrapper.
+
+You may access the `optuna.Study` like so: `optimizer.study`.
+
+Therefore, with `plotly` installed, you may create create Flyte Decks of the study like so:
+
+```python
+import plotly
+
+fig = optuna.visualization.plot_timeline(optimizer.study)
+fl.Deck(name, plotly.io.to_html(fig))
+```
+
+## Advanced Configuration
+
+### Custom Dictionary Inputs
+
+Suggestions may be defined in recursive dictionaries:
+
+```python
+import flytekit as fl
+from flytekitplugins.optuna import Optimizer, suggest
+
+image = fl.ImageSpec(packages=["flytekitplugins.optuna"])
+
+
+@fl.task(container_image=image)
+async def objective(params: dict[str, int | float | str]) -> float:
+    ...
+
+
+@fl.eager(container_image=image)
+async def train(concurrency: int, n_trials: int):
+
+    study = optuna.create_study(direction="maximize")
+
+    optimizer = Optimizer(objective=objective, concurrency=concurrency, n_trials=n_trials, study=study)
+
+    params = {
+        "lambda": suggest.float(1e-8, 1.0, log=True),
+        "alpha": suggest.float(1e-8, 1.0, log=True),
+        "subsample": suggest.float(0.2, 1.0),
+        "colsample_bytree": suggest.float(0.2, 1.0),
+        "max_depth": suggest.integer(3, 9, step=2),
+        "objective": "binary:logistic",
+        "tree_method": "exact",
+        "booster": "dart",
+    }
+
+    await optimizer(params=params)
+```
+
+### Custom Callbacks
+
+In some cases, you may need to define the suggestions programmatically. This may be done
+
+```python
+import flytekit as fl
+import optuna
+from flytekitplugins.optuna import optimize
+
+image = fl.ImageSpec(packages=["flytekitplugins.optuna"])
+
+@fl.task(container_image=image)
+async def objective(params: dict[str, int | float | str]) -> float:
+    ...
+
+@optimize
+def optimizer(trial: optuna.Trial, verbosity: int, tree_method: str):
+
+    params = {
+        "verbosity:": verbosity,
+        "tree_method": tree_method,
+        "objective": "binary:logistic",
+        # defines booster, gblinear for linear functions.
+        "booster": trial.suggest_categorical("booster", ["gbtree", "gblinear", "dart"]),
+        # sampling according to each tree.
+        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.2, 1.0),
+    }
+
+    if params["booster"] in ["gbtree", "dart"]:
+        # maximum depth of the tree, signifies complexity of the tree.
+        params["max_depth"] = trial.suggest_int("max_depth", 3, 9, step=2)
+
+    if params["booster"] == "dart":
+        params["sample_type"] = trial.suggest_categorical("sample_type", ["uniform", "weighted"])
+        params["normalize_type"] = trial.suggest_categorical("normalize_type", ["tree", "forest"])
+
+    return objective(params)
+
+
+@fl.eager(container_image=image)
+async def train(concurrency: int, n_trials: int):
+
+    optimizer.concurrency = concurrency
+    optimizer.n_trials = n_trials
+
+    study = optuna.create_study(direction="maximize")
+
+    await optimizer(verbosity=0, tree_method="exact")
+```
+
+## Troubleshooting
+
+Resource Constraints: Ensure sufficient compute resources are allocated for the number of parallel jobs specified.
+
+Flyte Errors: Refer to the Flyte logs and documentation to debug workflow execution issues.

flytekitplugins_optuna-1.15.0/flytekitplugins/optuna/__init__.py

@@ -0,0 +1,3 @@
+from .optimizer import Optimizer, optimize, suggest
+
+__all__ = ["Optimizer", "optimize", "suggest"]

flytekitplugins_optuna-1.15.0/flytekitplugins/optuna/optimizer.py

@@ -0,0 +1,212 @@
+import asyncio
+import inspect
+from copy import copy, deepcopy
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Any, Awaitable, Callable, Optional, Union
+
+from typing_extensions import Concatenate, ParamSpec
+
+import optuna
+from flytekit.core.python_function_task import AsyncPythonFunctionTask
+from flytekit.exceptions.eager import EagerException
+
+
+class Suggestion: ...
+
+
+class Number(Suggestion):
+    def __post_init__(self):
+        if self.low >= self.high:
+            raise ValueError("low must be less than high")
+
+        if self.step is not None and self.step > (self.high - self.low):
+            raise ValueError("step must be less than the range of the suggestion")
+
+
+@dataclass
+class Float(Number):
+    low: float
+    high: float
+    step: Optional[float] = None
+    log: bool = False
+
+
+@dataclass
+class Integer(Number):
+    low: int
+    high: int
+    step: int = 1
+    log: bool = False
+
+
+@dataclass
+class Category(Suggestion):
+    choices: list[Any]
+
+
+suggest = SimpleNamespace(float=Float, integer=Integer, category=Category)
+
+P = ParamSpec("P")
+
+Result = Union[float, tuple[float, ...]]
+
+CallbackType = Callable[Concatenate[optuna.Trial, P], Union[Awaitable[Result], Result]]
+
+
+@dataclass
+class Optimizer:
+    objective: Union[CallbackType, AsyncPythonFunctionTask]
+    concurrency: int
+    n_trials: int
+    study: Optional[optuna.Study] = None
+    delay: int = 0
+
+    """
+    Optimizer is a class that allows for the distributed optimization of a flytekit Task using Optuna.
+
+    Args:
+        objective: The objective function to be optimized. This can be a AsyncPythonFunctionTask or a callable.
+        concurrency: The number of trials to run concurrently.
+        n_trials: The number of trials to run in total.
+        study: The study to use for optimization. If None, a new study will be created.
+        delay: The delay in seconds between starting each trial. Default is 0.
+    """
+
+    @property
+    def is_imperative(self) -> bool:
+        return isinstance(self.objective, AsyncPythonFunctionTask)
+
+    def __post_init__(self):
+        if self.study is None:
+            self.study = optuna.create_study()
+
+        if (not isinstance(self.concurrency, int)) or (not self.concurrency > 0):
+            raise ValueError("concurrency must be an integer greater than 0")
+
+        if (not isinstance(self.n_trials, int)) or (not self.n_trials > 0):
+            raise ValueError("n_trials must be an integer greater than 0")
+
+        if not isinstance(self.study, optuna.Study):
+            raise ValueError("study must be an optuna.Study")
+
+        if not isinstance(self.delay, int) or (not self.delay >= 0):
+            raise ValueError("delay must be an integer greater than or equal to 0")
+
+        if self.is_imperative:
+            signature = inspect.signature(self.objective.task_function)
+
+            if signature.return_annotation is float:
+                if len(self.study.directions) != 1:
+                    raise ValueError("the study must have a single objective if objective returns a single float")
+
+            elif hasattr(signature.return_annotation, "__args__"):
+                args = signature.return_annotation.__args__
+                if len(args) != len(self.study.directions):
+                    raise ValueError("objective must return the same number of directions in the study")
+
+                if not all(arg is float for arg in args):
+                    raise ValueError("objective function must return a float or tuple of floats")
+
+            else:
+                raise ValueError("objective function must return a float or tuple of floats")
+
+        else:
+            if not callable(self.objective):
+                raise ValueError("objective must be a callable or a AsyncPythonFunctionTask")
+
+            signature = inspect.signature(self.objective)
+
+            if "trial" not in signature.parameters:
+                raise ValueError(
+                    "objective function must have a parameter called 'trial' if not a AsyncPythonFunctionTask"
+                )
+
+    async def __call__(self, **inputs: P.kwargs):
+        """
+        Asynchronously executes the objective function remotely.
+        Parameters:
+            **inputs: inputs to objective function
+        """
+
+        # create semaphore to manage concurrency
+        semaphore = asyncio.Semaphore(self.concurrency)
+
+        # create list of async trials
+        trials = [self.spawn(semaphore, deepcopy(inputs)) for _ in range(self.n_trials)]
+
+        # await all trials to complete
+        await asyncio.gather(*trials)
+
+    async def spawn(self, semaphore: asyncio.Semaphore, inputs: dict[str, Any]):
+        async with semaphore:
+            await asyncio.sleep(self.delay)
+
+            # ask for a new trial
+            trial: optuna.Trial = self.study.ask()
+
+            try:
+                result: Union[float, tuple[float, ...]]
+
+                # schedule the trial
+                if self.is_imperative:
+                    result = await self.objective(**process(trial, inputs))
+
+                else:
+                    out = self.objective(trial=trial, **inputs)
+                    result = out if not inspect.isawaitable(out) else await out
+
+                # tell the study the result
+                self.study.tell(trial, result, state=optuna.trial.TrialState.COMPLETE)
+
+            # if the trial fails, tell the study
+            except EagerException:
+                self.study.tell(trial, state=optuna.trial.TrialState.FAIL)
+
+
+def optimize(
+    objective: Optional[Union[CallbackType, AsyncPythonFunctionTask]] = None,
+    concurrency: int = 1,
+    n_trials: int = 1,
+    study: Optional[optuna.Study] = None,
+):
+    if objective is not None:
+        if callable(objective) or isinstance(objective, AsyncPythonFunctionTask):
+            return Optimizer(
+                objective=objective,
+                concurrency=concurrency,
+                n_trials=n_trials,
+                study=study,
+            )
+
+        else:
+            raise ValueError("This decorator must be called with a callable or a flyte Task")
+    else:
+
+        def decorator(objective):
+            return Optimizer(objective=objective, concurrency=concurrency, n_trials=n_trials, study=study)
+
+        return decorator
+
+
+def process(trial: optuna.Trial, inputs: dict[str, Any], root: Optional[list[str]] = None) -> dict[str, Any]:
+    if root is None:
+        root = []
+
+    suggesters = {
+        Float: trial.suggest_float,
+        Integer: trial.suggest_int,
+        Category: trial.suggest_categorical,
+    }
+
+    for key, value in inputs.items():
+        path = copy(root) + [key]
+
+        if isinstance(inputs[key], Suggestion):
+            suggester = suggesters[type(value)]
+            inputs[key] = suggester(name=(".").join(path), **vars(value))
+
+        elif isinstance(value, dict):
+            inputs[key] = process(trial=trial, inputs=value, root=path)
+
+    return inputs

flytekitplugins_optuna-1.15.0/flytekitplugins_optuna.egg-info/PKG-INFO

@@ -0,0 +1,28 @@
+Metadata-Version: 2.2
+Name: flytekitplugins-optuna
+Version: 1.15.0
+Summary: Optuna plugin for flytekit
+Author: flyteorg
+Author-email: admin@flyte.org
+License: apache2
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Requires-Dist: flytekit>=1.15.0
+Requires-Dist: optuna<5.0.0,>=4.0.0
+Requires-Dist: typing-extensions<5.0,>=4.10
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

flytekitplugins_optuna-1.15.0/flytekitplugins_optuna.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+README.md
+setup.py
+flytekitplugins/optuna/__init__.py
+flytekitplugins/optuna/optimizer.py
+flytekitplugins_optuna.egg-info/PKG-INFO
+flytekitplugins_optuna.egg-info/SOURCES.txt
+flytekitplugins_optuna.egg-info/dependency_links.txt
+flytekitplugins_optuna.egg-info/entry_points.txt
+flytekitplugins_optuna.egg-info/namespace_packages.txt
+flytekitplugins_optuna.egg-info/requires.txt
+flytekitplugins_optuna.egg-info/top_level.txt
+tests/test_callback.py
+tests/test_decorator.py
+tests/test_imperative.py
+tests/test_validation.py

flytekitplugins_optuna-1.15.0/flytekitplugins_optuna.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

flytekitplugins_optuna-1.15.0/flytekitplugins_optuna.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[flytekit.plugins]
+optuna = flytekitplugins.optuna

flytekitplugins_optuna-1.15.0/flytekitplugins_optuna.egg-info/namespace_packages.txt

@@ -0,0 +1 @@
+flytekitplugins

flytekitplugins_optuna-1.15.0/flytekitplugins_optuna.egg-info/requires.txt

@@ -0,0 +1,3 @@
+flytekit>=1.15.0
+optuna<5.0.0,>=4.0.0
+typing-extensions<5.0,>=4.10

flytekitplugins_optuna-1.15.0/flytekitplugins_optuna.egg-info/top_level.txt

@@ -0,0 +1 @@
+flytekitplugins

flytekitplugins_optuna-1.15.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

flytekitplugins_optuna-1.15.0/setup.py

@@ -0,0 +1,35 @@
+from setuptools import setup
+
+PLUGIN_NAME = "optuna"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+plugin_requires = ["flytekit>=1.15.0", "optuna>=4.0.0,<5.0.0", "typing-extensions>=4.10,<5.0"]
+
+__version__ = "1.15.0"
+
+setup(
+    name=microlib_name,
+    version=__version__,
+    author="flyteorg",
+    author_email="admin@flyte.org",
+    description="Optuna plugin for flytekit",
+    namespace_packages=["flytekitplugins"],
+    packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+    install_requires=plugin_requires,
+    license="apache2",
+    python_requires=">=3.9",
+    classifiers=[
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    entry_points={"flytekit.plugins": [f"{PLUGIN_NAME}=flytekitplugins.{PLUGIN_NAME}"]},
+)

flytekitplugins_optuna-1.15.0/tests/test_callback.py

@@ -0,0 +1,91 @@
+from typing import Union
+
+import asyncio
+from typing import Union
+import optuna
+
+import flytekit as fl
+from flytekitplugins.optuna import Optimizer
+
+
+def test_callback():
+
+
+    @fl.task
+    async def objective(letter: str, number: Union[float, int], other: str, fixed: str) -> float:
+
+        loss = len(letter) + number + len(other) + len(fixed)
+
+        return float(loss)
+
+    def callback(trial: optuna.Trial, fixed: str):
+
+        letter = trial.suggest_categorical("booster", ["A", "B", "BLAH"])
+
+        if letter == "A":
+            number = trial.suggest_int("number_A", 1, 10)
+        elif letter == "B":
+            number = trial.suggest_float("number_B", 10., 20.)
+        else:
+            number = 10
+
+        other = trial.suggest_categorical("other", ["Something", "another word", "a phrase"])
+
+        return objective(letter, number, other, fixed)
+
+
+    @fl.eager
+    async def train(concurrency: int, n_trials: int) -> float:
+
+        study = optuna.create_study(direction="maximize")
+
+        optimizer = Optimizer(callback, concurrency=concurrency, n_trials=n_trials, study=study)
+
+        await optimizer(fixed="hello!")
+
+        return float(optimizer.study.best_value)
+
+    loss = asyncio.run(train(concurrency=2, n_trials=10))
+
+    assert isinstance(loss, float)
+
+
+def test_async_callback():
+
+    @fl.task
+    async def objective(letter: str, number: Union[float, int], other: str, fixed: str) -> float:
+
+        loss = len(letter) + number + len(other) + len(fixed)
+
+        return float(loss)
+
+    async def callback(trial: optuna.Trial, fixed: str):
+
+        letter = trial.suggest_categorical("booster", ["A", "B", "BLAH"])
+
+        if letter == "A":
+            number = trial.suggest_int("number_A", 1, 10)
+        elif letter == "B":
+            number = trial.suggest_float("number_B", 10., 20.)
+        else:
+            number = 10
+
+        other = trial.suggest_categorical("other", ["Something", "another word", "a phrase"])
+
+        return await objective(letter, number, other, fixed)
+
+
+    @fl.eager
+    async def train(concurrency: int, n_trials: int) -> float:
+
+        study = optuna.create_study(direction="maximize")
+
+        optimizer = Optimizer(callback, concurrency=concurrency, n_trials=n_trials, study=study)
+
+        await optimizer(fixed="hello!")
+
+        return float(optimizer.study.best_value)
+
+    loss = asyncio.run(train(concurrency=2, n_trials=10))
+
+    assert isinstance(loss, float)

flytekitplugins_optuna-1.15.0/tests/test_decorator.py

@@ -0,0 +1,114 @@
+from typing import Union
+import math
+
+import asyncio
+from typing import Union
+import optuna
+
+import flytekit as fl
+from flytekitplugins.optuna import optimize, suggest
+
+
+def test_local_exec():
+
+    @fl.eager
+    async def train(concurrency: int, n_trials: int) -> float:
+
+        @optimize(concurrency=concurrency, n_trials=n_trials)
+        @fl.task
+        async def optimizer(x: float, y: int, z: int, power: int) -> float:
+            return (((x - 5) ** 2) + (y + 4) ** 4 + (3 * z - 3) ** 2) ** power
+
+        await optimizer(
+            x=suggest.float(low=-10, high=10),
+            y=suggest.integer(low=-10, high=10),
+            z=suggest.category([-5, 0, 3, 6, 9]),
+            power=2,
+        )
+
+        return optimizer.study.best_value
+
+    loss = asyncio.run(train(concurrency=2, n_trials=10))
+
+    assert isinstance(loss, float)
+
+
+def test_callback():
+
+
+    @fl.task
+    async def objective(letter: str, number: Union[float, int], other: str, fixed: str) -> float:
+
+        loss = len(letter) + number + len(other) + len(fixed)
+
+        return float(loss)
+
+    @optimize(n_trials=10, concurrency=2)
+    def optimizer(trial: optuna.Trial, fixed: str):
+
+        letter = trial.suggest_categorical("booster", ["A", "B", "BLAH"])
+
+        if letter == "A":
+            number = trial.suggest_int("number_A", 1, 10)
+        elif letter == "B":
+            number = trial.suggest_float("number_B", 10., 20.)
+        else:
+            number = 10
+
+        other = trial.suggest_categorical("other", ["Something", "another word", "a phrase"])
+
+        return objective(letter, number, other, fixed)
+
+
+    @fl.eager
+    async def train(concurrency: int, n_trials: int) -> float:
+
+        await optimizer(fixed="hello!")
+
+        return float(optimizer.study.best_value)
+
+    loss = asyncio.run(train(concurrency=2, n_trials=10))
+
+    assert isinstance(loss, float)
+
+
+def test_unparameterized_callback():
+
+
+    @fl.task
+    async def objective(letter: str, number: Union[float, int], other: str, fixed: str) -> float:
+
+        loss = len(letter) + number + len(other) + len(fixed)
+
+        return float(loss)
+
+    @optimize
+    def optimizer(trial: optuna.Trial, fixed: str):
+
+        letter = trial.suggest_categorical("booster", ["A", "B", "BLAH"])
+
+        if letter == "A":
+            number = trial.suggest_int("number_A", 1, 10)
+        elif letter == "B":
+            number = trial.suggest_float("number_B", 10., 20.)
+        else:
+            number = 10
+
+        other = trial.suggest_categorical("other", ["Something", "another word", "a phrase"])
+
+        return objective(letter, number, other, fixed)
+
+
+    @fl.eager
+    async def train(concurrency: int, n_trials: int) -> float:
+
+        optimizer.n_trials = n_trials
+        optimizer.concurrency = concurrency
+
+        await optimizer(fixed="hello!")
+
+        return float(optimizer.study.best_value)
+
+    loss = asyncio.run(train(concurrency=2, n_trials=10))
+
+    assert isinstance(loss, float)

flytekitplugins_optuna-1.15.0/tests/test_imperative.py

@@ -0,0 +1,97 @@
+from typing import Union
+
+import asyncio
+from typing import Union
+
+import optuna
+import flytekit as fl
+from flytekitplugins.optuna import Optimizer, suggest
+
+
+
+def test_local_exec():
+
+
+    @fl.task
+    async def objective(x: float, y: int, z: int, power: int) -> float:
+        return (((x - 5) ** 2) + (y + 4) ** 4 + (3 * z - 3) ** 2) ** power
+
+
+    @fl.eager
+    async def train(concurrency: int, n_trials: int) -> float:
+        optimizer = Optimizer(objective, concurrency=concurrency, n_trials=n_trials)
+
+        await optimizer(
+            x=suggest.float(low=-10, high=10),
+            y=suggest.integer(low=-10, high=10),
+            z=suggest.category([-5, 0, 3, 6, 9]),
+            power=2,
+        )
+
+        return optimizer.study.best_value
+
+    loss = asyncio.run(train(concurrency=2, n_trials=10))
+
+    assert isinstance(loss, float)
+
+def test_tuple_out():
+
+    @fl.task
+    async def objective(x: float, y: int, z: int, power: int) -> tuple[float, float]:
+
+        y0 = (((x - 5) ** 2) + (y + 4) ** 4 + (3 * z - 3) ** 2) ** power
+        y1 = (((x - 2) ** 4) + (y + 1) ** 2 + (4 * z - 1))
+
+        return y0, y1
+
+
+    @fl.eager
+    async def train(concurrency: int, n_trials: int):
+        optimizer = Optimizer(
+            objective=objective,
+            concurrency=concurrency,
+            n_trials=n_trials,
+            study=optuna.create_study(directions=["maximize", "maximize"])
+        )
+
+        await optimizer(
+            x=suggest.float(low=-10, high=10),
+            y=suggest.integer(low=-10, high=10),
+            z=suggest.category([-5, 0, 3, 6, 9]),
+            power=2,
+        )
+
+    asyncio.run(train(concurrency=2, n_trials=10))
+
+
+def test_bundled_local_exec():
+
+    @fl.task
+    async def objective(suggestions: dict[str, Union[int, float]], z: int, power: int) -> float:
+
+        # building out a large set of typed inputs is exhausting, so we can just use a dict
+
+        x, y = suggestions["x"], suggestions["y"]
+
+        return (((x - 5) ** 2) + (y + 4) ** 4) ** power
+
+
+    @fl.eager
+    async def train(concurrency: int, n_trials: int) -> float:
+        optimizer = Optimizer(objective, concurrency=concurrency, n_trials=n_trials)
+
+        suggestions = {
+            "x": suggest.float(low=-10, high=10),
+            "y": suggest.integer(low=-10, high=10),
+        }
+
+        await optimizer(
+            suggestions=suggestions,
+            z=suggest.category([-5, 0, 3, 6, 9]),
+            power=2,
+        )
+
+        return optimizer.study.best_value
+    loss = asyncio.run(train(concurrency=2, n_trials=10))
+
+    assert isinstance(loss, float)

flytekitplugins_optuna-1.15.0/tests/test_validation.py

@@ -0,0 +1,85 @@
+import flytekit as fl
+import pytest
+
+from flytekitplugins.optuna import Optimizer
+
+
+@fl.task
+async def objective(x: float, y: int, z: int, power: int) -> float:
+    return 1.0
+
+def test_concurrency():
+
+    with pytest.raises(ValueError):
+        Optimizer(objective, concurrency=-1, n_trials=10)
+
+    with pytest.raises(ValueError):
+        Optimizer(objective, concurrency=0, n_trials=10)
+
+    with pytest.raises(ValueError):
+        Optimizer(objective, concurrency="abc", n_trials=10)
+
+
+def test_n_trials():
+
+    with pytest.raises(ValueError):
+        Optimizer(objective, concurrency=3, n_trials=-10)
+
+    with pytest.raises(ValueError):
+        Optimizer(objective, concurrency=3, n_trials=0)
+
+    with pytest.raises(ValueError):
+        Optimizer(objective, concurrency=3, n_trials="abc")
+
+
+
+def test_study():
+
+    with pytest.raises(ValueError):
+        Optimizer(objective, concurrency=3, n_trials=10, study="abc")
+
+
+
+def test_delay():
+
+    with pytest.raises(ValueError):
+        Optimizer(objective, concurrency=3, n_trials=10, delay=-1)
+
+    with pytest.raises(ValueError):
+        Optimizer(objective, concurrency=3, n_trials=10, delay="abc")
+
+
+
+def test_objective():
+
+    @fl.workflow
+    def workflow(x: int, y: int, z: int, power: int) -> float:
+        return 1.0
+
+
+    with pytest.raises(ValueError):
+        Optimizer(workflow, concurrency=3, n_trials=10)
+
+    @fl.task
+    async def mistyped_objective(x: int, y: int, z: int, power: int) -> str:
+        return "abc"
+
+    with pytest.raises(ValueError):
+        Optimizer(mistyped_objective, concurrency=3, n_trials=10)
+
+    @fl.task
+    def synchronous_objective(x: int, y: int, z: int, power: int) -> float:
+        return 1.0
+
+    with pytest.raises(ValueError):
+        Optimizer(synchronous_objective, concurrency=3, n_trials=10)
+
+
+
+def test_callback():
+
+    def callback(value: float):
+        return 1.0
+
+    with pytest.raises(ValueError):
+        Optimizer(callback, concurrency=3, n_trials=10)