Repository: PrefectHQ/prefect

Issue:
Auto install packages for worker when worker type is not found
When a worker type is not found the users sees an error:

```
Unable to start worker of type 'docker'. Please ensure that you have installed this worker type on this machine. 
```

Using the `prefect-collection-registry` we can publish required packages for a worker and use that information to install the required packages on users' behalf.

Auto-installation will be handled either through a prompt or a flag that indicates that required packages should be auto-installed.

This feature will also be useful for running workers directly from the Prefect base Docker image.


Discussion:

Resolving patch:
diff --git a/src/prefect/cli/worker.py b/src/prefect/cli/worker.py
--- a/src/prefect/cli/worker.py
+++ b/src/prefect/cli/worker.py
@@ -1,14 +1,16 @@
 from functools import partial
-from typing import List, Optional
+from typing import List, Optional, Type
+from enum import Enum
 
 import os
+import sys
 import anyio
 import threading
 import typer
 
 from prefect.cli._types import PrefectTyper, SettingsOption
 from prefect.cli._utilities import exit_with_error
-from prefect.cli.root import app
+from prefect.cli.root import app, is_interactive
 from prefect.client.orchestration import get_client
 from prefect.exceptions import ObjectNotFound
 from prefect.settings import (
@@ -17,13 +19,15 @@
     PREFECT_WORKER_QUERY_SECONDS,
 )
 from prefect.utilities.dispatch import lookup_type
-from prefect.utilities.processutils import setup_signal_handlers_worker
+from prefect.utilities.processutils import run_process, setup_signal_handlers_worker
 from prefect.utilities.services import critical_service_loop
 from prefect.workers.base import BaseWorker
-from prefect.workers.process import ProcessWorker
 from prefect.workers.server import start_healthcheck_server
 from prefect.plugins import load_prefect_collections
 
