This repository has been archived by the owner on Feb 3, 2021. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 65
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Feature: New Toolkit configuration (#507)
- Loading branch information
1 parent
9bc7639
commit 7a7e63c
Showing
27 changed files
with
512 additions
and
106 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1 +1,2 @@ | ||
from .toolkit import Toolkit, TOOLKIT_MAP | ||
from .models import * |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,115 @@ | ||
from aztk.internal import ConfigurationBase | ||
from aztk.error import InvalidModelError | ||
from aztk.utils import constants, deprecate | ||
|
||
|
||
class ToolkitDefinition: | ||
def __init__(self, versions, environments): | ||
self.versions = versions | ||
self.environments = environments | ||
|
||
class ToolkitEnvironmentDefinition: | ||
def __init__(self, versions=None, default=""): | ||
self.versions = versions or [""] | ||
self.default = default | ||
|
||
TOOLKIT_MAP = dict( | ||
spark=ToolkitDefinition( | ||
versions=["1.6.3", "2.1.0", "2.2.0", "2.3.0"], | ||
environments=dict( | ||
base=ToolkitEnvironmentDefinition(), | ||
r=ToolkitEnvironmentDefinition(), | ||
miniconda=ToolkitEnvironmentDefinition(), | ||
anaconda=ToolkitEnvironmentDefinition(), | ||
) | ||
), | ||
) | ||
|
||
|
||
class Toolkit(ConfigurationBase): | ||
""" | ||
Toolkit for a cluster. | ||
This will help pick the docker image needed | ||
Args: | ||
software (str): Name of the toolkit(spark) | ||
version (str): Version of the toolkit | ||
environment (str): Which environment to use for this toolkit | ||
environment_version (str): If there is multiple version for an environment you can specify which one | ||
""" | ||
def __init__(self, | ||
software: str, | ||
version: str, | ||
environment: str = None, | ||
environment_version: str = None, | ||
docker_repo=None): | ||
|
||
self.software = software | ||
self.version = str(version) | ||
self.environment = environment | ||
self.environment_version = environment_version | ||
self.docker_repo = docker_repo | ||
|
||
|
||
def validate(self): | ||
self._validate_required(["software", "version"]) | ||
|
||
if self.software not in TOOLKIT_MAP: | ||
raise InvalidModelError("Toolkit '{0}' is not in the list of allowed toolkits {1}".format( | ||
self.software, list(TOOLKIT_MAP.keys()))) | ||
|
||
toolkit_def = TOOLKIT_MAP[self.software] | ||
|
||
if self.version not in toolkit_def.versions: | ||
raise InvalidModelError("Toolkit '{0}' with version '{1}' is not available. Use one of: {2}".format( | ||
self.software, self.version, toolkit_def.versions)) | ||
if self.version == "1.6": | ||
deprecate("Spark version 1.6 is being deprecated for Aztk. Please use 2.1 and above.") | ||
|
||
if self.environment: | ||
if self.environment not in toolkit_def.environments: | ||
raise InvalidModelError("Environment '{0}' for toolkit '{1}' is not available. Use one of: {2}".format( | ||
self.environment, self.software, list(toolkit_def.environments.keys()))) | ||
|
||
env_def = toolkit_def.environments[self.environment] | ||
|
||
if self.environment_version and self.environment_version not in env_def.versions: | ||
raise InvalidModelError( | ||
"Environment '{0}' version '{1}' for toolkit '{2}' is not available. Use one of: {3}".format( | ||
self.environment, self.environment_version, self.software, env_def.versions)) | ||
|
||
|
||
def get_docker_repo(self, gpu: bool): | ||
if self.docker_repo: | ||
return self.docker_repo | ||
|
||
repo = "aztk/{0}".format(self.software) | ||
|
||
return "{repo}:{tag}".format( | ||
repo=repo, | ||
tag=self._get_docker_tag(gpu), | ||
) | ||
|
||
def _get_docker_tag(self, gpu: bool): | ||
environment = self.environment or "base" | ||
environment_def = self._get_environent_definition() | ||
environment_version = self.environment_version or (environment_def and environment_def.default) | ||
|
||
array = [ | ||
"v{docker_image_version}".format(docker_image_version=constants.DOCKER_IMAGE_VERSION), | ||
"{toolkit}{version}".format(toolkit=self.software, version=self.version), | ||
] | ||
if self.environment: | ||
array.append("{0}{1}".format(environment, environment_version)) | ||
|
||
array.append("gpu" if gpu else "base") | ||
|
||
return '-'.join(array) | ||
|
||
|
||
def _get_environent_definition(self) -> ToolkitEnvironmentDefinition: | ||
toolkit = TOOLKIT_MAP.get(self.software) | ||
|
||
if toolkit: | ||
return toolkit.environments.get(self.environment or "base") | ||
return None |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.