Conda Pack build related enhancements (#908)

Author: qiuosier

ads/common/auth.py

@@ -73,7 +73,7 @@ def __post_init__(self):
         self.oci_key_profile = self.oci_key_profile or os.environ.get(
             "OCI_CONFIG_PROFILE", DEFAULT_PROFILE
         )
-        self.oci_config = self.oci_config or {}
+        self.oci_config = self.oci_config or {"region": os.environ["OCI_RESOURCE_REGION"]} if os.environ.get("OCI_RESOURCE_REGION") else {}
         self.oci_signer_kwargs = self.oci_signer_kwargs or {}
         self.oci_client_kwargs = self.oci_client_kwargs or {}
 
@@ -82,7 +82,7 @@ def set_auth(
     auth: Optional[str] = AuthType.API_KEY,
     oci_config_location: Optional[str] = DEFAULT_LOCATION,
     profile: Optional[str] = DEFAULT_PROFILE,
-    config: Optional[Dict] = {},
+    config: Optional[Dict] = {"region": os.environ["OCI_RESOURCE_REGION"]} if os.environ.get("OCI_RESOURCE_REGION") else {},
     signer: Optional[Any] = None,
     signer_callable: Optional[Callable] = None,
     signer_kwargs: Optional[Dict] = {},
@@ -678,7 +678,7 @@ def create_signer(self) -> Dict:
         >>> signer_generator = AuthFactory().signerGenerator(AuthType.RESOURCE_PRINCIPAL)
         >>> signer_generator(signer_args).create_signer()
         """
-        configuration = ads.telemetry.update_oci_client_config()
+        configuration = ads.telemetry.update_oci_client_config(AuthState().oci_config)
         signer_dict = {
             "config": configuration,
             "signer": oci.auth.signers.get_resource_principals_signer(),
@@ -739,7 +739,7 @@ def create_signer(self) -> Dict:
         >>> signer_generator = AuthFactory().signerGenerator(AuthType.INSTANCE_PRINCIPAL)
         >>> signer_generator(signer_args).create_signer()
         """
-        configuration = ads.telemetry.update_oci_client_config()
+        configuration = ads.telemetry.update_oci_client_config(AuthState().oci_config)
         signer_dict = {
             "config": configuration,
             "signer": oci.auth.signers.InstancePrincipalsSecurityTokenSigner(

ads/opctl/conda/cmds.py

@@ -80,7 +80,7 @@ def _check_job_image_exists(gpu: bool) -> None:
 def _get_name(name: str, env_file: str) -> str:
     if not name and env_file:
         with open(env_file) as f:
-            name = yaml.safe_load(f.read()).get("name", None)
+            name = yaml.safe_load(f.read()).get("manifest").get("name", None)
     if not name:
         raise ValueError(
             "Either specify environment name in environment yaml or with `--name`."
@@ -146,7 +146,14 @@ def _create(
     if not os.path.exists(env_file):
         raise FileNotFoundError(f"Environment file {env_file} is not found.")
 
-    slug = f"{name}_v{version}".replace(" ", "").replace(".", "_").lower()
+    conda_dep = None
+    with open(env_file) as mfile:
+        conda_dep = yaml.safe_load(mfile.read())
+    # If manifest exists in the environment.yaml file, use that
+    manifest = conda_dep.get("manifest", {})
+    slug = manifest.get(
+        "slug", f"{name}_v{version}".replace(" ", "").replace(".", "_").lower()
+    )
     pack_folder_path = os.path.join(
         os.path.abspath(os.path.expanduser(conda_pack_folder)), slug
     )
@@ -171,23 +178,38 @@ def _create(
 
     os.makedirs(pack_folder_path, exist_ok=True)
 
+    logger.info(
+        f"Preparing manifest. Manifest in the environment: {conda_dep.get('manifest')}"
+    )
     manifest = _fetch_manifest_template()
-    manifest["manifest"]["name"] = name
+    if "name" not in manifest:
+        manifest["manifest"]["name"] = name
     manifest["manifest"]["slug"] = slug
-    manifest["manifest"]["type"] = "published"
-    manifest["manifest"]["version"] = version
+    if "type" not in manifest:
+        logger.info("Setting manifest to published")
+        manifest["manifest"]["type"] = "published"
+    if "version" not in manifest:
+        manifest["manifest"]["version"] = version
     manifest["manifest"]["arch_type"] = "GPU" if gpu else "CPU"
 
     manifest["manifest"]["create_date"] = datetime.utcnow().strftime(
         "%a, %b %d, %Y, %H:%M:%S %Z UTC"
     )
-    manifest["manifest"]["manifest_version"] = "1.0"
+
+    if not "manifest_version" in manifest:
+        manifest["manifest"]["manifest_version"] = "1.0"
 
     logger.info(f"Creating conda environment {slug}")
-    conda_dep = None
-    with open(env_file) as mfile:
-        conda_dep = yaml.safe_load(mfile.read())
-    conda_dep["manifest"] = manifest["manifest"]
+    manifest_dict = {
+        k: manifest["manifest"][k]
+        for k in manifest["manifest"]
+        if manifest["manifest"][k]
+    }
+    if "manifest" in conda_dep:
+        conda_dep["manifest"].update(manifest_dict)
+    else:
+        conda_dep["manifest"] = manifest_dict
+    logger.info(f"Updated conda environment manifest: {conda_dep.get('manifest')}")
 
     if is_in_notebook_session() or NO_CONTAINER:
         command = f"conda env create --prefix {pack_folder_path} --file {os.path.abspath(os.path.expanduser(env_file))}"
@@ -200,13 +222,12 @@ def _create(
         )
 
         create_command = f"conda env create --prefix {docker_pack_folder_path} --file {docker_env_file_path}"
-            
+
         volumes = {
             pack_folder_path: {"bind": docker_pack_folder_path},
             os.path.abspath(os.path.expanduser(env_file)): {
                 "bind": docker_env_file_path
             },
-
         }
 
         if gpu:
@@ -217,26 +238,42 @@ def _create(
             if prepare_publish:
                 tmp_file = tempfile.NamedTemporaryFile(suffix=".yaml")
                 # Save the manifest in the temp file that can be mounted inside the container so that archiving will work
-                with open(tmp_file.name, 'w') as f:
-                    yaml.safe_dump(conda_dep, f)      
+                with open(tmp_file.name, "w") as f:
+                    yaml.safe_dump(conda_dep, f)
 
-                pack_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pack.py")
+                pack_script = os.path.join(
+                    os.path.dirname(os.path.abspath(__file__)), "pack.py"
+                )
                 pack_command = f"python {os.path.join(DEFAULT_IMAGE_HOME_DIR, 'pack.py')} --conda-path {docker_pack_folder_path} --manifest-location {os.path.join(DEFAULT_IMAGE_HOME_DIR, 'manifest.yaml')}"
 
                 # add pack script and manifest file to the mount so that archive can be created in the same container run
                 condapack_script = {
-                    pack_script: {"bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, "pack.py")},
-                    tmp_file.name: {"bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, "manifest.yaml")}
+                    pack_script: {
+                        "bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, "pack.py")
+                    },
+                    tmp_file.name: {
+                        "bind": os.path.join(DEFAULT_IMAGE_HOME_DIR, "manifest.yaml")
+                    },
                 }
-                volumes = {**volumes, **condapack_script} # | not supported in python 3.8
+                volumes = {
+                    **volumes,
+                    **condapack_script,
+                }  # | not supported in python 3.8
 
                 run_container(
-                    image=image, bind_volumes=volumes, entrypoint="/bin/bash -c ", env_vars={}, command=f" '{create_command} && {pack_command}'"
+                    image=image,
+                    bind_volumes=volumes,
+                    entrypoint="/bin/bash -c ",
+                    env_vars={},
+                    command=f" '{create_command} && {pack_command}'",
                 )
             else:
                 run_container(
-                    image=image, bind_volumes=volumes, env_vars={}, command=create_command
-                )                
+                    image=image,
+                    bind_volumes=volumes,
+                    env_vars={},
+                    command=create_command,
+                )
         except Exception:
             if os.path.exists(pack_folder_path):
                 shutil.rmtree(pack_folder_path)
@@ -507,9 +544,11 @@ def publish(**kwargs) -> None:
             conda_pack_folder=exec_config["conda_pack_folder"],
             gpu=exec_config.get("gpu", False),
             overwrite=exec_config["overwrite"],
-            prepare_publish=True
+            prepare_publish=True,
+        )
+        skip_archive = (
+            True  # The conda pack archive is already created during create process.
         )
-        skip_archive = True # The conda pack archive is already created during create process.
     else:
         slug = exec_config.get("slug")
     if not slug:
@@ -526,10 +565,10 @@ def publish(**kwargs) -> None:
         oci_profile=exec_config.get("oci_profile"),
         overwrite=exec_config["overwrite"],
         auth_type=exec_config["auth"],
-        skip_archive=skip_archive
+        skip_archive=skip_archive,
     )
 
-    
+
 def _publish(
     conda_slug: str,
     conda_uri_prefix: str,
@@ -538,7 +577,7 @@ def _publish(
     oci_profile: str,
     overwrite: bool,
     auth_type: str,
-    skip_archive: bool = False
+    skip_archive: bool = False,
 ) -> None:
     """Publish a local conda pack to object storage location
 
@@ -616,7 +655,11 @@ def _publish(
     pack_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pack.py")
     if not skip_archive:
         if is_in_notebook_session() or NO_CONTAINER:
+            # Set the CONDA_PUBLISH_TYPE environment variable so that the `type` attribute inside the manifest is not changed
+            publish_type = os.environ.get("CONDA_PUBLISH_TYPE")
             command = f"python {pack_script} --conda-path {pack_folder_path}"
+            if publish_type:
+                command = f"CONDA_PUBLISH_TYPE={publish_type} {command}"
             run_command(command, shell=True)
         else:
             volumes = {
@@ -641,7 +684,9 @@ def _publish(
     NOT_ALLOWED_CHARS = "@#$%^&*/"
 
     if any(chr in conda_slug for chr in NOT_ALLOWED_CHARS):
-        raise ValueError(f"Invalid conda_slug. Found {NOT_ALLOWED_CHARS} in slug name. Please use a different slug name.")
+        raise ValueError(
+            f"Invalid conda_slug. Found {NOT_ALLOWED_CHARS} in slug name. Please use a different slug name."
+        )
     pack_file = os.path.join(pack_folder_path, f"{conda_slug}.tar.gz")
     if not os.path.exists(pack_file):
         raise RuntimeError(f"Pack {pack_file} was not created.")
@@ -664,14 +709,19 @@ def _publish(
         str(manifest["version"]),
         publish_slug,
     )
-    manifest["pack_path"] = os.path.join(
-        prefix,
-        manifest.get("arch_type", "CPU").lower(),
-        manifest["name"],
-        str(manifest["version"]),
-        publish_slug,
-    )
-    manifest["pack_uri"] = pack_uri
+    if os.environ.get("CONDA_PUBLISH_TYPE") != "service":
+        # Set these values only for published conda pack
+        manifest["pack_path"] = os.path.join(
+            prefix,
+            manifest.get("arch_type", "CPU").lower(),
+            manifest["name"],
+            str(manifest["version"]),
+            publish_slug,
+        )
+        manifest["pack_uri"] = pack_uri
+    else:
+        manifest["type"] = "published"
+
     with open(manifest_location, "w") as f:
         yaml.safe_dump(env, f)
     if pack_size > 100:

ads/opctl/conda/pack.py

@@ -20,7 +20,6 @@
 
 
 def main(pack_folder_path, manifest_file=None):
-    slug = os.path.basename(pack_folder_path)
     manifest_path = (
         manifest_file or glob.glob(os.path.join(pack_folder_path, "*_manifest.yaml"))[0]
     )
@@ -47,7 +46,9 @@ def main(pack_folder_path, manifest_file=None):
             raise e
 
         manifest = env["manifest"]
-        manifest["type"] = "published"
+        slug = manifest.get("slug", os.path.basename(pack_folder_path))
+        if os.environ.get("CONDA_PUBLISH_TYPE") != "service":
+            manifest["type"] = "published"
         new_env_info["manifest"] = manifest
         with open(manifest_path, "w") as f:
             yaml.safe_dump(new_env_info, f)

tests/unitary/with_extras/ads_string/test_ads_string.py

@@ -48,6 +48,14 @@ def setup_class(self):
             "averaged_perceptron_tagger",
             download_dir=os.environ["CONDA_PREFIX"] + "/nltk",
         )
+        nltk.download(
+            "averaged_perceptron_tagger_eng",
+            download_dir=os.environ["CONDA_PREFIX"] + "/nltk",
+        )
+        nltk.download(
+            "punkt_tab",
+            download_dir=os.environ["CONDA_PREFIX"] + "/nltk",
+        )
 
     def test_basic_string_methods(self):
         s = ADSString("HELLO WORLD")

tests/unitary/with_extras/feature_types/test_feature_types.py

@@ -45,6 +45,14 @@ def setup_class(self):
             "averaged_perceptron_tagger",
             download_dir=os.environ["CONDA_PREFIX"] + "/nltk",
         )
+        nltk.download(
+            "averaged_perceptron_tagger_eng",
+            download_dir=os.environ["CONDA_PREFIX"] + "/nltk",
+        )
+        nltk.download(
+            "punkt_tab",
+            download_dir=os.environ["CONDA_PREFIX"] + "/nltk",
+        )
 
     def test_datatime_type(self):
         assert (