Esempio n. 1
0
def test_populate_job_yaml_multiple_containers():
    with tempfile.TemporaryDirectory() as directory:

        with open(os.path.join(directory, "job.yaml"), "w+") as file:
            file.write("job")

        environment = KubernetesJobEnvironment(
            job_spec_file=os.path.join(directory, "job.yaml"))

        file_path = os.path.dirname(
            prefect.environments.execution.dask.k8s.__file__)

        with open(path.join(file_path, "job.yaml")) as job_file:
            job = yaml.safe_load(job_file)
            # Generate yaml object with multiple containers
            job["spec"]["template"]["spec"]["containers"][0]["env"] = []
            job["spec"]["template"]["spec"]["containers"].append(
                copy.deepcopy(
                    job["spec"]["template"]["spec"]["containers"][0]))
            job["spec"]["template"]["spec"]["containers"][1]["env"] = []

        with set_temporary_config({
                "cloud.graphql": "gql_test",
                "cloud.auth_token": "auth_test"
        }):
            with prefect.context(flow_run_id="id_test",
                                 namespace="namespace_test"):
                yaml_obj = environment._populate_job_spec_yaml(
                    yaml_obj=job,
                    docker_name="test1/test2:test3",
                    flow_file_path="test4",
                )

        assert (yaml_obj["metadata"]["labels"]["identifier"] ==
                environment.identifier_label)
        assert yaml_obj["metadata"]["labels"]["flow_run_id"] == "id_test"
        assert (yaml_obj["spec"]["template"]["metadata"]["labels"]
                ["identifier"] == environment.identifier_label)

        # Assert First Container
        env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"]

        assert env[0]["value"] == "gql_test"
        assert env[1]["value"] == "auth_test"
        assert env[2]["value"] == "id_test"
        assert env[3]["value"] == "namespace_test"
        assert env[4]["value"] == "test1/test2:test3"
        assert env[5]["value"] == "test4"

        assert (yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"]
                == "test1/test2:test3")

        assert yaml_obj["spec"]["template"]["spec"]["containers"][0][
            "command"] == [
                "/bin/sh",
                "-c",
            ]
        assert yaml_obj["spec"]["template"]["spec"]["containers"][0]["args"] == [
            "python -c 'from prefect.environments import KubernetesJobEnvironment; KubernetesJobEnvironment().run_flow()'"
        ]

        # Assert Second Container
        env = yaml_obj["spec"]["template"]["spec"]["containers"][1]["env"]

        assert env[0]["value"] == "gql_test"
        assert env[1]["value"] == "auth_test"
        assert env[2]["value"] == "id_test"
        assert env[3]["value"] == "namespace_test"
        assert env[4]["value"] == "test1/test2:test3"
        assert env[5]["value"] == "test4"

        assert (yaml_obj["spec"]["template"]["spec"]["containers"][1]["image"]
                != "test1/test2:test3")

        assert yaml_obj["spec"]["template"]["spec"]["containers"][1]["args"] != [
            "python -c 'from prefect.environments import KubernetesJobEnvironment; KubernetesJobEnvironment().run_flow()'"
        ]
Esempio n. 2
0
def test_agent_config_options():
    with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
        agent = Agent()
        assert agent.client.get_auth_token() == "TEST_TOKEN"
        assert agent.logger
