Source code for afe.devkit.devkit_context_manager

#! /usr/bin/env python3
#########################################################
# Copyright (C) 2024 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Nevena Stojnic
#########################################################
"""
This file is used to store SetupConnection class used for
acquiring a connection to the DevKit.

"""
import subprocess
from typing import ContextManager, Optional

from sima_utils.logging import sima_logger

from devkit_inference_models.utils.ssh_utils import create_forward_tunnel_generic, scp_files_to_remote
from devkit_inference_models.apis.pipeline import Pipeline


[docs] class SetupConnection(ContextManager): """ Provides the context in which the user can connect to the DevKit Usage example: with SetupConnection(...): ... """
[docs] pipeline: Optional[Pipeline] = None
_ssh_connection: subprocess.Popen = None _port: int = 0 _username: str = '' _devkit: str = '' _password: str = '' _max_attempt: int = 0 _elf_file_path: str = '' _elf_folder: str = '' _zmq_port: int = 0 def __init__(self, port: int, username: str, devkit: str, password: str, max_attempts: int, elf_file_path: str, elf_folder_path: str, zmq_port: int): self._port = port self._username = username self._devkit = devkit self._password = password self._max_attempt = max_attempts self._elf_file_path = elf_file_path self._elf_folder = elf_folder_path self._zmq_port = zmq_port def __enter__(self): self._ssh_connection, local_port = create_forward_tunnel_generic(self._port, self._username, self._devkit, self._password, self._max_attempt) if self._ssh_connection is None: raise sima_logger.UserFacingException("Failed to forward local port after {self._max_attempt} attempts.") # we start to work with the local_port from now on dv_port = local_port # Copy the model file to the board. scp_file = scp_files_to_remote(None, self._elf_file_path, self._password, self._elf_folder, self._max_attempt, self._username, self._devkit) if scp_file is None: raise sima_logger.UserFacingException("Failed to scp the model file after {self._max_attempt} attempts.") # create tunnel for zmq client/server, use zmq for bigger size data # transfer such as ifms and ofms. self._zmq_ssh_connection, _ = create_forward_tunnel_generic(self._zmq_port, self._username, self._devkit, self._password, self._max_attempt) if self._zmq_ssh_connection is None: raise sima_logger.UserFacingException("Failed to scp the model file after {self._max_attempt} attempts.") self.pipeline = Pipeline(self._devkit, dv_port, self._elf_file_path) return self
[docs] def get_accel_pipeline(self) -> Pipeline: return self.pipeline
def __exit__(self, exc_type, exc_val, exc_tb): self._ssh_connection.kill() self._zmq_ssh_connection.kill() self.pipeline.release()