Support `Enum` task arguments

tcompa commented 1 month ago

I've started to test this in https://github.com/fractal-analytics-platform/fractal-tasks-core/blob/test-enum-argument/tests/dev/test_enum_arguments.py. Note that -depending on the final args-schema- this may also require additional work in fractal-web.

Ref:

jluethi commented 4 weeks ago

Here is the case that I'd want to use it for: Fractal-faim-ipa has a convert_ome_zarr function: https://github.com/jluethi/fractal-faim-ipa/blob/main/src/fractal_faim_ipa/convert_ome_zarr.py

3 of the inputs should be Enums:

mode: ModeEnum = "MD Stack Acquisition",
layout: PlateLayout = 96,
tile_alignment: TileAlignmentOptions = "GridAlignment",

But those currently fail to build a correct manifest and the error is:

2024-06-02 23:05:10,611; INFO; Start generating a new manifest
2024-06-02 23:05:10,679; INFO; [convert_ome_zarr.py] START
2024-06-02 23:05:10,679; INFO; [create_schema_for_single_task] START
2024-06-02 23:05:10,679; INFO; [create_schema_for_single_task] function_name='convert_ome_zarr'
2024-06-02 23:05:18,123; INFO; [create_schema_for_single_task] task_function=<cyfunction convert_ome_zarr at 0x17cfa5630>
2024-06-02 23:05:18,124; INFO; [_validate_function_signature] END
2024-06-02 23:05:18,132; INFO; [_remove_args_kwargs_properties] END
2024-06-02 23:05:18,132; INFO; [_remove_pydantic_internals] END
2024-06-02 23:05:18,139; INFO; [_remove_attributes_from_descriptions] END
Traceback (most recent call last):
  File "/Users/joel/Library/CloudStorage/Dropbox/Joel/BioVisionCenter/Code/fractal/fractal-faim-hcs/src/fractal_faim_ipa/dev/create_manifest.py", line 7, in <module>
    create_manifest(package=PACKAGE)
  File "/Users/joel/mambaforge/envs/fractal-faim-ipa/lib/python3.10/site-packages/fractal_tasks_core/dev/create_manifest.py", line 115, in create_manifest
    schema = create_schema_for_single_task(
  File "/Users/joel/mambaforge/envs/fractal-faim-ipa/lib/python3.10/site-packages/fractal_tasks_core/dev/lib_args_schemas.py", line 193, in create_schema_for_single_task
    schema = _include_titles(schema)
  File "/Users/joel/mambaforge/envs/fractal-faim-ipa/lib/python3.10/site-packages/fractal_tasks_core/dev/lib_titles.py", line 68, in _include_titles
    def_schema["properties"]
KeyError: 'properties'

tcompa commented 3 weeks ago

With #749, I am making the tools for generating JSON Schemas a bit more flexible, so that they they are easier to play with. This means that we now can quickly test any new task argument, as in:


class ColorA(Enum):
    RED = "this-is-red"
    GREEN = "this-is-green"

ColorB = Enum(
    "ColorB",
    {"RED": "this-is-red", "GREEN": "this-is-green"},
    type=str,
)

@validate_arguments
def task_function(
    arg_A: ColorA,
    arg_B: ColorB,
):
    """
    Short task description

    Args:
        arg_A: Description of arg_A.
        arg_B: Description of arg_B.
    """
    pass

def test_enum_argument():
    """
    This test only asserts that `create_schema_for_single_task` runs
    successfully. Its goal is also to offer a quick way to experiment
    with new task arguments and play with the generated JSON Schema,
    without re-building the whole fractal-tasks-core manifest.
    """
    schema = create_schema_for_single_task(
        task_function=task_function,
        executable=__file__,
        package=None,
        verbose=True,
    )
    debug(schema)
    print(json.dumps(schema, indent=2))

After a fix for supporting Enums, the test above now runs (within #749) successfully. And while working in the json-schema branch of fractal-web (where https://github.com/fractal-analytics-platform/fractal-web/issues/503 has been addressed), we get the correct rendering:

tcompa commented 3 weeks ago

This is in principle fixed with #749, but I want to make sure it doesn't break manifest-creation flow for other packages. I'll now make a 1.0.3a0 prerelease so that I can test it in fractal-tasks-templates.

jluethi commented 3 weeks ago

This is super neat, looking forward to it! :)

tcompa commented 3 weeks ago

Here how I checked:

I created a task package, based on the template
I installed it with fractal-tasks-core 1.0.2
I added an Enum argument to the task
I ran create_manifest.py, which failed with the KeyError mentioned above
I changed the required version of fractal-tasks-core to 1.0.3a0
I ran create_manifest.py, which ran successfully
I ran python -m build and created a wheel file
I triggered collection of the task package from fractal-web (in the json-schema current branch)
I added the task to a workflow
The argument appears correctly in the form

The relevant part of the args_schema property reads

Form entry

Any additional testing is welcome, but in the meanwhile I'm closing this issue and the next release (1.0.3 or 1.1.0) will include this feature. Depending on the status of the other PRs, we can either make a new release right away or wait a little bit longer and include other updates.

jluethi commented 3 weeks ago

That's awesome! Let's bundle some other PRs into the next release for 1.1.0 :)

Also in view of making sure Fractal web is ready for it & deployed before those task schemas come in

fractal-analytics-platform / fractal-tasks-core

Support `Enum` task arguments #745