Esempio n. 3
0
def test_entire_environment_process_together(monkeypatch):
    boto3_client = MagicMock()
    boto3_client.describe_task_definition.side_effect = ClientError({}, None)
    boto3_client.register_task_definition.return_value = {}
    boto3_client.run_task.return_value = {}
    monkeypatch.setattr("boto3.client", MagicMock(return_value=boto3_client))

    flow_runner = MagicMock()
    monkeypatch.setattr(
        "prefect.engine.get_default_flow_runner_class",
        MagicMock(return_value=flow_runner),
    )

    monkeypatch.setenv("AWS_ACCESS_KEY_ID", "id")
    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret")
    monkeypatch.setenv("REGION_NAME", "region")

    with prefect.context({"flow_run_id": "id"}):

        storage = Docker(registry_url="test",
                         image_name="image",
                         image_tag="tag")

        environment = FargateTaskEnvironment(
            containerDefinitions=[{
                "name": "flow",
                "image": "image",
                "command": [],
                "environment": [],
                "essential": True,
            }],
            cluster="test",
        )

        assert environment
        assert environment.aws_access_key_id == "id"
        assert environment.aws_secret_access_key == "secret"
        assert environment.region_name == "region"

        environment.setup(storage=storage)

        assert boto3_client.describe_task_definition.called
        assert boto3_client.register_task_definition.called
        assert (boto3_client.register_task_definition.call_args[1]["family"] ==
                "prefect-task-id-custom")
        assert boto3_client.register_task_definition.call_args[1][
            "containerDefinitions"] == [{
                "name":
                "flow",
                "image":
                "test/image:tag",
                "command": [
                    "/bin/sh",
                    "-c",
                    "python -c 'from prefect.environments import FargateTaskEnvironment; FargateTaskEnvironment().run_flow()'",
                ],
                "environment": [
                    {
                        "name": "PREFECT__CLOUD__GRAPHQL",
                        "value": prefect.config.cloud.graphql,
                    },
                    {
                        "name": "PREFECT__CLOUD__USE_LOCAL_SECRETS",
                        "value": "false"
                    },
                    {
                        "name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
                        "value": "prefect.engine.cloud.CloudFlowRunner",
                    },
                    {
                        "name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
                        "value": "prefect.engine.cloud.CloudTaskRunner",
                    },
                    {
                        "name": "PREFECT__LOGGING__LOG_TO_CLOUD",
                        "value": "true"
                    },
                ],
                "essential":
                True,
            }]

        environment.execute(storage=storage, flow_location=".prefect/flows")

        assert boto3_client.run_task.called
        assert (boto3_client.run_task.call_args[1]["taskDefinition"] ==
                "prefect-task-id-custom")
        assert boto3_client.run_task.call_args[1]["overrides"] == {
            "containerOverrides": [{
                "name":
                "flow",
                "environment": [
                    {
                        "name": "PREFECT__CLOUD__AUTH_TOKEN",
                        "value": prefect.config.cloud.agent.auth_token,
                    },
                    {
                        "name": "PREFECT__CONTEXT__FLOW_RUN_ID",
                        "value": "id"
                    },
                    {
                        "name": "PREFECT__CONTEXT__IMAGE",
                        "value": "test/image:tag"
                    },
                    {
                        "name": "PREFECT__CONTEXT__FLOW_FILE_PATH",
                        "value": ".prefect/flows",
                    },
                ],
            }]
        }
        assert boto3_client.run_task.call_args[1]["launchType"] == "FARGATE"
        assert boto3_client.run_task.call_args[1]["cluster"] == "test"

        with tempfile.TemporaryDirectory() as directory:
            with open(os.path.join(directory, "flow_env.prefect"), "w+"):
                flow = prefect.Flow("test")
                flow_path = os.path.join(directory, "flow_env.prefect")
                with open(flow_path, "wb") as f:
                    cloudpickle.dump(flow, f)

            with set_temporary_config({"cloud.auth_token": "test"}):
                with prefect.context(flow_file_path=os.path.join(
                        directory, "flow_env.prefect")):
                    environment.run_flow()

            assert flow_runner.call_args[1]["flow"].name == "test"
Esempio n. 4
0
def test_agent_env_vars(cloud_api):
    with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
        agent = Agent(env_vars=dict(AUTH_THING="foo"))
        assert agent.env_vars == dict(AUTH_THING="foo")
Esempio n. 5
0
def test_agent_labels(cloud_api):
    with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
        agent = Agent(labels=["test", "2"])
        assert agent.labels == ["test", "2"]
Esempio n. 6
0
def test_default_storage_ignores_bad_config():
    with set_temporary_config({"flows.defaults.storage.default_class": "FOOBAR"}):

        with pytest.warns(UserWarning):
            assert get_default_storage_class() is Local
Esempio n. 7
0
def test_agent_labels_from_config_var(cloud_api):
    with set_temporary_config({"cloud.agent.labels": ["test", "2"]}):
        agent = Agent()
        assert agent.labels == ["test", "2"]
Esempio n. 8
0
def test_validate_edges():
    with set_temporary_config({"flows.eager_edge_validation": True}):
        f = Flow(name="test")
        t1, t2 = Task(), Task()  # these tasks don't support keyed edges
        with pytest.raises(TypeError):
            f.add_edge(t1, t2, key="x")
