JupyterTask

@task
def init_mlflow_run(experiment_id=None):
    client = MlflowClient(DEFAULT_MLFLOW_TRACKING_URI)
    if experiment_id is None:
        flow_name = prefect.context.get('flow_name')
        artifact_location = os.path.join(DEFAULT_BUCKET_PATH, "mlflow-data", flow_name)
        experiment_id = get_or_create_mlflow_experiment_id(flow_name, artifact_location)
    run = client.create_run(
        experiment_id=experiment_id,
        tags=dict(prefect_run_id=prefect.context.get('flow_run_id'))
    )
    run_id = run.info.run_id
    for k, v in prefect.context.get('parameters', {}).items():
        client.log_param(run_id, k, v)
    return run_id

run_id

class IterablesAsListsJSONEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, frozenset):
            return list(o)
        return super().default(o)


@task(trigger=prefect.triggers.always_run)
def mlflow_logging(
        run_id: Optional[str] = None,
        params: Optional[dict] = None,
        metrics: Optional[dict] = None,
        artifacts: Optional[dict] = None,
        terminate=True
):
    if run_id is None:
        get_logger().warning("Got empty run_id. Skipping MLFlow logging")
        return None
    client = MlflowClient(DEFAULT_MLFLOW_TRACKING_URI)
    if params is not None and not isinstance(params, PrefectStateSignal):
        for k, v in params.items():
            client.log_param(run_id, k, v)
    if metrics is not None and not isinstance(metrics, PrefectStateSignal):
        for k, v in metrics.items():
            client.log_metric(run_id, k, v)
    if artifacts is not None and not isinstance(artifacts, PrefectStateSignal):
        shutil.rmtree(ARTIFACTS_PATH, ignore_errors=True)
        os.makedirs(ARTIFACTS_PATH)  # TODO: figure out how to upload stuff without actually saving to disk
        for filename, data in artifacts.items():
            try:
                with open(os.path.join(ARTIFACTS_PATH, f"{filename}.json"), "w") as json_f:
                    json.dump(data, json_f, cls=IterablesAsListsJSONEncoder)
            except TypeError as e:
                get_logger().warning("Failed to save data as JSON. Saving as pickle")
                with open(os.path.join(ARTIFACTS_PATH, filename), "wb") as f:
                    pickle.dump(data, f)
        client.log_artifacts(run_id, ARTIFACTS_PATH)
    if terminate:
        client.set_terminated(run_id)

    return run_id