dagster-cloud-cli 1.12.16__py2.py3-none-any.whl

dagster_cloud_cli/core/pex_builder/selftest.py

@@ -0,0 +1,14 @@
+# Simple environment test script
+
+import os
+
+from dagster_cloud_cli.core.pex_builder import util
+
+if __name__ == "__main__":
+    pex_environ = {name: value for name, value in os.environ.items() if name.startswith("PEX")}
+
+    proc = util.run_pex_command(["--version"])
+
+    proc = util.run_dagster_cloud_cli_command(["--version"])
+
+    proc = util.run_dagster_command(["--version"])

dagster_cloud_cli/core/pex_builder/source.py

@@ -0,0 +1,166 @@
+# Build source.pex, given a project root
+
+import os
+import os.path
+import shutil
+import subprocess
+import tempfile
+from uuid import uuid4
+
+import click
+
+from dagster_cloud_cli import ui
+from dagster_cloud_cli.core.pex_builder import deps, util
+
+# inspired by https://github.com/github/gitignore/blob/main/Python.gitignore
+# would be nice to just read the gitignore and use that
+IGNORED_PATTERNS = [
+    ".git",
+    "__pycache__",
+    ".pytest_cache",
+    ".tox",
+    ".nox",
+    "*.pyc",
+    ".mypy_cache",
+    ".cache",
+    ".venv",
+]
+
+
+def build_source_pex(
+    code_directory: str, local_package_paths: list[str], output_directory, python_version
+):
+    output_directory = os.path.abspath(output_directory)
+    os.makedirs(output_directory, exist_ok=True)
+    code_directory = os.path.abspath(code_directory)
+    tmp_pex_path = os.path.join(output_directory, f"source-tmp-{uuid4()}.pex")
+
+    build_pex_using_setup_py(code_directory, local_package_paths, tmp_pex_path, python_version)
+
+    pex_info = util.get_pex_info(tmp_pex_path)
+    pex_hash = pex_info["pex_hash"]
+    pex_name = f"source-{pex_hash}.pex"
+    final_pex_path = os.path.join(output_directory, pex_name)
+    os.rename(tmp_pex_path, final_pex_path)
+    ui.print(f"Wrote source pex: {final_pex_path}")
+    return final_pex_path
+
+
+def _build_local_package(local_dir: str, build_dir: str, python_interpreter: str):
+    if os.path.exists(os.path.join(local_dir, "setup.py")):
+        ui.print(f"Building package at {local_dir!r} using {python_interpreter} setup.py build")
+        command = [
+            python_interpreter,
+            "setup.py",
+            "build",
+            "--build-lib",
+            build_dir,
+        ]
+        subprocess.run(command, check=True, cwd=local_dir)
+    elif os.path.exists(os.path.join(local_dir, "pyproject.toml")):
+        ui.print(f"Building package at {local_dir!r} using {python_interpreter} -m pip install")
+        # Use pip install with --target to build and install the package to the build directory
+        # This handles pyproject.toml files properly using modern Python build standards
+        command = [
+            python_interpreter,
+            "-m",
+            "pip",
+            "install",
+            "--target",
+            build_dir,
+            "--no-deps",  # Don't install dependencies, just the package itself
+            ".",
+        ]
+        subprocess.run(command, check=True, cwd=local_dir)
+    else:
+        ui.warn(f"No setup.py or pyproject.toml found in {local_dir!r} - will not build.")
+
+
+def build_pex_using_setup_py(
+    code_directory: str,
+    local_package_paths: list[str],
+    tmp_pex_path,
+    python_version,
+):
+    """Builds package using setup.py and copies built output into PEX."""
+    python_interpreter = util.python_interpreter_for(python_version)
+    with tempfile.TemporaryDirectory() as build_dir, tempfile.TemporaryDirectory() as sources_dir:
+        included_dirs = [build_dir]
+
+        for local_dir in local_package_paths:
+            _build_local_package(
+                local_dir=local_dir,
+                build_dir=build_dir,
+                python_interpreter=python_interpreter,
+            )
+
+        _build_local_package(
+            local_dir=code_directory,
+            build_dir=build_dir,
+            python_interpreter=python_interpreter,
+        )
+
+        # We always include the code_directory source in a special package called working_directory
+        ui.print(
+            f"Bundling the source {code_directory} into the 'working_directory' package at {sources_dir}",
+        )
+        _prepare_working_directory(code_directory, sources_dir)
+        included_dirs.append(sources_dir)
+
+        # note this may include tests directories
+        proc = util.build_pex(
+            included_dirs,
+            requirements_filepaths=[],
+            pex_flags=util.get_pex_flags(python_version),
+            output_pex_path=tmp_pex_path,
+        )
+
+        if proc.returncode:
+            ui.error("build source pex stdout:")
+            ui.error(util.indent(proc.stdout.decode("utf-8")))
+            ui.error("build source pex stderr:")
+            ui.error(util.indent(proc.stderr.decode("utf-8")))
+            raise ui.error(
+                f"Failed to package source code into a source pex, error code: {proc.returncode}"
+            )
+        ui.print(f"Built source pex: {tmp_pex_path}")
+
+
+def _prepare_working_directory(code_directory, sources_directory):
+    # Copy code_directory contents into a package called working_directory under sources_directory
+    package_dir = os.path.join(sources_directory, "working_directory")
+    os.makedirs(package_dir, exist_ok=True)
+
+    with open(os.path.join(package_dir, "__init__.py"), "w", encoding="utf-8") as init_file:
+        init_file.write("# Auto generated package containing the original source at root/")
+
+    shutil.copytree(
+        code_directory,
+        os.path.join(package_dir, "root"),
+        dirs_exist_ok=True,
+        ignore=shutil.ignore_patterns(*IGNORED_PATTERNS),
+        # We copy over symlinks as symlinks rather than try and resolve them and copy the target.
+        # This prevents build failures if the target file doesn't exist locally. If the target
+        # doesn't exist at runtime, this will fail and that is acceptable.
+        symlinks=True,
+    )
+
+
+@click.command()
+@click.argument("project_dir", type=click.Path(exists=True))
+@click.argument("build_output_dir", type=click.Path(exists=False))
+@util.python_version_option()
+def source_main(project_dir, build_output_dir, python_version):
+    local_packages, _ = deps.get_deps_requirements(
+        code_directory=project_dir, python_version=python_version
+    )
+    build_source_pex(
+        project_dir,
+        local_packages.local_package_paths,
+        build_output_dir,
+        util.parse_python_version(python_version),
+    )
+
+
+if __name__ == "__main__":
+    source_main()