Esempio n. 9
0
def start(
    ctx,
    agent_option,
    token,
    name,
    verbose,
    label,
    env,
    namespace,
    no_pull,
    no_cloud_logs,
    base_url,
    import_path,
    show_flow_logs,
    volume,
    max_polls,
):
    """
    Start an agent.

    \b
    Arguments:
        agent-option    TEXT    The name of an agent to start (e.g. `docker`, `kubernetes`, `local`, `fargate`, `nomad`)
                                Defaults to `local`

    \b
    Options:
        --token, -t     TEXT    A Prefect Cloud API token with RUNNER scope
        --name, -n      TEXT    A name to use for the agent
        --verbose, -v           Enable verbose agent DEBUG logs
                                Defaults to INFO level logging
        --label, -l     TEXT    Labels the agent will use to query for flow runs
                                Multiple values supported e.g. `-l label1 -l label2`
        --env, -e       TEXT    Environment variables to set on each submitted flow run.
                                Note that equal signs in environment variable values are not currently supported from the CLI.
                                Multiple values supported e.g. `-e AUTH=token -e PKG_SETTING=true`
        --max-polls     INT     Maximum number of times the agent should poll the Prefect API for flow runs. Will run forever
                                if not specified.
        --no-cloud-logs         Turn off logging to the Prefect API for all flow runs
                                Defaults to `False`

    \b
    Local Agent Options:
        --import-path, -p   TEXT    Import paths which will be provided to each Flow's runtime environment.
                                    Used for Flows which might import from scripts or local packages.
                                    Multiple values supported e.g. `-p /root/my_scripts -p /utilities`
        --show-flow-logs, -f        Display logging output from flows run by the agent (available for Local and Docker agents only)

    \b
    Docker Agent Options:
        --base-url, -b  TEXT    A Docker daemon host URL for a DockerAgent
        --no-pull               Pull images for a DockerAgent
                                Defaults to pulling if not provided
        --volume        TEXT    Host paths for Docker bind mount volumes attached to each Flow runtime container.
                                Multiple values supported e.g. `--volume /some/path --volume /some/other/path`

    \b
    Kubernetes Agent Options:
        --namespace     TEXT    A Kubernetes namespace to create Prefect jobs in
                                Defaults to env var `NAMESPACE` or `default`

    \b
    Fargate Agent Options:
        Any of the configuration options outlined in the docs can be provided here
        https://docs.prefect.io/cloud/agents/fargate.html#configuration
    """

    # Split context
    kwargs = dict()
    for item in ctx.args:
        item = item.replace("--", "")
        kwargs.update([item.split("=")])

    tmp_config = {
        "cloud.agent.auth_token": token or config.cloud.agent.auth_token,
        "logging.log_to_cloud": False if no_cloud_logs else True,
    }
    if verbose:
        tmp_config["cloud.agent.level"] = "DEBUG"

    with set_temporary_config(tmp_config):
        retrieved_agent = _agents.get(agent_option, None)

        if not retrieved_agent:
            click.secho("{} is not a valid agent".format(agent_option),
                        fg="red")
            return

        env_vars = dict()
        for env_var in env:
            k, v = env_var.split("=")
            env_vars[k] = v

        if agent_option == "local":
            from_qualified_name(retrieved_agent)(
                name=name,
                labels=list(label),
                env_vars=env_vars,
                max_polls=max_polls,
                import_paths=list(import_path),
                show_flow_logs=show_flow_logs,
            ).start()
        elif agent_option == "docker":
            from_qualified_name(retrieved_agent)(
                name=name,
                labels=list(label),
                env_vars=env_vars,
                max_polls=max_polls,
                base_url=base_url,
                no_pull=no_pull,
                show_flow_logs=show_flow_logs,
                volumes=list(volume),
            ).start()
        elif agent_option == "fargate":
            from_qualified_name(retrieved_agent)(name=name,
                                                 labels=list(label),
                                                 env_vars=env_vars,
                                                 max_polls=max_polls,
                                                 **kwargs).start()
        elif agent_option == "kubernetes":
            from_qualified_name(retrieved_agent)(
                namespace=namespace,
                name=name,
                labels=list(label),
                env_vars=env_vars,
                max_polls=max_polls,
            ).start()
        else:
            from_qualified_name(retrieved_agent)(
                name=name,
                labels=list(label),
                env_vars=env_vars,
                max_polls=max_polls,
            ).start()
