Repository: PrefectHQ/prefect

Issue:
Add `--api` Flag Documentation to Prefect CLI and Docs
## Current behavior
*Please describe how the feature works today*
The `prefect agent start`/`prefect agent install` commands have an `--api` flag that a user can pass to specify the endpoint the agent should poll for Flow Runs. This flag used to be for internal use only, but now that Prefect Server has been open-sourced, our users find it useful as well!


## Proposed behavior
*Please describe your proposed change to the current behavior*

We should document this flag both in:

- [ ] Our Documentation
- [ ] The CLI help commands for `prefect agent start` and `prefect agent install`


## Example
*Please give an example of how the enhancement would be useful*
Addressed above.


Discussion:
Actually, looks like there’s an `--api` on the install but not on the start. Currently, start only reads from config. Should we consider updating the `start` command to respect this flag as well?
Yeah we should, I sometimes find myself updating the value in my config but the flag would be nice 👍 

Resolving patch:
diff --git a/src/prefect/agent/kubernetes/agent.py b/src/prefect/agent/kubernetes/agent.py
--- a/src/prefect/agent/kubernetes/agent.py
+++ b/src/prefect/agent/kubernetes/agent.py
@@ -175,6 +175,14 @@ def replace_job_spec_yaml(self, flow_run: GraphQLResult, image: str) -> dict:
             resources["requests"]["cpu"] = os.getenv("JOB_CPU_REQUEST")
         if os.getenv("JOB_CPU_LIMIT"):
             resources["limits"]["cpu"] = os.getenv("JOB_CPU_LIMIT")
+        if os.getenv("IMAGE_PULL_POLICY"):
+            job["spec"]["template"]["spec"]["containers"][0][
+                "imagePullPolicy"
+            ] = os.getenv("IMAGE_PULL_POLICY")
+        if os.getenv("SERVICE_ACCOUNT_NAME"):
+            job["spec"]["template"]["spec"]["serviceAccountName"] = os.getenv(
+                "SERVICE_ACCOUNT_NAME"
+            )
 
         return job
 
