Xiaoya update prefect3

[xiaoyachong]

This PR focus on Prefect upgrade. It is compatible with the upcoming changes in mlex_utils and mlex_prefect_worker and includes the following updates:

1.Add algorithm registry:
The application now reads the .json file and saves the algorithm details to MLflow, and subsequently retrieves them from MLflow when needed.
This enables the Prefect worker to access algorithm information directly from MLflow, rather than relying on parameters passed from the application.

2. Simplify parameters and remove flow_type in segmentation.py:
   The structure of parameters sent from the application to the Prefect worker has been simplified as follows:

{
    "model_name": model_name,
    "task_name": "train",
    "params": {
        "io_parameters": io_parameters,
        "model_parameters": model_parameters,
    },
}

Meanwhile, the following credentials have been removed from io_parameters:

data_tiled_api_key
mask_tiled_api_key
seg_tiled_api_key
mlflow_tracking_username
mlflow_tracking_password

These credentials are now stored on the Prefect worker side, eliminating the need to send them from the application to the Prefect worker and improving security and configuration consistency.

Companion PRs:
tomo: mlexchange/mlex_tomo_framework#16
prefect worker: mlexchange/mlex_prefect_worker#26
mlex_utils: mlexchange/mlex_utils#5
dlsia_proto: mlexchange/mlex_dlsia_segmentation_prototype#38

---

• To see the specific tasks where the Asana app for GitHub is being used, see below:
  • https://app.asana.com/0/0/1211925772191119

[gitnotebooks]

Review these changes at https://app.gitnotebooks.com/mlexchange/mlex_highres_segmentation/pull/201

[taxe10]

This is great! I have a couple of comments that should be addressed prior merging:

• Enhanced interface startup when no algorithms have been registered - I accidentally started the app without registering the algorithms first, and the application failed to initialize, so I added the following:

% git diff components/control_bar.py 
diff --git a/components/control_bar.py b/components/control_bar.py
index 80e0b59..c39a056 100644
--- a/components/control_bar.py
+++ b/components/control_bar.py
@@ -594,7 +594,7 @@ def layout():
                                         data=models.modelname_list,
                                         value=(
                                             models.modelname_list[0]
-                                            if models.modelname_list[0]
+                                            if len(models.modelname_list) > 0 and models.modelname_list[0]
                                             else None
                                         ),
                                         placeholder="Select a model...",

In the future - we could add a "refresh" button nearby the model selection controls to check if new algorithms have been made available in the application, but I don't think this is needed at this time.

Additionally, I was wondering if we could do a quick model registration at the application startup - something like:

  mlex_segmentation:
    build:
      context: ./mlex_highres_segmentation
      dockerfile: Dockerfile
    command: 'python scripts/save_mlflow_algorithm.py && gunicorn -b 0.0.0.0:8075 --reload app:server'

FYI - I have not been able to test the Prefect 3.x integration locally due to additional comments in the worker's PR

[xiaoyachong]

This is great! I have a couple of comments that should be addressed prior merging:

• Enhanced interface startup when no algorithms have been registered - I accidentally started the app without registering the algorithms first, and the application failed to initialize, so I added the following:

% git diff components/control_bar.py 
diff --git a/components/control_bar.py b/components/control_bar.py
index 80e0b59..c39a056 100644
--- a/components/control_bar.py
+++ b/components/control_bar.py
@@ -594,7 +594,7 @@ def layout():
                                         data=models.modelname_list,
                                         value=(
                                             models.modelname_list[0]
-                                            if models.modelname_list[0]
+                                            if len(models.modelname_list) > 0 and models.modelname_list[0]
                                             else None
                                         ),
                                         placeholder="Select a model...",

In the future - we could add a "refresh" button nearby the model selection controls to check if new algorithms have been made available in the application, but I don't think this is needed at this time.

Additionally, I was wondering if we could do a quick model registration at the application startup - something like:

  mlex_segmentation:
    build:
      context: ./mlex_highres_segmentation
      dockerfile: Dockerfile
    command: 'python scripts/save_mlflow_algorithm.py && gunicorn -b 0.0.0.0:8075 --reload app:server'

FYI - I have not been able to test the Prefect 3.x integration locally due to additional comments in the worker's PR

Thanks for your revision! I’ve made the changes accordingly in the latest commit. I also change the command in mlexchange/mlex_tomo_framework#16.

[taxe10]

These changes worked well at my end.

Just a couple of comments for follow-up in the next PR:

• The status check for the prefect worker switches to ready as soon as the parent flow pool becomes available, even if the individual child worker (conda/docker/etc.) is not actually ready. This can mislead users, so we should refine this logic or clarify the status reporting.
• For MLflow model registration, we're currently using the Prefect flow run ID. We should consider switching to human-readable names—possibly the job name—but we need to think through a long-term strategy that supports multi-tenant scenarios.

[Wiebke]

Thanks for the PR. I think this is pretty ready to merge. I left some minor comments.
Prior to merging we should:

• first merge the mlex_utils PR and update requirements.txt accordingly,
• [optionally] still introduce the guarding against no models being found

[Wiebke]

I am not sure if this is coming out of the options here, but with no models registered yet, I get this error:

Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/mlex_utils/mlflow_utils/mlflow_algorithm_client.py", line 267, in __getitem__
    return self.algorithms[key]
           ^^^^^^^^^^^^^^^^^^^^
KeyError: None

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/app/utils/data_utils.py", line 402, in __getitem__
    return self.mlflow_client[key]
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/mlex_utils/mlflow_utils/mlflow_algorithm_client.py", line 269, in __getitem__
    raise KeyError(f"An algorithm with name '{key}' does not exist.")
KeyError: "An algorithm with name 'None' does not exist."

During handling of the above exception, another exception occurred:

KeyError: 'A model with name None does not exist.'

Note that the reason for me starting the application without models registered is that I unintentionally overwrote the command in my docker-compose.override.yaml.
Some error handling for this might be useful though.

[Wiebke]

One idea could be to update the update_model_parameters callback

https://github.com/xiaoyachong/mlex_highres_segmentation/blob/f742fb97868606da4d7f471b8f1f138fc60832cd/callbacks/control_bar.py#L956-L969

with an additional check:

    if not model_name:
        return html.Div("No model available.")

[xiaoyachong]

Sounds good. I’ve added a safeguard in the update_model_parameters callback in control_bar.py to handle cases where no models are found.

[Wiebke]

In a future PR (in line with archiving mlex_tomo_framework, it might make sense to convert this script to use typer and read from environment variables.

[xiaoyachong]

Thanks for the kind reminder!

[xiaoyachong]

These changes worked well at my end.

Just a couple of comments for follow-up in the next PR:

• The status check for the prefect worker switches to ready as soon as the parent flow pool becomes available, even if the individual child worker (conda/docker/etc.) is not actually ready. This can mislead users, so we should refine this logic or clarify the status reporting.
• For MLflow model registration, we're currently using the Prefect flow run ID. We should consider switching to human-readable names—possibly the job name—but we need to think through a long-term strategy that supports multi-tenant scenarios.

Thanks for the kind reminder!

[Wiebke]

LGTM!