Esempio n. 10
0
def test_k8s_agent_replace_yaml_uses_user_env_vars(monkeypatch, cloud_api):
    get_jobs = MagicMock(return_value=[])
    monkeypatch.setattr(
        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
        get_jobs,
    )

    monkeypatch.setenv("IMAGE_PULL_SECRETS", "my-secret")
    monkeypatch.setenv("JOB_MEM_REQUEST", "mr")
    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(
        {
            "flow": GraphQLResult(
                {
                    "storage": Docker(
                        registry_url="test", image_name="name", image_tag="tag"
                    ).serialize(),
                    "environment": LocalEnvironment().serialize(),
                    "id": "new_id",
                    "core_version": "0.13.0",
                }
            ),
            "id": "id",
        }
    )

    with set_temporary_config(
        {"cloud.agent.auth_token": "token", "logging.log_to_cloud": True}
    ):
        agent = KubernetesAgent(env_vars=dict(AUTH_THING="foo", PKG_SETTING="bar"))
        job = agent.generate_job_spec_from_environment(flow_run, image="test/name:tag")

        assert job["metadata"]["labels"]["prefect.io/flow_run_id"] == "id"
        assert job["metadata"]["labels"]["prefect.io/flow_id"] == "new_id"
        assert (
            job["spec"]["template"]["metadata"]["labels"]["prefect.io/flow_run_id"]
            == "id"
        )
        assert (
            job["spec"]["template"]["spec"]["containers"][0]["image"] == "test/name:tag"
        )

        env = job["spec"]["template"]["spec"]["containers"][0]["env"]

        assert env[0]["value"] == "https://api.prefect.io"
        assert env[1]["value"] == "token"
        assert env[2]["value"] == "id"
        assert env[3]["value"] == "new_id"
        assert env[4]["value"] == "default"
        assert env[5]["value"] == "[]"
        assert env[6]["value"] == "true"

        user_vars = [
            dict(name="AUTH_THING", value="foo"),
            dict(name="PKG_SETTING", value="bar"),
        ]
        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"

        assert job["spec"]["template"]["spec"]["imagePullSecrets"] == [
            {"name": "my-secret"}
        ]
Esempio n. 11
0
def test_k8s_agent_replace_yaml(monkeypatch, cloud_api):
    get_jobs = MagicMock(return_value=[])
    monkeypatch.setattr(
        "prefect.agent.kubernetes.agent.KubernetesAgent.manage_jobs",
        get_jobs,
    )

    monkeypatch.setenv("IMAGE_PULL_SECRETS", "my-secret")
    monkeypatch.setenv("JOB_MEM_REQUEST", "mr")
    monkeypatch.setenv("JOB_MEM_LIMIT", "ml")
    monkeypatch.setenv("JOB_CPU_REQUEST", "cr")
    monkeypatch.setenv("JOB_CPU_LIMIT", "cl")

    flow_run = GraphQLResult(
        {
            "flow": GraphQLResult(
                {
                    "storage": Docker(
                        registry_url="test", image_name="name", image_tag="tag"
                    ).serialize(),
                    "environment": LocalEnvironment().serialize(),
                    "id": "new_id",
                    "core_version": "0.13.0",
                }
            ),
            "id": "id",
        }
    )

    with set_temporary_config(
        {"cloud.agent.auth_token": "token", "logging.log_to_cloud": True}
    ):
        volume_mounts = [{"name": "my-vol", "mountPath": "/mnt/my-mount"}]
        volumes = [{"name": "my-vol", "hostPath": "/host/folder"}]
        agent = KubernetesAgent(volume_mounts=volume_mounts, volumes=volumes)
        job = agent.generate_job_spec_from_environment(flow_run, image="test/name:tag")

        assert job["metadata"]["labels"]["prefect.io/flow_run_id"] == "id"
        assert job["metadata"]["labels"]["prefect.io/flow_id"] == "new_id"
        assert (
            job["spec"]["template"]["metadata"]["labels"]["prefect.io/flow_run_id"]
            == "id"
        )
        assert (
            job["spec"]["template"]["spec"]["containers"][0]["image"] == "test/name:tag"
        )

        env = job["spec"]["template"]["spec"]["containers"][0]["env"]

        assert env[0]["value"] == "https://api.prefect.io"
        assert env[1]["value"] == "token"
        assert env[2]["value"] == "id"
        assert env[3]["value"] == "new_id"
        assert env[4]["value"] == "default"
        assert env[5]["value"] == "[]"
        assert env[6]["value"] == "true"

        assert (
            job["spec"]["template"]["spec"]["imagePullSecrets"][0]["name"]
            == "my-secret"
        )

        resources = job["spec"]["template"]["spec"]["containers"][0]["resources"]
        assert resources["requests"]["memory"] == "mr"
        assert resources["limits"]["memory"] == "ml"
        assert resources["requests"]["cpu"] == "cr"
        assert resources["limits"]["cpu"] == "cl"

        volumeMounts = job["spec"]["template"]["spec"]["containers"][0]["volumeMounts"]
        assert volumeMounts[0]["name"] == "my-vol"
        assert volumeMounts[0]["mountPath"] == "/mnt/my-mount"

        assert (
            job["spec"]["template"]["spec"]["containers"][0]["imagePullPolicy"]
            == "IfNotPresent"
        )

        volumes = job["spec"]["template"]["spec"]["volumes"]
        assert volumes[0]["name"] == "my-vol"
        assert volumes[0]["hostPath"] == "/host/folder"

        assert job["spec"]["template"]["spec"].get("serviceAccountName", None) is None
