build: update isort known first party config (#50)

container/deps/vllm/data_plane/data_plane.py

@@ -27,6 +27,7 @@ import torch
 import torch.distributed
 import tritonserver
 import zmq
+
 from triton_distributed.icp.data_plane import (
     set_icp_data_type,
     set_icp_memory_type,

examples/hello_world/client/client.py

@@ -19,9 +19,10 @@ import sys
 import cupy
 import numpy
 from tqdm import tqdm
+from tritonserver import MemoryType
+
 from triton_distributed.icp import NatsRequestPlane, UcpDataPlane
 from triton_distributed.worker import RemoteOperator
-from tritonserver import MemoryType
 
 
 def _get_input_sizes(args):

examples/hello_world/operators/encoder_decoder.py

 import numpy
+
 from triton_distributed.worker import Operator, RemoteInferenceRequest, RemoteOperator
 
 

examples/hello_world/single_file.py

@@ -21,6 +21,8 @@ from pathlib import Path
 import cupy
 import numpy
 from tqdm import tqdm
+from tritonserver import MemoryType
+
 from triton_distributed.icp.nats_request_plane import NatsRequestPlane
 from triton_distributed.icp.ucp_data_plane import UcpDataPlane
 from triton_distributed.worker import (
@@ -32,7 +34,6 @@ from triton_distributed.worker import (
     TritonCoreOperator,
     WorkerConfig,
 )
-from tritonserver import MemoryType
 
 
 class EncodeDecodeOperator(Operator):

icp/src/python/triton_distributed/icp/data_plane.py

@@ -21,7 +21,6 @@ from typing import Optional, Sequence
 
 import cupy
 import numpy
-from triton_distributed.icp.protos.icp_pb2 import ModelInferRequest, ModelInferResponse
 from tritonserver import (
     DataType,
     InvalidArgumentError,
@@ -43,6 +42,8 @@ from tritonserver._c.triton_bindings import (
     TRITONSERVER_StringToDataType as StringToDataType,
 )
 
+from triton_distributed.icp.protos.icp_pb2 import ModelInferRequest, ModelInferResponse
+
 
 class DataPlaneError(Exception):
     pass

icp/src/python/triton_distributed/icp/nats_request_plane.py

@@ -25,6 +25,8 @@ from typing import Dict, Optional
 from urllib.parse import urlsplit, urlunsplit
 
 import nats
+from tritonserver import InvalidArgumentError
+
 from triton_distributed.icp.protos.icp_pb2 import ModelInferRequest, ModelInferResponse
 from triton_distributed.icp.request_plane import (
     RequestPlane,
@@ -37,7 +39,6 @@ from triton_distributed.icp.request_plane import (
     set_icp_request_to_uri,
     set_icp_response_to_uri,
 )
-from tritonserver import InvalidArgumentError
 
 
 class AsyncModelInferRequestIterator:

icp/src/python/triton_distributed/icp/request_plane.py

@@ -19,9 +19,10 @@ import abc
 import uuid
 from typing import AsyncIterator, Awaitable, Callable, Optional
 
-from triton_distributed.icp.protos.icp_pb2 import ModelInferRequest, ModelInferResponse
 from tritonserver import TritonError
 
+from triton_distributed.icp.protos.icp_pb2 import ModelInferRequest, ModelInferResponse
+
 ICP_REQUEST_ID = "icp_request_id"
 ICP_FINAL_RESPONSE = "icp_final_response"
 ICP_RESPONSE_FROM_URI = "icp_response_from_uri"

icp/src/python/triton_distributed/icp/ucp_data_plane.py

@@ -28,6 +28,8 @@ import numpy
 import tritonserver
 import ucp
 from cupy_backends.cuda.api.runtime import CUDARuntimeError
+from tritonserver import InvalidArgumentError, MemoryBuffer, MemoryType, Tensor
+
 from triton_distributed.icp.data_plane import (
     DataPlane,
     DataPlaneError,
@@ -47,7 +49,6 @@ from triton_distributed.icp.data_plane import (
     set_icp_tensor_uri,
 )
 from triton_distributed.icp.protos.icp_pb2 import ModelInferRequest, ModelInferResponse
-from tritonserver import InvalidArgumentError, MemoryBuffer, MemoryType, Tensor
 
 LOGGER = logging.getLogger(__name__)
 

icp/tests/python/test_data_plane.py

@@ -24,14 +24,15 @@ import numpy
 import pytest
 import ucp
 from cupy_backends.cuda.api.runtime import CUDARuntimeError
+from tritonserver import DataType, MemoryType, Tensor
+from tritonserver._api._datautils import TRITON_TO_NUMPY_DTYPE
+
 from triton_distributed.icp.data_plane import DataPlaneError
 from triton_distributed.icp.ucp_data_plane import (
     UcpDataPlane,
     get_icp_tensor_uri,
     set_icp_tensor_uri,
 )
-from tritonserver import DataType, MemoryType, Tensor
-from tritonserver._api._datautils import TRITON_TO_NUMPY_DTYPE
 
 # TODO decide if some tests should be removed
 # from pre_merge

icp/tests/python/test_request_plane.py

@@ -21,6 +21,7 @@ import uuid
 from multiprocessing import Process, Queue
 
 import pytest
+
 from triton_distributed.icp.nats_request_plane import NatsRequestPlane
 from triton_distributed.icp.protos.icp_pb2 import ModelInferRequest, ModelInferResponse
 from triton_distributed.icp.request_plane import (

pyproject.toml

@@ -43,6 +43,7 @@ line_length = 88
 balanced_wrapping = true
 indent = "    "
 skip = ["build"]
+known_first_party = ["triton_distributed"]
 
 [tool.pytest.ini_options]
 minversion = "8.0"

worker/src/python/triton_distributed/worker/deployment.py

@@ -16,6 +16,8 @@ import multiprocessing
 from pprint import pformat
 from typing import Optional, Type
 
+from tritonserver import InvalidArgumentError
+
 from triton_distributed.icp import (
     DataPlane,
     NatsRequestPlane,
@@ -25,7 +27,6 @@ from triton_distributed.icp import (
 )
 from triton_distributed.worker.log_formatter import setup_logger
 from triton_distributed.worker.worker import Worker, WorkerConfig
-from tritonserver import InvalidArgumentError
 
 LOGGER_NAME = __name__
 

worker/src/python/triton_distributed/worker/operator.py

@@ -19,10 +19,11 @@ import abc
 from dataclasses import dataclass, field
 from typing import Any, Optional, Type
 
+from tritonserver import Server
+
 from triton_distributed.icp.data_plane import DataPlane
 from triton_distributed.icp.request_plane import RequestPlane
 from triton_distributed.worker.remote_request import RemoteInferenceRequest
-from tritonserver import Server
 
 
 class Operator(abc.ABC):

worker/src/python/triton_distributed/worker/remote_operator.py

@@ -19,11 +19,12 @@ import asyncio
 import uuid
 from typing import Optional
 
+from tritonserver import InvalidArgumentError
+
 from triton_distributed.icp.data_plane import DataPlane
 from triton_distributed.icp.request_plane import RequestPlane
 from triton_distributed.worker.remote_request import RemoteInferenceRequest
 from triton_distributed.worker.remote_response import AsyncRemoteResponseIterator
-from tritonserver import InvalidArgumentError
 
 
 class RemoteOperator:

worker/src/python/triton_distributed/worker/remote_request.py

@@ -25,12 +25,13 @@ from dataclasses import dataclass, field
 from typing import Any, Optional
 
 import tritonserver
+from tritonserver import InferenceRequest, InvalidArgumentError, Tensor
+
 from triton_distributed.icp.data_plane import DataPlane
 from triton_distributed.icp.protos.icp_pb2 import ModelInferRequest
 from triton_distributed.icp.request_plane import RequestPlane, get_icp_component_id
 from triton_distributed.worker.remote_response import RemoteInferenceResponse
 from triton_distributed.worker.remote_tensor import RemoteTensor
-from tritonserver import InferenceRequest, InvalidArgumentError, Tensor
 
 
 @dataclass

worker/src/python/triton_distributed/worker/remote_response.py

@@ -29,6 +29,9 @@ if TYPE_CHECKING:
 
 import uuid
 
+from tritonserver import InternalError, Tensor, TritonError
+from tritonserver._api._response import InferenceResponse
+
 from triton_distributed.icp.request_plane import (
     get_icp_component_id,
     get_icp_final_response,
@@ -37,8 +40,6 @@ from triton_distributed.icp.request_plane import (
     set_icp_response_error,
 )
 from triton_distributed.worker.remote_tensor import RemoteTensor
-from tritonserver import InternalError, Tensor, TritonError
-from tritonserver._api._response import InferenceResponse
 
 
 class AsyncRemoteResponseIterator:

worker/src/python/triton_distributed/worker/remote_tensor.py

@@ -20,6 +20,13 @@ from typing import Optional, Sequence
 
 import cupy
 from cupy_backends.cuda.api.runtime import CUDARuntimeError
+from tritonserver import DataType, InvalidArgumentError, MemoryType, Tensor
+
+# TODO
+# Export from tritonserver
+from tritonserver._api._dlpack import DLDeviceType
+from tritonserver._api._tensor import DeviceOrMemoryType
+
 from triton_distributed.icp.data_plane import (
     DataPlane,
     get_icp_data_type,
@@ -28,12 +35,6 @@ from triton_distributed.icp.data_plane import (
     get_icp_tensor_size,
 )
 from triton_distributed.icp.protos.icp_pb2 import ModelInferRequest, ModelInferResponse
-from tritonserver import DataType, InvalidArgumentError, MemoryType, Tensor
-
-# TODO
-# Export from tritonserver
-from tritonserver._api._dlpack import DLDeviceType
-from tritonserver._api._tensor import DeviceOrMemoryType
 
 # Run cupy's cuda.is_available once to
 # avoid the exception hitting runtime code.

worker/src/python/triton_distributed/worker/triton_core_operator.py

@@ -21,13 +21,14 @@ import uuid
 from typing import Optional
 
 from google.protobuf import json_format, text_format
+from tritonclient.grpc import model_config_pb2
+from tritonserver import InvalidArgumentError, Server
+
 from triton_distributed.icp.data_plane import DataPlane
 from triton_distributed.icp.request_plane import RequestPlane
 from triton_distributed.worker.operator import Operator
 from triton_distributed.worker.remote_request import RemoteInferenceRequest
 from triton_distributed.worker.remote_response import RemoteInferenceResponse
-from tritonclient.grpc import model_config_pb2
-from tritonserver import InvalidArgumentError, Server
 
 
 class TritonCoreOperator(Operator):

worker/src/python/triton_distributed/worker/worker.py

@@ -26,6 +26,7 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Optional, Type
 
 import tritonserver
+
 from triton_distributed.icp.data_plane import DataPlane
 from triton_distributed.icp.nats_request_plane import NatsRequestPlane
 from triton_distributed.icp.request_plane import RequestPlane

worker/tests/python/conftest.py

@@ -21,6 +21,7 @@ import time
 
 import pytest
 import pytest_asyncio
+
 from triton_distributed.icp.nats_request_plane import NatsServer
 from triton_distributed.worker.log_formatter import LOGGER_NAME