flytekit.extras.pytorch.checkpoint

Directory

Classes

Protocols

flytekit.extras.pytorch.checkpoint.IsDataclass

Base class for protocol classes.

Protocol classes are defined as::

class Proto(Protocol):
    def meth(self) -> int:
        ...

Such classes are primarily used with static type checkers that recognize structural subtyping (static duck-typing).

For example::

class C:
    def meth(self) -> int:
        return 0

def func(x: Proto) -> int:
    return x.meth()

func(C())  # Passes static type check

See PEP 544 for details. Protocol classes decorated with @typing.runtime_checkable act as simple-minded runtime protocols that check only the presence of given attributes, ignoring their type signatures. Protocol classes can be generic, they are defined as::

class GenProto[T](Protocol):
    def meth(self) -> T:
        ...

protocol IsDataclass()

flytekit.extras.pytorch.checkpoint.PyTorchCheckpoint

This class is helpful to save a checkpoint.

class PyTorchCheckpoint(
    module: typing.Optional[torch.nn.modules.module.Module],
    hyperparameters: typing.Union[typing.Dict[str, typing.Any], NamedTuple, flytekit.extras.pytorch.checkpoint.IsDataclass, NoneType],
    optimizer: typing.Optional[torch.optim.optimizer.Optimizer],
)

Methods

from_dict()

def from_dict(
    kvs: typing.Union[dict, list, str, int, float, bool, NoneType],
    infer_missing,
) -> ~A

from_json()

def from_json(
    s: typing.Union[str, bytes, bytearray],
    parse_float,
    parse_int,
    parse_constant,
    infer_missing,
    kw,
) -> ~A

schema()

def schema(
    infer_missing: bool,
    only,
    exclude,
    many: bool,
    context,
    load_only,
    dump_only,
    partial: bool,
    unknown,
) -> SchemaType[A]

to_dict()

def to_dict(
    encode_json,
) -> typing.Dict[str, typing.Union[dict, list, str, int, float, bool, NoneType]]

to_json()

def to_json(
    skipkeys: bool,
    ensure_ascii: bool,
    check_circular: bool,
    allow_nan: bool,
    indent: typing.Union[int, str, NoneType],
    separators: typing.Tuple[str, str],
    default: typing.Callable,
    sort_keys: bool,
    kw,
) -> str

flytekit.extras.pytorch.checkpoint.PyTorchCheckpointTransformer

TypeTransformer that supports serializing and deserializing checkpoint.

def PyTorchCheckpointTransformer()

Methods

assert_type()

def assert_type(
    t: Type[T],
    v: T,
)

from_binary_idl()

def from_binary_idl(
    binary_idl_object: Binary,
    expected_python_type: Type[T],
) -> Optional[T]

This function primarily handles deserialization for untyped dicts, dataclasses, Pydantic BaseModels, and attribute access.｀

For untyped dict, dataclass, and pydantic basemodel: Life Cycle (Untyped Dict as example): python val -> msgpack bytes -> binary literal scalar -> msgpack bytes -> python val (to_literal) (from_binary_idl)

For attribute access: Life Cycle: python val -> msgpack bytes -> binary literal scalar -> resolved golang value -> binary literal scalar -> msgpack bytes -> python val (to_literal) (propeller attribute access) (from_binary_idl)

from_generic_idl()

def from_generic_idl(
    generic: Struct,
    expected_python_type: Type[T],
) -> Optional[T]

TODO: Support all Flyte Types. This is for dataclass attribute access from input created from the Flyte Console.

Note:

• This can be removed in the future when the Flyte Console support generate Binary IDL Scalar as input.

get_literal_type()

def get_literal_type(
    t: typing.Type[flytekit.extras.pytorch.checkpoint.PyTorchCheckpoint],
) -> flytekit.models.types.LiteralType

Converts the python type to a Flyte LiteralType

guess_python_type()

def guess_python_type(
    literal_type: flytekit.models.types.LiteralType,
) -> typing.Type[flytekit.extras.pytorch.checkpoint.PyTorchCheckpoint]

Converts the Flyte LiteralType to a python object type.

isinstance_generic()

def isinstance_generic(
    obj,
    generic_alias,
)

to_html()

def to_html(
    ctx: FlyteContext,
    python_val: T,
    expected_python_type: Type[T],
) -> str

Converts any python val (dataframe, int, float) to a html string, and it will be wrapped in the HTML div

to_literal()

def to_literal(
    ctx: flytekit.core.context_manager.FlyteContext,
    python_val: flytekit.extras.pytorch.checkpoint.PyTorchCheckpoint,
    python_type: typing.Type[flytekit.extras.pytorch.checkpoint.PyTorchCheckpoint],
    expected: flytekit.models.types.LiteralType,
) -> flytekit.models.literals.Literal

Converts a given python_val to a Flyte Literal, assuming the given python_val matches the declared python_type. Implementers should refrain from using type(python_val) instead rely on the passed in python_type. If these do not match (or are not allowed) the Transformer implementer should raise an AssertionError, clearly stating what was the mismatch

to_python_value()

def to_python_value(
    ctx: flytekit.core.context_manager.FlyteContext,
    lv: flytekit.models.literals.Literal,
    expected_python_type: typing.Type[flytekit.extras.pytorch.checkpoint.PyTorchCheckpoint],
) -> flytekit.extras.pytorch.checkpoint.PyTorchCheckpoint

Converts the given Literal to a Python Type. If the conversion cannot be done an AssertionError should be raised

Properties