🌻Workflow

Introduction

The workflow module aims to help users to create reusable and persistable ML assets. It allows users to -

• Automate and streamline complex tasks, making it easier to manage their assets.

• Simple code modification to change their processing and logical functions to transform these data into Dagster assets.

• Easily modify the workflows to accommodate new requirements as they emerge.

• Saving time and computing costs in the long run.

By building reusable workflows, users can achieve greater efficiency and productivity and reduce the risk of errors and omissions.

Storage

This is the storage utilities for define assets' IO Manager. It defines at which each asset result is persisted. Now, Fiat support the following storage services and data formats -

class PersistStorage(Enum):
    Local = "local"
    OSS = "oss" # Alicloud Object Storage Service


class DataFormat(Enum):
    Binary = ".bin"
    JSON = ".json"
    PlainText = ".txt"

If you are using an external object storage service to persist assets, make sure you include authentication information within the application.json config file -

{
    "oss": { # Storage type: oss, cos, obs, etc. You won't need this when using Local
        "accessKey": {
            "ID": "your_ak_id",
            "Secret": "your_ak_secret"
        },
        "bucket_name": "your-bucket",
        "endpoint": "oss-cn-beijing.aliyuncs.com"
    }
    # other config object
}

To get an IO Manager instance, the user only needs to call the get_io_manager util function -

def get_io_manager(
        medium: PersistStorage = PersistStorage.Local,
        base_path: str = "tmp",
        data_format: DataFormat = DataFormat.Binary
) -> IOManager

Local

from fiat_copilot.workflow.storage import PersistStorage, DataFormat, get_io_manager

local_bin_io_manager = get_io_manager(
        medium=PersistStorage.Local,
        base_path="path/to/store",
        data_format=DataFormat.Binary
)

OSS

Annotations

With annotation, you can easily create an asset -

def as_asset(
        name: str,
        description: str | None = None,
        io_manager_key: str = "default_io_manager"
) -> AssetDefinition

Use it as a decorator, modify your code with @as_asset -

from fiat_copilot.workflow.annotations import as_asset

@as_asset(
    name="get_user_info",
    description="get user information"
)
def get_user_info():
    with open("application.json", "r+", encoding='utf-8') as file_obj:
        conf: dict = json.load(file_obj)
        info: dict = conf['info']

    return info

What if I want to specify a dependent asset? Sure! Just add the asset function name as an input argument. Moreover, you can set the dedicated IOManger for each asset with io_manager_key -

import json
from fiat_copilot.workflow.annotations import as_asset

@as_asset(
    name="to_json",
    description="Convert dict to json",
    io_manager_key="json_io_manager"
)
def to_json(get_user_info):
    user_info = get_user_info # get the result of other asset
    print(user_info)
    
    return json.dumps(user_info)

After defining all of your assets, you need to explicitly state your assets "Definitions" -

def form_definitions(
        assets: list,
        resources: dict
) -> Definitions

# Same .py file with your assets
from fiat_copilot.workflow.annotations import as_asset, form_definitions

defs = form_definitions(
    assets=[get_baidu_html, get_user_info, to_json], # your assets functions
    resources={ # Define your io dedicated io manager here
        "default_io_manager": get_io_manager(
            medium=PersistStorage.OSS,
            base_path="test"
        ),
        "json_io_manager": get_io_manager( # key should match your 'io_manager_key'
            medium=PersistStorage.OSS,
            base_path="test",
            data_format=DataFormat.JSON
        ),
        "text_io_manager": get_io_manager(
            medium=PersistStorage.OSS,
            base_path="test",
            data_format=DataFormat.PlainText
        ),
    }
)

Ray Utils

We also provide a few Ray utilities to make you easy interact with your Ray computing cluster.

Adhoc Task On ray

class ResourceQuota(BaseModel):
    cpu: int
    gpu: int


@contextmanager
def adhoc_task_on_ray(
        cluster: str | None = None,
        resource: ResourceQuota = ResourceQuota(
            cpu=1,
            gpu=0
        )
) -> Callable[Callable[..., Any], Any]

