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

ads/opctl/distributed/cmds.py

@@ -0,0 +1,731 @@
+#!/usr/bin/env python
+# -*- coding: utf-8; -*-
+
+# Copyright (c) 2022, 2023 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+import json
+import os
+import re
+import subprocess
+from configparser import ConfigParser
+from urllib.parse import urlparse
+
+import fsspec
+import yaml
+
+from ads.common.auth import AuthContext, AuthType, create_signer
+from ads.common.decorator.runtime_dependency import (
+    OptionalDependency,
+    runtime_dependency,
+)
+from ads.jobs import Job
+from ads.jobs.builders.runtimes.artifact import Artifact
+from ads.opctl.utils import publish_image as publish_image_cmd
+from ads.opctl.utils import run_command
+
+ini_file = "config.ini"
+
+
+def _artifact_file_name(cluster_type, version):
+    return f"{cluster_type}_{version}.tar.gz".replace("-", "_")
+
+
+def initialize_workspace(cluster_type, version):
+    base_uri = None
+    local_file_name = f".tmp_{_artifact_file_name(cluster_type, version)}".replace(
+        "-", "_"
+    )
+    if os.environ.get("OCI_DISTRIBUTED_DEV_MODE"):
+        base_uri = dev_mode_base_uri(cluster_type, version)
+    else:
+        base_uri = production_mode_base_uri(cluster_type, version)
+    artificat_location = os.path.join(
+        base_uri, _artifact_file_name(cluster_type, version)
+    )
+    print(f"Downloading from {artificat_location} to {local_file_name}")
+    Artifact.copy_from_uri(uri=artificat_location, to_path=local_file_name)
+    os.makedirs("oci_dist_training_artifacts", exist_ok=True)
+    try:
+        subprocess.call(
+            ["tar", "-xvf", local_file_name, "-C", "oci_dist_training_artifacts"]
+        )
+        print(
+            f"Check oci_dist_training_artifacts/{cluster_type.replace('-', '_')}/{version}/README.md for build instructions "
+        )
+    finally:
+        if os.path.exists(local_file_name):
+            os.remove(local_file_name)
+
+
+def dev_mode_base_uri(cluster_type, version):
+    """
+    Temporary method to mock initialization using bucket.
+    """
+    bucket_name = os.environ.get("OCI_DISTRIBUTED_DEV_MODE_BUCKET")
+    namespace = os.environ.get("OCI_DISTRIBUTED_DEV_MODE_NAMESPACE")
+    prefix = os.environ.get("OCI_DISTRIBUTED_DEV_MODE_PREFIX")
+    base_uri = f"oci://{bucket_name}@{namespace}/{prefix}/"
+    return base_uri
+
+
+def production_mode_base_uri(cluster_type, version):
+    return "https://raw.githubusercontent.com/oracle-samples/oci-data-science-ai-samples/master/distributed_training/artifacts/"
+
+
+def cancel_distributed_run(job_id, cluster_file_name, **kwargs):
+    if cluster_file_name:
+        with open(cluster_file_name) as cf:
+            content = yaml.load(cf, yaml.SafeLoader)
+            job_id = content["jobId"]
+
+    with AuthContext():
+        import ads
+
+        ads.set_auth(auth="api_key", profile=kwargs.get("oci_profile", "DEFAULT"))
+
+        job = Job.from_datascience_job(job_id)
+        runs = job.run_list()
+        for job_run in runs:
+            print(f"Cancelling Job Run: {job_run.id}")
+            try:
+                job_run.cancel()
+            except Exception as e:
+                print(f"Error cancelling: {e}")
+        if not runs:
+            print(f"No Job runs found for : {job_id}")
+
+
+def show_config_info(job_id, work_dir, cluster_file_name, worker_info, **kwargs):
+    if cluster_file_name:
+        with open(cluster_file_name) as cf:
+            content = yaml.load(cf, yaml.SafeLoader)
+            job_id = content["jobId"]
+            work_dir = content["workDir"]
+    config_location = os.path.join(work_dir, job_id)
+    scheme = urlparse(
+        config_location,
+    ).scheme
+
+    oci_auth = (
+        create_signer(
+            AuthType.API_KEY,
+            kwargs.get("oci_config") or os.path.expanduser("~/.oci/config"),
+            kwargs.get("oci_profile") or "DEFAULT",
+        )
+        if scheme == "oci"
+        else {}
+    )
+    filesystem = fsspec.filesystem(
+        scheme, **oci_auth
+    )  # FileSystem class corresponding to the URI scheme.
+    main_file = os.path.join(config_location, "MAIN_config.json")
+    if filesystem.exists(main_file):
+        with fsspec.open(main_file, **oci_auth) as mfile:
+            print("Main Info:")
+            print(yaml.dump(json.loads(mfile.read())))
+
+        if worker_info:
+            files = fsspec.open_files(
+                os.path.join(config_location, "WORKER*.json"), **oci_auth
+            )
+            for fl in files:
+                with fsspec.open(
+                    f"{'oci://' if oci_auth else ''}{fl.path}", **oci_auth
+                ) as f:
+                    print(f"Worker Info from {os.path.basename(fl.path)}")
+                    print(yaml.dump(json.loads(f.read())))
+    else:
+        print(
+            f"MAIN_config file not found at location {config_location}. If you just started the cluster please wait until the main node is in `inprogress` state"
+        )
+
+
+@runtime_dependency(module="docker", install_from=OptionalDependency.OPCTL)
+def verify_image(img_name):
+    """
+    Verify if the input image exists in OCI registry
+
+    Parameters
+    ----------
+    img_name
+        Name of the docker image
+
+    Returns
+        None
+    -------
+
+    """
+    import docker
+
+    client = docker.from_env()
+    try:
+        client.images.get_registry_data(img_name)
+        return 1
+    except Exception:
+        return 0
+
+
+def update_ini(tag, registry, dockerfile, source_folder, config, nobuild):
+    """
+    Stores and updates input args in config.ini in Local
+
+    Parameters
+    ----------
+    tag
+        tag of image
+    registry
+        registry to push to
+    dockerfile
+        relative path to Dockerfile
+    source_folder
+        source scripts folder that will be mounted during a local run
+    config
+        Job config
+    nobuild
+        flag for building the docker
+
+    Returns
+        dict of config.ini file
+    -------
+
+    """
+    ini = ConfigParser(allow_no_value=True)
+    if config is None:
+        img_name = "@test"
+    else:
+        img_name = (
+            config.get("spec", {}).get("cluster", {}).get("spec", {}).get("image")
+        )
+
+    if tag is not None:
+        tag1 = tag
+    else:
+        if img_name.startswith("@"):
+            tag1 = tag
+        else:
+            if len(img_name.rsplit(":", 1)) == 1:
+                tag1 = "latest"
+            else:
+                tag1 = img_name.rsplit(":", 1)[1]
+    if registry is not None:
+        registry1 = registry
+    else:
+        if img_name.startswith("@"):
+            registry1 = registry
+        else:
+            registry1 = img_name.rsplit(":", 1)[0]
+
+    if os.path.isfile(ini_file):
+        ini.read(ini_file)
+        if tag1 is not None:
+            ini.set("main", "tag", tag1)
+        if registry1 is not None:
+            ini.set("main", "registry", registry1)
+        if dockerfile is not None:
+            ini.set("main", "dockerfile", dockerfile)
+        if source_folder is not None:
+            ini.set("main", "source_folder", source_folder)
+    else:
+        ini.add_section("main")
+        if tag1 is not None:
+            ini.set("main", "tag", tag1)
+        else:
+            raise ValueError("tag arg is missing")
+        if registry1 is not None:
+            ini.set("main", "registry", registry1)
+        else:
+            raise ValueError("registry arg is missing")
+        if dockerfile is not None:
+            ini.set("main", "dockerfile", dockerfile)
+        else:
+            if nobuild:
+                ini.set("main", "dockerfile", "DUMMY_PATH")
+            else:
+                raise ValueError("dockerfile arg is missing")
+        if source_folder is not None:
+            ini.set("main", "source_folder", source_folder)
+        else:
+            ini.set("main", "source_folder", ".")
+    ini.set("main", "; mount oci keys for local testing", None)
+    if ini.has_option("main", "oci_key_mnt"):
+        pass
+    else:
+        ini.set("main", "oci_key_mnt", "~/.oci:/home/oci_dist_training/.oci")
+    if ini.get("main", "dockerfile") == "DUMMY_PATH" and not nobuild:
+        raise ValueError("dockerfile arg is missing")
+    if os.path.exists(ini_file):
+        os.remove(ini_file)
+    with open(ini_file, "w") as f:
+        ini.write(f)
+    return ini
+
+
+def load_ini():
+    """
+    Loads config.ini from local in dict
+    Returns
+        None
+    -------
+
+    """
+    ini = ConfigParser(allow_no_value=True)
+    if os.path.isfile(ini_file):
+        ini.read(ini_file)
+        return ini
+    else:
+        raise RuntimeError(f"ini_file file {ini_file} not found !")
+
+
+def increment_tag_in_ini(ini):
+    """
+    increments tag of image and update latest in config.ini file
+
+    Parameters
+    ----------
+    ini
+        config.ini file dictionary
+    Returns
+        updated config.ini file
+    -------
+
+    """
+    ini = increment_tag(ini)
+    refresh_ini(ini, "config.ini")
+    return ini
+
+
+def increment_tag(ini):
+    """
+    Increments the tag of the image
+
+    Parameters
+    ----------
+    ini
+        config.ini file dictionary
+
+    Returns
+        updated ini dict
+    -------
+
+    """
+    tag_name = ini.get("main", "tag")
+    m = re.search(r"_v_\d+$", tag_name)
+    if m is not None:
+        updated_tag = re.sub(
+            r"[0-9]+$",
+            lambda x: f"{str(int(x.group()) + 1).zfill(len(x.group()))}",
+            tag_name,
+        )
+    else:
+        updated_tag = f"{tag_name}_v_1"
+    ini.set("main", "tag", updated_tag)
+    return ini
+
+
+def refresh_ini(ini, ini_file):
+    """
+    write updated config.ini in local
+
+    Parameters
+    ----------
+    ini
+        config.ini file dictionary
+    ini_file
+        Name of the config file
+
+    Returns
+    -------
+
+    """
+    if os.path.exists(ini_file):
+        os.remove(ini_file)
+    with open(ini_file, "w") as f:
+        ini.write(f)
+    return ini
+
+
+def docker_build_cmd(ini):
+    """
+    Builds the docker image
+
+    Parameters
+    ----------
+    ini
+        config.ini file dictionary
+
+    Returns
+        None
+    -------
+
+    """
+    cmd = get_cmd(ini)
+    return run_cmd(cmd)
+
+
+def get_cmd(ini):
+    """
+    Docker build command
+
+    Parameters
+    ----------
+    ini
+        config.ini file dictionary
+    Returns
+        Docker build command
+    -------
+
+    """
+    command = [
+        "docker",
+        "build",
+        "--build-arg",
+        "CODE_DIR=" + ini.get("main", "source_folder"),
+        "-t",
+        ini.get("main", "registry") + ":" + ini.get("main", "tag"),
+        "-f",
+        ini.get("main", "dockerfile"),
+    ]
+    if os.environ.get("no_proxy"):
+        command += ["--build-arg", f"no_proxy={os.environ['no_proxy']}"]
+    if os.environ.get("http_proxy"):
+        command += ["--build-arg", f"http_proxy={os.environ['http_proxy']}"]
+    if os.environ.get("https_proxy"):
+        command += ["--build-arg", f"https_proxy={os.environ['https_proxy']}"]
+    command += ["."]
+    return command
+
+
+def run_cmd(cmd):
+    """
+    Run cli command
+    Parameters
+    ----------
+    cmd
+        input cli command
+
+    Returns
+        None
+    -------
+
+    """
+    proc = run_command(cmd)
+    if proc.returncode != 0:
+        raise RuntimeError(f"Docker build failed.")
+    return 1
+
+
+def horovod_cmd(code_mount, oci_key_mount, config):
+    """
+    CLI command to run horovod distributed framework on local
+
+    Parameters
+    ----------
+    code_mount
+        source_folder to be mounted
+    oci_key_mount
+        oci keys to be mounted
+    config
+        Job config
+
+    Returns
+        List: CLI command
+    -------
+
+    """
+    command = [
+        "docker",
+        "run",
+        "-v",
+        code_mount,
+        "-v",
+        oci_key_mount,
+        "--env",
+        "OCI_IAM_TYPE=api_key",
+        "--rm",
+        "--entrypoint",
+        "/miniconda/envs/env/bin/horovodrun",
+        config["spec"]["cluster"]["spec"]["image"],
+        "--gloo",
+        "-np",
+        "2",
+        "-H",
+        "localhost:2",
+        "/miniconda/envs/env/bin/python",
+        config["spec"]["runtime"]["spec"]["entryPoint"],
+    ]
+    return command
+
+
+def pytorch_cmd(code_mount, oci_key_mount, config):
+    """
+    CLI command to run Pytorch distributed framework on local
+
+    Parameters
+    ----------
+    code_mount
+        source_folder to be mounted
+    oci_key_mount
+        oci keys to be mounted
+    config
+        Job config
+
+    Returns
+        List: CLI command
+    -------
+
+    """
+    command = [
+        "docker",
+        "run",
+        "-v",
+        code_mount,
+        "-v",
+        oci_key_mount,
+        "--env",
+        "OCI_IAM_TYPE=api_key",
+        "--env",
+        "WORLD_SIZE=1",
+        "--env",
+        "RANK=0",
+        "--env",
+        "LOCAL_RANK=0",
+        "--rm",
+        "--entrypoint",
+        "/opt/conda/bin/python",
+        config["spec"]["cluster"]["spec"]["image"],
+        config["spec"]["runtime"]["spec"]["entryPoint"],
+    ]
+    return command
+
+
+def dask_cmd(code_mount, oci_key_mount, config):
+    """
+    CLI command to run Dask distributed framework on local
+
+    Parameters
+    ----------
+    code_mount
+        source_folder to be mounted
+    oci_key_mount
+        oci keys to be mounted
+    config
+        Job config
+
+    Returns
+        List: CLI command
+    -------
+
+    """
+    command = [
+        "docker",
+        "run",
+        "-v",
+        code_mount,
+        "-v",
+        oci_key_mount,
+        "--env",
+        "OCI_IAM_TYPE=api_key",
+        "--env",
+        "SCHEDULER_IP=tcp://127.0.0.1",
+        "--rm",
+        "--entrypoint",
+        "/bin/sh",
+        config["spec"]["cluster"]["spec"]["image"],
+        "-c",
+        "(nohup dask-scheduler >scheduler.log &) && (nohup dask-worker localhost:8786 >worker.log &) && "
+        "/miniconda/envs/daskenv/bin/python "
+        + config["spec"]["runtime"]["spec"]["entryPoint"],
+    ]
+    return command
+
+
+def tensorflow_cmd(code_mount, oci_key_mount, config):
+    """
+    CLI command to run Tensorflow distributed framework on local
+
+    Parameters
+    ----------
+    code_mount
+        source_folder to be mounted
+    oci_key_mount
+        oci keys to be mounted
+    config
+        Job config
+
+    Returns
+        List: CLI command
+    -------
+
+    """
+    if "ps" in config["spec"]["cluster"]["spec"]:
+        # this is for the parameter server strategy.
+        command = [
+            "docker",
+            "run",
+            "-v",
+            code_mount,
+            "-v",
+            oci_key_mount,
+            "--env",
+            "OCI_IAM_TYPE=api_key",
+            "--rm",
+            "--entrypoint",
+            "/bin/sh",
+            config["spec"]["cluster"]["spec"]["image"],
+            "/etc/datascience/local_test.sh",
+            config["spec"]["runtime"]["spec"]["entryPoint"],
+        ]
+
+    else:
+        # this is for the mirror and multiworkermirror strategy.
+        command = [
+            "docker",
+            "run",
+            "-v",
+            code_mount,
+            "-v",
+            oci_key_mount,
+            "--env",
+            "OCI_IAM_TYPE=api_key",
+            "--env",
+            'TF_CONFIG={"cluster": {"worker": ["localhost:12345"]}, "task": {"type": "worker", "index": 0}}',
+            "--rm",
+            "--entrypoint",
+            "/miniconda/bin/python",
+            config["spec"]["cluster"]["spec"]["image"],
+            config["spec"]["runtime"]["spec"]["entryPoint"],
+        ]
+    return command
+
+
+def local_run(config, ini):
+    """
+
+    Local run distributed framework on local based on input args
+
+    Parameters
+    ----------
+    ini
+        config.ini file dictionary
+    config
+        Job run config
+
+    Returns
+        None
+    -------
+
+    """
+    cwd = os.path.join(os.getcwd(), ini.get("main", "source_folder"))
+    code_mount = os.path.join(cwd, ":/code/")
+    oci_key_mount = (
+        os.path.expanduser(ini.get("main", "oci_key_mnt").split(":")[0])
+        + ":"
+        + ini.get("main", "oci_key_mnt").split(":")[1]
+    )
+    if config["spec"]["cluster"]["kind"].lower() == "horovod":
+        command = horovod_cmd(code_mount, oci_key_mount, config)
+    elif config["spec"]["cluster"]["kind"].lower() == "pytorch":
+        command = pytorch_cmd(code_mount, oci_key_mount, config)
+    elif config["spec"]["cluster"]["kind"].lower() == "dask":
+        command = dask_cmd(code_mount, oci_key_mount, config)
+    elif config["spec"]["cluster"]["kind"].lower() == "tensorflow":
+        command = tensorflow_cmd(code_mount, oci_key_mount, config)
+    else:
+        raise RuntimeError(f"Framework not supported")
+    try:
+        command += [str(arg) for arg in config["spec"]["runtime"]["spec"]["args"]]
+    except KeyError:
+        pass
+    print("Running: ", " ".join(command))
+    proc = run_command(command)
+    if proc.returncode != 0:
+        raise RuntimeError(f"Failed to run local")
+    return 1
+
+
+def update_image(config, ini):
+    """
+    Update Image name in job config with latest build image
+
+    Parameters
+    ----------
+    config
+        Job run config
+    ini
+        config.ini file dictionary
+
+    Returns
+        Updated Job run config
+    -------
+
+    """
+    config["spec"]["cluster"]["spec"]["image"] = (
+        ini.get("main", "registry") + ":" + ini.get("main", "tag")
+    )
+    return config
+
+
+def verify_and_publish_image(nopush, config):
+    """
+    verify existence of docker image in OCI registry and publish it
+    Parameters
+    ----------
+    nopush
+        Doesn't push image to OCI registry
+    config
+        Job run config
+
+    Returns
+        None
+    -------
+
+    """
+    if not nopush:
+        publish_image_cmd(config["spec"]["cluster"]["spec"]["image"])
+    else:
+        if not verify_image(config["spec"]["cluster"]["spec"]["image"]):
+            print(
+                "\u26A0 Image: "
+                + config["spec"]["cluster"]["spec"]["image"]
+                + " does not exist in registry"
+            )
+            print("In order to push the image to registry enter Y else N ")
+            inp = input("[Y/N]\n")
+            if inp == "Y":
+                print("\u2705 pushing image to registry")
+                publish_image_cmd(config["spec"]["cluster"]["spec"]["image"])
+            else:
+                raise RuntimeError(
+                    "Stopping the execution as image doesn't exist in OCI registry"
+                )
+    return 1
+
+
+def update_config_image(config):
+    """
+    updates image name in config
+
+    Parameters
+    ----------
+    config
+        Job config
+
+    Returns
+        updated config dict
+    -------
+
+    """
+    ini = ConfigParser(allow_no_value=True)
+    img_name = config.get("spec", {}).get("cluster", {}).get("spec", {}).get("image")
+    if img_name.startswith("@"):
+        if os.path.isfile(ini_file):
+            ini.read(ini_file)
+            return update_image(config, ini)
+        else:
+            raise ValueError(
+                f"Invalid image name: {img_name} and also not able to locate config.ini file. "
+                f"Please update image name in Job config."
+            )
+    else:
+        return config