diff --git a/compute/protocol.py b/compute/protocol.py index 22ce3133..ae151c82 100644 --- a/compute/protocol.py +++ b/compute/protocol.py @@ -87,6 +87,7 @@ class Allocate(bt.Synapse): docker_action: dict = { "action": "", "ssh_key": "", + "key_type": "", } def deserialize(self) -> dict: diff --git a/neurons/Miner/container.py b/neurons/Miner/container.py index 2995c242..6752b726 100644 --- a/neurons/Miner/container.py +++ b/neurons/Miner/container.py @@ -432,7 +432,7 @@ def unpause_container(): bt.logging.info(f"Error unpausing container {e}") return {"status": False} -def exchange_key_container(new_ssh_key: str): +def exchange_key_container(new_ssh_key: str, key_type: str = "user"): try: client, containers = get_docker() running_container = None @@ -443,7 +443,22 @@ def exchange_key_container(new_ssh_key: str): if running_container: # stop and remove the container by using the SIGTERM signal to PID 1 (init) process in the container if running_container.status == "running": - running_container.exec_run(cmd=f"bash -c \"echo '{new_ssh_key}' > /root/.ssh/authorized_keys & sync & sleep 1\"") + exist_key = running_container.exec_run(cmd="cat /root/.ssh/authorized_keys") + exist_key = exist_key.output.decode("utf-8").split("\n") + user_key = exist_key[0] + terminal_key = "" + if len(exist_key) > 1: + terminal_key = exist_key[1] + if key_type == "terminal": + terminal_key = new_ssh_key + elif key_type == "user": + user_key = new_ssh_key + else: + bt.logging.debug("Invalid key type to swap the SSH key") + return {"status": False} + key_list = user_key + "\n" + terminal_key + # bt.logging.debug(f"New SSH key: {key_list}") + running_container.exec_run(cmd=f"bash -c \"echo '{key_list}' > /root/.ssh/authorized_keys & sync & sleep 1\"") running_container.exec_run(cmd="kill -15 1") running_container.wait() running_container.restart() diff --git a/neurons/miner.py b/neurons/miner.py index 57b0656a..e53de45a 100644 --- a/neurons/miner.py +++ b/neurons/miner.py @@ -438,7 +438,8 @@ def allocate(self, synapse: Allocate) -> Allocate: if docker_action["action"] == "exchange_key": public_key = synapse.public_key new_ssh_key = docker_action["ssh_key"] - result = exchange_key_container(new_ssh_key) + key_type = docker_action["key_type"] + result = exchange_key_container(new_ssh_key, key_type) synapse.output = result elif docker_action["action"] == "restart": public_key = synapse.public_key diff --git a/neurons/register_api.py b/neurons/register_api.py index 0daff4e7..484fc718 100644 --- a/neurons/register_api.py +++ b/neurons/register_api.py @@ -524,12 +524,12 @@ async def allocate_hotkey(hotkey: str, ssh_key: Optional[str] = None, uuid_key = str(uuid.uuid1()) private_key, public_key = rsa.generate_key_pair() - if ssh_key: - if docker_requirement is None: - docker_requirement = DockerRequirement() - docker_requirement.ssh_key = ssh_key - else: - docker_requirement.ssh_key = ssh_key + if docker_requirement is None: + docker_requirement = DockerRequirement() + if ssh_key is None: + docker_requirement.ssh_key = "" + else: + docker_requirement.ssh_key = ssh_key run_start = time.time() result = await run_in_threadpool(self._allocate_container_hotkey, requirements, hotkey, @@ -1142,7 +1142,7 @@ async def unpause_docker(hotkey: str, uuid_key: str) -> JSONResponse: "description": "An error occurred while exchanging docker key.", }, }) - async def exchange_docker_key(hotkey: str, uuid_key: str, ssh_key: str) -> JSONResponse: + async def exchange_docker_key(hotkey: str, uuid_key: str, ssh_key: str, key_type: str = "user") -> JSONResponse: # Instantiate the connection to the db db = ComputeDb() cursor = db.get_cursor() @@ -1170,13 +1170,14 @@ async def exchange_docker_key(hotkey: str, uuid_key: str, ssh_key: str) -> JSONR docker_action = { "action": "exchange_key", "ssh_key": ssh_key, + "key_type": key_type, } if uuid_key_db == uuid_key: index = self.metagraph.hotkeys.index(hotkey) axon = self.metagraph.axons[index] run_start = time.time() - allocate_class = Allocate(timeline=0, device_requirement={}, checking=False, public_key=regkey, + allocate_class = Allocate(timeline=1, device_requirement={}, checking=False, public_key=regkey, docker_change=True, docker_action=docker_action) response = await run_in_threadpool( self.dendrite.query, axon, allocate_class, timeout=60