[Need Help] Access data from one job (FedJob) in another job (FedAvgJob)

[aayush-kapoor]

Python version (python3 -V)

3.9.18

NVFlare version (python3 -m pip list | grep "nvflare")

2.5.2

NVFlare branch (if running examples, please use the branch that corresponds to the NVFlare version, git branch)

main

Operating system

Rocky Linux 9.3

Have you successfully run any of the following examples?

• hello-numpy-sag with simulator
• hello-pt with simulator
• hello-numpy-sag with POC
• hello-pt with POC

Please describe your question

I'm not exactly running into a problem per se, but need help from the developers in figuring out how I can send data between 2 different clients that are defined on separate jobs and are ran separately.

I've been using the streaming class that was setup in the recent version of the library in order to stream the dataset in batches for training my model. I'm using this class especially because even though we can use Pytorch to stream data, I wanted to decouple the compute layer of the model from the data layer (if that makes sense)

So the following are the 2 classes that I've setup for my use-case (Controller and Executor class):

class SimpleStreamingController(Controller):
    def __init__(self, dataset_retriever_id=None, task_timeout=60, task_check_period: float = 0.5):
        Controller.__init__(self, task_check_period=task_check_period)
        self.dataset_retriever_id = dataset_retriever_id
        self.dataset_retriever = None
        self.task_timeout = task_timeout
        self.container_name = "data_batch"
        self.initialize_dataloader()


    def start_controller(self, fl_ctx: FLContext):
        #model = self._get_test_model()
        data_batch = self.get_batch()
        self.dataset_retriever.add_container(self.container_name, data_batch)

    def stop_controller(self, fl_ctx: FLContext):
        pass

    def handle_event(self, event_type: str, fl_ctx: FLContext):
        if event_type == EventType.START_RUN:
            engine = fl_ctx.get_engine()
            if self.dataset_retriever_id:
                c = engine.get_component(self.dataset_retriever_id)
                if not isinstance(c, ContainerRetriever):
                    self.system_panic(
                        f"invalid dataset_retriever {self.dataset_retriever_id}, wrong type: {type(c)}",
                        fl_ctx,
                    )
                    return
                self.dataset_retriever = c

    def control_flow(self, abort_signal: Signal, fl_ctx: FLContext):
        s = Shareable()
        s["name"] = self.container_name
        task = Task(name="retrieve_dict", data=s, timeout=self.task_timeout)
        self.broadcast_and_wait(
            task=task,
            fl_ctx=fl_ctx,
            min_responses=1,
            abort_signal=abort_signal,
        )
        client_resps = {}
        for ct in task.client_tasks:
            assert isinstance(ct, ClientTask)
            resp = ct.result
            if resp is None:
                resp = "no answer"
            else:
                assert isinstance(resp, Shareable)
                self.log_info(fl_ctx, f"got resp {resp} from client {ct.client.name}")
                resp = resp.get_return_code()
            client_resps[ct.client.name] = resp
        return {"status": "OK", "data": client_resps}

    def process_result_of_unknown_task(
        self, client: Client, task_name: str, client_task_id: str, result: Shareable, fl_ctx: FLContext
    ):
        pass

    @staticmethod
    def _get_test_model() -> dict:
        model = {}
        for i in range(10):
            key = f"layer-{i}"
            model[key] = randbytes(1024)

        return model
    
    @staticmethod
    def get_batch():
        global data_iterator

        if data_iterator is None:
            return {"error": "Dataloader not initialized"}, 400

        with data_lock:  # Synchronize access to the iterator
            try:
                # Fetch the next batch
                batch = next(data_iterator)
                inputs, labels = batch
                return {
                    "inputs": inputs.tolist(),  # Convert tensors to lists for JSON serialization
                    "labels": labels.tolist()
                }
            except StopIteration:
                # Reset the iterator if the end of the dataset is reached
                data_iterator = iter(dataloader)
                return {"error": "End of dataset reached"}, 400
            except Exception as e:
                # Catch other unexpected errors
                return {"error": f"Unexpected error: {str(e)}"}, 500


    @staticmethod
    def initialize_dataloader():
        print("Initializing dataloader...")
        global dataloader, data_iterator

        dataset_name = "CIFAR10"
        data_path = "~/dataset"
        batch_size = 32

        # Load dataset
        dataset_wrapper = DatasetWrapper(dataset_name, data_path)
        trainset = dataset_wrapper.trainset
        dataloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)

        # Create an iterator for streaming
        data_iterator = iter(dataloader)

        return {"message": "Dataloader initialized", "batch_size": batch_size}