When using adhoc_task_on_ray, the context manager will return a helper function. It accepts a function to be executed on the remote Ray cluster. The Fiat Copilot will wrap it into a Ray Remote Task and send it to the cluster.

EnsureMake sure the function can get its dependencies and data on the remote cluster.

from fiat_copilot.workflow.ray_utils import adhoc_task_on_ray, ResourceQuota

@as_asset(
    name="some_computing_intensive_task",
    description="Compute on Ray",
    io_manager_key="json_io_manager"
)
def some_computing_intensive_task(get_dataset, get_model):
    ray_conf = get_application_conf()['ray']
    
    with adhoc_task_on_ray(
        cluster=ray_conf['cluster'],
        resource=ResourceQuota(
            cpu=ray_conf['resource']['cpu'],
            gpu=ray_conf['resource']['gpu']
        )
    ) as remote_util:
        dataset = get_dataset
        model = get_model

        def train():
            model.fit(dataset)
            
            return model.metrics

        model_metrics = remote_util(train)

    return model_metrics

Long Running Job Submission

Sometimes if your project has lots of dependencies or is fairly complicated. Submitting a running job may be a good choice. For this purpose, Fiat Copilot provides a programmatic approach to submit your job onto a remote Ray computing cluster.

class JobDescription(BaseModel):
    entrypoint: str  # "python script.py"
    working_dir: str  # "examples/"
    pip_packages: list[str]
    

def long_run_job_client(
        cluster: str,
        description: JobDescription
) -> str

You need to define the job description object and provide the target cluster URL -

from fiat_copilot.trainer.torch_util import build_lightning_config, \
    form_lightning_tuner
from fiat_copilot.workflow.annotations import as_asset
from fiat_copilot.workflow.ray_utils import long_run_job_client, JobDescription

@as_asset(
    name="tuning_with_ray",
    description="Fit tuner and get best results.",
    io_manager_key="json_io_manager"
)
def tuning_with_ray(define_lightning_tuner): # Manually define dependency
    cluster_url = "http://10.112.67.227:8265"
    job_description = JobDescription(
        entrypoint="python script.py",
        working_dir="examples/",
        pip_packages=[]
    )
    job_id = long_run_job_client(
        cluster=cluster_url,
        description=JobDescription(
            entrypoint="python script.py",
            working_dir="tune_test/",
            pip_packages=[]
        )
    )

    return {
        "job_id": job_id,
        "cluster": cluster_url,
        "job_descrption": job_description.dict()
    }

Track and manage your assets in DagitUI

Start local Dagster dev server with the Fiat CLI dev run command -

fiat dev run -f my_asset.py

2023-05-17 14:51:15 +0800 - dagster - INFO - Using temporary directory /Users/ken/WorkSpaces/PycharmProjects/fiat-copilot/tmpfw9ig33x for storage. This will be removed when dagster dev exits.
2023-05-17 14:51:15 +0800 - dagster - INFO - To persist information across sessions, set the environment variable DAGSTER_HOME to a directory to use.
2023-05-17 14:51:15 +0800 - dagster - INFO - Launching Dagster services...
[2023-05-17 14:51:18] INFO     Started Dagster code server for file my_assets.py in process 5909       api.py:732
[2023-05-17 14:51:18] INFO     Started Dagster code server for file my_assets.py in process 5912       api.py:732
2023-05-17 14:51:18 +0800 - dagster.daemon - INFO - Instance is configured with the following daemons: ['AssetDaemon', 'BackfillDaemon', 'SchedulerDaemon', 'SensorDaemon']
2023-05-17 14:51:18 +0800 - dagit - INFO - Serving dagit on http://127.0.0.1:3000 in process 5902
2023-05-17 14:51:18 +0800 - dagster.daemon.SensorDaemon - INFO - Not checking for any runs since no sensors have been started.

Open your browser and visit http://127.0.0.1:3000 -

DagitUI

Then, you can materialize your assets and check your results at your storage -

Materialization

Check persisted assets' results