from google.protobuf import struct_pb2 as _struct_pb2
from google.protobuf.internal import containers as _containers
from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
from google.protobuf import descriptor as _descriptor
from google.protobuf import message as _message
from collections.abc import Mapping as _Mapping
from typing import ClassVar as _ClassVar, Optional as _Optional, Union as _Union

DESCRIPTOR: _descriptor.FileDescriptor

class SessionState(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
    __slots__ = ()
    STATE_UNSPECIFIED: _ClassVar[SessionState]
    QUEUED: _ClassVar[SessionState]
    PROVISIONING: _ClassVar[SessionState]
    READY: _ClassVar[SessionState]
    IN_PROGRESS: _ClassVar[SessionState]
    COMPLETED: _ClassVar[SessionState]
    ERRORED: _ClassVar[SessionState]
    TIMEOUT: _ClassVar[SessionState]

class SessionType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
    __slots__ = ()
    TYPE_UNSPECIFIED: _ClassVar[SessionType]
    INTERACTIVE: _ClassVar[SessionType]
    INTERNAL_BENCHMARK: _ClassVar[SessionType]
    EXTERNAL_BENCHMARK: _ClassVar[SessionType]
    TRAIN: _ClassVar[SessionType]
STATE_UNSPECIFIED: SessionState
QUEUED: SessionState
PROVISIONING: SessionState
READY: SessionState
IN_PROGRESS: SessionState
COMPLETED: SessionState
ERRORED: SessionState
TIMEOUT: SessionState
TYPE_UNSPECIFIED: SessionType
INTERACTIVE: SessionType
INTERNAL_BENCHMARK: SessionType
EXTERNAL_BENCHMARK: SessionType
TRAIN: SessionType

class InteractiveSessionSpec(_message.Message):
    __slots__ = ("gym_id", "num_envs", "kwargs", "locations")
    class LocationsEntry(_message.Message):
        __slots__ = ("key", "value")
        KEY_FIELD_NUMBER: _ClassVar[int]
        VALUE_FIELD_NUMBER: _ClassVar[int]
        key: str
        value: float
        def __init__(self, key: _Optional[str] = ..., value: _Optional[float] = ...) -> None: ...
    GYM_ID_FIELD_NUMBER: _ClassVar[int]
    NUM_ENVS_FIELD_NUMBER: _ClassVar[int]
    KWARGS_FIELD_NUMBER: _ClassVar[int]
    LOCATIONS_FIELD_NUMBER: _ClassVar[int]
    gym_id: str
    num_envs: int
    kwargs: _struct_pb2.Struct
    locations: _containers.ScalarMap[str, float]
    def __init__(self, gym_id: _Optional[str] = ..., num_envs: _Optional[int] = ..., kwargs: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ..., locations: _Optional[_Mapping[str, float]] = ...) -> None: ...

class InternalBenchmarkSessionSpec(_message.Message):
    __slots__ = ()
    def __init__(self) -> None: ...

class ExternalBenchmarkSessionSpec(_message.Message):
    __slots__ = ("comp_id", "user_token", "locations")
    class LocationsEntry(_message.Message):
        __slots__ = ("key", "value")
        KEY_FIELD_NUMBER: _ClassVar[int]
        VALUE_FIELD_NUMBER: _ClassVar[int]
        key: str
        value: float
        def __init__(self, key: _Optional[str] = ..., value: _Optional[float] = ...) -> None: ...
    COMP_ID_FIELD_NUMBER: _ClassVar[int]
    USER_TOKEN_FIELD_NUMBER: _ClassVar[int]
    LOCATIONS_FIELD_NUMBER: _ClassVar[int]
    comp_id: str
    user_token: str
    locations: _containers.ScalarMap[str, float]
    def __init__(self, comp_id: _Optional[str] = ..., user_token: _Optional[str] = ..., locations: _Optional[_Mapping[str, float]] = ...) -> None: ...

class TrainSessionSpec(_message.Message):
    __slots__ = ()
    def __init__(self) -> None: ...

class Session(_message.Message):
    __slots__ = ("id", "type", "state", "queue_position", "address")
    ID_FIELD_NUMBER: _ClassVar[int]
    TYPE_FIELD_NUMBER: _ClassVar[int]
    STATE_FIELD_NUMBER: _ClassVar[int]
    QUEUE_POSITION_FIELD_NUMBER: _ClassVar[int]
    ADDRESS_FIELD_NUMBER: _ClassVar[int]
    id: str
    type: SessionType
    state: SessionState
    queue_position: int
    address: str
    def __init__(self, id: _Optional[str] = ..., type: _Optional[_Union[SessionType, str]] = ..., state: _Optional[_Union[SessionState, str]] = ..., queue_position: _Optional[int] = ..., address: _Optional[str] = ...) -> None: ...