Esempio n. 12
0
 def test_client_settings_path_depends_on_home_dir(self, cloud_api):
     with set_temporary_config(dict(home_dir="abc/def")):
         path = Client(api_server="xyz")._api_token_settings_path
         expected = os.path.join("abc", "def", "client", "xyz", "settings.toml")
         assert str(path) == os.path.expanduser(expected)
Esempio n. 13
0
 def test_client_token_priotizes_arg_over_config(self, cloud_api):
     with set_temporary_config({"cloud.auth_token": "CONFIG_TOKEN"}):
         client = Client(api_token="ARG_TOKEN")
     assert client._api_token == "ARG_TOKEN"
Esempio n. 14
0
# import pendulum
#
# from prefect.schedules import Schedule
# from prefect.schedules.clocks import DatesClock
#
# schedule = Schedule(
#     clocks=[DatesClock([pendulum.now().add(seconds=1)])])

temp_config = {
    "cloud.agent.auth_token": prefect.config.cloud.agent.auth_token,
}


with Flow("Hydat-ETL") as flow:

    with set_temporary_config(temp_config):
        if flow.run_config is not None:
            labels = list(flow.run_config.labels or ())
        elif flow.environment is not None:
            labels = list(flow.environment.labels or ())
        else:
            labels = []
        agent = agent.local.LocalAgent(
            labels=labels, max_polls=50
        )

    url = 'https://collaboration.cmc.ec.gc.ca/cmc/hydrometrics/www/'
    ext = 'zip'

    path = extract_hydat_path(url, ext)
    cond = verify_if_to_date(path)
Esempio n. 15
0
 def test_uses_config_value_if_not_provided(self, sub, config_value):
     with set_temporary_config(
         {"flows.defaults.storage.add_default_labels": config_value}
     ):
         inst = sub()
         assert inst.add_default_labels == config_value
Esempio n. 16
0
def test_secret_get_none():
    secret = Secret(name="test")
    with set_temporary_config({"cloud.use_local_secrets": True}):
        assert secret.get() is None
Esempio n. 17
0
def test_default_storage_responds_to_config():
    with set_temporary_config(
        {"flows.defaults.storage.default_class": "prefect.storage.Docker"}
    ):
        assert get_default_storage_class() is Docker
Esempio n. 18
0
def test_secret_value_pulled_from_context():
    secret = Secret(name="test")
    with set_temporary_config({"cloud.use_local_secrets": True}):
        with prefect.context(secrets=dict(test=42)):
            assert secret.get() == 42
        assert secret.get() is None
Esempio n. 19
0
def mock_cloud_config(cloud_api):
    with set_temporary_config({
            "cloud.agent.auth_token": "TEST_TOKEN",
            "logging.log_to_cloud": True
    }):
        yield
Esempio n. 20
0
def test_local_secrets_remain_plain_dictionaries():
    secret = Secret(name="test")
    with set_temporary_config({"cloud.use_local_secrets": True}):
        with prefect.context(secrets=dict(test={"x": 42})):
            assert isinstance(prefect.context.secrets["test"], dict)
            assert secret.get() == {"x": 42}
