#! /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()