https://outerbounds.com/ logo
Join Slack
Powered by
hi, is there a way to override the serializer betw...
# ask-metaflow
w
hi, is there a way to override the serializer between steps from 
pickle
to 
dill
?
✅ 1
v
not directly, but you can use whatever serializer by yourself like this:
Copy code
self.my_special_artifact = dill.dumps(my_object)
and on the other side
Copy code
my_object = dill.loads(self.my_special_artifact)
technically you could do all this is in a custom 
@dill
decorator automatically for all artifacts
out of curiosity, what kind of object gives you trouble with 
pickle
?
ah, I guess this relates to your previous question 🙂 https://outerboundsco.slack.com/archives/C02116BBNTU/p1674837578435779
w
it's just a regular pytorch nn.Module that uses pytorch lightning. But for whatever reason, after calling pickle.dump() on it, the module gets garbage collected 😕
👍 1
v
interesting and unfortunate. We can do some testing on our side too to see if we can reproduce 😕
meanwhile you can use their native serializer or 
dill
as shown above
w
i found this ticket which seems to indicate a problem with the lightning trainer: https://github.com/Lightning-AI/lightning/issues/12233 but i upgraded to lightning 1.8.1 and am still having the problem
at what point does the datastore actually serialize the objects? is it right after self.next() is called?
or is it every time something is written to self?
v
after the 
@step
function finishes
that's why you could override the serializer with a custom decorator that executes after your code but before the whole thing finishes from Metaflow's point of view
c
+1 Ran into the same problem today. My hypothesis is that the Trainer object alters the Model and the references to the model cause the issue. Here's a minimal representation of my code:
Copy code
@step
    def setup_data_module(self):
        from rnn import MyDataModule

        self.data_module = MyDataModule(...)
        self.next(self.setup_model)

    @step
    def setup_model(self):
        from rnn import LSTMModule, LightningModel

        lstm_module = LSTMModule(...)
        self.model = LightningModel(lstm_module, ...)

        self.next(self.train_model)

    @step
    def train_model(self):
        from pytorch_lightning import Trainer, seed_everything
        from pytorch_lightning.loggers import WandbLogger

        seed_everything(0, workers=True)

        trainer = Trainer(
            max_epochs=200,
            accelerator="auto",
            devices="auto",
            deterministic=True,
        )

        # .tune() and .fit() are called on self.model, then trainer is garbage collected
        # storing the trainer object via self.trainer object = trainer doesn't fix the issue
        trainer.tune(self.model, datamodule=self.data_module)
        trainer.fit(self.model, self.data_module)

        self.next(self.end)
It seems that removing Trainer().tune() fixes the issue for me
It seems that you can keep `Trainer().tune()`which will set values on the model, but then remove the reference to the trainer itself via 
model._trainer = None
Copy code
@step
def train_model(self):
    from pytorch_lightning import Trainer

    self.trainer_config = dict(...)
        
    trainer = Trainer(**self.trainer_config)

    trainer.tune(self.model, datamodule=self.data_module)
    trainer.fit(self.model, self.data_module)

    self.model._trainer = None

    self.next(self.end)
4 Views
Open in Slack
PreviousNext
Powered by