oracle-ads 2.13.9rc0__py3-none-any.whl → 2.13.10rc0__py3-none-any.whl

ads/jobs/builders/runtimes/base.py

@@ -0,0 +1,325 @@
+#!/usr/bin/env python
+# -*- coding: utf-8; -*-
+
+# Copyright (c) 2022, 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+from __future__ import annotations
+import re
+import time
+import traceback
+
+from typing import Dict, TypeVar
+from ads.jobs.builders.base import Builder
+from ads.jobs import env_var_parser
+
+
+Self = TypeVar("Self", bound="Runtime")
+
+
+class Runtime(Builder):
+    """Base class for job runtime"""
+
+    # Constant strings
+    CONST_ENV_VAR = "env"
+    CONST_ARGS = "args"
+    CONST_MAXIMUM_RUNTIME_IN_MINUTES = "maximumRuntimeInMinutes"
+    CONST_FREEFORM_TAGS = "freeformTags"
+    CONST_DEFINED_TAGS = "definedTags"
+
+    attribute_map = {
+        CONST_FREEFORM_TAGS: "freeform_tags",
+        CONST_DEFINED_TAGS: "defined_tags",
+        CONST_ENV_VAR: CONST_ENV_VAR,
+        CONST_ARGS: CONST_ARGS,
+    }
+
+    def __init__(self, spec: Dict = None, **kwargs) -> None:
+        env = {}
+        if spec and "env" in spec and isinstance(spec["env"], dict):
+            env = spec.pop("env", {})
+        if "env" in kwargs:
+            env = kwargs.pop("env")
+        super().__init__(spec, **kwargs)
+        if isinstance(env, dict):
+            self.with_environment_variable(**env)
+
+    @property
+    def kind(self) -> str:
+        """Kind of the object to be stored in YAML.
+        All runtime implementations will have "runtime" as kind.
+        Subclass will have different types.
+        """
+        return "runtime"
+
+    @property
+    def type(self) -> str:
+        """The type of the object as showing in YAML"""
+        class_name = self.__class__.__name__
+        if class_name.endswith("Runtime") and class_name != "Runtime":
+            class_name = re.sub(r"Runtime$", "", class_name)
+        return class_name[0].lower() + (class_name[1:] if len(class_name) > 1 else "")
+
+    def with_argument(self: Self, *args, **kwargs) -> Self:
+        """Adds command line arguments to the runtime.
+
+        This method can be called (chained) multiple times to add various arguments.
+
+        Parameters
+        ----------
+        args:
+            Positional arguments.
+            In a single method call, positional arguments are always added before keyword arguments.
+            You can call with_argument() to add positional arguments after keyword arguments.
+
+        kwargs:
+            Keyword arguments.
+            To add a keyword argument without value, set the value to None.
+
+        Returns
+        -------
+        Self
+            This method returns self to support chaining methods.
+
+        Raises
+        ------
+        ValueError
+            Keyword arguments with space in a key.
+
+        Examples
+        --------
+        >>> runtime = Runtime().with_argument(key1="val1", key2="val2").with_argument("pos1")
+        >>> print(runtime.args)
+        ["--key1", "val1", "--key2", "val2", "pos1"]
+
+        >>> runtime = Runtime()
+        >>> runtime.with_argument("pos1")
+        >>> runtime.with_argument(key1="val1", key2="val2.1 val2.2")
+        >>> runtime.with_argument("pos2")
+        >>> print(runtime.args)
+        ['pos1', '--key1', 'val1', '--key2', 'val2.1 val2.2', 'pos2']
+
+        >>> runtime = Runtime()
+        >>> runtime.with_argument("pos1")
+        >>> runtime.with_argument(key1=None, key2="val2")
+        >>> runtime.with_argument("pos2")
+        >>> print(runtime.args)
+        ["pos1", "--key1", "--key2", "val2", "pos2"]
+
+        """
+        arg_values = self.get_spec(self.CONST_ARGS, [])
+        args = [str(arg) for arg in args]
+        arg_values.extend(args)
+        for k, v in kwargs.items():
+            if " " in k:
+                raise ValueError("Argument key cannot contain space.")
+            arg_values.append(f"--{str(k)}")
+            # Ignore None value
+            if v is None:
+                continue
+            arg_values.append(str(v))
+        self.set_spec(self.CONST_ARGS, arg_values)
+        return self
+
+    def with_environment_variable(self: Self, **kwargs) -> Self:
+        """Sets environment variables
+
+        Environment variables enclosed by ``${...}`` will be substituted.
+
+        * You can use ``$$`` to escape the substitution.
+        * Undefined variable enclosed by ``${}`` will be ignored.
+        * Double dollar signs ``$$`` will be substituted by a single one ``$``.
+
+        Returns
+        -------
+        Self
+            This method returns self to support chaining methods.
+
+        Examples
+        --------
+        >>> runtime = (
+        ...     PythonRuntime()
+        ...     .with_environment_variable(
+        ...         HOST="10.0.0.1",
+        ...         PORT="443",
+        ...         URL="http://${HOST}:${PORT}/path/",
+        ...         ESCAPED_URL="http://$${HOST}:$${PORT}/path/",
+        ...         MISSING_VAR="This is ${UNDEFINED}",
+        ...         VAR_WITH_DOLLAR="$10",
+        ...         DOUBLE_DOLLAR="$$10"
+        ...     )
+        ... )
+        >>> for k, v in runtime.environment_variables.items():
+        ...     print(f"{k}: {v}")
+        HOST: 10.0.0.1
+        PORT: 443
+        URL: http://10.0.0.1:443/path/
+        ESCAPED_URL: http://${HOST}:${PORT}/path/
+        MISSING_VAR: This is ${UNDEFINED}
+        VAR_WITH_DOLLAR: $10
+        DOUBLE_DOLLAR: $10
+
+
+        """
+        if not kwargs:
+            return self
+        envs = [{"name": k, "value": v} for k, v in kwargs.items()]
+        return self.set_spec(self.CONST_ENV_VAR, envs)
+
+    def with_freeform_tag(self: Self, **kwargs) -> Self:
+        """Sets freeform tags
+
+        Returns
+        -------
+        Self
+            This method returns self to support chaining methods.
+        """
+        return self.set_spec(self.CONST_FREEFORM_TAGS, kwargs)
+
+    def with_defined_tag(self: Self, **kwargs) -> Self:
+        """Sets defined tags
+
+        Returns
+        -------
+        Self
+            This method returns self to support chaining methods.
+        """
+        return self.set_spec(self.CONST_DEFINED_TAGS, kwargs)
+
+    def with_maximum_runtime_in_minutes(
+        self: Self, maximum_runtime_in_minutes: int
+    ) -> Self:
+        """Sets maximum runtime in minutes
+
+        Returns
+        -------
+        Self
+            This method returns self to support chaining methods.
+        """
+        return self.set_spec(
+            self.CONST_MAXIMUM_RUNTIME_IN_MINUTES, maximum_runtime_in_minutes
+        )
+
+    @property
+    def environment_variables(self) -> dict:
+        """Environment variables
+
+        Returns
+        -------
+        dict
+            The runtime environment variables.
+            The returned dictionary is a copy.
+        """
+        env_var_list = self.get_spec(self.CONST_ENV_VAR)
+        if env_var_list:
+            return env_var_parser.parse(env_var_list)
+        return {}
+
+    @property
+    def envs(self) -> dict:
+        """Environment variables"""
+        return self.environment_variables
+
+    @property
+    def freeform_tags(self) -> dict:
+        """Freeform tags"""
+        return self.get_spec(self.CONST_FREEFORM_TAGS, {})
+
+    @property
+    def defined_tags(self) -> dict:
+        """Defined tags"""
+        return self.get_spec(self.CONST_DEFINED_TAGS, {})
+
+    @property
+    def args(self) -> list:
+        """Command line arguments"""
+        return self.get_spec(self.CONST_ARGS, [])
+
+    @property
+    def maximum_runtime_in_minutes(self) -> int:
+        """Maximum runtime in minutes"""
+        return self.get_spec(self.CONST_MAXIMUM_RUNTIME_IN_MINUTES)
+
+    def init(self, **kwargs) -> Self:
+        """Initializes a starter specification for the runtime.
+
+        Returns
+        -------
+        Self
+            This method returns self to support chaining methods.
+        """
+        return (
+            self.with_environment_variable(
+                **kwargs.get(self.attribute_map[self.CONST_ENV_VAR], {})
+            )
+            .with_freeform_tag(
+                **kwargs.get(self.attribute_map[self.CONST_FREEFORM_TAGS], {})
+            )
+            .with_argument(**kwargs.get(self.attribute_map[self.CONST_ARGS], {}))
+        )
+
+
+class MultiNodeRuntime(Runtime):
+    """Represents runtime supporting multi-node jobs."""
+
+    CONST_REPLICA = "replicas"
+
+    def with_replica(self, count: int):
+        """Specifies the number of nodes (job runs) for the job.
+
+        Parameters
+        ----------
+        count : int
+            Number of nodes (job runs)
+
+        Returns
+        -------
+        self
+            The runtime instance.
+        """
+        return self.set_spec(self.CONST_REPLICA, count)
+
+    @property
+    def replica(self) -> int:
+        """The number of nodes (job runs)."""
+        return self.get_spec(self.CONST_REPLICA)
+
+    def run(self, dsc_job, **kwargs):
+        """Starts the job runs"""
+        replicas = self.replica if self.replica else 1
+        main_run = None
+        job_runs = []
+        try:
+            for i in range(replicas):
+                replica_kwargs = kwargs.copy()
+
+                # Only update display name and env vars if replica is specified (not None).
+                if self.replica is not None:
+                    envs = replica_kwargs.get("environment_variables")
+                    if not envs:
+                        envs = {}
+                    # HuggingFace accelerate requires machine rank
+                    # Here we use NODE_RANK to store the machine rank
+                    envs["NODE_RANK"] = str(i)
+                    envs["NODE_COUNT"] = str(replicas)
+                    if main_run:
+                        envs["MAIN_JOB_RUN_OCID"] = main_run.id
+                    name = replica_kwargs.get("display_name")
+                    if not name:
+                        name = dsc_job.display_name
+
+                    replica_kwargs["display_name"] = f"{name}-{str(i)}"
+                    replica_kwargs["environment_variables"] = envs
+                run = dsc_job.run(**replica_kwargs)
+                job_runs.append(run)
+                if i == 0:
+                    main_run = run
+        except Exception as ex:
+            traceback.print_exc()
+            # Wait a few second to avoid the job run being in a transient state.
+            time.sleep(2)
+            # If there is any error when creating the job runs
+            # cancel all the job runs.
+            for run in job_runs:
+                run.cancel()
+            raise ex
+        return main_run