@@ -191,6 +199,8 @@ def generate_deployment_yaml(
         mem_limit: str = None,
         cpu_request: str = None,
         cpu_limit: str = None,
+        image_pull_policy: str = None,
+        service_account_name: str = None,
         labels: Iterable[str] = None,
         env_vars: dict = None,
         backend: str = None,
@@ -216,6 +226,10 @@ def generate_deployment_yaml(
             - mem_limit (str, optional): Limit memory for Prefect init job.
             - cpu_request (str, optional): Requested CPU for Prefect init job.
             - cpu_limit (str, optional): Limit CPU for Prefect init job.
+            - image_pull_policy (str, optional): imagePullPolicy to use for Prefect init job.
+                Job defaults to `IfNotPresent`.
+            - service_account_name (str, optional): Name of a service account to use for
+                Prefect init job. Job defaults to using `default` service account.
             - labels (List[str], optional): a list of labels, which are arbitrary string
                 identifiers used by Prefect Agents when polling for work
             - env_vars (dict, optional): additional environment variables to attach to all
@@ -236,6 +250,8 @@ def generate_deployment_yaml(
         mem_limit = mem_limit or ""
         cpu_request = cpu_request or ""
         cpu_limit = cpu_limit or ""
+        image_pull_policy = image_pull_policy or ""
+        service_account_name = service_account_name or ""
         backend = backend or config.backend
 
         version = prefect.__version__.split("+")
@@ -256,13 +272,15 @@ def generate_deployment_yaml(
         agent_env[2]["value"] = namespace
         agent_env[3]["value"] = image_pull_secrets or ""
         agent_env[4]["value"] = str(labels)
-        agent_env[9]["value"] = backend
+        agent_env[11]["value"] = backend
 
         # Populate job resource env vars
         agent_env[5]["value"] = mem_request
         agent_env[6]["value"] = mem_limit
         agent_env[7]["value"] = cpu_request
         agent_env[8]["value"] = cpu_limit
+        agent_env[9]["value"] = image_pull_policy
+        agent_env[10]["value"] = service_account_name
 
         if env_vars:
             for k, v in env_vars.items():
diff --git a/src/prefect/cli/agent.py b/src/prefect/cli/agent.py
--- a/src/prefect/cli/agent.py
+++ b/src/prefect/cli/agent.py
@@ -49,6 +49,7 @@ def agent():
 @click.option(
     "--token", "-t", required=False, help="A Prefect Cloud API token.", hidden=True
 )
+@click.option("--api", "-a", required=False, help="A Prefect API URL.", hidden=True)
 @click.option(
     "--name",
     "-n",
@@ -137,6 +138,7 @@ def start(
     ctx,
     agent_option,
     token,
+    api,
     name,
     verbose,
     label,
@@ -164,6 +166,7 @@ def start(
     \b
     Options:
         --token, -t     TEXT    A Prefect Cloud API token with RUNNER scope
+        --api, -a       TEXT    A Prefect API URL
         --name, -n      TEXT    A name to use for the agent
         --verbose, -v           Enable verbose agent DEBUG logs
                                 Defaults to INFO level logging
@@ -221,6 +224,8 @@ def start(
     }
     if verbose:
         tmp_config["cloud.agent.level"] = "DEBUG"
+    if api:
+        tmp_config["cloud.api"] = api
 
     with set_temporary_config(tmp_config):
         retrieved_agent = _agents.get(agent_option, None)
@@ -335,6 +340,18 @@ def start(
 @click.option(
     "--cpu-limit", required=False, help="Limit CPU for Prefect init job.", hidden=True
 )
+@click.option(
+    "--image-pull-policy",
+    required=False,
+    help="imagePullPolicy for Prefect init job",
+    hidden=True,
+)
+@click.option(
+    "--service-account-name",
+    required=False,
+    help="Name of Service Account for Prefect init job",
+    hidden=True,
+)
 @click.option(
     "--label",
     "-l",
@@ -383,6 +400,8 @@ def install(
     mem_limit,
     cpu_request,
     cpu_limit,
+    image_pull_policy,
+    service_account_name,
     label,
     env,
     import_path,
@@ -418,6 +437,8 @@ def install(
         --mem-limit                 TEXT    Limit memory for Prefect init job
         --cpu-request               TEXT    Requested CPU for Prefect init job
         --cpu-limit                 TEXT    Limit CPU for Prefect init job
+        --image-pull-policy         TEXT    imagePullPolicy for Prefect init job
+        --service-account-name      TEXT    Name of Service Account for Prefect init job
         --backend                   TEST    Prefect backend to use for this agent
                                             Defaults to the backend currently set in config.
 
@@ -457,6 +478,8 @@ def install(
             mem_limit=mem_limit,
             cpu_request=cpu_request,
             cpu_limit=cpu_limit,
+            image_pull_policy=image_pull_policy,
+            service_account_name=service_account_name,
             labels=list(label),
             env_vars=env_vars,
             backend=backend,


Test patch:
diff --git a/tests/agent/test_k8s_agent.py b/tests/agent/test_k8s_agent.py
--- a/tests/agent/test_k8s_agent.py
+++ b/tests/agent/test_k8s_agent.py
@@ -156,6 +156,8 @@ def test_k8s_agent_replace_yaml_uses_user_env_vars(
     monkeypatch.setenv("JOB_MEM_LIMIT", "ml")
     monkeypatch.setenv("JOB_CPU_REQUEST", "cr")
     monkeypatch.setenv("JOB_CPU_LIMIT", "cl")
+    monkeypatch.setenv("IMAGE_PULL_POLICY", "custom_policy")
+    monkeypatch.setenv("SERVICE_ACCOUNT_NAME", "svc_name")
 
     flow_run = GraphQLResult(
         {
@@ -205,6 +207,12 @@ def test_k8s_agent_replace_yaml_uses_user_env_vars(
         assert env[-1] in user_vars
         assert env[-2] in user_vars
 
+        assert (
+            job["spec"]["template"]["spec"]["containers"][0]["imagePullPolicy"]
+            == "custom_policy"
+        )
+        assert job["spec"]["template"]["spec"]["serviceAccountName"] == "svc_name"
+
 
 def test_k8s_agent_replace_yaml(monkeypatch, runner_token, cloud_api):
     k8s_config = MagicMock()
@@ -268,6 +276,12 @@ def test_k8s_agent_replace_yaml(monkeypatch, runner_token, cloud_api):
         assert resources["requests"]["cpu"] == "cr"
         assert resources["limits"]["cpu"] == "cl"
 
+        assert (
+            job["spec"]["template"]["spec"]["containers"][0]["imagePullPolicy"]
+            == "IfNotPresent"
+        )
+        assert job["spec"]["template"]["spec"].get("serviceAccountName", None) is None
+
 
 @pytest.mark.parametrize("flag", [True, False])
 def test_k8s_agent_replace_yaml_responds_to_logging_config(
@@ -371,7 +385,7 @@ def test_k8s_agent_generate_deployment_yaml(monkeypatch, runner_token):
     assert agent_env[0]["value"] == "test_token"
     assert agent_env[1]["value"] == "test_api"
     assert agent_env[2]["value"] == "test_namespace"
-    assert agent_env[9]["value"] == "backend-test"
+    assert agent_env[11]["value"] == "backend-test"
 
     assert resource_manager_env[0]["value"] == "test_token"
     assert resource_manager_env[1]["value"] == "test_api"
@@ -391,10 +405,10 @@ def test_k8s_agent_generate_deployment_yaml_env_vars(monkeypatch, runner_token):
 
     agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
 
-    assert agent_env[11]["name"] == "PREFECT__CLOUD__AGENT__ENV_VARS__test1"
-    assert agent_env[11]["value"] == "test2"
-    assert agent_env[12]["name"] == "PREFECT__CLOUD__AGENT__ENV_VARS__test3"
-    assert agent_env[12]["value"] == "test4"
+    assert agent_env[13]["name"] == "PREFECT__CLOUD__AGENT__ENV_VARS__test1"
+    assert agent_env[13]["value"] == "test2"
+    assert agent_env[14]["name"] == "PREFECT__CLOUD__AGENT__ENV_VARS__test3"
+    assert agent_env[14]["value"] == "test4"
 
 
 def test_k8s_agent_generate_deployment_yaml_backend_default(monkeypatch, server_api):
@@ -408,7 +422,7 @@ def test_k8s_agent_generate_deployment_yaml_backend_default(monkeypatch, server_
 
     agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
 
-    assert agent_env[9]["value"] == "server"
+    assert agent_env[11]["value"] == "server"
 
 
 @pytest.mark.parametrize(
@@ -569,6 +583,8 @@ def test_k8s_agent_generate_deployment_yaml_contains_resources(
         mem_limit="ml",
         cpu_request="cr",
         cpu_limit="cl",
+        image_pull_policy="custom_policy",
+        service_account_name="svc",
     )
 
     deployment = yaml.safe_load(deployment)
@@ -579,6 +595,8 @@ def test_k8s_agent_generate_deployment_yaml_contains_resources(
     assert env[6]["value"] == "ml"
     assert env[7]["value"] == "cr"
     assert env[8]["value"] == "cl"
+    assert env[9]["value"] == "custom_policy"
+    assert env[10]["value"] == "svc"
 
 
 def test_k8s_agent_generate_deployment_yaml_rbac(monkeypatch, runner_token):
diff --git a/tests/cli/test_agent.py b/tests/cli/test_agent.py
--- a/tests/cli/test_agent.py
+++ b/tests/cli/test_agent.py
@@ -60,6 +60,15 @@ def test_docker_agent_start_verbose(monkeypatch, runner_token):
     assert result.exit_code == 0
 
 
+def test_agent_start_api(monkeypatch, runner_token):
+    start = MagicMock()
+    monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start)
+
+    runner = CliRunner()
+    result = runner.invoke(agent, ["start", "local", "--api", "test_api"])
+    assert result.exit_code == 0
+
+
 def test_agent_start_local(monkeypatch, runner_token):
     start = MagicMock()
     monkeypatch.setattr("prefect.agent.local.LocalAgent.start", start)
@@ -357,6 +366,10 @@ def test_agent_install_k8s_asses_args():
             "cpu_req",
             "--cpu-limit",
             "cpu_limt",
+            "--image-pull-policy",
+            "custom_policy",
+            "--service-account-name",
+            "svc_name",
             "--label",
             "test_label1",
             "-l",
@@ -380,6 +393,8 @@ def test_agent_install_k8s_asses_args():
     assert "mem_lim" in result.output
     assert "cpu_req" in result.output
     assert "cpu_lim" in result.output
+    assert "custom_policy" in result.output
+    assert "svc_name" in result.output
     assert "secret-test" in result.output
     assert "test_label1" in result.output
     assert "test_label2" in result.output