dagster_cloud_cli/core/pex_builder/util.py

@@ -0,0 +1,165 @@
+import json
+import logging
+import os
+import subprocess
+import sys
+from subprocess import CompletedProcess
+from zipfile import ZipFile
+
+import click
+from packaging import version
+
+from dagster_cloud_cli.core.pex_builder.platforms import COMPLETE_PLATFORMS
+from dagster_cloud_cli.utils import DEFAULT_PYTHON_VERSION, SUPPORTED_PYTHON_VERSIONS
+
+TARGET_PYTHON_VERSIONS = [
+    version.Version(python_version) for python_version in SUPPORTED_PYTHON_VERSIONS
+]
+
+
+def run_python_subprocess(args: list[str], env=None) -> CompletedProcess:
+    """Invoke python with given args, using an environment identical to current environment."""
+    # If running a pex file directly, we invoke the executable pex file again.
+    # Otherwise we assume we're running in a pex generated venv and use the python executable.
+    cmd = os.getenv("PEX", sys.executable)
+
+    args = [cmd] + args
+    logging.debug(f"Running {args} in {os.path.abspath(os.curdir)}")
+    proc = subprocess.run(args, capture_output=True, env=env, check=False)
+    return proc
+
+
+def run_self_module(module_name, args: list[str], env=None) -> CompletedProcess:
+    """Invoke this executable again with -m {module}."""
+    return run_python_subprocess(["-m", module_name, *args], env=env)
+
+
+def get_pex_flags(python_version: version.Version, build_sdists: bool = True) -> list[str]:
+    """python_version should includes the major and minor version only, eg Version('3.11')."""
+    if python_version not in TARGET_PYTHON_VERSIONS:
+        raise ValueError(
+            f"Unsupported python version {python_version}. Supported: {TARGET_PYTHON_VERSIONS}."
+        )
+    version_tag = f"{python_version.major}{python_version.minor}"  # eg '38'
+    # Resolves dependencies using the local interpreter, effectively allowing source distributions
+    # to work (since they get built by the local interpreter).
+    # If build_sdists is false and a binary distribution matching the platform below is not
+    # available for a dependency, then the build will fail.
+    # see also https://linear.app/elementl/issue/CLOUD-2023/pex-builds-fail-for-dbt-core-dependency
+    resolve_local = ["--resolve-local-platforms"] if build_sdists else []
+    # This is mainly useful in local mac test environments
+    include_current = ["--platform=current"] if os.getenv("PEX_INCLUDE_CURRENT_PLATFORM") else []
+
+    complete_platform = os.getenv(
+        "PEX_COMPLETE_PLATFORM", json.dumps(COMPLETE_PLATFORMS[version_tag])
+    )
+
+    return [
+        # this complete platform matches what can run on our serverless base images
+        f"--complete-platform={complete_platform}",
+        *include_current,
+        # this ensures PEX_PATH is not cleared and any subprocess invoked can also use this.
+        # this is important for running console scripts that use the pex environment (eg dbt)
+        "--no-strip-pex-env",
+        # use the latest version of pip bundled with pex - this will typically provide
+        # the best dependency resolution logic. added in
+        # https://github.com/pantsbuild/pex/releases/tag/v2.1.132
+        "--pip-version=latest",
+        "-v",  # verbose logging, level 3
+        "-v",
+        "-v",
+        *resolve_local,
+    ]
+
+
+def build_pex(
+    sources_directories: list[str],
+    requirements_filepaths: list[str],
+    pex_flags: list[str],
+    output_pex_path: str,
+    pex_root: str | None = None,
+) -> CompletedProcess:
+    """Invoke pex with common build flags and pass parameters through to specific pex flags.
+
+    sources_directories: passed to the pex -D flag (aka --source-directories)
+    requirements_filepaths: passed to the pex -r flag (aka --requirement)
+
+    Platform:
+    For the pex --platform tag used below, see pex --help and
+    https://peps.python.org/pep-0425/
+    https://peps.python.org/pep-0427/
+
+    Packages for the current platform are only included if requested with PEX_INCLUDE_CURRENT_PLATFORM
+    The manylinux platform ensures pexes built on local machines (macos, windows) are compatible
+    with linux on cloud.
+
+    The python_version tuple eg ('3', '9') determines the target runtime environment. In theory
+    we can build a pex in an environment without the target python version present.
+    """
+    flags = pex_flags.copy()
+    if not sources_directories and not requirements_filepaths:
+        raise ValueError("At least one of sources_directories or requirements_filepath required.")
+    for src_dir in sources_directories:
+        flags.extend(["-D", src_dir])
+    for req_file in requirements_filepaths:
+        flags.extend(["-r", req_file])
+    pex_args = [*flags, "-o", output_pex_path]
+    if pex_root:
+        pex_args.extend(["--pex-root", pex_root])
+    return run_pex_command(pex_args)
+
+
+def run_pex_command(args: list[str]) -> CompletedProcess:
+    # https://github.com/pantsbuild/pex/issues/1969#issuecomment-1336105021
+    env = {**os.environ, "_PEX_FILE_LOCK_STYLE": "bsd"}
+
+    return run_self_module("pex", args, env=env)
+
+
+def run_dagster_cloud_cli_command(args: list[str]) -> CompletedProcess:
+    return run_self_module("dagster_cloud_cli.entrypoint", args)
+
+
+def run_dagster_command(args: list[str]) -> CompletedProcess:
+    return run_self_module("dagster", args)
+
+
+def get_pex_info(pex_filepath):
+    with ZipFile(pex_filepath) as pex_zip:
+        return json.load(pex_zip.open("PEX-INFO"))
+
+
+def build_pex_tag(filepaths: list[str]) -> str:
+    return "files=" + ":".join(sorted(os.path.basename(filepath) for filepath in filepaths))
+
+
+def python_interpreter_for(python_version: version.Version) -> str:
+    return "python" + str(python_version)  # eg 'python3.11'
+
+
+def python_version_option():
+    """Reusable click.option."""
+    return click.option(
+        "--python-version",
+        type=click.Choice(SUPPORTED_PYTHON_VERSIONS),
+        default=DEFAULT_PYTHON_VERSION,
+        show_default=True,
+        help="Target Python version.",
+    )
+
+
+def parse_python_version(python_version: str) -> version.Version:
+    return version.Version(python_version)
+
+
+def parse_kv(ctx, param: str, value: str | None):
+    if not value:
+        return {}
+    try:
+        return dict(part.split("=", 1) for part in value.split(","))
+    except ValueError as err:
+        raise ValueError(f"Value {value!r} could not be parsed: {err}")
+
+
+def indent(text: str, prefix="| "):
+    return "".join([prefix + line for line in text.splitlines(keepends=True)])