Esempio n. 21
0
def test_agent_log_level(cloud_api):
    with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
        agent = Agent()
        assert agent.logger.level == 20
Esempio n. 22
0
def test_get_flow_run_info_with_nontrivial_payloads(patch_post):
    response = {
        "flow_run_by_pk": {
            "id":
            "da344768-5f5d-4eaf-9bca-83815617f713",
            "flow_id":
            "da344768-5f5d-4eaf-9bca-83815617f713",
            "name":
            "flow-run-name",
            "version":
            0,
            "parameters": {
                "x": {
                    "deep": {
                        "nested": 5
                    }
                }
            },
            "context": {
                "my_val": "test"
            },
            "scheduled_start_time":
            "2019-01-25T19:15:58.632412+00:00",
            "serialized_state": {
                "type": "Pending",
                "_result": {
                    "type": "SafeResult",
                    "value": "42",
                    "result_handler": {
                        "type": "JSONResultHandler"
                    },
                },
                "message": None,
                "__version__": "0.3.3+309.gf1db024",
                "cached_inputs": None,
            },
            "task_runs": [{
                "id": "da344768-5f5d-4eaf-9bca-83815617f713",
                "task": {
                    "id": "da344768-5f5d-4eaf-9bca-83815617f713",
                    "slug": "da344768-5f5d-4eaf-9bca-83815617f713",
                },
                "version": 0,
                "serialized_state": {
                    "type": "Pending",
                    "result": None,
                    "message": None,
                    "__version__": "0.3.3+309.gf1db024",
                    "cached_inputs": None,
                },
            }],
        }
    }
    post = patch_post(dict(data=response))

    with set_temporary_config({
            "cloud.graphql": "http://my-cloud.foo",
            "cloud.auth_token": "secret_token"
    }):
        client = Client()
    result = client.get_flow_run_info(flow_run_id="74-salt")
    assert isinstance(result, FlowRunInfoResult)
    assert isinstance(result.scheduled_start_time, datetime.datetime)
    assert result.scheduled_start_time.minute == 15
    assert result.scheduled_start_time.year == 2019
    assert isinstance(result.state, Pending)
    assert result.state.result == "42"
    assert result.state.message is None
    assert result.version == 0
    assert isinstance(result.parameters, dict)
    assert result.parameters["x"]["deep"]["nested"] == 5
    # ensures all sub-dictionaries are actually dictionaries
    assert json.loads(json.dumps(result.parameters)) == result.parameters
    assert isinstance(result.context, dict)
    assert result.context["my_val"] == "test"
Esempio n. 23
0
def test_agent_max_polls(cloud_api):
    with set_temporary_config({"cloud.agent.auth_token": "TEST_TOKEN"}):
        agent = Agent(max_polls=10)
        assert agent.max_polls == 10
Esempio n. 24
0
def test_secrets_dont_raise_just_because_flow_key_is_populated():
    secret = Secret(name="test")
    with set_temporary_config({"cloud.use_local_secrets": True}):
        with prefect.context(secrets=dict(test=42), flow="not None"):
            assert secret.get() == 42
Esempio n. 25
0
def test_deploy_flows_passes_base_agent():
    with set_temporary_config({"cloud.agent.auth_token": "token"}):
        agent = Agent()
        assert not agent.deploy_flows([])
Esempio n. 26
0
def test_secret_raises_if_doesnt_exist():
    secret = Secret(name="test")
    with set_temporary_config({"cloud.use_local_secrets": True}):
        with pytest.raises(ValueError, match="not found"):
            secret.get()
Esempio n. 27
0
def test_agent_fails_no_auth_token():
    with set_temporary_config({"cloud.agent.auth_token": None}):
        with pytest.raises(AuthorizationError):
            agent = Agent()
            agent.query_tenant_id()