class SimpleStreamingExecutor(Executor):
    def __init__(self, dataset_retriever_id=None):
        Executor.__init__(self)
        self.dataset_retriever_id = dataset_retriever_id
        self.dataset_retriever = None

    def handle_event(self, event_type: str, fl_ctx: FLContext):
        if event_type == EventType.START_RUN:
            engine = fl_ctx.get_engine()
            if self.dataset_retriever_id:
                c = engine.get_component(self.dataset_retriever_id)
                if not isinstance(c, ContainerRetriever):
                    self.system_panic(
                        f"invalid dataset_retriever {self.dataset_retriever_id}, wrong type: {type(c)}",
                        fl_ctx,
                    )
                    return
                self.dataset_retriever = c

    def execute(self, task_name: str, shareable: Shareable, fl_ctx: FLContext, abort_signal: Signal) -> Shareable:
        self.log_info(fl_ctx, f"got task {task_name}")
        if task_name == "retrieve_dict":
            name = shareable.get("name")
            if not name:
                self.log_error(fl_ctx, "missing name in request")
                return make_reply(ReturnCode.BAD_TASK_DATA)
            if not self.dataset_retriever:
                self.log_error(fl_ctx, "no container retriever")
                return make_reply(ReturnCode.SERVICE_UNAVAILABLE)

            assert isinstance(self.dataset_retriever, ContainerRetriever)
            rc, data_batch = self.dataset_retriever.retrieve_container(
                from_site="server",
                fl_ctx=fl_ctx,
                timeout=10.0,
                name=name,
            )
            if rc != ReturnCode.OK:
                self.log_error(fl_ctx, f"failed to retrieve dict {name}: {rc}")
                return make_reply(rc)

            self.log_info(fl_ctx, f"received container type: {type(data_batch)} size: {len(data_batch)}")
            return make_reply(ReturnCode.OK)
        else:
            self.log_error(fl_ctx, f"got unknown task {task_name}")
            return make_reply(ReturnCode.TASK_UNKNOWN)

So the above streaming classes when ran with the following commands:

job = FedJob(name="simple_dataset_streaming", min_clients=1)

    dataset_retriever = ContainerRetriever()
    job.to_server(dataset_retriever, id="dataset_retriever")
    job.to_clients(dataset_retriever, id="dataset_retriever")

    controller = SimpleStreamingController(dataset_retriever_id="dataset_retriever")
    job.to_server(controller)

    executor = SimpleStreamingExecutor(dataset_retriever_id="dataset_retriever")
    job.to_clients(executor, tasks=["*"])

    job_dir = "~/streaming/workspace/jobs/simple_dataset_streaming"
    print("job_dir=", job_dir)
    job.export_job(job_dir)

    work_dir = "~/streaming/workspace/works/simple_dataset_streaming"
    print("workspace_dir=", work_dir)

    job.simulator_run(work_dir, n_clients=1, threads=1)

are able to fetch the batches of the dataset.

I run my training script for a simple federated average job in the following way:

train_script = "src/pt_fl.py"

    job = FedAvgJob(
        name=f"pt_{dataset_name}_{model_name}_fedavg",
         n_clients=n_clients, 
         num_rounds=num_rounds, 
         initial_model=model
    )

    # Add clients
    for i in range(n_clients):
        executor = ScriptRunner(
            script=train_script, script_args=""
        )
        job.to(executor, f"site-{i + 1}")

    job.simulator_run("/tmp/nvflare/jobs/workdir", gpu="0")