ads/jobs/builders/runtimes/container_runtime.py

@@ -0,0 +1,242 @@
+#!/usr/bin/env python
+# -*- coding: utf-8; -*-
+
+# Copyright (c) 2021, 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+import logging
+from typing import Union
+from ads.jobs.builders.runtimes.base import MultiNodeRuntime
+
+logger = logging.getLogger(__name__)
+
+
+class ContainerRuntime(MultiNodeRuntime):
+    """Represents a container job runtime
+
+    To define container runtime:
+
+    >>> ContainerRuntime()
+    >>> .with_image("iad.ocir.io/<your_tenancy>/<your_image>:<tag>")
+    >>> .with_cmd("sleep 5 && echo Hello World")
+    >>> .with_entrypoint(["/bin/sh", "-c"])
+    >>> .with_image_digest("<image_digest>")
+    >>> .with_image_signature_id("<image_signature_id>")
+    >>> .with_environment_variable(MY_ENV="MY_VALUE")
+
+    Alternatively, you can define the ``entrypoint``, ``cmd``,
+    ``image_digest``and ``image_signature_id`` along with the image.
+
+    >>> ContainerRuntime()
+    >>> .with_image(
+    >>>     "iad.ocir.io/<your_tenancy>/<your_image>:<tag>",
+    >>>     entrypoint=["/bin/sh", "-c"],
+    >>>     cmd="sleep 5 && echo Hello World",
+    >>>     image_digest="<image_digest>",
+    >>>     image_signature_id="<image_signature_id>",
+    >>> )
+    >>> .with_environment_variable(MY_ENV="MY_VALUE")
+
+    The entrypoint and cmd can be either "exec form" or "shell form" (See references).
+    The exec form is used when a list is passed in.
+    The shell form is used when a space separated string is passed in.
+
+    When using the ContainerRuntime with OCI Data Science Job, the exec form is recommended.
+    For most images, when the entrypoint is set to ``["/bin/sh", "-c"]``,
+    ``cmd`` can be a string as if you are running shell command.
+
+    References
+    ----------
+    https://docs.docker.com/engine/reference/builder/#entrypoint
+    https://docs.docker.com/engine/reference/builder/#cmd
+
+    """
+
+    CONST_IMAGE = "image"
+    CONST_ENTRYPOINT = "entrypoint"
+    CONST_CMD = "cmd"
+    CONST_IMAGE_DIGEST = "imageDigest"
+    CONST_IMAGE_SIGNATURE_ID = "imageSignatureId"
+    CONST_SCRIPT_PATH = "scriptPathURI"
+    attribute_map = {
+        CONST_IMAGE: CONST_IMAGE,
+        CONST_ENTRYPOINT: CONST_ENTRYPOINT,
+        CONST_CMD: CONST_CMD,
+        CONST_IMAGE_DIGEST: "image_digest",
+        CONST_IMAGE_SIGNATURE_ID: "image_signature_id",
+    }
+    attribute_map.update(MultiNodeRuntime.attribute_map)
+
+    @property
+    def job_env_type(self) -> str:
+        """The container type"""
+        return "OCIR_CONTAINER"
+
+    @property
+    def image(self) -> str:
+        """The container image"""
+        return self.get_spec(self.CONST_IMAGE)
+
+    def with_image(
+        self,
+        image: str,
+        entrypoint: Union[str, list, None] = None, 
+        cmd: str = None,
+        image_digest: str = None,
+        image_signature_id: str = None,
+    ) -> "ContainerRuntime":
+        """Specify the image for the container job.
+
+        Parameters
+        ----------
+        image : str
+            The container image, e.g. iad.ocir.io/<your_tenancy>/<your_image>:<your_tag>
+        entrypoint : str or list, optional
+            Entrypoint for the job, by default None (the entrypoint defined in the image will be used).
+        cmd : str, optional
+            Command for the job, by default None.
+        image_digest: str, optional
+            The image digest, by default None.
+        image_signature_id: str, optional
+            The image signature id, by default None.
+
+        Returns
+        -------
+        ContainerRuntime
+            The runtime instance.
+        """
+        if not isinstance(image, str):
+            raise ValueError(
+                "Custom image must be provided as a string."
+            )
+        if image.find(":") < 0:
+            logger.warning(
+                "Tag is required for custom image. Accepted format: iad.ocir.io/<tenancy>/<image>:<tag>."
+            )
+        self.with_entrypoint(entrypoint)
+        self.set_spec(self.CONST_CMD, cmd)
+        self.with_image_digest(image_digest)
+        self.with_image_signature_id(image_signature_id)
+        return self.set_spec(self.CONST_IMAGE, image)
+
+    @property
+    def image_digest(self) -> str:
+        """The container image digest."""
+        return self.get_spec(self.CONST_IMAGE_DIGEST)
+
+    def with_image_digest(self, image_digest: str) -> "ContainerRuntime":
+        """Sets the digest of custom image.
+
+        Parameters
+        ----------
+        image_digest: str
+            The image digest.
+
+        Returns
+        -------
+        ContainerRuntime
+            The runtime instance.
+        """
+        return self.set_spec(self.CONST_IMAGE_DIGEST, image_digest)
+
+    @property
+    def image_signature_id(self) -> str:
+        """The container image signature id."""
+        return self.get_spec(self.CONST_IMAGE_SIGNATURE_ID)
+
+    def with_image_signature_id(self, image_signature_id: str) -> "ContainerRuntime":
+        """Sets the signature id of custom image.
+
+        Parameters
+        ----------
+        image_signature_id: str
+            The image signature id.
+
+        Returns
+        -------
+        ContainerRuntime
+            The runtime instance.
+        """
+        return self.set_spec(
+            self.CONST_IMAGE_SIGNATURE_ID,
+            image_signature_id
+        )
+
+    @property
+    def entrypoint(self) -> str:
+        """Entrypoint of the container job"""
+        return self.get_spec(self.CONST_ENTRYPOINT)
+
+    def with_entrypoint(self, entrypoint: Union[str, list]) -> "ContainerRuntime":
+        """Specifies the entrypoint for the container job.
+
+        Parameters
+        ----------
+        entrypoint : str or list
+            Entrypoint for the container job
+
+        Returns
+        -------
+        ContainerRuntime
+            The runtime instance.
+        """
+        self._spec[self.CONST_ENTRYPOINT] = entrypoint
+        return self
+
+    @property
+    def cmd(self) -> str:
+        """Command of the container job"""
+        return self.get_spec(self.CONST_CMD)
+
+    def with_cmd(self, cmd: str) -> "ContainerRuntime":
+        """Specifies the command for the container job.
+
+        Parameters
+        ----------
+        cmd : str
+            Command for the container job
+
+        Returns
+        -------
+        ContainerRuntime
+            The runtime instance.
+        """
+        self._spec[self.CONST_CMD] = cmd
+        return self
+
+    def init(self, **kwargs) -> "ContainerRuntime":
+        """Initializes a starter specification for the runtime.
+
+        Returns
+        -------
+        ContainerRuntime
+            The runtime instance.
+        """
+        super().init(**kwargs)
+
+        return self.with_image(
+            image=kwargs.get("image", "iad.ocir.io/namespace/image:tag"),
+            entrypoint=["bash", "--login", "-c"],
+            cmd="{Container CMD. For MLflow and Operator will be auto generated}",
+        )
+
+    @property
+    def artifact_uri(self) -> str:
+        """The URI of the source code"""
+        return self.get_spec(self.CONST_SCRIPT_PATH)
+
+    def with_artifact(self, uri: str):
+        """Specifies the artifact to be added to the container.
+
+        Parameters
+        ----------
+        uri : str
+            URI to the source code script, which can be any URI supported by fsspec,
+            including http://, https:// and OCI object storage.
+            For example: oci://your_bucket@your_namespace/path/to/script.py
+
+        Returns
+        -------
+        self
+            The runtime instance.
+        """
+        return self.set_spec(self.CONST_SCRIPT_PATH, uri)