Esempio n. 28
0
def start(
    ctx,
    agent_option,
    token,
    api,
    agent_config_id,
    name,
    verbose,
    label,
    env,
    namespace,
    job_template,
    no_pull,
    no_cloud_logs,
    base_url,
    import_path,
    show_flow_logs,
    volume,
    network,
    no_docker_interface,
    max_polls,
    agent_address,
    hostname_label,
    storage_labels,
    docker_client_timeout,
):
    """
    Start an agent.

    DEPRECATED: use `prefect agent <agent-type> start` instead.

    \b
    Arguments:
        agent-option    TEXT    The name of an agent to start (e.g. `docker`, `kubernetes`,
                                `local`, `fargate`). Defaults to `local`

    \b
    Options:
        --token, -t             TEXT    A Prefect Cloud API token with RUNNER scope
        --api, -a               TEXT    A Prefect API URL
        --agent-config--id      TEXT    An agent ID to link this agent instance with
        --name, -n              TEXT    A name to use for the agent
        --verbose, -v                   Enable verbose agent DEBUG logs
                                        Defaults to INFO level logging
        --label, -l             TEXT    Labels the agent will use to query for flow runs
                                        Multiple values supported e.g. `-l label1 -l label2`
        --env, -e               TEXT    Environment variables to set on each submitted flow
                                        run.
                                        Note that equal signs in environment variable values
                                        are not currently supported from the CLI.  Multiple
                                        values supported.
                                            e.g. `-e AUTH=token -e PKG_SETTING=true`
        --max-polls             INT     Maximum number of times the agent should poll the
                                        Prefect API for flow runs. Will run forever if not
                                        specified.
        --no-cloud-logs                 Turn off logging to the Prefect API for all flow runs
                                        Defaults to `False`
        --agent-address         TEXT    The address to server internal api at. Currently this
                                        is just health checks for use by an orchestration layer
                                        (e.g. kubernetes). Leave blank for no api server (default).

    \b
    Local Agent:
        --import-path, -p   TEXT    Import paths which will be provided to each Flow's
                                    runtime environment.  Used for Flows which might
                                    import from scripts or local packages.  Multiple values
                                    supported.
                                        e.g. `-p /root/my_scripts -p /utilities`
        --show-flow-logs, -f        Display logging output from flows run by the agent
                                    (available for Local and Docker agents only)
        --hostname-label            Add hostname to the Agent's labels
                                        (Default to True. Disable with --no-hostname-label option)

    \b
    Docker Agent:
        --base-url, -b      TEXT    A Docker daemon host URL for a DockerAgent
        --no-pull                   Pull images for a DockerAgent
                                    Defaults to pulling if not provided
        --volume            TEXT    Host paths for Docker bind mount volumes attached to
                                    each Flow runtime container. Multiple values supported.
                                        e.g. `--volume /some/path`
        --network           TEXT    Add containers to existing Docker networks.
                                    Multiple values supported.
                                        e.g. `--network network1 --network network2`
        --no-docker-interface       Disable the check of a Docker interface on this machine.
                                    Note: This is mostly relevant for some Docker-in-Docker
                                    setups that users may be running their agent with.
        --docker-client-timeout     Timeout for docker api requests

    \b
    Kubernetes Agent:
        --namespace     TEXT    A Kubernetes namespace to create Prefect jobs in
                                Defaults to env var `NAMESPACE` or `default`
        --job-template  TEXT    Path to a job template to use instead of the default.

    \b
    Fargate Agent Options:
        Any of the configuration options outlined in the docs can be provided here
        https://docs.prefect.io/orchestration/agents/fargate.html#configuration
    """
    # Split context
    kwargs = dict()
    for item in ctx.args:
        item = item.replace("--", "")
        kwargs.update([item.split("=", 1)])

    tmp_config = {
        "cloud.agent.auth_token": token or config.cloud.agent.auth_token,
    }
    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)

        if not retrieved_agent:
            click.secho("{} is not a valid agent".format(agent_option),
                        fg="red",
                        err=True)
            return

        if agent_option == "fargate":
            warn_fargate_deprecated()
        else:
            click.secho(
                f"Warning: `prefect agent start {agent_option}` is deprecated, use "
                f"`prefect agent {agent_option} start` instead",
                fg="yellow",
                err=True,
            )

        env_vars = dict()
        for env_var in env:
            k, v = env_var.split("=", 1)
            env_vars[k] = v

        labels = sorted(set(label))

        if agent_option == "local":
            from_qualified_name(retrieved_agent)(
                agent_config_id=agent_config_id,
                name=name,
                labels=labels,
                env_vars=env_vars,
                max_polls=max_polls,
                agent_address=agent_address,
                import_paths=list(import_path),
                show_flow_logs=show_flow_logs,
                no_cloud_logs=no_cloud_logs,
                hostname_label=hostname_label,
                storage_labels=storage_labels,
            ).start()
        elif agent_option == "docker":
            from_qualified_name(retrieved_agent)(
                agent_config_id=agent_config_id,
                name=name,
                labels=labels,
                env_vars=env_vars,
                max_polls=max_polls,
                no_cloud_logs=no_cloud_logs,
                agent_address=agent_address,
                base_url=base_url,
                no_pull=no_pull,
                show_flow_logs=show_flow_logs,
                volumes=list(volume),
                networks=tuple(network),
                docker_interface=not no_docker_interface,
                docker_client_timeout=docker_client_timeout,
            ).start()
        elif agent_option == "fargate":
            from_qualified_name(retrieved_agent)(
                agent_config_id=agent_config_id,
                name=name,
                labels=labels,
                env_vars=env_vars,
                max_polls=max_polls,
                no_cloud_logs=no_cloud_logs,
                agent_address=agent_address,
                _called_from_cli=True,
                **kwargs,
            ).start()
        elif agent_option == "kubernetes":
            from_qualified_name(retrieved_agent)(
                agent_config_id=agent_config_id,
                namespace=namespace,
                job_template_path=job_template,
                name=name,
                labels=labels,
                env_vars=env_vars,
                max_polls=max_polls,
                no_cloud_logs=no_cloud_logs,
                agent_address=agent_address,
            ).start()
        else:
            from_qualified_name(retrieved_agent)(
                agent_config_id=agent_config_id,
                name=name,
                labels=labels,
                env_vars=env_vars,
                max_polls=max_polls,
                no_cloud_logs=no_cloud_logs,
                agent_address=agent_address,
            ).start()