As you can see, I define a separate job to train the model using a different class. And so in this train_script, particularly in the training loop, I want to access the batches that are being streamed. Here's where exactly

flare.init()
    sys_info = flare.system_info()
    client_name = sys_info["site_name"]

    summary_writer = SummaryWriter()
    while flare.is_running():
        input_model = flare.receive()
        print(f"current_round={input_model.current_round}")
        print(f"input model={input_model.meta}")

        model.load_state_dict(input_model.params)
        model.to(device)

        steps = epochs * len(trainloader)
        for epoch in range(epochs):
            running_loss = 0.0
            for i, batch in enumerate(trainloader):
                inputs, labels = batch_fetched_from_streaming_class() # THIS IS WHERE THEY SHOULD BE FETCHED 
                # Rest of taining loop.....

I hope the question makes sense, and if not, kindly let me know what further context I can provide. I've been stuck on figuring out a solution for this for days now so I hope this forum will be able to help me.

[aayush-kapoor]

@YuanTingHsieh @holgerroth sorry to bother you guys, but do you mind taking a look at this? Would the object_retriever module be applicable in this case?

[chesterxgchen]

There are a several things in this which doesn't make sense to me.

1. The code is trying to stream dataset from server to clients to do training. is this what intended? One key principal for FL is not share data. Unless you just want to use FLARE as a distributed mechanism, not for FL
2. You want to stream data between clients, but you are streaming between server to client
3. Streaming functionality is generic for streaming anything, but primarily used for model streaming
4. Why do you want streaming dataset ( not recommend to do ) ?
5. if you are looking for p2p type of interaction, look at swarming learning example, and distributed optimization example.

[aayush-kapoor]

Hi sorry I'm just seeing this now @chesterxgchen

1. Yes that is what is intended. That is not the sole purpose of why I'm using FLARE but one of the reasons.
2. I mean my team was initially thinking of a nested-job type approach (where client would then act as a server) but technically a server will stream the dataset to the client
3. But that still means we can stream objects right? (ie datasets in my case)
4. My team and I are working on a heterogenous federated learning architecture where the training layer is decoupled from the data layer (if that makes sense) and that's why I thought this streaming class could come in handy for the dataset. What exactly are the downsides if you don't mind me asking
5. alright, i'm going to discuss that with the team, see if it's applicable

Thanks for your help!

[chesterxgchen]

From pure technical point of view: stream any thing to any where ( client ==> server, server ==> client), FLARE has the API to do that. That being said.

We need to look the requirements and end goal to see if you want to do that.

Based on what you said above, this is still a FL project.
"Heterogenous" could mean many things

• heterogenous data set from different location or
• heterogenous compute resources ( different GPU, compute powers)
  Not sure which one you are referring.

Before we get to the implementation approach, lets analyze 2nd requirement you mentioned, "the training layer is decoupled from data layer"

I am not sure what does means

• ML training code has no direct access to data ?
  ML training needs data, there must be a data loader that will provide the data to the training code, via File loader, some kind of API ( streaming or not)
• ML Infrastructure is different Data Engineering Infrastructure ? or
• FL infrastructure can not be directly part of existing Data engineering infrastructure ?

In either of these cases, it doesn't make sense to stream data from client to server or server to client. Remember the FL means don't move data.

With above, couldn't make sense of

• Nested Job Type
• Client acts as server
• streaming dataset

if by "server", you mean aggregator, which aggregation weights. For whatever reason, you want to do this on the client side rather on FL server side. You could do that, checkout the Swarm learning workflow. Where in each training round, the workflow select one of the client as aggregator and receive from other clients' weight (peer-to-peer) and perform aggregation.
@aayush-kapoor

[aayush-kapoor]

Thank you for your reply again @chesterxgchen! This was helpful, and we realised we need to rethink our strategy.