from typing import Any, Mapping, Optional, Sequence
import docker
from dagster import Field, In, Nothing, OpExecutionContext, StringSource, op
from dagster._annotations import experimental
from dagster._core.utils import parse_env_var
from dagster._serdes.utils import hash_str
from ..container_context import DockerContainerContext
from ..docker_run_launcher import DockerRunLauncher
from ..utils import DOCKER_CONFIG_SCHEMA, validate_docker_image
DOCKER_CONTAINER_OP_CONFIG = {
**DOCKER_CONFIG_SCHEMA,
"image": Field(
StringSource,
is_required=True,
description="The image in which to run the Docker container.",
),
"entrypoint": Field(
[str],
is_required=False,
description="The ENTRYPOINT for the Docker container",
),
"command": Field(
[str],
is_required=False,
description="The command to run in the container within the launched Docker container.",
),
}
def _get_client(docker_container_context: DockerContainerContext):
client = docker.client.from_env()
if docker_container_context.registry:
client.login(
registry=docker_container_context.registry["url"],
username=docker_container_context.registry["username"],
password=docker_container_context.registry["password"],
)
return client
def _get_container_name(run_id, op_name, retry_number):
container_name = hash_str(run_id + op_name)
retry_number = retry_number
if retry_number > 0:
container_name = f"{container_name}-{retry_number}"
return container_name
def _create_container(
op_context: OpExecutionContext,
client,
container_context: DockerContainerContext,
image: str,
entrypoint: Optional[Sequence[str]],
command: Optional[Sequence[str]],
):
env_vars = dict([parse_env_var(env_var) for env_var in container_context.env_vars])
return client.containers.create(
image,
name=_get_container_name(op_context.run_id, op_context.op.name, op_context.retry_number),
detach=True,
network=container_context.networks[0] if len(container_context.networks) else None,
entrypoint=entrypoint,
command=command,
environment=env_vars,
**container_context.container_kwargs,
)
[docs]@experimental
def execute_docker_container(
context: OpExecutionContext,
image: str,
entrypoint: Optional[Sequence[str]] = None,
command: Optional[Sequence[str]] = None,
networks: Optional[Sequence[str]] = None,
registry: Optional[Mapping[str, str]] = None,
env_vars: Optional[Sequence[str]] = None,
container_kwargs: Optional[Mapping[str, Any]] = None,
):
"""This function is a utility for executing a Docker container from within a Dagster op.
Args:
image (str): The image to use for the launched Docker container.
entrypoint (Optional[Sequence[str]]): The ENTRYPOINT to run in the launched Docker
container. Default: None.
command (Optional[Sequence[str]]): The CMD to run in the launched Docker container.
Default: None.
networks (Optional[Sequence[str]]): Names of the Docker networks to which to connect the
launched container. Default: None.
registry: (Optional[Mapping[str, str]]): Information for using a non local/public Docker
registry. Can have "url", "username", or "password" keys.
env_vars (Optional[Sequence[str]]): List of environemnt variables to include in the launched
container. ach can be of the form KEY=VALUE or just KEY (in which case the value will be
pulled from the calling environment.
container_kwargs (Optional[Dict[str[Any]]]): key-value pairs that can be passed into
containers.create in the Docker Python API. See
https://docker-py.readthedocs.io/en/stable/containers.html for the full list
of available options.
"""
run_container_context = DockerContainerContext.create_for_run(
context.dagster_run,
(
context.instance.run_launcher
if isinstance(context.instance.run_launcher, DockerRunLauncher)
else None
),
)
validate_docker_image(image)
op_container_context = DockerContainerContext(
registry=registry, env_vars=env_vars, networks=networks, container_kwargs=container_kwargs
)
container_context = run_container_context.merge(op_container_context)
client = _get_client(container_context)
try:
container = _create_container(
context, client, container_context, image, entrypoint, command
)
except docker.errors.ImageNotFound:
client.images.pull(image)
container = _create_container(
context, client, container_context, image, entrypoint, command
)
if len(container_context.networks) > 1:
for network_name in container_context.networks[1:]:
network = client.networks.get(network_name)
network.connect(container)
container.start()
for line in container.logs(stdout=True, stderr=True, stream=True, follow=True):
print(line) # noqa: T201
exit_status = container.wait()["StatusCode"]
if exit_status != 0:
raise Exception(f"Docker container returned exit code {exit_status}")
[docs]@op(ins={"start_after": In(Nothing)}, config_schema=DOCKER_CONTAINER_OP_CONFIG)
@experimental
def docker_container_op(context):
"""An op that runs a Docker container using the docker Python API.
Contrast with the `docker_executor`, which runs each Dagster op in a Dagster job in its
own Docker container.
This op may be useful when:
- You need to orchestrate a command that isn't a Dagster op (or isn't written in Python)
- You want to run the rest of a Dagster job using a specific executor, and only a single
op in docker.
For example:
.. literalinclude:: ../../../../../../python_modules/libraries/dagster-docker/dagster_docker_tests/test_example_docker_container_op.py
:start-after: start_marker
:end-before: end_marker
:language: python
You can create your own op with the same implementation by calling the `execute_docker_container` function
inside your own op.
"""
execute_docker_container(context, **context.op_config)