<gh_stars>0
import traceback
from pyignite import Client
from pyignite.utils import is_hinted
from pyignite.exceptions import ReconnectError
import grpc
from ..extensions.context_vars import fabric_service, fabric_execution
from ..protocol import FabricService, FabricExecution, TaskCollection
from .connection import establish_connection, configure_instance
from .context import StartupContext, startup_context, run_init_delegate, register_delegate


def run_notebook(*, agent="notebook", instance="notebook-Main"):  # It is important that this function is not async, as it sets context
    fabric_service.set(establish_connection())
    startup_context.set(StartupContext(agent, instance, TaskCollection()))
    fabric_execution.set(FabricExecution(agent, instance, "Main", fabric_service.get().generate_name(agent)))


async def run(agent, delegates={}, *, use_instances=False):
    fabric_service_instance = establish_connection()
    fabric_service.set(fabric_service_instance)
    try:
        instance = configure_instance() if use_instances else None

        task_collection = TaskCollection()
        task_collection.add(fabric_service.get().task_collection.wait(False))
        startup_context.set(StartupContext(agent, instance, task_collection))

        if "Init" in delegates:
            run_init_delegate(delegates.pop("Init"))

        for (delegate, function) in delegates.items():
            register_delegate(delegate, function)

        await task_collection.wait()
    finally:
        try:
            task_collection.cancel()
            await fabric_service_instance.stop()
        except Exception as ex:
            print(-1)
            raise

import traceback
from pyignite import Client
from pyignite.utils import is_hinted
from pyignite.exceptions import ReconnectError
import grpc
from ..extensions.context_vars import fabric_service, fabric_execution
from ..protocol import FabricService, FabricExecution, TaskCollection
from .connection import establish_connection, configure_instance
from .context import StartupContext, startup_context, run_init_delegate, register_delegate


def run_notebook(*, agent="notebook", instance="notebook-Main"):  # It is important that this function is not async, as it sets context
    fabric_service.set(establish_connection())
    startup_context.set(StartupContext(agent, instance, TaskCollection()))
    fabric_execution.set(FabricExecution(agent, instance, "Main", fabric_service.get().generate_name(agent)))


async def run(agent, delegates={}, *, use_instances=False):
    fabric_service_instance = establish_connection()
    fabric_service.set(fabric_service_instance)
    try:
        instance = configure_instance() if use_instances else None

        task_collection = TaskCollection()
        task_collection.add(fabric_service.get().task_collection.wait(False))
        startup_context.set(StartupContext(agent, instance, task_collection))

        if "Init" in delegates:
            run_init_delegate(delegates.pop("Init"))

        for (delegate, function) in delegates.items():
            register_delegate(delegate, function)

        await task_collection.wait()
    finally:
        try:
            task_collection.cancel()
            await fabric_service_instance.stop()
        except Exception as ex:
            print(-1)
            raise
# It is important that this function is not async, as it sets context