dagster_cloud_cli/core/workspace.py

@@ -0,0 +1,145 @@
+from collections.abc import Mapping
+from typing import Any
+
+from dagster_shared import check
+from dagster_shared.record import IHaveNew, LegacyNamedTupleMixin, copy, record, record_custom
+from dagster_shared.serdes import whitelist_for_serdes
+from dagster_shared.serdes.objects.models.defs_state_info import DefsStateInfo
+from dagster_shared.serdes.serdes import serialize_value
+
+from dagster_cloud_cli.core.agent_queue import AgentQueue
+
+
+@whitelist_for_serdes
+@record(kw_only=False)
+class GitMetadata:
+    commit_hash: str | None = None
+    url: str | None = None
+
+
+@whitelist_for_serdes
+@record(kw_only=False)
+class PexMetadata:
+    # pex_tag is a string like 'deps-234y4384.pex:source-39y3474.pex' that idenfies
+    # the pex files to execute
+    pex_tag: str
+    # python_version determines which pex base docker image to use
+    # only one of PexMetadata.python_version or CodeLocationDeployData.image should be specified
+    python_version: str | None = None
+
+
+# History of CodeLocationDeployData
+# 1. Removal of `enable_metrics` field
+# 2. Renamed from `CodeDeploymentMetadata` to `CodeLocationDeployData``
+@whitelist_for_serdes(storage_name="CodeDeploymentMetadata")
+@record_custom
+class CodeLocationDeployData(IHaveNew, LegacyNamedTupleMixin):
+    image: str | None
+    python_file: str | None
+    package_name: str | None
+    module_name: str | None
+    working_directory: str | None
+    executable_path: str | None
+    attribute: str | None
+    git_metadata: GitMetadata | None
+    container_context: Mapping[str, Any]
+    cloud_context_env: Mapping[str, Any]
+    pex_metadata: PexMetadata | None
+    agent_queue: AgentQueue | None
+    autoload_defs_module_name: str | None
+    defs_state_info: DefsStateInfo | None
+
+    def __new__(
+        cls,
+        image: str | None = None,
+        python_file: str | None = None,
+        package_name: str | None = None,
+        module_name: str | None = None,
+        working_directory: str | None = None,
+        executable_path: str | None = None,
+        attribute: str | None = None,
+        git_metadata: GitMetadata | None = None,
+        container_context: Mapping[str, Any] | None = None,
+        cloud_context_env: Mapping[str, Any] | None = None,
+        pex_metadata: PexMetadata | None = None,
+        agent_queue: AgentQueue | None = None,
+        autoload_defs_module_name: str | None = None,
+        defs_state_info: DefsStateInfo | None = None,
+    ):
+        check.invariant(
+            len(
+                [
+                    val
+                    for val in [
+                        python_file,
+                        package_name,
+                        module_name,
+                        autoload_defs_module_name,
+                    ]
+                    if val
+                ]
+            )
+            == 1,
+            "Must supply exactly one of python_file, package_name, module_name, or autoload_defs_module_name.",
+        )
+
+        return super().__new__(
+            cls,
+            image=image,
+            python_file=python_file,
+            package_name=package_name,
+            module_name=module_name,
+            working_directory=working_directory,
+            executable_path=executable_path,
+            attribute=attribute,
+            git_metadata=git_metadata,
+            container_context=container_context or {},
+            cloud_context_env=cloud_context_env or {},
+            pex_metadata=pex_metadata,
+            agent_queue=agent_queue,
+            autoload_defs_module_name=autoload_defs_module_name,
+            defs_state_info=defs_state_info,
+        )
+
+    def with_cloud_context_env(
+        self,
+        cloud_context_env: Mapping[str, Any],
+    ) -> "CodeLocationDeployData":
+        return copy(self, cloud_context_env=cloud_context_env)
+
+    def get_multipex_server_command(
+        self,
+        port: int | None,
+        socket: str | None = None,
+        metrics_enabled: bool = False,
+    ) -> list[str]:
+        return (
+            ["dagster-cloud", "pex", "grpc", "--host", "0.0.0.0"]
+            + (["--port", str(port)] if port else [])
+            + (["--socket", str(socket)] if socket else [])
+            + (["--enable-metrics"] if metrics_enabled else [])
+            + (
+                ["--defs-state-info", serialize_value(self.defs_state_info)]
+                if self.defs_state_info
+                else []
+            )
+        )
+
+    def get_multipex_server_env(self) -> dict[str, str]:
+        return {"DAGSTER_CURRENT_IMAGE": self.image} if self.image else {}
+
+    def get_grpc_server_command(self, metrics_enabled: bool = False) -> list[str]:
+        return (
+            ([self.executable_path, "-m"] if self.executable_path else [])
+            + [
+                "dagster",
+                "api",
+                "grpc",
+            ]
+            + (["--enable-metrics"] if metrics_enabled else [])
+            + (
+                ["--defs-state-info", serialize_value(self.defs_state_info)]
+                if self.defs_state_info
+                else []
+            )
+        )

