def __post_init__(self) -> None:
self.state_sub = distributed.Sub(TaskStateEvent.topic_name(),
client=self.client)
self.progress_sub = distributed.Sub(TaskProgressEvent.topic_name(),
client=self.client)
self.logs_sub = distributed.Sub(TaskLogEvent.topic_name(),
client=self.client)
async def _task_progress_change_handler(self, event: str) -> None:
task_progress_event = TaskProgressEvent.parse_raw(event)
logger.debug("received task progress update: %s", task_progress_event)
*_, user_id, project_id, node_id = parse_dask_job_id(task_progress_event.job_id)
message = ProgressRabbitMessage(
user_id=user_id,
project_id=project_id,
node_id=node_id,
progress=task_progress_event.progress,
)
await self.rabbitmq_client.publish_message(message)
def publish_task_logs(
progress_pub: distributed.Pub,
logs_pub: distributed.Pub,
log_type: LogType,
message_prefix: str,
message: str,
) -> None:
logger.info("[%s - %s]: %s", message_prefix, log_type.name, message)
if log_type == LogType.PROGRESS:
publish_event(
progress_pub,
TaskProgressEvent.from_dask_worker(progress=float(message)),
)
else:
publish_event(logs_pub, TaskLogEvent.from_dask_worker(log=message))
def fake_remote_fct(
docker_auth: DockerBasicAuth,
service_key: str,
service_version: str,
input_data: TaskInputData,
output_data_keys: TaskOutputDataSchema,
log_file_url: AnyUrl,
command: List[str],
) -> TaskOutputData:
state_pub = distributed.Pub(TaskStateEvent.topic_name())
progress_pub = distributed.Pub(TaskProgressEvent.topic_name())
logs_pub = distributed.Pub(TaskLogEvent.topic_name())
state_pub.put("my name is state")
progress_pub.put("my name is progress")
logs_pub.put("my name is logs")
# tell the client we are done
published_event = Event(name=_DASK_START_EVENT)
published_event.set()
return TaskOutputData.parse_obj({"some_output_key": 123})
def test_task_progress_from_worker(mocked_dask_worker_job_id: str):
event = TaskProgressEvent.from_dask_worker(0.7)
assert event.job_id == mocked_dask_worker_job_id
assert event.progress == 0.7
def __post_init__(self):
self.state = distributed.Pub(TaskStateEvent.topic_name())
self.progress = distributed.Pub(TaskProgressEvent.topic_name())
self.logs = distributed.Pub(TaskLogEvent.topic_name())