gen-worker 0.1.4__py3-none-any.whl

gen_worker/pb/worker_scheduler_pb2.py

@@ -0,0 +1,69 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: worker_scheduler.proto
+# Protobuf Python Version: 5.27.4
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    5,
+    27,
+    4,
+    '',
+    'worker_scheduler.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16worker_scheduler.proto\x12\tscheduler\"\x82\x05\n\x0fWorkerResources\x12\x11\n\tworker_id\x18\x01 \x01(\t\x12\x11\n\tcpu_cores\x18\x02 \x01(\x05\x12\x14\n\x0cmemory_bytes\x18\x03 \x01(\x03\x12\x11\n\tgpu_count\x18\x04 \x01(\x05\x12\x18\n\x10gpu_memory_bytes\x18\x05 \x01(\x03\x12\x1b\n\x13\x61vailable_functions\x18\x06 \x03(\t\x12\x18\n\x10\x61vailable_models\x18\x07 \x03(\t\x12\x1e\n\x16supports_model_loading\x18\x08 \x01(\x08\x12\x15\n\rdeployment_id\x18\t \x01(\t\x12\x15\n\rrunpod_pod_id\x18\n \x01(\t\x12\x13\n\x0bgpu_is_busy\x18\x0b \x01(\x08\x12\x1d\n\x15gpu_memory_used_bytes\x18\x0c \x01(\x03\x12\x10\n\x08gpu_name\x18\r \x01(\t\x12\x12\n\ngpu_driver\x18\x0e \x01(\t\x12\x1d\n\x15gpu_memory_free_bytes\x18\x0f \x01(\x03\x12\x17\n\x0fmax_concurrency\x18\x10 \x01(\x05\x12Q\n\x14\x66unction_concurrency\x18\x11 \x03(\x0b\x32\x33.scheduler.WorkerResources.FunctionConcurrencyEntry\x12\x14\n\x0c\x63uda_version\x18\x12 \x01(\t\x12\x15\n\rtorch_version\x18\x13 \x01(\t\x12\x33\n\x10\x66unction_schemas\x18\x14 \x03(\x0b\x32\x19.scheduler.FunctionSchema\x1a:\n\x18\x46unctionConcurrencyEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x05:\x02\x38\x01\"U\n\x0e\x46unctionSchema\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x19\n\x11input_schema_json\x18\x02 \x01(\x0c\x12\x1a\n\x12output_schema_json\x18\x03 \x01(\x0c\"Y\n\x12WorkerRegistration\x12-\n\tresources\x18\x01 \x01(\x0b\x32\x1a.scheduler.WorkerResources\x12\x14\n\x0cis_heartbeat\x18\x02 \x01(\x08\"$\n\x10LoadModelCommand\x12\x10\n\x08model_id\x18\x01 \x01(\t\"&\n\x12UnloadModelCommand\x12\x10\n\x08model_id\x18\x01 \x01(\t\"&\n\x14InterruptTaskCommand\x12\x0e\n\x06run_id\x18\x01 \x01(\t\"\xa5\x01\n\x14TaskExecutionRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x15\n\rfunction_name\x18\x02 \x01(\t\x12\x15\n\rinput_payload\x18\x03 \x01(\x0c\x12\x17\n\x0frequired_models\x18\x04 \x03(\t\x12\x12\n\ntimeout_ms\x18\x05 \x01(\x03\x12\x11\n\ttenant_id\x18\x06 \x01(\t\x12\x0f\n\x07user_id\x18\x07 \x01(\t\"e\n\x13TaskExecutionResult\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x0f\n\x07success\x18\x02 \x01(\x08\x12\x16\n\x0eoutput_payload\x18\x03 \x01(\x0c\x12\x15\n\rerror_message\x18\x04 \x01(\t\"`\n\x11SpawnTasksRequest\x12\x15\n\rparent_run_id\x18\x01 \x01(\t\x12\x34\n\x0b\x63hild_tasks\x18\x02 \x03(\x0b\x32\x1f.scheduler.TaskExecutionRequest\"G\n\x0bWorkerEvent\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x12\n\nevent_type\x18\x02 \x01(\t\x12\x14\n\x0cpayload_json\x18\x03 \x01(\x0c\"K\n\x0fLoadModelResult\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x0f\n\x07success\x18\x02 \x01(\x08\x12\x15\n\rerror_message\x18\x03 \x01(\t\"M\n\x11UnloadModelResult\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x0f\n\x07success\x18\x02 \x01(\x08\x12\x15\n\rerror_message\x18\x03 \x01(\t\"4\n\x15\x44\x65ploymentModelConfig\x12\x1b\n\x13supported_model_ids\x18\x01 \x03(\t\"\x9b\x05\n\x16WorkerSchedulerMessage\x12<\n\x13worker_registration\x18\x01 \x01(\x0b\x32\x1d.scheduler.WorkerRegistrationH\x00\x12\x34\n\nrun_result\x18\x02 \x01(\x0b\x32\x1e.scheduler.TaskExecutionResultH\x00\x12\x33\n\x0bspawn_tasks\x18\x03 \x01(\x0b\x32\x1c.scheduler.SpawnTasksRequestH\x00\x12\x37\n\x11load_model_result\x18\x04 \x01(\x0b\x32\x1a.scheduler.LoadModelResultH\x00\x12;\n\x13unload_model_result\x18\x05 \x01(\x0b\x32\x1c.scheduler.UnloadModelResultH\x00\x12.\n\x0cworker_event\x18\x06 \x01(\x0b\x32\x16.scheduler.WorkerEventH\x00\x12\x36\n\x0brun_request\x18\n \x01(\x0b\x32\x1f.scheduler.TaskExecutionRequestH\x00\x12\x35\n\x0eload_model_cmd\x18\x0b \x01(\x0b\x32\x1b.scheduler.LoadModelCommandH\x00\x12\x39\n\x10unload_model_cmd\x18\x0c \x01(\x0b\x32\x1d.scheduler.UnloadModelCommandH\x00\x12<\n\x11interrupt_run_cmd\x18\r \x01(\x0b\x32\x1f.scheduler.InterruptTaskCommandH\x00\x12\x43\n\x17\x64\x65ployment_model_config\x18\x0e \x01(\x0b\x32 .scheduler.DeploymentModelConfigH\x00\x42\x05\n\x03msg2s\n\x16SchedulerWorkerService\x12Y\n\rConnectWorker\x12!.scheduler.WorkerSchedulerMessage\x1a!.scheduler.WorkerSchedulerMessage(\x01\x30\x01\x42\x41Z?github.com/cozy-creator/gen-orchestrator/pkg/pb/workerSchedulerb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'worker_scheduler_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'Z?github.com/cozy-creator/gen-orchestrator/pkg/pb/workerScheduler'
+  _globals['_WORKERRESOURCES_FUNCTIONCONCURRENCYENTRY']._loaded_options = None
+  _globals['_WORKERRESOURCES_FUNCTIONCONCURRENCYENTRY']._serialized_options = b'8\001'
+  _globals['_WORKERRESOURCES']._serialized_start=38
+  _globals['_WORKERRESOURCES']._serialized_end=680
+  _globals['_WORKERRESOURCES_FUNCTIONCONCURRENCYENTRY']._serialized_start=622
+  _globals['_WORKERRESOURCES_FUNCTIONCONCURRENCYENTRY']._serialized_end=680
+  _globals['_FUNCTIONSCHEMA']._serialized_start=682
+  _globals['_FUNCTIONSCHEMA']._serialized_end=767
+  _globals['_WORKERREGISTRATION']._serialized_start=769
+  _globals['_WORKERREGISTRATION']._serialized_end=858
+  _globals['_LOADMODELCOMMAND']._serialized_start=860
+  _globals['_LOADMODELCOMMAND']._serialized_end=896
+  _globals['_UNLOADMODELCOMMAND']._serialized_start=898
+  _globals['_UNLOADMODELCOMMAND']._serialized_end=936
+  _globals['_INTERRUPTTASKCOMMAND']._serialized_start=938
+  _globals['_INTERRUPTTASKCOMMAND']._serialized_end=976
+  _globals['_TASKEXECUTIONREQUEST']._serialized_start=979
+  _globals['_TASKEXECUTIONREQUEST']._serialized_end=1144
+  _globals['_TASKEXECUTIONRESULT']._serialized_start=1146
+  _globals['_TASKEXECUTIONRESULT']._serialized_end=1247
+  _globals['_SPAWNTASKSREQUEST']._serialized_start=1249
+  _globals['_SPAWNTASKSREQUEST']._serialized_end=1345
+  _globals['_WORKEREVENT']._serialized_start=1347
+  _globals['_WORKEREVENT']._serialized_end=1418
+  _globals['_LOADMODELRESULT']._serialized_start=1420
+  _globals['_LOADMODELRESULT']._serialized_end=1495
+  _globals['_UNLOADMODELRESULT']._serialized_start=1497
+  _globals['_UNLOADMODELRESULT']._serialized_end=1574
+  _globals['_DEPLOYMENTMODELCONFIG']._serialized_start=1576
+  _globals['_DEPLOYMENTMODELCONFIG']._serialized_end=1628
+  _globals['_WORKERSCHEDULERMESSAGE']._serialized_start=1631
+  _globals['_WORKERSCHEDULERMESSAGE']._serialized_end=2298
+  _globals['_SCHEDULERWORKERSERVICE']._serialized_start=2300
+  _globals['_SCHEDULERWORKERSERVICE']._serialized_end=2415
+# @@protoc_insertion_point(module_scope)

