gen-worker 0.1.4__py3-none-any.whl

gen_worker/__init__.py

@@ -0,0 +1,19 @@
+# Make src/gen_worker a Python package
+from .decorators import worker_function, ResourceRequirements
+from .worker import ActionContext
+from .errors import RetryableError, FatalError
+from .types import Asset
+from .model_interface import ModelManager
+from .downloader import ModelDownloader, CozyHubDownloader
+
+__all__ = [
+    "worker_function",
+    "ResourceRequirements",
+    "ActionContext",
+    "RetryableError",
+    "FatalError",
+    "Asset",
+    "ModelManager",
+    "ModelDownloader",
+    "CozyHubDownloader",
+]

gen_worker/decorators.py

@@ -0,0 +1,66 @@
+import functools
+from typing import Callable, Any, Dict, Optional, TypeVar, cast
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+class ResourceRequirements:
+    """
+    Specifies the resource requirements for a worker function.
+    """
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        model_family: Optional[str] = None,
+        min_vram_gb: Optional[float] = None,
+        recommended_vram_gb: Optional[float] = None,
+        requires_gpu: bool = False,
+        expects_pipeline_arg: bool = False,
+        max_concurrency: Optional[int] = None
+        # Add other potential requirements here:
+        # e.g., cpu_cores: Optional[int] = None,
+        # specific_accelerators: Optional[list[str]] = None,
+        # etc.
+    ) -> None:
+        self.model_name = model_name
+        self.model_family = model_family
+        self.min_vram_gb = min_vram_gb
+        self.recommended_vram_gb = recommended_vram_gb
+        self.requires_gpu = requires_gpu
+        self.expects_pipeline_arg = expects_pipeline_arg
+        self.max_concurrency = max_concurrency
+        # Store all defined attributes for easy access
+        self._requirements = {k: v for k, v in locals().items() if k != 'self' and v is not None}
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Returns a dictionary representation of the defined requirements."""
+        return self._requirements
+
+    def __repr__(self) -> str:
+        return f"ResourceRequirements({self._requirements})"
+
+
+def worker_function(
+    resources: Optional[ResourceRequirements] = None,
+) -> Callable[[F], F]:
+    """
+    Decorator to mark a function as a worker task and associate resource requirements.
+
+    Args:
+        resources: An optional ResourceRequirements object describing the function's needs.
+    """
+    if resources is None:
+        resources = ResourceRequirements() # Default empty requirements
+
+    def decorator(func: F) -> F:
+        # Attach metadata directly to the function object.
+        # The SDK's runner component will look for these attributes.
+        setattr(func, '_is_worker_function', True)
+        setattr(func, '_worker_resources', resources)
+
+        # Return the original function, now marked with attributes.
+        # Use functools.wraps to preserve original function metadata (like __name__, __doc__)
+        # even though we are returning the function itself.
+        wrapped = functools.wraps(func)(func)
+        return cast(F, wrapped)
+
+    return decorator

gen_worker/default_model_manager/__init__.py

@@ -0,0 +1,5 @@
+"""Compatibility shim for legacy imports.
+
+Use gen_worker.torch_manager instead.
+"""
+from ..torch_manager import DefaultModelManager, load_config, set_config  # noqa: F401

gen_worker/downloader.py

@@ -0,0 +1,84 @@
+import asyncio
+import hashlib
+import os
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+from urllib.parse import urlparse
+
+import aiohttp
+import backoff
+from tqdm import tqdm
+
+
+class ModelDownloader(ABC):
+    @abstractmethod
+    async def download(self, model_ref: str, dest_dir: str, filename: Optional[str] = None) -> str:
+        """Download a model artifact and return the local file path."""
+        raise NotImplementedError
+
+
+@dataclass
+class DownloadResult:
+    path: str
+    sha256: Optional[str] = None
+    bytes_written: int = 0
+
+
+class CozyHubDownloader(ModelDownloader):
+    """
+    Simple async downloader for Cozy hub model artifacts.
+
+    If model_ref is a full URL (http/https), it is fetched directly.
+    Otherwise, it is appended to base_url as a path segment.
+    """
+
+    def __init__(self, base_url: str, token: Optional[str] = None, timeout_seconds: int = 120) -> None:
+        self.base_url = base_url.rstrip("/")
+        self.token = token
+        self.timeout_seconds = timeout_seconds
+
+    def _resolve_url(self, model_ref: str) -> str:
+        ref = model_ref.strip()
+        if ref.startswith("http://") or ref.startswith("https://"):
+            return ref
+        if not self.base_url:
+            raise ValueError("COZY_HUB_URL is required for non-URL model_ref")
+        return f"{self.base_url}/{ref.lstrip('/')}"
+
+    def _default_filename(self, url: str) -> str:
+        path = urlparse(url).path
+        name = os.path.basename(path)
+        return name or "model.bin"
+
+    @backoff.on_exception(backoff.expo, (aiohttp.ClientError, asyncio.TimeoutError), max_tries=5)
+    async def download(self, model_ref: str, dest_dir: str, filename: Optional[str] = None) -> str:
+        url = self._resolve_url(model_ref)
+        os.makedirs(dest_dir, exist_ok=True)
+        target_name = filename or self._default_filename(url)
+        target_path = os.path.join(dest_dir, target_name)
+
+        headers = {}
+        if self.token:
+            headers["Authorization"] = f"Bearer {self.token}"
+
+        timeout = aiohttp.ClientTimeout(total=self.timeout_seconds)
+        sha256 = hashlib.sha256()
+        bytes_written = 0
+
+        async with aiohttp.ClientSession(timeout=timeout, headers=headers) as session:
+            async with session.get(url) as resp:
+                resp.raise_for_status()
+                total = resp.content_length or 0
+                progress = tqdm(total=total, unit="B", unit_scale=True, desc=f"download {target_name}")
+                with open(target_path, "wb") as f:
+                    async for chunk in resp.content.iter_chunked(1 << 20):
+                        if not chunk:
+                            continue
+                        f.write(chunk)
+                        sha256.update(chunk)
+                        bytes_written += len(chunk)
+                        progress.update(len(chunk))
+                progress.close()
+
+        return target_path

gen_worker/entrypoint.py

@@ -0,0 +1,135 @@
+import logging
+import os
+import sys
+import  asyncio
+from typing import cast
+
+# Ensure the package source is potentially discoverable if running locally
+# In a proper install, this might not be strictly necessary
+# but helps during development if the current dir is the repo root.
+# script_dir = os.path.dirname(os.path.abspath(__file__))
+# sys.path.insert(0, os.path.dirname(script_dir))
+
+try:
+    from .worker import Worker
+    from .model_interface import ModelManagementInterface
+except ImportError as e:
+    print(f"Error importing Worker: {e}", file=sys.stderr)
+    print("Please ensure the gen_worker package is installed or accessible in PYTHONPATH.", file=sys.stderr)
+    sys.exit(1)
+
+# Optional Default Model Management Components
+DMM_AVAILABLE = False
+DefaultModelManager_cls = None
+dmm_load_config_func = None
+dmm_set_config_func = None
+
+try:
+    from .torch_manager import (
+        DefaultModelManager, # The class itself
+        load_config,         # The config loading utility
+        set_config           # The config setting utility
+        # ModelManager as DefaultModelDownloader # If you export your downloader too
+    )
+    DefaultModelManager_cls = DefaultModelManager
+    dmm_load_config_func = load_config
+    dmm_set_config_func = set_config
+    # DefaultModelDownloader_cls = DefaultModelDownloader
+    DMM_AVAILABLE = True
+except ImportError:
+    # This is not necessarily an error if user doesn't intend to use default MMM
+    pass # logging will happen in main
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger('WorkerEntrypoint')
+
+if DMM_AVAILABLE:
+    logger.info("DefaultModelManager components are available in this gen-worker installation.")
+else:
+    logger.info("DefaultModelManager components not found. Built-in dynamic model management will be unavailable "
+                "unless ENABLE_DEFAULT_MODEL_MANAGER is explicitly set to true.")
+
+# --- Configuration ---
+# Read from environment variables or set defaults
+SCHEDULER_ADDR = os.getenv('SCHEDULER_ADDR', 'localhost:8080')
+SCHEDULER_ADDRS = os.getenv('SCHEDULER_ADDRS', '')
+SEED_ADDRS = [addr.strip() for addr in SCHEDULER_ADDRS.split(',') if addr.strip()]
+
+# Default user module name, can be overridden by environment variable
+default_user_modules = 'functions' # A sensible default
+user_modules_str = os.getenv('USER_MODULES', default_user_modules)
+USER_MODULES = [mod.strip() for mod in user_modules_str.split(',') if mod.strip()]
+
+WORKER_ID = os.getenv('WORKER_ID', "worker-1") # Optional, will be generated if None
+AUTH_TOKEN = os.getenv('AUTH_TOKEN') or os.getenv('WORKER_JWT') # Optional
+USE_TLS = os.getenv('USE_TLS', 'false').lower() in ('true', '1', 't')
+RECONNECT_DELAY = int(os.getenv('RECONNECT_DELAY', '5'))
+MAX_RECONNECT_ATTEMPTS = int(os.getenv('MAX_RECONNECT_ATTEMPTS', '0'))
+ENABLE_DEFAULT_MODEL_MANAGER = os.getenv('ENABLE_DEFAULT_MODEL_MANAGER', 'false').lower() in ('true', '1', 't')
+
+
+if __name__ == '__main__':
+
+    # if sys.platform == "win32":
+    #     asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy()) # Or SelectorEventLoopPolicy
+
+    logger.info(f'Starting worker...')
+    logger.info(f'  Scheduler Address: {SCHEDULER_ADDR}')
+    if SEED_ADDRS:
+        logger.info(f'  Scheduler Seeds: {SEED_ADDRS}')
+    logger.info(f'  User Function Modules: {USER_MODULES}')
+    logger.info(f'  Worker ID: {WORKER_ID or "(generated)"}')
+    logger.info(f'  Use TLS: {USE_TLS}')
+    logger.info(f'  Reconnect Delay: {RECONNECT_DELAY}s')
+    logger.info(f'  Max Reconnect Attempts: {MAX_RECONNECT_ATTEMPTS or "Infinite"}')
+    logger.info(f'  Enable Default Model Manager: {ENABLE_DEFAULT_MODEL_MANAGER}')
+
+    if not USER_MODULES:
+        logger.error("No user function modules specified. Set the USER_MODULES environment variable.")
+        sys.exit(1)
+
+    model_manager_instance_to_pass = None
+
+    if ENABLE_DEFAULT_MODEL_MANAGER:
+        if DMM_AVAILABLE and DefaultModelManager_cls and dmm_load_config_func and dmm_set_config_func:
+            logger.info("DefaultModelManager is ENABLED and AVAILABLE. Initializing...")
+            try:
+                # Load config (e.g., from DB/YAML) needed by DefaultModelManager and its utils
+                app_cfg = dmm_load_config_func()
+                dmm_set_config_func(app_cfg) # Set it globally for utils used by DMM
+                logger.info("Application configuration loaded for DefaultModelManager.")
+                
+                model_manager_instance_to_pass = cast(ModelManagementInterface, DefaultModelManager_cls())
+                logger.info("DefaultModelManager instance created.")
+            except Exception as e_dmm_init:
+                logger.exception(f"Failed to initialize DefaultModelManager: {e_dmm_init}. "
+                                 "Proceeding without dynamic model management.")
+                model_manager_instance_to_pass = None
+        else:
+            logger.warning("ENABLE_DEFAULT_MODEL_MANAGER is true, but DefaultModelManager components "
+                           "are not fully available/imported. Proceeding without dynamic model management.")
+    else:
+        logger.info("ENABLE_DEFAULT_MODEL_MANAGER is false. Worker will run without built-in dynamic model management.")
+
+    try:
+        worker = Worker(
+            scheduler_addr=SCHEDULER_ADDR,
+            scheduler_addrs=SEED_ADDRS,
+            user_module_names=USER_MODULES,
+            worker_id=WORKER_ID,
+            auth_token=AUTH_TOKEN,
+            use_tls=USE_TLS,
+            reconnect_delay=RECONNECT_DELAY,
+            max_reconnect_attempts=MAX_RECONNECT_ATTEMPTS,
+            model_manager=model_manager_instance_to_pass
+        )
+        # This blocks until the worker stops
+        worker.run()
+        logger.info('Worker process finished gracefully.')
+        sys.exit(0)
+    except ImportError as e:
+        logger.exception(f"Failed to import user module(s) or dependencies: {e}. Make sure modules '{USER_MODULES}' and their requirements are installed.")
+        sys.exit(1)
+    except Exception as e:
+        logger.exception(f"Worker failed unexpectedly: {e}")
+        sys.exit(1) 

gen_worker/errors.py

@@ -0,0 +1,10 @@
+class WorkerError(Exception):
+    """Base class for worker execution errors."""
+
+
+class RetryableError(WorkerError):
+    """Indicates the job can be retried safely."""
+
+
+class FatalError(WorkerError):
+    """Indicates the job should not be retried."""

gen_worker/model_interface.py

@@ -0,0 +1,48 @@
+from abc import ABC, abstractmethod
+from typing import List, Any, Optional, Dict
+
+DownloaderType = Any 
+
+class ModelManagementInterface(ABC):
+    @abstractmethod
+    async def process_supported_models_config(
+        self, 
+        supported_model_ids: List[str], 
+        downloader_instance: Optional[DownloaderType] 
+    ) -> None:
+        pass
+
+    @abstractmethod
+    async def load_model_into_vram(self, model_id: str) -> bool:
+        pass
+    
+    @abstractmethod
+    async def get_active_pipeline(self, model_id: str) -> Optional[Any]:
+        pass
+
+    @abstractmethod
+    def get_vram_loaded_models(self) -> List[str]:
+        pass
+
+
+class ModelManager(ABC):
+    """
+    Core model manager interface (no torch imports).
+    Implementations are responsible for loading/unloading models into memory.
+    """
+
+    @abstractmethod
+    def load(self, model_ref: str, local_path: Optional[str] = None, **opts: Any) -> Any:
+        pass
+
+    @abstractmethod
+    def get(self, model_ref: str) -> Optional[Any]:
+        pass
+
+    @abstractmethod
+    def unload(self, model_ref: str) -> None:
+        pass
+
+    @abstractmethod
+    def stats(self) -> Dict[str, Any]:
+        pass

gen_worker/pb/__init__.py

@@ -0,0 +1,27 @@
+"""Protocol buffer module for worker communication."""
+
+from __future__ import annotations
+
+import sys
+from typing import Any
+
+from . import frontend_pb2, worker_scheduler_pb2
+
+__all__ = [
+    "worker_scheduler_pb2",
+    "worker_scheduler_pb2_grpc",
+    "frontend_pb2",
+    "frontend_pb2_grpc",
+]
+
+# Compatibility: generated grpc stubs use absolute imports (e.g., worker_scheduler_pb2).
+# Register module aliases before importing *_grpc stubs.
+sys.modules.setdefault("worker_scheduler_pb2", worker_scheduler_pb2)  # type: ignore[arg-type]
+sys.modules.setdefault("frontend_pb2", frontend_pb2)  # type: ignore[arg-type]
+
+# Import grpc stubs after aliases are registered.
+from . import frontend_pb2_grpc, worker_scheduler_pb2_grpc
+
+# Also expose grpc stubs for convenience.
+sys.modules.setdefault("worker_scheduler_pb2_grpc", worker_scheduler_pb2_grpc)  # type: ignore[arg-type]
+sys.modules.setdefault("frontend_pb2_grpc", frontend_pb2_grpc)  # type: ignore[arg-type]

gen_worker/pb/frontend_pb2.py

@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: frontend.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,
+    '',
+    'frontend.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0e\x66rontend.proto\x12\x08\x66rontend\"\\\n\rActionOptions\x12\n\n\x02id\x18\x01 \x01(\t\x12\x12\n\ntimeout_ms\x18\x02 \x01(\x03\x12+\n\x0cretry_policy\x18\x03 \x01(\x0b\x32\x15.frontend.RetryPolicy\"\xa1\x01\n\x0bRetryPolicy\x12\x1b\n\x13initial_interval_ms\x18\x01 \x01(\x03\x12\x1b\n\x13\x62\x61\x63koff_coefficient\x18\x02 \x01(\x02\x12\x1b\n\x13maximum_interval_ms\x18\x03 \x01(\x03\x12\x18\n\x10maximum_attempts\x18\x04 \x01(\x05\x12!\n\x19non_retryable_error_types\x18\x05 \x03(\t\"\xb1\x01\n\x14\x45xecuteActionRequest\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x15\n\rinput_payload\x18\x02 \x01(\x0c\x12(\n\x07options\x18\x03 \x01(\x0b\x32\x17.frontend.ActionOptions\x12\x15\n\rdeployment_id\x18\x04 \x01(\t\x12\x19\n\x11required_model_id\x18\x05 \x01(\t\x12\x0f\n\x07user_id\x18\x06 \x01(\t\"\'\n\x15\x45xecuteActionResponse\x12\x0e\n\x06run_id\x18\x01 \x01(\t\"\x1f\n\rGetRunRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\t\"P\n\x0eGetRunResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x16\n\x0eoutput_payload\x18\x02 \x01(\x0c\x12\x15\n\rerror_message\x18\x03 \x01(\t\"\"\n\x10\x43\x61ncelRunRequest\x12\x0e\n\x06run_id\x18\x01 \x01(\t\"\x1f\n\x11\x43\x61ncelRunResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x32\xe6\x01\n\x0f\x46rontendService\x12P\n\rExecuteAction\x12\x1e.frontend.ExecuteActionRequest\x1a\x1f.frontend.ExecuteActionResponse\x12;\n\x06GetRun\x12\x17.frontend.GetRunRequest\x1a\x18.frontend.GetRunResponse\x12\x44\n\tCancelRun\x12\x1a.frontend.CancelRunRequest\x1a\x1b.frontend.CancelRunResponseB:Z8github.com/cozy-creator/gen-orchestrator/pkg/pb/frontendb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'frontend_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
+  _globals['DESCRIPTOR']._serialized_options = b'Z8github.com/cozy-creator/gen-orchestrator/pkg/pb/frontend'
+  _globals['_ACTIONOPTIONS']._serialized_start=28
+  _globals['_ACTIONOPTIONS']._serialized_end=120
+  _globals['_RETRYPOLICY']._serialized_start=123
+  _globals['_RETRYPOLICY']._serialized_end=284
+  _globals['_EXECUTEACTIONREQUEST']._serialized_start=287
+  _globals['_EXECUTEACTIONREQUEST']._serialized_end=464
+  _globals['_EXECUTEACTIONRESPONSE']._serialized_start=466
+  _globals['_EXECUTEACTIONRESPONSE']._serialized_end=505
+  _globals['_GETRUNREQUEST']._serialized_start=507
+  _globals['_GETRUNREQUEST']._serialized_end=538
+  _globals['_GETRUNRESPONSE']._serialized_start=540
+  _globals['_GETRUNRESPONSE']._serialized_end=620
+  _globals['_CANCELRUNREQUEST']._serialized_start=622
+  _globals['_CANCELRUNREQUEST']._serialized_end=656
+  _globals['_CANCELRUNRESPONSE']._serialized_start=658
+  _globals['_CANCELRUNRESPONSE']._serialized_end=689
+  _globals['_FRONTENDSERVICE']._serialized_start=692
+  _globals['_FRONTENDSERVICE']._serialized_end=922
+# @@protoc_insertion_point(module_scope)

gen_worker/pb/frontend_pb2_grpc.py

@@ -0,0 +1,189 @@
+# 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 frontend_pb2 as frontend__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 frontend_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 FrontendServiceStub(object):
+    """The gRPC service for client <-> scheduler.
+    """
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.ExecuteAction = channel.unary_unary(
+                '/frontend.FrontendService/ExecuteAction',
+                request_serializer=frontend__pb2.ExecuteActionRequest.SerializeToString,
+                response_deserializer=frontend__pb2.ExecuteActionResponse.FromString,
+                _registered_method=True)
+        self.GetRun = channel.unary_unary(
+                '/frontend.FrontendService/GetRun',
+                request_serializer=frontend__pb2.GetRunRequest.SerializeToString,
+                response_deserializer=frontend__pb2.GetRunResponse.FromString,
+                _registered_method=True)
+        self.CancelRun = channel.unary_unary(
+                '/frontend.FrontendService/CancelRun',
+                request_serializer=frontend__pb2.CancelRunRequest.SerializeToString,
+                response_deserializer=frontend__pb2.CancelRunResponse.FromString,
+                _registered_method=True)
+
+
+class FrontendServiceServicer(object):
+    """The gRPC service for client <-> scheduler.
+    """
+
+    def ExecuteAction(self, request, context):
+        """1) Submit a new action/job to the scheduler.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def GetRun(self, request, context):
+        """2) Wait/await the final result of an existing action. Blocks until completed or error.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def CancelRun(self, request, context):
+        """3) Cancel an in-flight action/job.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_FrontendServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'ExecuteAction': grpc.unary_unary_rpc_method_handler(
+                    servicer.ExecuteAction,
+                    request_deserializer=frontend__pb2.ExecuteActionRequest.FromString,
+                    response_serializer=frontend__pb2.ExecuteActionResponse.SerializeToString,
+            ),
+            'GetRun': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetRun,
+                    request_deserializer=frontend__pb2.GetRunRequest.FromString,
+                    response_serializer=frontend__pb2.GetRunResponse.SerializeToString,
+            ),
+            'CancelRun': grpc.unary_unary_rpc_method_handler(
+                    servicer.CancelRun,
+                    request_deserializer=frontend__pb2.CancelRunRequest.FromString,
+                    response_serializer=frontend__pb2.CancelRunResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'frontend.FrontendService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('frontend.FrontendService', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class FrontendService(object):
+    """The gRPC service for client <-> scheduler.
+    """
+
+    @staticmethod
+    def ExecuteAction(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/frontend.FrontendService/ExecuteAction',
+            frontend__pb2.ExecuteActionRequest.SerializeToString,
+            frontend__pb2.ExecuteActionResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def GetRun(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/frontend.FrontendService/GetRun',
+            frontend__pb2.GetRunRequest.SerializeToString,
+            frontend__pb2.GetRunResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def CancelRun(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/frontend.FrontendService/CancelRun',
+            frontend__pb2.CancelRunRequest.SerializeToString,
+            frontend__pb2.CancelRunResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)