dagster_cloud_cli/docker_utils.py

@@ -0,0 +1,98 @@
+import base64
+import importlib.resources
+import os
+import subprocess
+import sys
+import uuid
+from collections.abc import Generator
+from contextlib import ExitStack, contextmanager
+
+from dagster_cloud_cli import ui
+from dagster_cloud_cli.utils import DEFAULT_PYTHON_VERSION
+
+
+def verify_docker() -> None:
+    if subprocess.call("docker -v", shell=True) != 0:
+        raise ui.error("Docker must be installed locally to deploy to Dagster Cloud Serverless")
+
+
+@contextmanager
+def _template_dockerfile(env_vars, custom_base_image=None) -> Generator[bytes, None, None]:
+    DOCKERFILE_TEMPLATE = str(
+        importlib.resources.files("dagster_cloud_cli") / "commands/serverless/Dockerfile"
+    )
+    base_image_command = (
+        f"FROM {custom_base_image}"
+        if custom_base_image
+        else f"FROM public.ecr.aws/docker/library/python:{DEFAULT_PYTHON_VERSION}-slim"
+    )
+    with open(DOCKERFILE_TEMPLATE, encoding="utf-8") as template:
+        dockerfile_content = "\n".join(
+            [base_image_command, template.read(), *[f"ENV {env_var}" for env_var in env_vars]]
+        )
+
+        yield bytes(dockerfile_content, "utf-8")
+
+
+def build_image(
+    source_directory,
+    image: str,
+    registry_info,
+    env_vars: list[str],
+    base_image: str | None,
+    use_editable_dagster: bool,
+    dockerfile_path: str | None = None,
+) -> int:
+    registry = registry_info["registry_url"]
+
+    with ExitStack() as stack:
+        if dockerfile_path:
+            file_path = dockerfile_path
+            build_input = None
+        else:
+            file_path = "-"
+            build_input = stack.enter_context(_template_dockerfile(env_vars, base_image))
+
+        cmd = [
+            "docker",
+            "build",
+            source_directory,
+            "-t",
+            f"{registry}:{image}",
+            "-f",
+            file_path,
+            "--platform",
+            "linux/amd64",
+        ]
+        if use_editable_dagster:
+            cmd += [
+                "--build-context",
+                f"oss={os.environ['DAGSTER_GIT_REPO_DIR']}",
+                "--build-context",
+                f"internal={os.environ['DAGSTER_INTERNAL_GIT_REPO_DIR']}",
+            ]
+        return subprocess.run(cmd, input=build_input, check=True).returncode
+
+
+def upload_image(image, registry_info) -> int:
+    registry = registry_info["registry_url"]
+    aws_token = registry_info["aws_auth_token"]
+    if not registry or not aws_token:
+        raise ui.error(
+            "No registry found. You may need to wait for your Dagster serverless deployment to"
+            " activate."
+        )
+
+    username, password = base64.b64decode(aws_token).decode("utf-8").split(":")
+    subprocess.check_output(
+        f"echo {password} | docker login --username {username} --password-stdin {registry}",
+        shell=True,
+    )
+    return subprocess.call(
+        ["docker", "push", f"{registry}:{image}"], stderr=sys.stderr, stdout=sys.stdout
+    )
+
+
+def default_image_tag(deployment: str, location_name: str, commit_hash: str | None) -> str:
+    commit_hash = commit_hash if commit_hash else str(uuid.uuid4().hex)
+    return f"{deployment}-{location_name}-{commit_hash}"