+from prefect.client.collections import get_collections_metadata_client
+from prefect.cli._prompts import confirm
+
 
 worker_app = PrefectTyper(
     name="worker", help="Commands for starting and interacting with workers."
@@ -31,6 +35,13 @@
 app.add_typer(worker_app)
 
 
+class InstallPolicy(str, Enum):
+    ALWAYS = "always"
+    IF_NOT_PRESENT = "if-not-present"
+    NEVER = "never"
+    PROMPT = "prompt"
+
+
 @worker_app.command()
 async def start(
     worker_name: str = typer.Option(
@@ -43,7 +54,11 @@ async def start(
         ),
     ),
     work_pool_name: str = typer.Option(
-        ..., "-p", "--pool", help="The work pool the started worker should poll."
+        ...,
+        "-p",
+        "--pool",
+        help="The work pool the started worker should poll.",
+        prompt=True,
     ),
     work_queues: List[str] = typer.Option(
         None,
@@ -79,37 +94,22 @@ async def start(
     with_healthcheck: bool = typer.Option(
         False, help="Start a healthcheck server for the worker."
     ),
+    install_policy: InstallPolicy = typer.Option(
+        InstallPolicy.PROMPT.value,
+        "--install-policy",
+        help="Install policy to use workers from Prefect integration packages.",
+        case_sensitive=False,
+    ),
 ):
     """
     Start a worker process to poll a work pool for flow runs.
     """
-    try:
-        if worker_type is None:
-            async with get_client() as client:
-                work_pool = await client.read_work_pool(work_pool_name=work_pool_name)
-            worker_type = work_pool.type
-            app.console.print(
-                f"Discovered worker type {worker_type!r} for work pool"
-                f" {work_pool.name!r}."
-            )
-
-        load_prefect_collections()
-        worker_cls = lookup_type(BaseWorker, worker_type)
-    except KeyError:
-        # TODO: Use collection registry info to direct users on how to install the worker type
+    worker_cls = await _get_worker_class(worker_type, work_pool_name, install_policy)
+    if worker_cls is None:
         exit_with_error(
-            f"Unable to start worker of type {worker_type!r}. "
-            "Please ensure that you have installed this worker type on this machine."
-        )
-    except ObjectNotFound:
-        app.console.print(
-            (
-                f"Work pool {work_pool_name!r} does not exist and no worker type was"
-                " provided. Starting a process worker..."
-            ),
-            style="yellow",
+            "Unable to start worker. Please ensure you have the necessary dependencies"
+            " installed to run your desired worker type."
         )
-        worker_cls = ProcessWorker
 
     worker_process_id = os.getpid()
     setup_signal_handlers_worker(
@@ -123,7 +123,7 @@ async def start(
         limit=limit,
         prefetch_seconds=prefetch_seconds,
     ) as worker:
-        app.console.print(f"Worker {worker.name!r} started!")
+        app.console.print(f"Worker {worker.name!r} started!", style="green")
         async with anyio.create_task_group() as tg:
             # wait for an initial heartbeat to configure the worker
             await worker.sync_with_backend()
@@ -179,3 +179,102 @@ async def start(
 
     await worker._emit_worker_stopped_event(started_event)
     app.console.print(f"Worker {worker.name!r} stopped!")
+
+
+async def _retrieve_worker_type_from_pool(work_pool_name: Optional[str] = None) -> str:
+    try:
+        async with get_client() as client:
+            work_pool = await client.read_work_pool(work_pool_name=work_pool_name)
+        worker_type = work_pool.type
+        app.console.print(
+            f"Discovered worker type {worker_type!r} for work pool {work_pool.name!r}."
+        )
+    except ObjectNotFound:
+        app.console.print(
+            (
+                f"Work pool {work_pool_name!r} does not exist and no worker type was"
+                " provided. Starting a process worker..."
+            ),
+            style="yellow",
+        )
+        worker_type = "process"
+    return worker_type
+
+
+def _load_worker_class(worker_type: str) -> Optional[Type[BaseWorker]]:
+    try:
+        load_prefect_collections()
+        return lookup_type(BaseWorker, worker_type)
+    except KeyError:
+        return None
+
+
+async def _install_package(
+    package: str, upgrade: bool = False
+) -> Optional[Type[BaseWorker]]:
+    app.console.print(f"Installing {package}...")
+    command = [sys.executable, "-m", "pip", "install", package]
+    if upgrade:
+        command.append("--upgrade")
+    await run_process(command, stream_output=True)
+
+
+async def _find_package_for_worker_type(worker_type: str) -> Optional[str]:
+    async with get_collections_metadata_client() as client:
+        worker_metadata = await client.read_worker_metadata()
+
+    worker_types_with_packages = {
+        worker_type: package_name
+        for package_name, worker_dict in worker_metadata.items()
+        for worker_type in worker_dict
+        if worker_type != "prefect-agent"
+    }
+
+    return worker_types_with_packages[worker_type]
+
+
+async def _get_worker_class(
+    worker_type: Optional[str] = None,
+    work_pool_name: Optional[str] = None,
+    install_policy: InstallPolicy = InstallPolicy.PROMPT,
+) -> Optional[Type[BaseWorker]]:
+    if worker_type is None and work_pool_name is None:
+        raise ValueError("Must provide either worker_type or work_pool_name.")
+
+    if worker_type is None:
+        worker_type = await _retrieve_worker_type_from_pool(work_pool_name)
+
+    if install_policy == InstallPolicy.ALWAYS:
+        package = await _find_package_for_worker_type(worker_type)
+        if package:
+            await _install_package(package, upgrade=True)
+            worker_cls = _load_worker_class(worker_type)
+
+    worker_cls = _load_worker_class(worker_type)
+
+    if worker_cls is None:
+        package = await _find_package_for_worker_type(worker_type)
+        # Check if the package exists
+        if package:
+            # Prompt to install if the package is not present
+            if install_policy == InstallPolicy.IF_NOT_PRESENT:
+                should_install = True
+
+            # Confirm with the user for installation in an interactive session
+            elif install_policy == InstallPolicy.PROMPT and is_interactive():
+                message = (
+                    f"Could not find a {worker_type} worker in the current"
+                    " environment. Install it now?"
+                )
+                should_install = confirm(message, default=True)
+
+            # If none of the conditions met, don't install the package
+            else:
+                should_install = False
+
+            # If should_install is True, install the package
+            if should_install:
+                await _install_package(package)
+                worker_cls = _load_worker_class(worker_type)
+
+    return worker_cls


Test patch:
diff --git a/tests/cli/test_worker.py b/tests/cli/test_worker.py
--- a/tests/cli/test_worker.py
+++ b/tests/cli/test_worker.py
@@ -5,6 +5,7 @@
 from unittest.mock import ANY
 
 import anyio
+import httpx
 import pytest
 
 import prefect
@@ -15,9 +16,82 @@
     get_current_settings,
 )
 from prefect.testing.cli import invoke_and_assert
-from prefect.testing.utilities import MagicMock
+from prefect.testing.utilities import AsyncMock, MagicMock
 from prefect.utilities.asyncutils import run_sync_in_worker_thread
 from prefect.utilities.processutils import open_process
+from prefect.settings import PREFECT_API_URL
+
+from prefect.client.schemas.actions import WorkPoolCreate
+import readchar
+import respx
+
+from typer import Exit
+
+from prefect.workers.base import BaseJobConfiguration, BaseWorker
+
+
+class MockKubernetesWorker(BaseWorker):
+    type = "kubernetes"
+    job_configuration = BaseJobConfiguration
+
+    async def run(self):
+        pass
+
+    async def kill_infrastructure(self, *args, **kwargs):
+        pass
+
+
+@pytest.fixture
+def interactive_console(monkeypatch):
+    monkeypatch.setattr("prefect.cli.worker.is_interactive", lambda: True)
+
+    # `readchar` does not like the fake stdin provided by typer isolation so we provide
+    # a version that does not require a fd to be attached
+    def readchar():
+        sys.stdin.flush()
+        position = sys.stdin.tell()
+        if not sys.stdin.read():
+            print("TEST ERROR: CLI is attempting to read input but stdin is empty.")
+            raise Exit(-2)
+        else:
+            sys.stdin.seek(position)
+        return sys.stdin.read(1)
+
+    monkeypatch.setattr("readchar._posix_read.readchar", readchar)
+
+
+@pytest.fixture
+async def kubernetes_work_pool(prefect_client: PrefectClient):
+    work_pool = await prefect_client.create_work_pool(
+        work_pool=WorkPoolCreate(name="test-k8s-work-pool", type="kubernetes")
+    )
+
+    with respx.mock(
+        assert_all_mocked=False, base_url=PREFECT_API_URL.value()
+    ) as respx_mock:
+        respx_mock.route(path__startswith="/work_pools/").pass_through()
+        respx_mock.route(path__startswith="/flow_runs/").pass_through()
+        respx_mock.get("/collections/views/aggregate-worker-metadata").mock(
+            return_value=httpx.Response(
+                200,
+                json={
+                    "prefect": {
+                        "prefect-agent": {
+                            "type": "prefect-agent",
+                            "default_base_job_configuration": {},
+                        }
+                    },
+                    "prefect-kubernetes": {
+                        "kubernetes": {
+                            "type": "kubernetes",
+                            "default_base_job_configuration": {},
+                        }
+                    },
+                },
+            )
+        )
+
+        yield work_pool
 
 
 @pytest.mark.usefixtures("use_hosted_api_server")
@@ -262,6 +336,204 @@ async def test_start_worker_without_type_creates_process_work_pool(
     assert workers[0].name == "test-worker"
 
 
+@pytest.mark.usefixtures("use_hosted_api_server")
+class TestInstallPolicyOption:
+    async def test_install_policy_if_not_present(
+        self, kubernetes_work_pool, monkeypatch
+    ):
+        run_process_mock = AsyncMock()
+        lookup_type_mock = MagicMock()
+        lookup_type_mock.side_effect = [KeyError, MockKubernetesWorker]
+        monkeypatch.setattr("prefect.cli.worker.run_process", run_process_mock)
+        monkeypatch.setattr("prefect.cli.worker.lookup_type", lookup_type_mock)
+        await run_sync_in_worker_thread(
+            invoke_and_assert,
+            command=[
+                "worker",
+                "start",
+                "--run-once",
+                "-p",
+                kubernetes_work_pool.name,
+                "-n",
+                "test-worker",
+                "--install-policy=if-not-present",
+            ],
+            expected_output_contains=[
+                "Installing prefect-kubernetes...",
+                "Worker 'test-worker' started!",
+                "Worker 'test-worker' stopped!",
+            ],
+        )
+
+        run_process_mock.assert_called_once_with(
+            [sys.executable, "-m", "pip", "install", "prefect-kubernetes"],
+            stream_output=True,
+        )
+
+    @pytest.mark.usefixtures("interactive_console")
+    async def test_install_policy_prompt(self, kubernetes_work_pool, monkeypatch):
+        run_process_mock = AsyncMock()
+        lookup_type_mock = MagicMock()
+        lookup_type_mock.side_effect = [KeyError, MockKubernetesWorker]
+        monkeypatch.setattr("prefect.cli.worker.run_process", run_process_mock)
+        monkeypatch.setattr("prefect.cli.worker.lookup_type", lookup_type_mock)
+        await run_sync_in_worker_thread(
+            invoke_and_assert,
+            command=[
+                "worker",
+                "start",
+                "--run-once",
+                "-p",
+                kubernetes_work_pool.name,
+                "-n",
+                "test-worker",
+            ],
+            user_input=readchar.key.ENTER,
+            expected_output_contains=[
+                (
+                    "Could not find a kubernetes worker in the current"
+                    " environment. Install it now?"
+                ),
+                "Installing prefect-kubernetes...",
+                "Worker 'test-worker' started!",
+                "Worker 'test-worker' stopped!",
+            ],
+        )
+
+        run_process_mock.assert_called_once_with(
+            [sys.executable, "-m", "pip", "install", "prefect-kubernetes"],
+            stream_output=True,
+        )
+
+    @pytest.mark.usefixtures("interactive_console")
+    async def test_install_policy_prompt_decline(self, monkeypatch, prefect_client):
+        run_process_mock = AsyncMock()
+        lookup_type_mock = MagicMock()
+        lookup_type_mock.side_effect = [KeyError, MockKubernetesWorker]
+        monkeypatch.setattr("prefect.cli.worker.run_process", run_process_mock)
+        monkeypatch.setattr("prefect.cli.worker.lookup_type", lookup_type_mock)
+        kubernetes_work_pool = await prefect_client.create_work_pool(
+            work_pool=WorkPoolCreate(name="test-k8s-work-pool", type="kubernetes")
+        )
+
+        await run_sync_in_worker_thread(
+            invoke_and_assert,
+            command=[
+                "worker",
+                "start",
+                "--run-once",
+                "-p",
+                kubernetes_work_pool.name,
+                "-n",
+                "test-worker",
+            ],
+            expected_code=1,
+            user_input="n" + readchar.key.ENTER,
+            expected_output_contains=[
+                "Unable to start worker. Please ensure you have the necessary"
+                " dependencies installed to run your desired worker type."
+            ],
+        )
+
+        run_process_mock.assert_not_called()
+
+    @pytest.mark.usefixtures("interactive_console")
+    async def test_install_policy_if_not_present_overrides_prompt(
+        self, kubernetes_work_pool, monkeypatch
+    ):
+        run_process_mock = AsyncMock()
+        lookup_type_mock = MagicMock()
+        lookup_type_mock.side_effect = [KeyError, MockKubernetesWorker]
+        monkeypatch.setattr("prefect.cli.worker.run_process", run_process_mock)
+        monkeypatch.setattr("prefect.cli.worker.lookup_type", lookup_type_mock)
+        await run_sync_in_worker_thread(
+            invoke_and_assert,
+            command=[
+                "worker",
+                "start",
+                "--run-once",
+                "-p",
+                kubernetes_work_pool.name,
+                "-n",
+                "test-worker",
+                "--install-policy=if-not-present",
+            ],
+            expected_output_contains=[
+                "Installing prefect-kubernetes...",
+                "Worker 'test-worker' started!",
+                "Worker 'test-worker' stopped!",
+            ],
+        )
+
+        run_process_mock.assert_called_once_with(
+            [sys.executable, "-m", "pip", "install", "prefect-kubernetes"],
+            stream_output=True,
+        )
+
+    @pytest.mark.usefixtures("interactive_console")
+    async def test_install_policy_always(self, kubernetes_work_pool, monkeypatch):
+        run_process_mock = AsyncMock()
+        lookup_type_mock = MagicMock()
+        lookup_type_mock.return_value = MockKubernetesWorker
+        monkeypatch.setattr("prefect.cli.worker.run_process", run_process_mock)
+        monkeypatch.setattr("prefect.cli.worker.lookup_type", lookup_type_mock)
+        await run_sync_in_worker_thread(
+            invoke_and_assert,
+            command=[
+                "worker",
+                "start",
+                "--run-once",
+                "-p",
+                kubernetes_work_pool.name,
+                "-n",
+                "test-worker",
+                "--install-policy=always",
+            ],
+            expected_output_contains=[
+                "Installing prefect-kubernetes...",
+                "Worker 'test-worker' started!",
+                "Worker 'test-worker' stopped!",
+            ],
+        )
+
+        run_process_mock.assert_called_once_with(
+            [sys.executable, "-m", "pip", "install", "prefect-kubernetes", "--upgrade"],
+            stream_output=True,
+        )
+
+    @pytest.mark.usefixtures("interactive_console")
+    async def test_install_policy_never(self, monkeypatch, prefect_client):
+        kubernetes_work_pool = await prefect_client.create_work_pool(
+            work_pool=WorkPoolCreate(name="test-k8s-work-pool", type="kubernetes")
+        )
+
+        run_process_mock = AsyncMock()
+        lookup_type_mock = MagicMock()
+        lookup_type_mock.side_effect = KeyError
+        monkeypatch.setattr("prefect.cli.worker.run_process", run_process_mock)
+        monkeypatch.setattr("prefect.cli.worker.lookup_type", lookup_type_mock)
+        await run_sync_in_worker_thread(
+            invoke_and_assert,
+            command=[
+                "worker",
+                "start",
+                "--run-once",
+                "-p",
+                kubernetes_work_pool.name,
+                "-n",
+                "test-worker",
+                "--install-policy=never",
+            ],
+            expected_code=1,
+            expected_output_contains=[
+                "Unable to start worker. Please ensure you have the necessary"
+                " dependencies installed to run your desired worker type."
+            ],
+        )
+
+        run_process_mock.assert_not_called()
+
+
 POLL_INTERVAL = 0.5
 STARTUP_TIMEOUT = 20
 SHUTDOWN_TIMEOUT = 5