Remove redundant get_user_key()

Author: un-def

src/dstack/_internal/server/routers/runs.py

@@ -118,7 +118,7 @@ async def get_plan(
     """
     user, project = user_project
     if not user.ssh_public_key and not body.run_spec.ssh_key_pub:
-        await users.refresh_ssh_key(session=session, user=user, username=user.name)
+        await users.refresh_ssh_key(session=session, user=user)
     run_plan = await runs.get_plan(
         session=session,
         project=project,
@@ -148,7 +148,7 @@ async def apply_plan(
     """
     user, project = user_project
     if not user.ssh_public_key and not body.plan.run_spec.ssh_key_pub:
-        await users.refresh_ssh_key(session=session, user=user, username=user.name)
+        await users.refresh_ssh_key(session=session, user=user)
     return CustomORJSONResponse(
         await runs.apply_plan(
             session=session,

src/dstack/_internal/server/routers/users.py

@@ -43,10 +43,7 @@ async def get_my_user(
 ):
     if user.ssh_private_key is None or user.ssh_public_key is None:
         # Generate keys for pre-0.19.33 users
-        updated_user = await users.refresh_ssh_key(session=session, user=user, username=user.name)
-        if updated_user is None:
-            raise ResourceNotExistsError()
-        user = updated_user
+        await users.refresh_ssh_key(session=session, user=user)
     return CustomORJSONResponse(users.user_model_to_user_with_creds(user))
 
 

src/dstack/_internal/server/services/users.py

@@ -147,8 +147,10 @@ async def update_user(
 async def refresh_ssh_key(
     session: AsyncSession,
     user: UserModel,
-    username: str,
+    username: Optional[str] = None,
 ) -> Optional[UserModel]:
+    if username is None:
+        username = user.name
     logger.debug("Refreshing SSH key for user [code]%s[/code]", username)
     if user.global_role != GlobalRole.ADMIN and user.name != username:
         raise error_forbidden()

src/dstack/_internal/server/testing/common.py

@@ -277,6 +277,7 @@ def get_run_spec(
     configuration_path: str = "dstack.yaml",
     profile: Union[Profile, Callable[[], Profile], None] = lambda: Profile(name="default"),
     configuration: Optional[AnyRunConfiguration] = None,
+    ssh_key_pub: Optional[str] = "user_ssh_key",
 ) -> RunSpec:
     if callable(profile):
         profile = profile()
@@ -288,7 +289,7 @@ def get_run_spec(
         configuration_path=configuration_path,
         configuration=configuration or DevEnvironmentConfiguration(ide="vscode"),
         profile=profile,
-        ssh_key_pub="user_ssh_key",
+        ssh_key_pub=ssh_key_pub,
     )
 
 

src/dstack/api/_public/runs.py

@@ -504,10 +504,6 @@ def get_run_plan(
             ssh_key_pub = Path(ssh_identity_file).with_suffix(".pub").read_text()
         else:
             ssh_key_pub = None  # using the server-managed user key
-            config_manager = ConfigManager()
-            key_manager = UserSSHKeyManager(self._api_client, config_manager.dstack_ssh_dir)
-            # Ensure we have a fresh key locally
-            key_manager.get_user_key()
         run_spec = RunSpec(
             run_name=configuration.name,
             repo_id=repo.repo_id,

src/tests/_internal/server/routers/test_runs.py

@@ -1364,6 +1364,32 @@ async def test_returns_update_or_create_action_on_conf_change(
         assert response_json["action"] == action
         assert response_json["current_resource"] == json.loads(run.json())
 
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("test_db")
+    async def test_generates_user_ssh_key(self, session: AsyncSession, client: AsyncClient):
+        user = await create_user(
+            session=session, global_role=GlobalRole.USER, ssh_public_key=None, ssh_private_key=None
+        )
+        project = await create_project(session=session, owner=user)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        repo = await create_repo(session=session, project_id=project.id)
+        run_spec = get_run_spec(run_name="test-run", repo_id=repo.name, ssh_key_pub=None)
+
+        response = await client.post(
+            f"/api/project/{project.name}/runs/get_plan",
+            headers=get_auth_headers(user.token),
+            json={"run_spec": run_spec.dict()},
+        )
+
+        assert response.status_code == 200, response.json()
+        run_spec_ssh_public_key = response.json()["effective_run_spec"]["ssh_key_pub"]
+        assert run_spec_ssh_public_key is not None
+        await session.refresh(user)
+        assert user.ssh_public_key == run_spec_ssh_public_key
+        assert user.ssh_private_key is not None
+
 
 class TestApplyPlan:
     @pytest.mark.asyncio
@@ -1517,6 +1543,38 @@ async def test_creates_pending_run_if_run_is_scheduled(
         assert run.status == RunStatus.PENDING
         assert run.next_triggered_at == datetime(2023, 1, 2, 3, 10, tzinfo=timezone.utc)
 
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("test_db")
+    async def test_generates_user_ssh_key(self, session: AsyncSession, client: AsyncClient):
+        user = await create_user(
+            session=session, global_role=GlobalRole.USER, ssh_public_key=None, ssh_private_key=None
+        )
+        project = await create_project(session=session, owner=user)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        repo = await create_repo(session=session, project_id=project.id)
+        run_spec = get_run_spec(run_name="test-run", repo_id=repo.name, ssh_key_pub=None)
+
+        response = await client.post(
+            f"/api/project/{project.name}/runs/apply",
+            headers=get_auth_headers(user.token),
+            json={
+                "plan": {
+                    "run_spec": run_spec.dict(),
+                    "current_resource": None,
+                },
+                "force": False,
+            },
+        )
+
+        assert response.status_code == 200, response.json()
+        run_spec_ssh_public_key = response.json()["run_spec"]["ssh_key_pub"]
+        assert run_spec_ssh_public_key is not None
+        await session.refresh(user)
+        assert user.ssh_public_key == run_spec_ssh_public_key
+        assert user.ssh_private_key is not None
+
 
 class TestSubmitRun:
     @pytest.mark.asyncio