Skip to content

Commit

Permalink
Moving Dask interceptor instance to scheduler_plugin.start dask/distr…
Browse files Browse the repository at this point in the history
  • Loading branch information
renan-souza committed Feb 4, 2025
1 parent 492cca9 commit e8ad034
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions src/flowcept/flowceptor/adapters/dask/dask_plugins.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from uuid import uuid4

from dask.distributed import WorkerPlugin, SchedulerPlugin
from distributed import Client
from distributed import Client, Scheduler

from flowcept import WorkflowObject
from flowcept.configs import INSTRUMENTATION
Expand Down Expand Up @@ -69,8 +69,11 @@ def register_dask_workflow(
class FlowceptDaskSchedulerAdapter(SchedulerPlugin):
"""Dask schedule adapter."""

def __init__(self, scheduler):
self.address = scheduler.address
def __init__(self):
self.interceptor = None

def start(self, scheduler: Scheduler) -> None:
"""Run this when scheduler starts"""
self.interceptor = DaskSchedulerInterceptor(scheduler)

def transition(self, key, start, finish, *args, **kwargs):
Expand Down

0 comments on commit e8ad034

Please sign in to comment.