[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.