File size: 3,762 Bytes
b84549f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
import os
import sys
import json
import tempfile
import time
import socket
import string
import random
import glob
from colorama import Fore
import filelock
import psutil
import yaml
from .constants import ERROR_INFO, NORMAL_INFO, WARNING_INFO
def get_yml_content(file_path):
'''Load yaml file content'''
try:
with open(file_path, 'r') as file:
return yaml.safe_load(file)
except yaml.scanner.ScannerError as err:
print_error('yaml file format error!')
print_error(err)
exit(1)
except Exception as exception:
print_error(exception)
exit(1)
def get_json_content(file_path):
'''Load json file content'''
try:
with open(file_path, 'r') as file:
return json.load(file)
except TypeError as err:
print_error('json file format error!')
print_error(err)
return None
def print_error(*content):
'''Print error information to screen'''
print(Fore.RED + ERROR_INFO + ' '.join([str(c) for c in content]) + Fore.RESET)
def print_green(*content):
'''Print information to screen in green'''
print(Fore.GREEN + ' '.join([str(c) for c in content]) + Fore.RESET)
def print_normal(*content):
'''Print error information to screen'''
print(NORMAL_INFO, *content)
def print_warning(*content):
'''Print warning information to screen'''
print(Fore.YELLOW + WARNING_INFO + ' '.join([str(c) for c in content]) + Fore.RESET)
def detect_process(pid):
'''Detect if a process is alive'''
try:
process = psutil.Process(pid)
return process.is_running()
except:
return False
def detect_port(port):
'''Detect if the port is used'''
socket_test = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
socket_test.connect(('127.0.0.1', int(port)))
socket_test.close()
return True
except:
return False
def get_user():
if sys.platform == 'win32':
return os.environ['USERNAME']
else:
return os.environ['USER']
def generate_temp_dir():
'''generate a temp folder'''
def generate_folder_name():
return os.path.join(tempfile.gettempdir(), 'nni', ''.join(random.sample(string.ascii_letters + string.digits, 8)))
temp_dir = generate_folder_name()
while os.path.exists(temp_dir):
temp_dir = generate_folder_name()
os.makedirs(temp_dir)
return temp_dir
class SimplePreemptiveLock(filelock.SoftFileLock):
'''this is a lock support check lock expiration, if you do not need check expiration, you can use SoftFileLock'''
def __init__(self, lock_file, stale=-1):
super(__class__, self).__init__(lock_file, timeout=-1)
self._lock_file_name = '{}.{}'.format(self._lock_file, os.getpid())
self._stale = stale
def _acquire(self):
open_mode = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_TRUNC
try:
lock_file_names = glob.glob(self._lock_file + '.*')
for file_name in lock_file_names:
if os.path.exists(file_name) and (self._stale < 0 or time.time() - os.stat(file_name).st_mtime < self._stale):
return None
fd = os.open(self._lock_file_name, open_mode)
except (IOError, OSError):
pass
else:
self._lock_file_fd = fd
return None
def _release(self):
os.close(self._lock_file_fd)
self._lock_file_fd = None
try:
os.remove(self._lock_file_name)
except OSError:
pass
return None
def get_file_lock(path: string, stale=-1):
return SimplePreemptiveLock(path + '.lock', stale=stale)
|