gen_worker/pb/worker_scheduler_pb2_grpc.py

@@ -0,0 +1,100 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+import worker_scheduler_pb2 as worker__scheduler__pb2
+
+GRPC_GENERATED_VERSION = '1.76.0'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + ' but the generated code in worker_scheduler_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )
+
+
+class SchedulerWorkerServiceStub(object):
+    """The single gRPC service with a bidirectional streaming method.
+    """
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.ConnectWorker = channel.stream_stream(
+                '/scheduler.SchedulerWorkerService/ConnectWorker',
+                request_serializer=worker__scheduler__pb2.WorkerSchedulerMessage.SerializeToString,
+                response_deserializer=worker__scheduler__pb2.WorkerSchedulerMessage.FromString,
+                _registered_method=True)
+
+
+class SchedulerWorkerServiceServicer(object):
+    """The single gRPC service with a bidirectional streaming method.
+    """
+
+    def ConnectWorker(self, request_iterator, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_SchedulerWorkerServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'ConnectWorker': grpc.stream_stream_rpc_method_handler(
+                    servicer.ConnectWorker,
+                    request_deserializer=worker__scheduler__pb2.WorkerSchedulerMessage.FromString,
+                    response_serializer=worker__scheduler__pb2.WorkerSchedulerMessage.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'scheduler.SchedulerWorkerService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('scheduler.SchedulerWorkerService', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class SchedulerWorkerService(object):
+    """The single gRPC service with a bidirectional streaming method.
+    """
+
+    @staticmethod
+    def ConnectWorker(request_iterator,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.stream_stream(
+            request_iterator,
+            target,
+            '/scheduler.SchedulerWorkerService/ConnectWorker',
+            worker__scheduler__pb2.WorkerSchedulerMessage.SerializeToString,
+            worker__scheduler__pb2.WorkerSchedulerMessage.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

gen_worker/testing/__init__.py

@@ -0,0 +1 @@
+"""Testing helpers for gen_worker."""

gen_worker/testing/stub_manager.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import asyncio
+import os
+from typing import Any, Dict, List, Optional
+
+from gen_worker.model_interface import DownloaderType, ModelManagementInterface
+
+
+class StubModelManager(ModelManagementInterface):
+    """
+    Minimal model manager for E2E testing.
+
+    Downloads model artifacts via the provided downloader and returns a
+    lightweight "pipeline" dict containing the local path.
+    """
+
+    def __init__(self) -> None:
+        self._cache_dir = os.getenv("MODEL_CACHE_DIR", "/models")
+        self._downloader: Optional[DownloaderType] = None
+        self._models: Dict[str, str] = {}
+        self._vram_loaded: List[str] = []
+        self._lock = asyncio.Lock()
+
+    async def process_supported_models_config(
+        self,
+        supported_model_ids: List[str],
+        downloader_instance: Optional[DownloaderType],
+    ) -> None:
+        self._downloader = downloader_instance
+        if not supported_model_ids:
+            return
+        for model_id in supported_model_ids:
+            await self._ensure_download(model_id)
+            if model_id not in self._vram_loaded:
+                self._vram_loaded.append(model_id)
+
+    async def load_model_into_vram(self, model_id: str) -> bool:
+        if not model_id:
+            return False
+        await self._ensure_download(model_id)
+        if model_id not in self._vram_loaded:
+            self._vram_loaded.append(model_id)
+        return True
+
+    async def get_active_pipeline(self, model_id: str) -> Optional[Any]:
+        if not model_id:
+            return None
+        path = await self._ensure_download(model_id)
+        return {
+            "model_id": model_id,
+            "model_path": path,
+        }
+
+    def get_vram_loaded_models(self) -> List[str]:
+        return list(self._vram_loaded)
+
+    async def _ensure_download(self, model_id: str) -> str:
+        if model_id in self._models:
+            return self._models[model_id]
+        if not self._downloader:
+            raise RuntimeError("Model downloader is not configured")
+        os.makedirs(self._cache_dir, exist_ok=True)
+        async with self._lock:
+            if model_id in self._models:
+                return self._models[model_id]
+            local_path = await self._downloader.download(model_id, self._cache_dir)
+            self._models[model_id] = local_path
+            return local_path

gen_worker/torch_manager/__init__.py

@@ -0,0 +1,4 @@
+from .manager import DefaultModelManager
+from .utils.config import load_config, set_config
+
+__all__ = ["DefaultModelManager", "load_config", "set_config"]