from __future__ import annotations

import logging

import torch
from transformers import OPTForCausalLM

from proxystore.connectors.file import FileConnector
from proxystore.store import Store


def main() -> None:
    logging.basicConfig(level=logging.DEBUG)

    model = OPTForCausalLM.from_pretrained(
        'facebook/opt-1.3b',
        cache_dir='/tmp/proxystore-issue-704/cache',
        torch_dtype=torch.float32,
    )

    connector = FileConnector('/tmp/proxystore-issue-704/store')
    with Store(name='issue-704', connector=connector) as store:
        key = store.put(model)
        store.get(key)


if __name__ == '__main__':
    main()

from __future__ import annotations

import logging
import io

import torch
from transformers import AutoConfig, OPTForCausalLM

from proxystore.connectors.file import FileConnector
from proxystore.serialize import serialize, deserialize, BytesLike, SerializationError
from proxystore.store import Store


def serialize_torch_model(obj: Any) -> bytes:
    if isinstance(obj, torch.nn.Module):
        buffer = io.BytesIO()
        buffer.write(b'PT\n')
        torch.save(obj, buffer)
        return buffer.getvalue()
    else:
        return serialize(obj)


def deserialize_torch_model(raw: BytesLike) -> Any:
    try:
        return deserialize(raw)
    except SerializationError:
        buffer = io.BytesIO(raw)
        assert buffer.readline() == b'PT\n'
        return torch.load(buffer, weights_only=False)


def main() -> None:
    logging.basicConfig(level=logging.DEBUG)

    model = OPTForCausalLM.from_pretrained(
        'facebook/opt-1.3b',
        cache_dir='/tmp/proxystore-issue-704/cache',
        torch_dtype=torch.float32,
    )

    connector = FileConnector('/tmp/proxystore-issue-704/store')
    with Store(name='issue-704', connector=connector) as store:
        key = store.put(model, serializer=serialize_torch_model)
        store.get(key, deserializer=deserialize_torch_model)


if __name__ == '__main__':
    main()