Esempio n. 29
0
def test_populate_job_yaml():
    with tempfile.TemporaryDirectory() as directory:

        with open(os.path.join(directory, "job.yaml"), "w+") as file:
            file.write("job")

        environment = KubernetesJobEnvironment(job_spec_file=os.path.join(
            directory, "job.yaml"),
                                               unique_job_name=True)

        file_path = os.path.dirname(
            prefect.environments.execution.dask.k8s.__file__)

        with open(path.join(file_path, "job.yaml")) as job_file:
            job = yaml.safe_load(job_file)
            job["spec"]["template"]["spec"]["containers"][0]["env"] = []

        with set_temporary_config({
                "cloud.graphql": "gql_test",
                "cloud.auth_token": "auth_test",
                "logging.extra_loggers": "['test_logger']",
        }):
            with prefect.context(flow_run_id="id_test",
                                 namespace="namespace_test"):
                yaml_obj = environment._populate_job_spec_yaml(
                    yaml_obj=job,
                    docker_name="test1/test2:test3",
                    flow_file_path="test4",
                )

        assert "prefect-dask-job-" in yaml_obj["metadata"]["name"]
        assert len(yaml_obj["metadata"]["name"]) == 25

        assert (yaml_obj["metadata"]["labels"]["prefect.io/identifier"] ==
                environment.identifier_label)
        assert yaml_obj["metadata"]["labels"][
            "prefect.io/flow_run_id"] == "id_test"
        assert (yaml_obj["spec"]["template"]["metadata"]["labels"]
                ["prefect.io/identifier"] == environment.identifier_label)

        env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"]

        assert env[0]["value"] == "gql_test"
        assert env[1]["value"] == "auth_test"
        assert env[2]["value"] == "id_test"
        assert env[3]["value"] == "namespace_test"
        assert env[4]["value"] == "test1/test2:test3"
        assert env[5]["value"] == "test4"
        assert env[10]["value"] == "['test_logger']"

        assert (yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"]
                == "test1/test2:test3")

        assert yaml_obj["spec"]["template"]["spec"]["containers"][0][
            "command"] == [
                "/bin/sh",
                "-c",
            ]
        assert yaml_obj["spec"]["template"]["spec"]["containers"][0]["args"] == [
            "python -c 'import prefect; prefect.Flow.load(prefect.context.flow_file_path).environment.run_flow()'"
        ]
 def test_secret_handler_writes_and_reads(self, res, secret_task):
     handler = SecretResultHandler(secret_task)
     with set_temporary_config({"cloud.use_local_secrets": True}):
         with prefect.context(secrets=dict(test=res)):
             final = handler.read(handler.write(res))
     assert final == res