Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

hotfix import torch #15849

Merged
merged 13 commits into from
Nov 28, 2022
Merged

hotfix import torch #15849

Changes from 2 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
2aad4e8
fix import torch
Borda Nov 28, 2022
a769c28
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Nov 28, 2022
3bda463
plugin
Borda Nov 28, 2022
dbbb502
Merge branch 'fix/import-torch' of https://github.com/PyTorchLightnin…
Borda Nov 28, 2022
c90cc2e
fix
Borda Nov 28, 2022
fa3e961
skip
Borda Nov 28, 2022
d03eea2
patch require
Borda Nov 28, 2022
828d0d2
seed
Borda Nov 28, 2022
c73c88e
warn
Borda Nov 28, 2022
dccb55e
.
Borda Nov 28, 2022
c0c0e97
..
Borda Nov 28, 2022
ccdf442
skip True
Borda Nov 28, 2022
ae2f4a7
0.0.3
Borda Nov 28, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions src/lightning_app/components/serve/python_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
from pathlib import Path
from typing import Any, Dict, Optional

import torch
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
Expand All @@ -13,10 +12,15 @@
from lightning_app.core.queues import MultiProcessQueue
from lightning_app.core.work import LightningWork
from lightning_app.utilities.app_helpers import Logger
from lightning_app.utilities.imports import _is_torch_available, requires
from lightning_app.utilities.proxies import _proxy_setattr, unwrap, WorkRunExecutor, WorkStateObserver

logger = Logger(__name__)

# Skip doctests if requirements aren't available
if _is_torch_available():
__doctest_skip__ = ["PythonServer", "PythonServer.*"]


class _PyTorchSpawnRunExecutor(WorkRunExecutor):

Expand Down Expand Up @@ -86,6 +90,7 @@ def _get_sample_data() -> Dict[Any, Any]:


class PythonServer(LightningWork, abc.ABC):
@requires("torch")
def __init__( # type: ignore
self,
host: str = "127.0.0.1",
Expand Down Expand Up @@ -199,11 +204,13 @@ def _get_sample_dict_from_datatype(datatype: Any) -> dict:
return out

def _attach_predict_fn(self, fastapi_app: FastAPI) -> None:
from torch import inference_mode

input_type: type = self.configure_input_type()
output_type: type = self.configure_output_type()

def predict_fn(request: input_type): # type: ignore
with torch.inference_mode():
with inference_mode():
return self.predict(request)

fastapi_app.post("/predict", response_model=output_type)(predict_fn)
Expand Down