.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples_apps/lightning_classy_vision/component.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        Click :ref:`here <sphx_glr_download_examples_apps_lightning_classy_vision_component.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_apps_lightning_classy_vision_component.py:


Trainer Component Example
=========================

This is a component definition that runs the example lightning_classy_vision app.

.. GENERATED FROM PYTHON SOURCE LINES 13-121

.. code-block:: default


    from typing import Optional, Dict

    import torchx.specs.api as torchx
    from torchx.specs import named_resources


    def trainer(
        image: str,
        output_path: str,
        data_path: Optional[str] = None,
        entrypoint: str = "examples/apps/lightning_classy_vision/train.py",
        load_path: str = "",
        log_path: str = "/logs",
        resource: Optional[str] = None,
        env: Optional[Dict[str, str]] = None,
        skip_export: bool = False,
        epochs: int = 1,
    ) -> torchx.AppDef:
        """Runs the example lightning_classy_vision app.

        Args:
            image: image to run (e.g. foobar:latest)
            output_path: output path for model checkpoints (e.g. file:///foo/bar)
            load_path: path to load pretrained model from
            data_path: path to the data to load, if data_path is not provided,
                auto generated test data will be used
            entrypoint: user script to launch.
            log_path: path to save tensorboard logs to
            resource: the resources to use
            env: env variables for the app
            skip_export: disable model export
            epochs: number of epochs to run
        """
        env = env or {}
        args = [
            "--output_path",
            output_path,
            "--load_path",
            load_path,
            "--log_path",
            log_path,
            "--epochs",
            str(epochs),
        ]
        if data_path:
            args += ["--data_path", data_path]
        else:
            args.append("--test")
        if skip_export:
            args.append("--skip_export")
        return torchx.AppDef(
            name="cv-trainer",
            roles=[
                torchx.Role(
                    name="worker",
                    entrypoint=entrypoint,
                    args=args,
                    env=env,
                    image=image,
                    resource=named_resources[resource]
                    if resource
                    else torchx.Resource(cpu=1, gpu=0, memMB=1500),
                )
            ],
        )


    def interpret(
        image: str,
        load_path: str,
        data_path: str,
        output_path: str,
        resource: Optional[str] = None,
        entrypoint: str = "examples/apps/lightning_classy_vision/interpret.py",
    ) -> torchx.AppDef:
        """Runs the model interpretability app on the model outputted by the training
        component.

        Args:
            image: image to run (e.g. foobar:latest)
            load_path: path to load pretrained model from
            data_path: path to the data to load
            output_path: output path for model checkpoints (e.g. file:///foo/bar)
            resource: the resources to use
            entrypoint: user script to launch.
        """
        return torchx.AppDef(
            name="cv-interpret",
            roles=[
                torchx.Role(
                    name="worker",
                    entrypoint=entrypoint,
                    args=[
                        "--load_path",
                        load_path,
                        "--data_path",
                        data_path,
                        "--output_path",
                        output_path,
                    ],
                    image=image,
                    resource=named_resources[resource]
                    if resource
                    else torchx.Resource(cpu=1, gpu=0, memMB=1024),
                )
            ],
        )


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.000 seconds)


.. _sphx_glr_download_examples_apps_lightning_classy_vision_component.py:


.. only :: html

 .. container:: sphx-glr-footer
    :class: sphx-glr-footer-example



  .. container:: sphx-glr-download sphx-glr-download-python

     :download:`Download Python source code: component.py <component.py>`



  .. container:: sphx-glr-download sphx-glr-download-jupyter

     :download:`Download Jupyter notebook: component.ipynb <component.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_