From 031fded07dbee95a1addeab2d76dd8536e2e3c7e Mon Sep 17 00:00:00 2001 From: Kevin REMY Date: Fri, 20 Dec 2024 10:17:59 +0100 Subject: [PATCH 1/3] Sync sources --- dataikuapi/apinode_admin/auth.py | 5 +- dataikuapi/dss/jupyternotebook.py | 7 + dataikuapi/dss/langchain/embeddings.py | 69 +++-- dataikuapi/dss/langchain/llm.py | 108 +++++-- dataikuapi/dss/langchain/utils.py | 7 + dataikuapi/dss/llm.py | 324 +++++++++++++++++--- dataikuapi/dss/llm_tracing/__init__.py | 156 ++++++++++ dataikuapi/dss/messaging_channel.py | 11 + dataikuapi/dss/ml.py | 45 ++- dataikuapi/dss/plugin.py | 2 +- dataikuapi/dss/project.py | 56 ++++ dataikuapi/dss/projectdeployer.py | 78 +++-- dataikuapi/dss/scenario.py | 67 +++++ dataikuapi/fm/cloudaccounts.py | 338 +++++++++++++++++++++ dataikuapi/fm/loadbalancers.py | 398 +++++++++++++++++++++++++ dataikuapi/fm/tenant.py | 4 +- dataikuapi/fm/virtualnetworks.py | 45 ++- dataikuapi/fmclient.py | 143 ++++++++- 18 files changed, 1728 insertions(+), 135 deletions(-) create mode 100644 dataikuapi/dss/langchain/utils.py create mode 100644 dataikuapi/dss/llm_tracing/__init__.py create mode 100644 dataikuapi/fm/cloudaccounts.py create mode 100644 dataikuapi/fm/loadbalancers.py diff --git a/dataikuapi/apinode_admin/auth.py b/dataikuapi/apinode_admin/auth.py index 00717426..093abbf6 100644 --- a/dataikuapi/apinode_admin/auth.py +++ b/dataikuapi/apinode_admin/auth.py @@ -10,12 +10,13 @@ def list_keys(self): """Lists the Admin API keys""" return self.client._perform_json("GET", "keys") - def add_key(self, label=None, description=None, created_by=None): + def add_key(self, label=None, description=None, created_by=None, expiry=None): """Add an Admin API key. Returns the key details""" key = { "label" : label, "description" : description, - "createdBy" : created_by + "createdBy" : created_by, + "expiry" : expiry } return self.client._perform_json("POST", "keys", body=key) diff --git a/dataikuapi/dss/jupyternotebook.py b/dataikuapi/dss/jupyternotebook.py index 0e217cf6..3a47e339 100644 --- a/dataikuapi/dss/jupyternotebook.py +++ b/dataikuapi/dss/jupyternotebook.py @@ -148,6 +148,13 @@ def delete(self): return self.client._perform_json("DELETE", "/projects/%s/jupyter-notebooks/%s" % (self.project_key, self.notebook_name)) + def clear_outputs(self): + """ + Clear this Jupyter notebook's outputs. + """ + return self.client._perform_json("DELETE", + "/projects/%s/jupyter-notebooks/%s/outputs" % (self.project_key, self.notebook_name)) + ######################################################## # Discussions ######################################################## diff --git a/dataikuapi/dss/langchain/embeddings.py b/dataikuapi/dss/langchain/embeddings.py index 7b5379c4..d7aaa99a 100644 --- a/dataikuapi/dss/langchain/embeddings.py +++ b/dataikuapi/dss/langchain/embeddings.py @@ -2,18 +2,33 @@ import asyncio import concurrent import logging -from typing import List, Any +import threading -from pydantic import BaseModel, Extra +from typing import List, Any +import pydantic from langchain.embeddings.base import Embeddings +from dataikuapi.dss.llm_tracing import new_trace +from dataikuapi.dss.langchain.utils import must_use_deprecated_pydantic_config logger = logging.getLogger(__name__) CHUNK_SIZE = 1000 -class DKUEmbeddings(BaseModel, Embeddings): +if must_use_deprecated_pydantic_config(): + class LockedDownBaseModel(pydantic.BaseModel): + class Config: + extra = pydantic.Extra.forbid + underscore_attrs_are_private = True +else: + class LockedDownBaseModel(pydantic.BaseModel): + model_config = { + 'extra': 'forbid', + } + + +class DKUEmbeddings(LockedDownBaseModel, Embeddings): """ Langchain-compatible wrapper around Dataiku-mediated embedding LLMs @@ -27,9 +42,12 @@ class DKUEmbeddings(BaseModel, Embeddings): _llm_handle = None """:class:`dataikuapi.dss.llm.DSSLLM` object to wrap.""" - class Config: - extra = Extra.forbid - underscore_attrs_are_private = True + # The embeddings class of LangChain can only return raw embedding, without any additional information + # (unlike ChatModel, which supports additional information), so we cannot use this to return the last + # trace to the caller. + # So, instead, we keep a thread local with the last trace, and the caller can get it from here + # (at the moment, it's mostly done by rag_query_server.py) + _last_trace = None def __init__(self, llm_handle=None, **data: Any): if llm_handle is None: @@ -45,6 +63,7 @@ def __init__(self, llm_handle=None, **data: Any): super().__init__(**data) self._llm_handle = llm_handle + self._last_trace = threading.local() def embed_documents(self, texts: List[str]) -> List[List[float]]: """Call out to Dataiku-mediated LLM @@ -55,29 +74,37 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]: Returns: List of embeddings, one for each text. """ - logging.info("Performing embedding of {num_texts} texts".format(num_texts=len(texts))) - embeddings = [] - for i in range(0, len(texts), CHUNK_SIZE): - query = self._llm_handle.new_embeddings(text_overflow_mode="FAIL") + with new_trace("DKUEmbeddings") as trace: + self._last_trace.trace = trace + + logger.info("Performing embedding of {num_texts} texts".format(num_texts=len(texts))) + + embeddings = [] + for i in range(0, len(texts), CHUNK_SIZE): + query = self._llm_handle.new_embeddings(text_overflow_mode="FAIL") + + for text in texts[i:i+CHUNK_SIZE]: + query.add_text(text) - for text in texts[i:i+CHUNK_SIZE]: - query.add_text(text) + resp = query.execute() - resp = query.execute() + # TODO + #if not resp.success: + # raise Exception("LLM call failed: %s" % resp._raw.get("errorMessage", "Unknown error")) - # TODO - #if not resp.success: - # raise Exception("LLM call failed: %s" % resp._raw.get("errorMessage", "Unknown error")) + if "responses" in resp._raw and len(resp._raw["responses"]) == 1: + if "trace" in resp._raw["responses"][0]: + trace.append_trace(resp._raw["responses"][0]["trace"]) - embeddings.extend(resp.get_embeddings()) + embeddings.extend(resp.get_embeddings()) - logging.info("Finished a chunk. Embedded {num_embedded} of {num_texts} texts".format( - num_embedded=min(i + CHUNK_SIZE, len(texts)), num_texts=len(texts))) + logger.info("Finished a chunk. Embedded {num_embedded} of {num_texts} texts".format( + num_embedded=min(i + CHUNK_SIZE, len(texts)), num_texts=len(texts))) - logging.info("Done performing embedding of {num_texts} texts".format(num_texts=len(texts))) + logger.info("Done performing embedding of {num_texts} texts".format(num_texts=len(texts))) - return embeddings + return embeddings async def aembed_documents(self, texts: List[str]) -> List[List[float]]: loop = asyncio.get_event_loop() diff --git a/dataikuapi/dss/langchain/llm.py b/dataikuapi/dss/langchain/llm.py index c307f02b..8e7123f5 100644 --- a/dataikuapi/dss/langchain/llm.py +++ b/dataikuapi/dss/langchain/llm.py @@ -17,7 +17,7 @@ Union, ) -from langchain.callbacks.manager import CallbackManagerForLLMRun, AsyncCallbackManagerForLLMRun +from langchain.callbacks.manager import CallbackManagerForLLMRun from langchain.llms.base import BaseLLM from langchain_core.messages import ( AIMessage, @@ -30,20 +30,21 @@ ToolCall, ToolMessage, ) +from langchain_core.messages.ai import UsageMetadata from langchain_core.outputs import Generation, GenerationChunk, ChatGenerationChunk, LLMResult, ChatResult from langchain_core.output_parsers.openai_tools import ( make_invalid_tool_call, parse_tool_call, ) from langchain_core.language_models import BaseChatModel -from langchain_core.pydantic_v1 import BaseModel from langchain_core.tools import BaseTool from langchain_core.utils.function_calling import convert_to_openai_tool from langchain.schema import ChatGeneration -from pydantic import Extra +import pydantic +from dataikuapi.dss.langchain.utils import must_use_deprecated_pydantic_config from dataikuapi.dss.tools.langchain import StopSequencesAwareStreamer - +from dataikuapi.dss.llm import DSSLLMStreamedCompletionFooter logger = logging.getLogger(__name__) @@ -131,7 +132,19 @@ def _enforce_stop_sequences(text: str, stop: List[str]) -> str: return re.split("|".join(escaped_stop), text, maxsplit=1)[0] -class DKULLM(BaseLLM): +if must_use_deprecated_pydantic_config(): + class LockedDownBaseLLM(BaseLLM): + class Config: + extra = pydantic.Extra.forbid + underscore_attrs_are_private = True +else: + class LockedDownBaseLLM(BaseLLM): + model_config = { + 'extra': 'forbid', + } + + +class DKULLM(LockedDownBaseLLM): """ Langchain-compatible wrapper around Dataiku-mediated LLMs @@ -174,10 +187,6 @@ class DKULLM(BaseLLM): _llm_handle = None """:class:`dataikuapi.dss.llm.DSSLLM` object to wrap.""" - class Config: - extra = Extra.forbid - underscore_attrs_are_private = True - def __init__(self, llm_handle=None, **data: Any): if llm_handle is None: if data.get("llm_id") is None: @@ -223,6 +232,9 @@ def _generate( total_estimated_cost = 0.0 generations = [] + + trace_of_last_response = None + for prompt_resp in resp.responses: if not prompt_resp.success: raise Exception("LLM call failed: %s" % prompt_resp._raw.get("errorMessage", "Unknown error")) @@ -236,12 +248,16 @@ def _generate( text = _enforce_stop_sequences(prompt_resp.text, stop) generations.append([Generation(text=text)]) + if prompt_resp._raw.get("trace", None) is not None: + trace_of_last_response = prompt_resp._raw.get("trace") + llm_output = { 'promptTokens': total_prompt_tokens, 'completionTokens': total_completion_tokens, 'totalTokens': total_total_tokens, 'tokenCountsAreEstimated': token_counts_are_estimated, - 'estimatedCost': total_estimated_cost + 'estimatedCost': total_estimated_cost, + 'lastTrace': trace_of_last_response } return LLMResult(generations=generations, llm_output=llm_output) @@ -280,7 +296,19 @@ def produce_chunk(chunk: GenerationChunk): yield streamer.yield_(produce_chunk) -class DKUChatModel(BaseChatModel): +if must_use_deprecated_pydantic_config(): + class LockedDownBaseChatModel(BaseChatModel): + class Config: + extra = pydantic.Extra.forbid + underscore_attrs_are_private = True +else: + class LockedDownBaseChatModel(BaseChatModel): + model_config = { + 'extra': 'forbid', + } + + +class DKUChatModel(LockedDownBaseChatModel): """ Langchain-compatible wrapper around Dataiku-mediated chat LLMs @@ -318,10 +346,6 @@ class DKUChatModel(BaseChatModel): _llm_handle = None """:class:`dataikuapi.dss.llm.DSSLLM` object to wrap.""" - class Config: - extra = Extra.forbid - underscore_attrs_are_private = True - def __init__(self, llm_handle=None, **data: Any): if llm_handle is None: if data.get("llm_id") is None: @@ -365,6 +389,9 @@ def _generate( token_counts_are_estimated = False total_estimated_cost = 0.0 generations = [] + + trace_of_last_response = None + for prompt_resp in resp.responses: if not prompt_resp.success: raise Exception("LLM call failed: %s" % prompt_resp._raw.get("errorMessage", "Unknown error")) @@ -400,12 +427,22 @@ def _generate( make_invalid_tool_call(raw_tool_call, str(e)) ) + if prompt_resp._raw.get("trace", None) is not None: + trace_of_last_response = prompt_resp._raw.get("trace") + + usage_metadata = UsageMetadata( + input_tokens=prompt_resp._raw.get("promptTokens", 0), + output_tokens=prompt_resp._raw.get("completionTokens", 0), + total_tokens=prompt_resp._raw.get("totalTokens", 0), + ) + generations.append( ChatGeneration(message=AIMessage( content=text, additional_kwargs=additional_kwargs, tool_calls=tool_calls, - invalid_tool_calls=invalid_tool_calls + invalid_tool_calls=invalid_tool_calls, + usage_metadata = usage_metadata )) ) @@ -414,7 +451,8 @@ def _generate( 'completionTokens': total_completion_tokens, 'totalTokens': total_total_tokens, 'tokenCountsAreEstimated': token_counts_are_estimated, - 'estimatedCost': total_estimated_cost + 'estimatedCost': total_estimated_cost, + 'lastTrace': trace_of_last_response } return ChatResult(generations=generations, llm_output=llm_output) @@ -458,6 +496,15 @@ def produce_chunk(chunk: ChatGenerationChunk): additional_kwargs["tool_calls"] = raw_tool_calls tool_call_chunks = _parse_tool_call_chunks(raw_tool_calls) + if type(raw_chunk) == DSSLLMStreamedCompletionFooter: + usage_metadata = UsageMetadata( + input_tokens=raw_chunk.data.get("promptTokens", 0), + output_tokens=raw_chunk.data.get("completionTokens", 0), + total_tokens=raw_chunk.data.get("totalTokens", 0), + ) + else: + usage_metadata = None + new_chunk = ChatGenerationChunk( message=AIMessageChunk( content=text, @@ -465,6 +512,7 @@ def produce_chunk(chunk: ChatGenerationChunk): additional_kwargs=additional_kwargs, ), generation_info=raw_chunk.data, + usage_metadata = usage_metadata ) streamer.append(new_chunk) @@ -479,10 +527,12 @@ def produce_chunk(chunk: ChatGenerationChunk): def bind_tools( self, - tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]], + tools: Sequence[Union[Dict[str, Any], Type[pydantic.BaseModel], Callable, BaseTool]], tool_choice: Optional[ Union[dict, str, Literal["auto", "none", "required", "any"], bool] ] = None, + strict: Optional[bool] = None, + compatible: Optional[bool] = None, **kwargs: Any, ): """ @@ -504,10 +554,16 @@ def bind_tools( - a dict of the form: {"type": "tool_name", "name": "<>"}, or {"type": "required"}, or {"type": "any"} or {"type": "none"}, or {"type": "auto"}; + strict: If specified, request the model to produce a JSON tool call that adheres to the provided schema. Support varies across models/providers. + compatible: Allow DSS to modify the schema in order to increase compatibility, depending on known limitations of the model/provider. Defaults to automatic. kwargs: Any additional parameters to bind. """ - formatted_tools = [convert_to_openai_tool(tool) for tool in tools] + formatted_tools = [ + _convert_to_llm_mesh_tool(tool, strict, compatible) + for tool in tools + ] + if tool_choice: kwargs["tool_choice"] = _convert_to_llm_mesh_tool_choice(tool_choice, formatted_tools) @@ -556,6 +612,20 @@ def _parse_tool_call_chunks(raw_tool_calls): return [] +def _convert_to_llm_mesh_tool( + tool: Union[Dict[str, Any], Type[pydantic.BaseModel], Callable, BaseTool], + strict: Optional[bool] = None, + compatible: Optional[bool] = None +) -> Dict[str, Any]: + """Map the tool to its LLM mesh compatible representation""" + mesh_tool = convert_to_openai_tool(tool) + if isinstance(strict, bool): + mesh_tool['function']['strict'] = strict + if isinstance(compatible, bool): + mesh_tool['function']['compatible'] = compatible + return mesh_tool + + def _convert_to_llm_mesh_tool_choice( tool_choice: Union[dict, str, Literal["auto", "none", "required", "any"], bool], formatted_tools: Sequence[Dict[str, Any]] diff --git a/dataikuapi/dss/langchain/utils.py b/dataikuapi/dss/langchain/utils.py new file mode 100644 index 00000000..f85cb315 --- /dev/null +++ b/dataikuapi/dss/langchain/utils.py @@ -0,0 +1,7 @@ +import langchain +import pydantic + + +def must_use_deprecated_pydantic_config(): + # Pydantic 2 models are supported starting with Pydantic 2 and Langchain 0.3 + return pydantic.__version__[0] < '2' or langchain.__version__[0] == '0' and langchain.__version__[2] < '3' diff --git a/dataikuapi/dss/llm.py b/dataikuapi/dss/llm.py index 3a512cef..36023dd8 100644 --- a/dataikuapi/dss/llm.py +++ b/dataikuapi/dss/llm.py @@ -266,17 +266,74 @@ def with_tool_output(self, tool_output, tool_call_id, role="tool"): class SettingsMixin(object): - def with_json_output(self): + def with_json_output(self, schema=None, strict=None, compatible=None): """ Request the model to generate a valid JSON response, for models that support it. - + Note that some models may require you to also explicitly request this in the user or system prompt to use this. + + .. caution:: + JSON output support is experimental for locally-running Hugging Face models. + + :param dict schema: (optional) If specified, request the model to produce a JSON response that adheres to the provided schema. Support varies across models/providers. + :param bool strict: (optional) If a schema is provided, whether to strictly enforce it. Support varies across models/providers. + :param bool compatible: (optional) Allow DSS to modify the schema in order to increase compatibility, depending on known limitations of the model/provider. Defaults to automatic. """ self._settings["responseFormat"] = { - "type": "json" + "type": "json", + "schema": schema, + "strict": strict, + "compatible": compatible, } return self + def with_structured_output(self, model_type, strict=None, compatible=None): + """ + Instruct the model to generate a response as an instance of a specified Pydantic model. + + This functionality depends on `with_json_output` and necessitates that the model supports JSON output with a schema. + + .. caution:: + Structured output support is experimental for locally-running Hugging Face models. + + :param pydantic.BaseModel model_type: A Pydantic model class used for structuring the response. + :param bool strict: (optional) see :func:`with_json_output` + :param bool compatible: (optional) see :func:`with_json_output` + """ + if hasattr(model_type, "model_json_schema") and hasattr(model_type, "model_validate_json"): + schema = model_type.model_json_schema() # Pydantic 2 BaseModel + self._response_parser = model_type.model_validate_json + elif hasattr(model_type, "schema") and hasattr(model_type, "parse_raw"): + schema = model_type.schema() # Pydantic 1 BaseModel + self._response_parser = model_type.parse_raw + else: + # 'model_type' is not a Pydantic BaseModel. Derive schema Python type hints. + try: + import pydantic + except ImportError: + raise Exception("Pydantic is required to use Python's type hints with structured output") + + if hasattr(pydantic, "TypeAdapter"): + # Pydantic 2 provides a TypeAdapter to work with regular Python classes / type hints + from pydantic import TypeAdapter + adapter = TypeAdapter(model_type) + schema = adapter.json_schema() + self._response_parser = adapter.validate_json + elif hasattr(pydantic, "schema_of") and hasattr(pydantic, "parse_obj_as"): + # Pydantic 1 had similar functionality via 'schema_of' and 'parse_obj_as' + schema = pydantic.schema_of(model_type) + + def response_parser(json_response): + parsed_json = json.loads(json_response) + return pydantic.parse_obj_as(model_type, parsed_json) + + self._response_parser = response_parser + else: + # Unsupported Pydantic version + raise Exception("Incompatible Pydantic version") + self.with_json_output(schema=schema, strict=strict, compatible=compatible) + return self + class DSSLLMCompletionQuery(DSSLLMCompletionsQuerySingleQuery, SettingsMixin): """ @@ -291,6 +348,7 @@ def __init__(self, llm): super().__init__() self.llm = llm self._settings = {} + self._response_parser = None @property def settings(self): @@ -310,7 +368,7 @@ def execute(self): queries = {"queries": [self.cq], "settings": self._settings, "llmId": self.llm.llm_id} ret = self.llm.client._perform_json("POST", "/projects/%s/llms/completions" % (self.llm.project_key), body=queries) - return DSSLLMCompletionResponse(ret["responses"][0]) + return DSSLLMCompletionResponse(raw_resp=ret["responses"][0], response_parser=self._response_parser) def execute_streamed(self): """ @@ -344,6 +402,7 @@ def __init__(self, llm): self.llm = llm self.queries = [] self._settings = {} + self._response_parser = None @property def settings(self): @@ -368,7 +427,7 @@ def execute(self): queries = {"queries": [q.cq for q in self.queries], "settings": self._settings, "llmId": self.llm.llm_id} ret = self.llm.client._perform_json("POST", "/projects/%s/llms/completions" % (self.llm.project_key), body=queries) - return DSSLLMCompletionsResponse(ret["responses"]) + return DSSLLMCompletionsResponse(ret["responses"], response_parser=self._response_parser) class DSSLLMCompletionQueryMultipartMessage(object): @@ -393,8 +452,8 @@ def with_inline_image(self, image, mime_type=None): """ Add an image part to the multipart message - :param image: bytes or str (base64) - :param mime_type str: None for default + :param Union[str, bytes] image: The image + :param str mime_type: None for default """ img_b64 = None if isinstance(image, str): @@ -424,6 +483,21 @@ class DSSLLMStreamedCompletionChunk(object): def __init__(self, data): self.data = data + @property + def type(self): + """Type of this chunk, either "content" or "event" """ + return self.data.get("type", "content") + + @property + def text(self): + """If this chunk is content and has text, the (partial) text""" + return self.data.get("text", None) + + @property + def event_kind(self): + """If this chunk is an event, its kind""" + return self.data.get("eventKind", None) + def __repr__(self): return "" % self.data @@ -432,6 +506,15 @@ class DSSLLMStreamedCompletionFooter(object): def __init__(self, data): self.data = data + # Compatibility for code that just checks for "type"" + @property + def type(self): + return "footer" + + @property + def trace(self): + return self.data.get("trace", None) + def __repr__(self): return "" % self.data @@ -490,28 +573,40 @@ def iterevents(self): class DSSLLMCompletionResponse(object): """ - A handle to interact with a completion response. - - .. important:: - Do not create this class directly, use :meth:`dataikuapi.dss.llm.DSSLLMCompletionQuery.execute` instead. + Response to a completion """ - def __init__(self, raw_resp): - self._raw = raw_resp + def __init__(self, raw_resp=None, text=None, finish_reason=None, response_parser=None, trace=None): + if raw_resp is not None: + self._raw = raw_resp + else: + self._raw = {} + self._raw["text"] = text + self._raw["finishReason"] = finish_reason + self._raw["trace"] = trace + self._json = None + self._response_parser = response_parser + self._parsed = None @property def json(self): """ :return: LLM response parsed as a JSON object """ - if not self.success: - error_message = self._raw.get("errorMessage", "An unknown error occurred") - raise Exception(error_message) - - if self._json is None: - self._json = json.loads(self._raw["text"]) + self._fail_unless_success() + if self._json is None and self.text is not None: + self._json = json.loads(self.text) return self._json + @property + def parsed(self): + self._fail_unless_success() + if self._parsed is None and self.text is not None: + if not self._response_parser: + raise Exception("Structured output is not enabled for this completion query") + self._parsed = self._response_parser(self.text) + return self._parsed + @property def success(self): """ @@ -526,6 +621,7 @@ def text(self): :return: The raw text of the LLM response. :rtype: Union[str, None] """ + self._fail_unless_success() return self._raw.get("text") @property @@ -534,6 +630,7 @@ def tool_calls(self): :return: The tool calls of the LLM response. :rtype: Union[list, None] """ + self._fail_unless_success() return self._raw.get("toolCalls") @property @@ -542,8 +639,18 @@ def log_probs(self): :return: The log probs of the LLM response. :rtype: Union[list, None] """ + self._fail_unless_success() return self._raw.get("logProbs") + @property + def trace(self): + return self._raw.get("trace", None) + + def _fail_unless_success(self): + if not self.success: + error_message = self._raw.get("errorMessage", "An unknown error occurred") + raise Exception(error_message) + class DSSLLMCompletionsResponse(object): """ A handle to interact with a multi-completion response. @@ -551,13 +658,14 @@ class DSSLLMCompletionsResponse(object): .. important:: Do not create this class directly, use :meth:`dataikuapi.dss.llm.DSSLLMCompletionsQuery.execute` instead. """ - def __init__(self, raw_resp): + def __init__(self, raw_resp, response_parser=None): self._raw = raw_resp + self._response_parser = response_parser @property def responses(self): """The array of responses""" - return [DSSLLMCompletionResponse(x) for x in self._raw] + return [DSSLLMCompletionResponse(raw_resp=x, response_parser=self._response_parser) for x in self._raw] class DSSLLMImageGenerationQuery(object): @@ -577,22 +685,49 @@ def __init__(self, llm): def with_prompt(self, prompt, weight=None): """ + Add a prompt to the image generation query. + + :param str prompt: The prompt text. + :param float weight: Optional weight between 0 and 1 for the prompt. """ self.gq["prompts"].append({"prompt": prompt, "weight": weight}) return self def with_negative_prompt(self, prompt, weight=None): """ + Add a negative prompt to the image generation query. + + :param str prompt: The prompt text. + :param float weight: Optional weight between 0 and 1 for the negative prompt. """ self.gq["negativePrompts"].append({"prompt": prompt, "weight": weight}) return self def with_original_image(self, image, mode=None, weight=None): + """ + Add an image to the generation query. + + To edit specific pixels of the original image. A mask can be applied by calling `with_mask()`: + + >>> query.with_original_image(image, mode="INPAINTING") # replace the pixels using a mask + + To edit an image: + + >>> query.with_original_image(image, mode="MASK_FREE") # edit the original image according to the prompt + + >>> query.with_original_image(image, mode="VARY") # generates a variation of the original image + + :param Union[str, bytes] image: The original image as `str` in base 64 or `bytes`. + :param str mode: The edition mode. Modes support varies across models/providers. + :param float weight: The original image weight between 0 and 1. + """ if isinstance(image, str): self.gq["originalImage"] = image elif isinstance(image, bytes): import base64 self.gq["originalImage"] = base64.b64encode(image).decode("utf8") + else: + raise Exception(u"The `image` parameter has to be of type `str` in base 64 or `bytes`. Got {} instead.".format(type(image))) if mode is not None: self.gq["originalImageEditionMode"] = mode @@ -601,7 +736,21 @@ def with_original_image(self, image, mode=None, weight=None): self.gq["originalImageWeight"] = weight return self - def with_mask(self, mode, image=None, text=None): + def with_mask(self, mode, image=None): + """ + Add a mask for edition to the generation query. Call this method alongside `with_original_image()`. + + To edit parts of the image using a black mask (replace the black pixels): + + >>> query.with_mask("MASK_IMAGE_BLACK", image=black_mask) + + To edit parts of the image that are transparent (replace the transparent pixels): + + >>> query.with_mask("ORIGINAL_IMAGE_ALPHA") + + :param str mode: The mask mode. Modes support varies across models/providers. + :param Union[str, bytes] image: The mask image to apply to the image edition. As `str` in base 64 or `bytes`. + """ self.gq["maskMode"] = mode if image is not None: @@ -610,82 +759,151 @@ def with_mask(self, mode, image=None, text=None): elif isinstance(image, bytes): import base64 self.gq["maskImage"] = base64.b64encode(image).decode("utf8") + else: + raise Exception(u"When specified, the mask `image` parameter has to be of type `str` in base 64 or `bytes`. Got type {} instead.".format(type(image))) return self - @property - def inference_steps(self): - return self.gq.get("nbInferenceSteps", None) - @inference_steps.setter - def inference_steps(self, new_value): - self.gq["nbInferenceSteps"] = new_value - - @property - def refiner_strength(self): - return self.gq.get("refinerStrength", None) - @refiner_strength.setter - def refiner_strength(self, new_value): - self.gq["refinerStrength"] = new_value - @property def height(self): + """ + :return: The generated image height in pixels. + :rtype: Optional[int] + """ return self.gq.get("height", None) @height.setter def height(self, new_value): - self.gq["height"] = new_value + """ + The generated image height in pixels. + + :param Optional[int] new_value: The generated image height in pixels. + """ + self.gq["height"] = int(new_value) if new_value is not None else None @property def width(self): + """ + :return: The generated image width in pixels. + :rtype: Optional[int] + """ return self.gq.get("width", None) @width.setter def width(self, new_value): - self.gq["width"] = new_value + """ + The generated image width in pixels. + + :param Optional[int] new_value: The generated image width in pixels. + """ + self.gq["width"] = int(new_value) if new_value is not None else None @property def fidelity(self): + """ + :return: From 0.0 to 1.0, how strongly to adhere to prompt. + :rtype: Optional[float] + """ return self.gq.get("fidelity", None) @fidelity.setter def fidelity(self, new_value): + """ + Quality of the image to generate. Valid values depend on the targeted model. + + :param Optional[float] new_value: From 0.0 to 1.0, how strongly to adhere to prompt. + """ self.gq["fidelity"] = new_value @property def quality(self): + """ + :return: Quality of the image to generate. Valid values depend on the targeted model. + :rtype: Optional[str] + """ return self.gq.get("quality", None) @quality.setter def quality(self, new_value): + """ + Quality of the image to generate. Valid values depend on the targeted model. + + :param str new_value: Quality of the image to generate. + """ self.gq["quality"] = new_value @property def seed(self): + """ + :return: Seed of the image to generate, gives deterministic results when set. + :rtype: Optional[int] + """ return self.gq.get("seed", None) @seed.setter def seed(self, new_value): + """ + Seed of the image to generate, gives deterministic results when set. + + :param str new_value: Seed of the image to generate. + """ self.gq["seed"] = new_value @property def style(self): - """Style of the image to generate. Valid values depend on the targeted model""" + """ + :return: Style of the image to generate. Valid values depend on the targeted model. + :rtype: Optional[str] + """ return self.gq.get("style", None) @style.setter def style(self, new_value): + """ + Style of the image to generate. Valid values depend on the targeted model. + + :param str new_value: Style of the image to generate. + """ self.gq["style"] = new_value @property def images_to_generate(self): + """ + :return: Number of images to generate per query. Valid values depend on the targeted model. + :rtype: Optional[int] + """ return self.gq.get("nbImagesToGenerate", None) @images_to_generate.setter def images_to_generate(self, new_value): + """ + Number of images to generate per query. Valid values depend on the targeted model. + + :param int new_value: Number of images to generate. Valid values depend on the targeted model. + """ self.gq["nbImagesToGenerate"] = new_value - def with_aspect_ratio(self, ar): - self.gq["height"] = 1024 - self.gq["width"] = int(1024 * ar) - return self + @property + def aspect_ratio(self): + """ + :return: The width/height aspect ratio or `None` if either is not set. + :rtype: Optional[float] + """ + if self.width is not None and self.width > 0 and self.height is not None and self.height > 0: + return self.width / self.height + return None + @aspect_ratio.setter + def aspect_ratio(self, ar): + """ + Aspect ratio of the image to generate. Valid values depend on the targeted model. Set/update the width or height, or both if none are set. + + :param float ar: The width/height aspect ratio. + """ + if self.height is not None and self.height > 0: + self.width = self.height * ar + elif self.width is not None and self.width > 0: + self.height = self.width / ar + else: + self.height = 1024 + self.width = 1024 * ar def execute(self): """ Executes the image generation - :rtype: :class:DSSLLMImageGenerationResponse + :rtype: :class:`DSSLLMImageGenerationResponse` """ ret = self.llm.client._perform_json("POST", "/projects/%s/llms/images" % (self.llm.project_key), body=self.gq) @@ -712,8 +930,9 @@ def success(self): def first_image(self, as_type="bytes"): """ - :return: The first generated image. - :rtype: str + :param str as_type: The type of image to return, 'bytes' for `bytes` otherwise 'str' for base 64 `str`. + :return: The first generated image as `bytes` or `str` depending on the `as_type` parameter. + :rtype: Union[bytes,str] """ if not self.success: @@ -731,8 +950,9 @@ def first_image(self, as_type="bytes"): def get_images(self, as_type="bytes"): """ - :return: The generated images. - :rtype: List[str] + :param str as_type: The type of images to return, 'bytes' for `bytes` otherwise 'str' for base 64 `str`. + :return: The generated images as `bytes` or `str` depending on the `as_type` parameter. + :rtype: Union[List[bytes], List[str]] """ if not self.success: @@ -746,3 +966,11 @@ def get_images(self, as_type="bytes"): return [base64.b64decode(image["data"]) for image in self._raw["images"]] else: return [image["data"] for image in self._raw["images"]] + + @property + def images(self): + """ + :return: The generated images in bytes format. + :rtype: List[bytes] + """ + return self.get_images(as_type="bytes") diff --git a/dataikuapi/dss/llm_tracing/__init__.py b/dataikuapi/dss/llm_tracing/__init__.py new file mode 100644 index 00000000..c848589a --- /dev/null +++ b/dataikuapi/dss/llm_tracing/__init__.py @@ -0,0 +1,156 @@ +import datetime +import time +import threading + +dku_tracing_ls = threading.local() + + +def new_trace(name): + return SpanBuilder(name) + + +def current_span_builder(): + if hasattr(dku_tracing_ls, "current_span_builder"): + return dku_tracing_ls.current_span_builder + else: + return None + + +def current_span_builder_or_noop(): + if hasattr(dku_tracing_ls, "current_span_builder"): + return dku_tracing_ls.current_span_builder + else: + return SpanBuilder("noop") + + +class SpanReader(object): + def __init__(self, data): + if isinstance(data, SpanBuilder): + self.span_data = data.span + else: + self.span_data = data + + @property + def name(self): + return self.span_data["name"] + + @property + def type(self): + return self.span_data["type"] + + @property + def attributes(self): + return self.span_data["attributes"] + + @property + def raw_children(self): + return self.span_data["children"] + + @property + def children(self): + for child in self.span_data["children"]: + child_sr = SpanReader(child) + yield child_sr + + @property + def inputs(self): + return self.span_data.get("inputs", {}) + + @property + def outputs(self): + return self.span_data.get("outputs", {}) + + @property + def begin(self): + return self.span_data.get("begin", None) + + @property + def end(self): + return self.span_data.get("end", None) + + @property + def duration(self): + return self.span_data.get("duration", None) + + @property + def begin_ts(self): + return int(datetime.datetime.strptime(self.span_data["begin"], "%Y-%m-%dT%H:%M:%S.%f%z").timestamp() * 1000) + + @property + def end_ts(self): + return int(datetime.datetime.strptime(self.span_data["end"], "%Y-%m-%dT%H:%M:%S.%f%z").timestamp() * 1000) + + +class SpanBuilder: + def __init__(self,name): + self.span = { + "type": "span", + "name": name, + "children": [], + "attributes": {}, + "inputs": {}, + "outputs": {} + } + + def to_dict(self): + if "begin" in self.span and not "end" in self.span: + self.end(int(time.time() * 1000)) + return self.span + + @property + def inputs(self): + if self.span.get("inputs", None) is None: + self.span["inputs"] = {} + return self.span["inputs"] + + @property + def outputs(self): + if self.span.get("outputs", None) is None: + self.span["outputs"] = {} + return self.span["outputs"] + + @property + def attributes(self): + return self.span["attributes"] + + def subspan(self, name): + sub = SpanBuilder(name) + self.span["children"].append(sub.span) + return sub + + def append_trace(self, trace_to_append): + if isinstance(trace_to_append, dict): + self.span["children"].append(trace_to_append) + elif isinstance(trace_to_append, SpanBuilder): + self.span["children"].append(trace_to_append.to_dict()) + else: + raise Exception("Cannot happen trace of type %s" % type(trace_to_append)) + + def begin(self, begin_time): + self._begin_ts = begin_time + self.span["begin"] = datetime.datetime.utcfromtimestamp(begin_time / 1000).strftime("%Y-%m-%dT%H:%M:%S.%fZ") + + def end(self, end_time): + self.span["end"] = datetime.datetime.utcfromtimestamp(end_time / 1000).strftime("%Y-%m-%dT%H:%M:%S.%fZ") + #print("Ending span: %s -> %s" % (self.span["begin"], self.span["end"])) + self.span["duration"] = int(end_time - self._begin_ts) + + def __enter__(self,): + self.previous_sb_on_thread = current_span_builder() + dku_tracing_ls.current_span_builder = self + self.begin(int(time.time() * 1000)) + return self + + def __exit__(self, type, value, traceback): + dku_tracing_ls.current_span_builder = self.previous_sb_on_thread + self.previous_sb_on_thread = None + self.end(int(time.time() * 1000)) + + +def mini_trace_dump(trace): + def _rec(span, level): + print("%s %s (%s -> %s: %s)" % (" "* (2 * level), span.name, span.begin, span.end, span.duration)) + for child in span.children: + _rec(child, level +1) + + _rec(SpanReader(trace), 0) \ No newline at end of file diff --git a/dataikuapi/dss/messaging_channel.py b/dataikuapi/dss/messaging_channel.py index 8c337fa5..f003d050 100644 --- a/dataikuapi/dss/messaging_channel.py +++ b/dataikuapi/dss/messaging_channel.py @@ -133,6 +133,7 @@ class DSSMailMessagingChannel(DSSMessagingChannel): def __init__(self, client, data): super().__init__(client, data) self._sender = data.get("sender", None) + self._use_current_user_as_sender = data.get("useCurrentUserAsSender", None) @property def sender(self): @@ -143,6 +144,16 @@ def sender(self): """ return self._sender + @property + def use_current_user_as_sender(self): + """ + Indicates whether the messaging channel will use the address of the current user as sender. + If True and the current user has no associated email address, the sender property is used instead. + + :rtype: bool + """ + return self._use_current_user_as_sender + def send(self, project_key, to, subject, body, attachments=None, plain_text=False, sender=None, cc=None, bcc=None): """ Send an email with or without attachments to a list of recipients diff --git a/dataikuapi/dss/ml.py b/dataikuapi/dss/ml.py index ffebbfff..8407a386 100644 --- a/dataikuapi/dss/ml.py +++ b/dataikuapi/dss/ml.py @@ -1897,6 +1897,43 @@ def __init__(self, raw_settings, hyperparameter_search_params): self.seed = self._register_single_value_hyperparameter("seed", accepted_types=[int]) +class GluonTSTorchSimpleFeedForwardSettings(PredictionAlgorithmSettings): + + def __init__(self, raw_settings, hyperparameter_search_params): + super(GluonTSTorchSimpleFeedForwardSettings, self).__init__(raw_settings, hyperparameter_search_params) + self.context_length = self._register_numerical_hyperparameter("context_length") + self.distr_output = self._register_categorical_hyperparameter("distr_output") + self.batch_normalization = self._register_categorical_hyperparameter("batch_normalization") + self.mean_scaling = self._register_categorical_hyperparameter("mean_scaling") + self.num_hidden_dimensions = self._register_single_value_hyperparameter("num_hidden_dimensions", accepted_types=[list]) + self.full_context = self._register_single_value_hyperparameter("full_context", accepted_types=[bool]) + self.batch_size = self._register_single_value_hyperparameter("batch_size", accepted_types=[int]) + self.epochs = self._register_single_value_hyperparameter("epochs", accepted_types=[int]) + self.auto_num_batches_per_epoch = self._register_single_value_hyperparameter("auto_num_batches_per_epoch", accepted_types=[bool]) + self.num_batches_per_epoch = self._register_single_value_hyperparameter("num_batches_per_epoch", accepted_types=[int]) + self.seed = self._register_single_value_hyperparameter("seed", accepted_types=[int]) + + +class GluonTSTorchDeepARSettings(PredictionAlgorithmSettings): + + def __init__(self, raw_settings, hyperparameter_search_params): + super(GluonTSTorchDeepARSettings, self).__init__(raw_settings, hyperparameter_search_params) + self.context_length = self._register_numerical_hyperparameter("context_length") + self.num_layers = self._register_numerical_hyperparameter("num_layers") + self.num_cells = self._register_numerical_hyperparameter("num_cells") + self.dropout_rate = self._register_numerical_hyperparameter("dropout_rate") + self.distr_output = self._register_categorical_hyperparameter("distr_output") + self.full_context = self._register_single_value_hyperparameter("full_context", accepted_types=[bool]) + self.scaling = self._register_single_value_hyperparameter("scaling", accepted_types=[bool]) + self.num_parallel_samples = self._register_single_value_hyperparameter("num_parallel_samples", accepted_types=[int]) + self.minimum_scale = self._register_single_value_hyperparameter("minimum_scale", accepted_types=[float]) + self.batch_size = self._register_single_value_hyperparameter("batch_size", accepted_types=[int]) + self.epochs = self._register_single_value_hyperparameter("epochs", accepted_types=[int]) + self.auto_num_batches_per_epoch = self._register_single_value_hyperparameter("auto_num_batches_per_epoch", accepted_types=[bool]) + self.num_batches_per_epoch = self._register_single_value_hyperparameter("num_batches_per_epoch", accepted_types=[int]) + self.seed = self._register_single_value_hyperparameter("seed", accepted_types=[int]) + + class GluonTSSimpleFeedForwardSettings(PredictionAlgorithmSettings): def __init__(self, raw_settings, hyperparameter_search_params): @@ -2325,6 +2362,10 @@ class DSSTimeseriesForecastingMLTaskSettings(AbstractTabularPredictionMLTaskSett "SEASONAL_LOESS": PredictionAlgorithmMeta("seasonal_loess_timeseries", SeasonalLoessSettings), "PROPHET": PredictionAlgorithmMeta("prophet_timeseries", ProphetSettings), "GLUONTS_NPTS_FORECASTER": PredictionAlgorithmMeta("gluonts_npts_timeseries", GluonTSNPTSForecasterSettings), + + "GLUONTS_TORCH_SIMPLE_FEEDFORWARD": PredictionAlgorithmMeta("gluonts_torch_simple_feed_forward_timeseries", GluonTSTorchSimpleFeedForwardSettings), + "GLUONTS_TORCH_DEEPAR": PredictionAlgorithmMeta("gluonts_torch_deepar_timeseries", GluonTSTorchDeepARSettings), + "GLUONTS_SIMPLE_FEEDFORWARD": PredictionAlgorithmMeta("gluonts_simple_feed_forward_timeseries", GluonTSSimpleFeedForwardSettings), "GLUONTS_DEEPAR": PredictionAlgorithmMeta("gluonts_deepar_timeseries", GluonTSDeepARSettings), "GLUONTS_TRANSFORMER": PredictionAlgorithmMeta("gluonts_transformer_timeseries", GluonTSTransformerSettings), @@ -2332,7 +2373,7 @@ class DSSTimeseriesForecastingMLTaskSettings(AbstractTabularPredictionMLTaskSett } _TIME_UNITS = {"MILLISECOND", "SECOND", "MINUTE", "HOUR", "DAY", "BUSINESS_DAY", "WEEK", "MONTH", "QUARTER", "HALF_YEAR", "YEAR"} - _INTERPOLATION_METHODS = {"NEAREST", "PREVIOUS", "NEXT", "LINEAR", "QUADRATIC", "CUBIC", "CONSTANT"} + _INTERPOLATION_METHODS = {"NEAREST", "PREVIOUS", "NEXT", "LINEAR", "QUADRATIC", "CUBIC", "CONSTANT","STAIRCASE"} _EXTRAPOLATION_METHODS = {"PREVIOUS_NEXT", "NO_EXTRAPOLATION", "CONSTANT", "LINEAR", "QUADRATIC", "CUBIC"} _CATEGORICAL_IMPUTATION_METHODS = {"MOST_COMMON", "NULL", "CONSTANT", "PREVIOUS_NEXT", "PREVIOUS", "NEXT"} _DUPLICATE_TIMESTAMPS_HANDLING_METHODS = {"FAIL_IF_CONFLICTING", "DROP_IF_CONFLICTING", "MEAN_MODE"} @@ -2429,7 +2470,7 @@ def set_numerical_interpolation(self, method=None, constant=None): Sets the time series resampling numerical interpolation parameters :param method: Interpolation method. Valid values are: NEAREST, PREVIOUS, NEXT, LINEAR, QUADRATIC, - CUBIC, CONSTANT (defaults to **None**, i.e. don't change) + CUBIC, CONSTANT, STAIRCASE (defaults to **None**, i.e. don't change) :type method: str, optional :param constant: Value for the CONSTANT interpolation method (defaults to **None**, i.e. don't change) :type constant: float, optional diff --git a/dataikuapi/dss/plugin.py b/dataikuapi/dss/plugin.py index 8a30a9d0..db734038 100644 --- a/dataikuapi/dss/plugin.py +++ b/dataikuapi/dss/plugin.py @@ -600,7 +600,7 @@ def create_code_env(self, python_interpreter=None, conda=False): settings.set_code_env(env_name) settings.save() - :param string python_interpreter: which version of python to use. Possible values: PYTHON27, PYTHON34, PYTHON35, PYTHON36, PYTHON37, PYTHON38, PYTHON39, PYTHON310, PYTHON311 + :param string python_interpreter: which version of python to use. Possible values: PYTHON27, PYTHON34, PYTHON35, PYTHON36, PYTHON37, PYTHON38, PYTHON39, PYTHON310, PYTHON311, PYTHON312 :param boolean conda: if True use conda to create the code env, if False use virtualenv and pip. :return: a handle on the operation diff --git a/dataikuapi/dss/project.py b/dataikuapi/dss/project.py index 0c4712ea..e5a4cdcd 100644 --- a/dataikuapi/dss/project.py +++ b/dataikuapi/dss/project.py @@ -1782,6 +1782,42 @@ def preload_bundle(self, bundle_id): "/projects/%s/bundles/imported/%s/actions/preload" % ( self.project_key, bundle_id)) + ######################################################## + # Testing with DSS test scenarios report + ######################################################## + + def get_last_test_scenario_runs_report(self, bundle_id): + """ + Download a report describing the outcome of the latest test scenario runs performed in this project, on an + Automation node, under a specified active bundle + + :param str bundle_id: bundle id tag + + :return: the test scenarios report, in JUnit XML format + :rtype: file-like + """ + return self.client._perform_raw( + "GET", + "/projects/%s/scenarios/last-test-scenario-runs-report" % self.project_key, + params={"bundleId": bundle_id} + ) + + def get_last_test_scenario_runs_html_report(self, bundle_id): + """ + Download a report describing the outcome of the latest test scenario runs performed in this project, on an + Automation node, under a specified active bundle + + :param str bundle_id: bundle id tag + + :return: the test scenarios report, in HTML format + :rtype: file-like + """ + return self.client._perform_raw( + "GET", + "/projects/%s/scenarios/last-test-scenario-runs-html-report" % self.project_key, + params={"bundleId": bundle_id} + ) + ######################################################## # Scenarios ######################################################## @@ -2493,6 +2529,26 @@ def get_data_quality_timeline(self, min_timestamp=None, max_timestamp=None): """ return self.client._perform_json("GET", "/projects/%s/data-quality/timeline" % self.project_key, params={"minTimestamp": min_timestamp, "maxTimestamp": max_timestamp}) + def list_test_scenarios(self): + """ + Lists all the test scenarios of a DSS Project + :return: list all test scenarios of a project + :rtype: list of :class:`DSSScenario` + """ + scenarios = self.client._perform_json("GET", "/projects/%s/scenarios/" % self.project_key) + return [DSSScenario(self.client, self.project_key, scenario["id"]) for scenario in scenarios if scenario.get("markedAsTest")] + + def get_testing_status(self, bundle_id=None): + """ + Get the testing status of a DSS Project. It combines the last run outcomes of all the test scenarios defined on the project, considering the worst outcome as a final result. + :param (optional) string bundle_id : if the project is on automation node, you can specify a bundle_id to filter only on the last + scenario runs when this bundle was active + :return: returns a dict with the keys 'nbTotalRanScenarios' and 'nbScenariosPerOutcome' + :rtype: dict + """ + return self.client._perform_json("GET", "/projects/%s/scenarios/testing-status" % self.project_key, params={"bundleId": bundle_id}) + + class TablesImportDefinition(object): """ Temporary structure holding the list of tables to import diff --git a/dataikuapi/dss/projectdeployer.py b/dataikuapi/dss/projectdeployer.py index 84590eca..91077f6a 100644 --- a/dataikuapi/dss/projectdeployer.py +++ b/dataikuapi/dss/projectdeployer.py @@ -23,11 +23,11 @@ def list_deployments(self, as_objects=True): # list all deployments with their current state for deployment in deployer.list_deployments(): status = deployment.get_status() - print("Deployment %s is %s" % (deployment.id, status.get_health())) + print("Deployment %s is %s" % (deployment.id, status.get_health())) :param boolean as_objects: if True, returns a list of :class:`DSSProjectDeployerDeployment`, else returns a list of dict. - :returns: list of deployments, either as :class:`DSSProjectDeployerDeployment` or as dict (with fields + :returns: list of deployments, either as :class:`DSSProjectDeployerDeployment` or as dict (with fields as in :meth:`DSSProjectDeployerDeploymentStatus.get_light()`) :rtype: list """ @@ -50,7 +50,7 @@ def get_deployment(self, deployment_id): def create_deployment(self, deployment_id, project_key, infra_id, bundle_id, deployed_project_key=None, project_folder_id=None, ignore_warnings=False): """ - Create a deployment and return the handle to interact with it. + Create a deployment and return the handle to interact with it. The returned deployment is not yet started and you need to call :meth:`~DSSProjectDeployerDeployment.start_update` @@ -117,11 +117,11 @@ def list_infras(self, as_objects=True): # list infrastructures that the user can deploy to for infrastructure in deployer.list_infras(as_objects=False): if infrastructure.get("canDeploy", False): - print("User can deploy to %s" % infrastructure["infraBasicInfo"]["id"]) + print("User can deploy to %s" % infrastructure["infraBasicInfo"]["id"]) :param boolean as_objects: if True, returns a list of :class:`DSSProjectDeployerInfra`, else returns a list of dict. - :return: list of infrastructures, either as :class:`DSSProjectDeployerInfra` or as dict (with fields + :return: list of infrastructures, either as :class:`DSSProjectDeployerInfra` or as dict (with fields as in :meth:`DSSProjectDeployerInfraStatus.get_raw()`) :rtype: list """ @@ -144,7 +144,7 @@ def create_infra(self, infra_id, stage, govern_check_policy="NO_CHECK"): """ settings = { "id": infra_id, - "stage": stage, + "stage": stage, "governCheckPolicy": govern_check_policy, } self.client._perform_json("POST", "/project-deployer/infras", body=settings) @@ -171,12 +171,12 @@ def list_projects(self, as_objects=True): # list project that the user can deploy bundles from for project in deployer.list_projects(as_objects=False): if project.get("canDeploy", False): - print("User can deploy to %s" % project["projectBasicInfo"]["id"]) + print("User can deploy to %s" % project["projectBasicInfo"]["id"]) :param boolean as_objects: if True, returns a list of :class:`DSSProjectDeployerProject`, else returns a list of dict. - :return: list of published projects, either as :class:`DSSProjectDeployerProject` or as dict (with fields + :return: list of published projects, either as :class:`DSSProjectDeployerProject` or as dict (with fields as in :meth:`DSSProjectDeployerProjectStatus.get_raw()`) :rtype: list """ @@ -215,7 +215,7 @@ def upload_bundle(self, fp, project_key=None): Upload a bundle archive for a project. :param file-like fp: a bundle archive (should be a zip) - :param string project_key: key of the published project where the bundle will be uploaded. If the project does not + :param string project_key: key of the published project where the bundle will be uploaded. If the project does not exist, it is created. If not set, the key of the bundle's source project is used. """ if project_key is None: @@ -266,7 +266,7 @@ def get_status(self): def get_settings(self): """ - Get the settings of this infrastructure. + Get the settings of this infrastructure. :rtype: :class:`DSSProjectDeployerInfraSettings` """ @@ -278,7 +278,7 @@ def get_settings(self): def delete(self): """ Delete this infra. - + .. note:: You may only delete an infra if there are no deployments using it. @@ -304,9 +304,9 @@ def __init__(self, client, infra_id, settings): def get_raw(self): """ - Get the raw settings of this infrastructure. + Get the raw settings of this infrastructure. - This returns a reference to the raw settings, not a copy, so changes made to the returned + This returns a reference to the raw settings, not a copy, so changes made to the returned object will be reflected when saving. :return: the settings, as a dict. @@ -347,7 +347,7 @@ def get_deployments(self): def get_raw(self): """ - Get the raw status information. + Get the raw status information. :return: the status, as a dict. The dict contains a list of the bundles currently deployed on the infrastructure as a **deployments** field. @@ -401,7 +401,7 @@ def get_governance_status(self, bundle_id=""): :param string bundle_id: (Optional) The ID of a specific bundle of the published project to get status from. If empty, the bundle currently used in the deployment. - :return: messages about the governance status, as a dict with a **messages** field, itself a list of meassage + :return: messages about the governance status, as a dict with a **messages** field, itself a list of meassage information, each one a dict of: * **severity** : severity of the error in the message. Possible values are SUCCESS, INFO, WARNING, ERROR @@ -411,13 +411,13 @@ def get_governance_status(self, bundle_id=""): * **message** : the error message * **details** : a more detailed error description - :rtype: dict + :rtype: dict """ return self.client._perform_json("GET", "/project-deployer/deployments/%s/governance-status" % (self.deployment_id), params={ "bundleId": bundle_id }) def get_settings(self): """ - Get the settings of this deployment. + Get the settings of this deployment. :rtype: :class:`DSSProjectDeployerDeploymentSettings` """ @@ -451,6 +451,20 @@ def delete(self): self.client._perform_empty( "DELETE", "/project-deployer/deployments/%s" % (self.deployment_id)) + def get_testing_status(self, bundle_id=None, automation_node_id=None): + """ + Get the testing status of a project deployment. + + :param (optional) string bundle_id: filters the scenario runs done on a specific bundle + :param (optional) automation_node_id: for multi-node deployments only, you need to specify the automation node id on which you want to retrieve + the testing status + """ + + return self.client._perform_json("GET", "/project-deployer/deployments/%s/testing-status" % self.deployment_id, params={ + "bundleId": bundle_id, + "automationNodeId": automation_node_id + }) + class DSSProjectDeployerDeploymentSettings(object): """ @@ -470,16 +484,16 @@ def __init__(self, client, deployment_id, settings): def get_raw(self): """ - Get the raw settings of this deployment. + Get the raw settings of this deployment. - This returns a reference to the raw settings, not a copy, so changes made to the returned + This returns a reference to the raw settings, not a copy, so changes made to the returned object will be reflected when saving. :return: the settings, as a dict. Notable fields are: * **id** : identifier of the deployment * **infraId** : identifier of the infrastructure on which the deployment is done - * **bundleId** : identifier of the bundle of the published project being deployed + * **bundleId** : identifier of the bundle of the published project being deployed :rtype: dict """ @@ -488,7 +502,7 @@ def get_raw(self): @property def bundle_id(self): """ - Get or set the identifier of the bundle currently used by this deployment. + Get or set the identifier of the bundle currently used by this deployment. If setting the value, you need to call :meth:`save()` afterward for the change to be effective. """ @@ -526,7 +540,7 @@ def __init__(self, client, deployment_id, light_status, heavy_status): def get_light(self): """ - Get the 'light' (summary) status. + Get the 'light' (summary) status. This returns a dictionary with various information about the deployment, but not the actual health of the deployment @@ -538,11 +552,11 @@ def get_light(self): def get_heavy(self): """ - Get the 'heavy' (full) status. + Get the 'heavy' (full) status. This returns various information about the deployment, notably its health. - :return: a status, as a dict. The overall status of the deployment is in a **health** field (possible values: UNKNOWN, ERROR, + :return: a status, as a dict. The overall status of the deployment is in a **health** field (possible values: UNKNOWN, ERROR, WARNING, HEALTHY, UNHEALTHY, OUT_OF_SYNC). :rtype: dict """ @@ -601,9 +615,9 @@ def id(self): def get_status(self): """ - Get status information about this published project. + Get status information about this published project. - This is used mostly to get information about which versions are available and which + This is used mostly to get information about which versions are available and which deployments are exposing this project :rtype: :class:`DSSProjectDeployerProjectStatus` @@ -613,7 +627,7 @@ def get_status(self): def get_settings(self): """ - Get the settings of this published project. + Get the settings of this published project. The main things that can be modified in a project settings are permissions @@ -662,9 +676,9 @@ def __init__(self, client, project_key, settings): def get_raw(self): """ - Get the raw settings of this published project. + Get the raw settings of this published project. - This returns a reference to the raw settings, not a copy, so changes made to the returned + This returns a reference to the raw settings, not a copy, so changes made to the returned object will be reflected when saving. :return: the settings, as a dict. @@ -698,8 +712,8 @@ def get_deployments(self, infra_id=None): """ Get the deployments that have been created from this published project. - :param string infra_id: (optional) identifier of an infrastructure. When set, only get the deployments deployed on - this infrastructure. When not set, the list contains all the deployments using this published project, + :param string infra_id: (optional) identifier of an infrastructure. When set, only get the deployments deployed on + this infrastructure. When not set, the list contains all the deployments using this published project, across every infrastructure of the Project Deployer. :returns: a list of deployments, each a :class:`DSSProjectDeployerDeployment` @@ -731,7 +745,7 @@ def get_infras(self): def get_raw(self): """ - Gets the raw status information. + Gets the raw status information. :return: the status, as a dict. A **deployments** sub-field contains a list of the deployments of bundles of this projects. :rtype: dict diff --git a/dataikuapi/dss/scenario.py b/dataikuapi/dss/scenario.py index 69cab537..b0bf9d2c 100644 --- a/dataikuapi/dss/scenario.py +++ b/dataikuapi/dss/scenario.py @@ -831,6 +831,73 @@ def get_duration(self): duration = property(get_duration) + def get_report(self): + """ + Download a report describing the outcome of a test scenario run, in JUnit XML format. + + :return: the scenario run report, in JUnit XML format + :rtype: file-like + """ + if not self.run["scenario"].get("markedAsTest"): + raise DataikuException( + "When run %s was performed, scenario %s was not marked as a test scenario. Reports are only available for test scenarios." + % (self.run["runId"], self.run["scenario"]["id"]) + ) + + return self.client._perform_raw( + "GET", "/projects/%s/scenarios/%s/%s/scenario-run-report" % ( + self.run["scenario"]["projectKey"], + self.run["scenario"]["id"], + self.run["runId"] + ), + ) + + def get_step_run_report(self, step_id): + """ + Download a report describing the outcome of a test scenario step run, in JUnit XML format. + + :param string step_id: identifier of the step + + :return: the step run report, in JUnit XML format + :rtype: file-like + """ + if not self.run["scenario"].get("markedAsTest"): + raise DataikuException( + "When run %s was performed, scenario %s was not marked as a test scenario. Reports are only available for test scenarios." + % (self.run["runId"], self.run["scenario"]["id"]) + ) + + return self.client._perform_raw( + "GET", + "/projects/%s/scenarios/%s/%s/step-run-report" + % ( + self.run["scenario"]["projectKey"], + self.run["scenario"]["id"], + self.run["runId"], + ), + params={"stepId": step_id}, + ) + + def get_log(self, step_id=None): + """ + Gets the logs of the scenario run. If a step_id is passed in the parameters + the logs will be scoped to that step. + + :param string step_id: (optional) the id of the step in the run whose log is requested (defaults to **None**) + + :returns: the scenario run logs + :rtype: string + """ + return self.client._perform_text( + "GET", "/projects/%s/scenarios/%s/%s/log" % ( + self.run["scenario"]["projectKey"], + self.run["scenario"]["id"], + self.run["runId"], + ), + params={ "stepId" : step_id } + ) + + class DSSScenarioRunDetails(dict): """ Details of a scenario run, notably the outcome of its steps. diff --git a/dataikuapi/fm/cloudaccounts.py b/dataikuapi/fm/cloudaccounts.py new file mode 100644 index 00000000..b6c9c553 --- /dev/null +++ b/dataikuapi/fm/cloudaccounts.py @@ -0,0 +1,338 @@ +from .future import FMFuture + + +class FMCloudAccountCreator(object): + def __init__(self, client, label): + """ + A builder class to create cloud accounts + + :param str label: The label of the cloud account + """ + self.client = client + self.data = {"label": label} + + def with_description(self, description): + """ + Set the description of this cloud account + + :param str description: the description of the cloud account + :rtype: :class:`dataikuapi.fm.cloudaccounts.FMCloudAccountCreator` + """ + self.data["description"] = description + return self + + +class FMAWSCloudAccountCreator(FMCloudAccountCreator): + def same_as_fm(self): + """ + Use the same authentication as Fleet Manager + """ + self.data["awsAuthenticationMode"] = "DEFAULT_INSTANCE_CREDENTIALS" + return self + + def with_iam_role(self, role_arn): + """ + Use an IAM Role + + :param str role_arn: ARN of the IAM Role + """ + self.data["awsAuthenticationMode"] = "IAM_ROLE" + self.data["awsIAMRoleARN"] = role_arn + return self + + def with_keypair(self, access_key_id, secret_access_key): + """ + Use an AWS Access Key + + :param str access_key_id: AWS Access Key ID + :param str secret_access_key: AWS Secret Access Key + """ + self.data["awsAuthenticationMode"] = "KEYPAIR" + self.data["awsAccessKeyId"] = access_key_id + self.data["awsSecretAccessKey"] = secret_access_key + return self + + def create(self): + """ + Create a new cloud account + + :return: a newly created cloud account + :rtype: :class:`dataikuapi.fm.cloudaccounts.FMAWSCloudAccount` + """ + account = self.client._perform_tenant_json( + "POST", "/cloud-accounts", body=self.data + ) + return FMAWSCloudAccount(self.client, account) + + +class FMAzureCloudAccountCreator(FMCloudAccountCreator): + def same_as_fm(self, tenant_id, image_resource_group): + """ + Use the same authentication as Fleet Manager + + :param str tenant_id: Azure Tenant Id + :param str image_resource_group: Azure image cached resource group + """ + self.data["azureAuthenticationMode"] = "DEFAULT_INSTANCE_CREDENTIALS" + self.data["azureTenantId"] = tenant_id + self.data["azureImageResourceGroup"] = image_resource_group + return self + + def with_secret(self, subscription, tenant_id, environment, image_resource_group, client_id, secret): + """ + Use a Secret based authentication + + :param str subscription: Azure Subscription + :param str tenant_id: Azure Tenant Id + :param str environment: Azure Environment + :param str client_id: Azure Client Id + :param str image_resource_group: Azure image cached resource group + :param str secret: Azure Secret + """ + + self.data["azureAuthenticationMode"] = "OAUTH2_AUTHENTICATION_METHOD_CLIENT_SECRET" + self.data["azureSubscription"] = subscription + self.data["azureTenantId"] = tenant_id + self.data["azureEnvironment"] = environment + self.data["azureClientId"] = client_id + self.data["azureSecret"] = secret + self.data["azureImageResourceGroup"] = image_resource_group + return self + + def with_certificate(self, subscription, tenant_id, environment, client_id, image_resource_group, certificate_path, certificate_password): + """ + Use a Secret based authentication + + :param str subscription: Azure Subscription + :param str tenant_id: Azure Tenant Id + :param str environment: Azure Environment + :param str client_id: Azure Client Id + :param str image_resource_group: Azure image cached resource group + :param str certificate_path: Azure certificate path + :param str certificate_password: Azure certificate password + """ + + self.data["azureAuthenticationMode"] = "OAUTH2_AUTHENTICATION_METHOD_CERTIFICATE" + self.data["azureSubscription"] = subscription + self.data["azureTenantId"] = tenant_id + self.data["azureEnvironment"] = environment + self.data["azureClientId"] = client_id + self.data["azureImageResourceGroup"] = image_resource_group + self.data["azureCertificatePath"] = certificate_path + self.data["azureCertificatePassword"] = certificate_password + return self + + def with_managed_identity(self, tenant_id, environment, image_resource_group, managed_identity): + """ + Use a Managed Identity based authentication + + :param str tenant_id: Azure Tenant identifier + :param str environment: Azure Environment + :param str image_resource_group: Azure image cached resource group + :param str managed_identity: Azure Managed Identity + """ + + self.data["azureAuthenticationMode"] = "MANAGED_IDENTITY" + self.data["azureTenantId"] = tenant_id + self.data["azureEnvironment"] = environment + self.data["azureManagedIdentityId"] = managed_identity + self.data["azureImageResourceGroup"] = image_resource_group + return self + + def create(self): + """ + Create a new cloud account + + :return: a newly created cloud account + :rtype: :class:`dataikuapi.fm.cloudaccounts.FMAzureCloudAccount` + """ + account = self.client._perform_tenant_json( + "POST", "/cloud-accounts", body=self.data + ) + return FMAzureCloudAccount(self.client, account) + + +class FMGCPCloudAccountCreator(FMCloudAccountCreator): + def same_as_fm(self): + """ + Use the same authentication as Fleet Manager + """ + self.data["gcpAuthenticationMode"] = "DEFAULT_INSTANCE_CREDENTIALS" + return self + + def with_service_account_key(self, project_id, service_account_key): + """ + Use a Service Account JSON key based authentication + + :param str project_id: GCP project + :param str service_account_key: Optional, service account key (JSON) + """ + self.data["gcpAuthenticationMode"] = "JSON_KEY" + self.data["gcpProjectId"] = project_id + self.data["gcpServiceAccountKey"] = service_account_key + return self + + def create(self): + """ + Create a new cloud account + + :return: a newly created cloud account + :rtype: :class:`dataikuapi.fm.cloudaccounts.FMGCPCloudAccount` + """ + account = self.client._perform_tenant_json( + "POST", "/cloud-accounts", body=self.data + ) + return FMGCPCloudAccount(self.client, account) + + +class FMCloudAccount(object): + + def __init__(self, client, account_data): + self.client = client + self.account_data = account_data + self.id = self.account_data["id"] + + def delete(self): + """ + Delete this cloud account. + + :return: the `Future` object representing the deletion process + :rtype: :class:`dataikuapi.fm.future.FMFuture` + """ + if (self.account_data["nbNetworks"] > 0): + raise Exception("This account is in use by some networks, you cannot delete it") + + future = self.client._perform_tenant_json( + "DELETE", "/cloud-accounts/%s" % self.id + ) + return FMFuture.from_resp(self.client, future) + + def save(self): + """ + Save this cloud account. + """ + self.client._perform_tenant_empty( + "PUT", "/cloud-accounts/%s" % self.id, body=self.account_data + ) + + +class FMAWSCloudAccount(FMCloudAccount): + def set_same_as_fm(self): + """ + Use the same authentication as Fleet Manager + """ + self.account_data["awsAuthenticationMode"] = "DEFAULT_INSTANCE_CREDENTIALS" + + def set_iam_role(self, role_arn): + """ + Use an IAM Role + + :param str role_arn: ARN of the IAM Role + """ + self.account_data["awsAuthenticationMode"] = "IAM_ROLE" + self.account_data["awsIAMRoleARN"] = role_arn + + def set_keypair(self, access_key_id, secret_access_key): + """ + Use an AWS Access Key + + :param str access_key_id: AWS Access Key ID + :param str secret_access_key: AWS Secret Access Key + """ + self.account_data["awsAuthenticationMode"] = "KEYPAIR" + self.account_data["awsAccessKeyId"] = access_key_id + self.account_data["awsSecretAccessKey"] = secret_access_key + + +class FMAzureCloudAccount(FMCloudAccount): + def set_same_as_fm(self, tenant_id, image_resource_group): + """ + Use the same authentication as Fleet Manager + + :param str tenant_id: Azure Tenant Id + :param str image_resource_group: Azure image cached resource group + """ + self.account_data["azureAuthenticationMode"] = "DEFAULT_INSTANCE_CREDENTIALS" + self.account_data["azureTenantId"] = tenant_id + self.account_data["azureImageResourceGroup"] = image_resource_group + + def set_secret(self, subscription, tenant_id, environment, client_id, image_resource_group, secret): + """ + Use a Secret based authentication + + :param str subscription: Azure Subscription + :param str tenant_id: Azure Tenant Id + :param str environment: Azure Environment + :param str client_id: Azure Client Id + :param str image_resource_group: Azure image cached resource group + :param str secret: Azure Secret + """ + + self.account_data["azureAuthenticationMode"] = "OAUTH2_AUTHENTICATION_METHOD_CLIENT_SECRET" + self.account_data["azureSubscription"] = subscription + self.account_data["azureTenantId"] = tenant_id + self.account_data["azureEnvironment"] = environment + self.account_data["azureClientId"] = client_id + self.account_data["azureSecret"] = secret + self.account_data["azureImageResourceGroup"] = image_resource_group + return self + + def set_certificate(self, subscription, tenant_id, environment, client_id, image_resource_group, certificate_path, certificate_password): + """ + Use a Secret based authentication + + :param str subscription: Azure Subscription + :param str tenant_id: Azure Tenant Id + :param str environment: Azure Environment + :param str client_id: Azure Client Id + :param str image_resource_group: Azure image cached resource group + :param str certificate_path: Azure certificate path + :param str certificate_password: Azure certificate password + """ + + self.account_data["azureAuthenticationMode"] = "OAUTH2_AUTHENTICATION_METHOD_CLIENT_SECRET" + self.account_data["azureSubscription"] = subscription + self.account_data["azureTenantId"] = tenant_id + self.account_data["azureEnvironment"] = environment + self.account_data["azureClientId"] = client_id + self.account_data["azureCertificatePath"] = certificate_path + self.account_data["azureCertificatePassword"] = certificate_password + self.account_data["azureImageResourceGroup"] = image_resource_group + return self + + def set_managed_identity(self, tenant_id, environment, managed_identity, image_resource_group): + """ + Use a Managed Identity based authentication + + :param str tenant_id: Azure Tenant identifier + :param str environment: Azure Environment + :param str image_resource_group: Azure image cached resource group + :param str managed_identity: Azure Managed Identity + """ + + self.account_data["azureAuthenticationMode"] = "MANAGED_IDENTITY" + self.account_data["azureTenantId"] = tenant_id + self.account_data["azureEnvironment"] = environment + self.account_data["azureManagedIdentityId"] = managed_identity + self.account_data["azureImageResourceGroup"] = image_resource_group + return self + + +class FMGCPCloudAccount(FMCloudAccount): + def set_same_as_fm(self): + """ + Use the same authentication as Fleet Manager + """ + self.account_data["gcpAuthenticationMode"] = "DEFAULT_INSTANCE_CREDENTIALS" + + def set_service_account_key(self, project_id, service_account_key): + """ + Use a Service Account JSON key based authentication + + :param str project_id: GCP project + :param str service_account_key: Optional, service account key (JSON) + """ + self.account_data["gcpAuthenticationMode"] = "JSON_KEY" + self.account_data["gcpProjectId"] = project_id + self.account_data["gcpServiceAccountKey"] = service_account_key + diff --git a/dataikuapi/fm/loadbalancers.py b/dataikuapi/fm/loadbalancers.py new file mode 100644 index 00000000..e7ac98e0 --- /dev/null +++ b/dataikuapi/fm/loadbalancers.py @@ -0,0 +1,398 @@ +from .future import FMFuture + +import sys + +if sys.version_info > (3, 4): + from enum import Enum +else: + class Enum(object): + pass + + +class FMLoadBalancerCreator(object): + def __init__(self, client, name, virtual_network_id): + """ + A builder class to create load balancer + + :param str name: The Name of the load balancer + :param str virtual_network_id: The id of the virtual network + """ + self.client = client + self.lb_data = {} + self.lb_data["name"] = name + self.lb_data["virtualNetworkId"] = virtual_network_id + self.lb_data["azureAssignPublicIP"] = False + self.lb_data["nodes"] = [] + + def with_description(self, description): + """ + Set the load balancer description + + :param str description + + """ + self.lb_data["description"] = description + return self + + def with_cloud_tags(self, cloud_tags): + """ + Set the tags to be applied to the cloud resources created for this load balancer + + :param dict cloud_tags: a key value dictionary of tags to be applied on the cloud resources + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancerCreator` + """ + self.lb_data["cloudTags"] = cloud_tags + return self + + def with_fm_tags(self, fm_tags): + """ + A list of tags to add on the load balancer in Fleet Manager + + :param list fm_tags: Optional, list of tags to be applied on the load balancer in the Fleet Manager + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancerCreator` + """ + self.lb_data["fmTags"] = fm_tags + return self + + def with_private_scheme(self): + """ + Setup the load balancer private + + """ + self.lb_data["publicIpMode"] = 'NO_PUBLIC_IP' + return self + + def add_node(self, hostname, instance): + """ + The node mapping to add on the load balancer in Fleet Manager + + :param str hostname: the hostname for the instance + :param :class:`dataikuapi.fm.instances.FMInstance` instance: the instance to assign to the load balancer + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancerCreator` + """ + self.lb_data["nodes"].append({ + "hostname": hostname, + "instanceId": instance.id + }) + return self + + def remove_node(self, hostname): + """ + The node mapping to remove on the load balancer in Fleet Manager + + :param str hostname: the hostname for the instance + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancerCreator` + """ + + self.lb_data["nodes"] = [x for x in self.lb_data["nodes"] if x["hostname"] != hostname] + return self + + +class FMAWSLoadBalancerCreator(FMLoadBalancerCreator): + def with_certificate_arn(self, aws_certificate_arn): + """ + Setup the certificate ARN to be used by the load balancer + + :param str aws_certificate_arn: certificate ARN + """ + self.lb_data["certificateMode"] = 'AWS_ARN' + self.lb_data["awsCertificateARN"] = aws_certificate_arn + return self + + def create(self): + """ + Create a new load balancer + + :return: a newly created load balancer + :rtype: :class:`dataikuapi.fm.loadbalancers.FMAWSLoadBalancer` + """ + lb = self.client._perform_tenant_json( + "POST", "/load-balancers", body=self.lb_data + ) + return FMAWSLoadBalancer(self.client, lb) + + +class FMAzureLoadBalancerCreator(FMLoadBalancerCreator): + + def with_tier(self, tier): + """ + Setup the tier of the load balancer + + :param :class:`dataikuapi.fm.loadbalancers.FMAzureLoadBalancerTier` tier: lb tier + """ + self.lb_data["tier"] = tier + return self + + def with_certificate_secret_id(self, azure_certificate_secret_id): + """ + Setup the certificate secret id to be used by the load balancer + + :param str azure_certificate_secret_id: certificate secret id + """ + self.lb_data["certificateMode"] = 'AZURE_SECRET_ID' + self.lb_data["azureCertificateSecretId"] = azure_certificate_secret_id + return self + + def with_public_ip(self, azure_public_ip): + """ + Setup the public IP to be used by the load balancer + + :param str azure_public_ip: public ip ID + """ + self.lb_data["publicIpMode"] = 'STATIC_PUBLIC_IP' + self.lb_data["azurePublicIPID"] = azure_public_ip + return self + + + def with_dynamic_public_ip(self): + """ + Setup the public IP to be dynamic for the load balancer + """ + self.lb_data["publicIpMode"] = 'DYNAMIC_PUBLIC_IP' + return self + + def create(self): + """ + Create a new load balancer + + :return: a newly created load balancer + :rtype: :class:`dataikuapi.fm.loadbalancers.FMAzureLoadBalancer` + """ + lb = self.client._perform_tenant_json( + "POST", "/load-balancers", body=self.lb_data + ) + return FMAzureLoadBalancer(self.client, lb) + +class FMLoadBalancer(object): + def __init__(self, client, lb_data): + self.client = client + self.lb_data = lb_data + self.id = self.lb_data["id"] + + def provision(self): + """ + Provision the physical load balancer + + :return: the `Future` object representing the reprovision process + :rtype: :class:`dataikuapi.fm.future.FMFuture` + """ + future = self.client._perform_tenant_json( + "POST", "/load-balancers/%s/actions/provision" % self.id + ) + return FMFuture.from_resp(self.client, future) + + def update(self): + """ + Update the physical load balancer + + :return: the `Future` object representing the update process + :rtype: :class:`dataikuapi.fm.future.FMFuture` + """ + future = self.client._perform_tenant_json( + "POST", "/load-balancers/%s/actions/update" % self.id + ) + return FMFuture.from_resp(self.client, future) + + def reprovision(self): + """ + Reprovision the physical load balancer + + :return: the `Future` object representing the reprovision process + :rtype: :class:`dataikuapi.fm.future.FMFuture` + """ + future = self.client._perform_tenant_json( + "POST", "/load-balancers/%s/actions/reprovision" % self.id + ) + return FMFuture.from_resp(self.client, future) + + def deprovision(self): + """ + Deprovision the physical load balancer + + :return: the `Future` object representing the deprovision process + :rtype: :class:`dataikuapi.fm.future.FMFuture` + """ + future = self.client._perform_tenant_json( + "POST", "/load-balancers/%s/actions/deprovision" % self.id + ) + return FMFuture.from_resp(self.client, future) + + def get_physical_status(self): + """ + Get the physical load balancer's status + + :return: the load balancer status + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancerPhysicalStatus` + """ + status = self.client._perform_tenant_json( + "GET", "/load-balancers/%s/physical/status" % self.id + ) + return FMLoadBalancerPhysicalStatus(status) + + + def set_description(self, description): + """ + Set the load balancer description + + :param str description" + + """ + self.lb_data["description"] = description + return self + + + def set_cloud_tags(self, cloud_tags): + """ + Set the tags to be applied to the cloud resources created for this load balancer + + :param dict cloud_tags: a key value dictionary of tags to be applied on the cloud resources + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancerCreator` + """ + self.lb_data["cloudTags"] = cloud_tags + return self + + def set_fm_tags(self, fm_tags): + """ + A list of tags to add on the load balancer in Fleet Manager + + :param list fm_tags: Optional, list of tags to be applied on the load balancer in the Fleet Manager + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancerCreator` + """ + self.lb_data["fmTags"] = fm_tags + return self + + def set_private(self): + """ + Setup the load balancer private + + """ + self.lb_data["publicIpMode"] = 'NO_PUBLIC_IP' + return self + + def add_node(self, hostname, instance): + """ + The node mapping to add on the load balancer in Fleet Manager + + :param str hostname: the hostname for the instance + :param :class:`dataikuapi.fm.instances.FMInstance` instance: the instance to assign to the load balancer + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancerCreator` + """ + self.lb_data["nodes"].append({ + "hostname": hostname, + "instanceId": instance.id + }) + return self + + def remove_node(self, hostname): + """ + The node mapping to remove on the load balancer in Fleet Manager + + :param str hostname: the hostname for the instance + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancerCreator` + """ + + self.lb_data["nodes"] = [x for x in self.lb_data["nodes"] if x["hostname"] != hostname] + return self + + def save(self): + """ + Update this load balancers. + """ + self.client._perform_tenant_empty( + "PUT", "/load-balancers/%s" % self.id, body=self.lb_data + ) + self.lb_data = self.client._perform_tenant_json( + "GET", "/load-balancers/%s" % self.id + ) + + def delete(self): + """ + Delete this load balancer. + + :return: the `Future` object representing the deletion process + :rtype: :class:`dataikuapi.fm.future.FMFuture` + """ + future = self.client._perform_tenant_json( + "DELETE", "/load-balancers/%s" % self.id + ) + return FMFuture.from_resp(self.client, future) + + +class FMAWSLoadBalancer(FMLoadBalancer): + def set_certificate_arn( + self, + certificate_arn + ): + """ + Set the certificate ARN for this load balancer + + :param str certificate_arn: the certificate ARN + """ + + self.lb_data["awsCertificateARN"] = certificate_arn + return self + +class FMAzureLoadBalancer(FMLoadBalancer): + def set_tier(self, tier): + """ + Setup the tier of the load balancer + + :param :class:`dataikuapi.fm.loadbalancers.FMAzureLoadBalancerTier` tier: lb tier + """ + self.lb_data["tier"] = tier + return self + + def set_certificate_secret_id( + self, + certificate_secret_id + ): + """ + Set the certificate secret id for this load balancer + + :param str certificate_secret_id: the certificate secret ID + """ + + self.lb_data["azureCertificateSecretId"] = certificate_secret_id + return self + + def set_public_ip(self, azure_public_ip): + """ + Setup the public IP to be used by the load balancer + + :param str azure_public_ip: public ip ID + """ + self.lb_data["publicIpMode"] = 'STATIC_PUBLIC_IP' + self.lb_data["azurePublicIPID"] = azure_public_ip + return self + + def set_dynamic_public_ip(self): + """ + Setup the public IP to be dynamic for the load balancer + """ + self.lb_data["publicIpMode"] = 'DYNAMIC_PUBLIC_IP' + return self + + +class FMLoadBalancerPhysicalStatus(dict): + """ + A class holding read-only information about an load balancer. + This class should not be created directly. Instead, use :meth:`FMLoadBalancer.get_physical_status` + """ + + def __init__(self, data): + """ + Do not call this directly, use :meth:`FMLoadBalancer.get_physical_status` + """ + super(FMLoadBalancerPhysicalStatus, self).__init__(data) + + +class FMAzureLoadBalancerTier(Enum): + WAF_V2 = "WAF_V2" + STANDARD_V2 = "STANDARD_V2" + + # Python2 emulated enum. to be removed on Python2 support removal + @staticmethod + def get_from_string(s): + if s == "WAF_V2": return FMAzureLoadBalancerTier.WAF_V2 + if s == "STANDARD_V2": return FMAzureLoadBalancerTier.STANDARD_V2 + raise Exception("Invalid load balancer tier " + s) diff --git a/dataikuapi/fm/tenant.py b/dataikuapi/fm/tenant.py index 760bc523..af619384 100644 --- a/dataikuapi/fm/tenant.py +++ b/dataikuapi/fm/tenant.py @@ -4,6 +4,7 @@ class FMCloudCredentials(object): """ A Tenant Cloud Credentials in the FM instance + """ def __init__(self, client, cloud_credentials): @@ -97,7 +98,6 @@ def save(self): """ Saves back the settings to the project """ - self.client._perform_tenant_empty( "PUT", "/cloud-credentials", body=self.cloud_credentials ) @@ -120,7 +120,6 @@ def save(self): """ Saves the tags on FM """ - self.client._perform_tenant_empty("PUT", "/cloud-tags", body=self.cloud_tags) @@ -133,6 +132,7 @@ def __init__(self, data): - :meth:`dataikuapi.fm.tenant.FMCloudAuthentication.aws_same_as_fm` to use the same authentication as Fleet Manager - :meth:`dataikuapi.fm.tenant.FMCloudAuthentication.aws_iam_role` to use a custom IAM Role - :meth:`dataikuapi.fm.tenant.FMCloudAuthentication.aws_keypair` to use a AWS Access key ID and AWS Secret Access Key pair + """ super(FMCloudAuthentication, self).__init__(data) diff --git a/dataikuapi/fm/virtualnetworks.py b/dataikuapi/fm/virtualnetworks.py index f8a77afb..838050f6 100644 --- a/dataikuapi/fm/virtualnetworks.py +++ b/dataikuapi/fm/virtualnetworks.py @@ -38,17 +38,52 @@ def with_default_values(self): self.use_default_values = True return self + def with_account(self, cloud_account=None, cloud_account_id=None): + """ + Set the Cloud Account for this virtual network + + :param cloud_account: The cloud account + :type cloud_account: :class:`dataikuapi.fm.cloudaccounts.FMCloudAccount` + :param str cloud_account_id: The cloud account identifier + + """ + if cloud_account_id is not None: + self.data["accountId"] = cloud_account_id + else: + if cloud_account is not None: + self.data["accountId"] = cloud_account.id + else: + raise ValueError("You must specify a Cloud Account or a Cloud Account identifier") + return self + + def with_auto_create_peering(self): + """ + Automatically create the network peering when creating this virtual network + """ + self.data["autoCreatePeerings"] = True + return self class FMAWSVirtualNetworkCreator(FMVirtualNetworkCreator): - def with_vpc(self, aws_vpc_id, aws_subnet_id): + def with_vpc(self, aws_vpc_id, aws_subnet_id, aws_second_subnet_id=None): """ - Setup the VPC and Subnet to be used by the virtual network + Set the VPC and Subnet to be used by the virtual network :param str aws_vpc_id: ID of the VPC to use :param str aws_subnet_id: ID of the subnet to use + :param str aws_second_subnet_id: ID of the second subnet to use """ self.data["awsVpcId"] = aws_vpc_id self.data["awsSubnetId"] = aws_subnet_id + self.data["awsSecondSubnetId"] = aws_second_subnet_id + return self + + def with_region(self, aws_region): + """ + Set the region where the VPC should be found + + :param str aws_region: the region of the VPC to use + """ + self.data["awsRegion"] = aws_region return self def with_auto_create_security_groups(self): @@ -82,15 +117,17 @@ def create(self): class FMAzureVirtualNetworkCreator(FMVirtualNetworkCreator): - def with_azure_virtual_network(self, azure_vn_id, azure_subnet_id): + def with_azure_virtual_network(self, azure_vn_id, azure_subnet_id, azure_second_subnet_id=None): """ Setup the Azure Virtual Network and Subnet to be used by the virtual network :param str azure_vn_id: Resource ID of the Azure Virtual Network to use - :param str azure_subnet_id: Resource ID of the subnet to use + :param str azure_subnet_id: Subnet name of the first subnet + :param str azure_second_subnet_id: Subnet name of the second subnet """ self.data["azureVnId"] = azure_vn_id self.data["azureSubnetId"] = azure_subnet_id + self.data["azureSecondSubnetId"] = azure_second_subnet_id return self def with_auto_update_security_groups(self, auto_update_security_groups=True): diff --git a/dataikuapi/fmclient.py b/dataikuapi/fmclient.py index 9bd8e9b3..8d9fa04a 100644 --- a/dataikuapi/fmclient.py +++ b/dataikuapi/fmclient.py @@ -9,6 +9,15 @@ from .iam.settings import FMSSOSettings, FMLDAPSettings, FMAzureADSettings from .fm.tenant import FMCloudCredentials, FMCloudTags +from .fm.cloudaccounts import ( + FMCloudAccount, + FMAWSCloudAccountCreator, + FMAzureCloudAccountCreator, + FMGCPCloudAccountCreator, + FMAWSCloudAccount, + FMAzureCloudAccount, + FMGCPCloudAccount +) from .fm.virtualnetworks import ( FMVirtualNetwork, FMAWSVirtualNetworkCreator, @@ -18,6 +27,12 @@ FMAzureVirtualNetwork, FMGCPVirtualNetwork ) +from .fm.loadbalancers import ( + FMAWSLoadBalancerCreator, + FMAzureLoadBalancerCreator, + FMAWSLoadBalancer, + FMAzureLoadBalancer +) from .fm.instances import ( FMInstance, FMInstanceEncryptionMode, @@ -85,17 +100,16 @@ def __init__( def get_tenant_id(self): return self.__tenant_id - + def get_cloud_credentials(self): """ Get the cloud credentials - :return: cloud credentials :rtype: :class:`dataikuapi.fm.tenant.FMCloudCredentials` """ creds = self._perform_tenant_json("GET", "/cloud-credentials") return FMCloudCredentials(self, creds) - + def get_sso_settings(self): """ Get the Single Sign-On (SSO) settings @@ -125,7 +139,7 @@ def get_azure_ad_settings(self): """ ldap = self._perform_tenant_json("GET", "/iam/azure-ad-settings") return FMAzureADSettings(self, ldap) - + def get_cloud_tags(self): """ Get the tenant's cloud tags @@ -136,6 +150,44 @@ def get_cloud_tags(self): tags = self._perform_tenant_json("GET", "/cloud-tags") return FMCloudTags(self, tags) + ######################################################## + # CloudAccount + ######################################################## + + def _make_cloud_account(self, account): + if self.cloud == "AWS": + return FMAWSCloudAccount(self, account) + elif self.cloud == "Azure": + return FMAzureCloudAccount(self, account) + elif self.cloud == "GCP": + return FMGCPCloudAccount(self, account) + else: + raise Exception("Unknown cloud type %s" % self.cloud) + + def list_cloud_accounts(self): + """ + List all cloud accounts + + :return: list of cloud accounts + :rtype: list of :class:`dataikuapi.fm.cloudaccounts.FMCloudAccount` + """ + vns = self._perform_tenant_json("GET", "/cloud-accounts") + return [self._make_cloud_account(x) for x in vns] + + def get_cloud_account(self, cloud_account_id): + """ + Get a cloud account by its id + + :param str cloud_account_id: the id of the cloud account to retrieve + + :return: the requested cloud account + :rtype: :class:`dataikuapi.fm.cloudaccounts.FMCloudAccount` + """ + vn = self._perform_tenant_json( + "GET", "/cloud-accounts/%s" % cloud_account_id + ) + return self._make_cloud_account(vn) + ######################################################## # VirtualNetwork ######################################################## @@ -174,6 +226,42 @@ def get_virtual_network(self, virtual_network_id): ) return self._make_virtual_network(vn) + ######################################################## + # Load balancers + ######################################################## + + def _make_load_balancer(self, vn): + if self.cloud == "AWS": + return FMAWSLoadBalancer(self, vn) + elif self.cloud == "Azure": + return FMAzureLoadBalancer(self, vn) + else: + raise Exception("Unknown cloud type %s" % self.cloud) + + def list_load_balancers(self): + """ + List all load balancers + + :return: list of load balancers + :rtype: list of :class:`dataikuapi.fm.loadbalancers.FMLoadBalancer` + """ + vns = self._perform_tenant_json("GET", "/load-balancers") + return [self._make_load_balancer(x) for x in vns] + + def get_load_balancer(self, load_balancer_id): + """ + Get a load balancer by its id + + :param str load_balancer_id: the id of the load balancer to retrieve + + :return: the requested load balancer + :rtype: :class:`dataikuapi.fm.loadbalancers.FMLoadBalancer` + """ + vn = self._perform_tenant_json( + "GET", "/load-balancers/%s" % load_balancer_id + ) + return self._make_load_balancer(vn) + ######################################################## # Instance settings template ######################################################## @@ -362,6 +450,15 @@ def __init__( host, api_key_id, api_key_secret, tenant_id, extra_headers, insecure_tls ) + def new_cloud_account_creator(self, label): + """ + Instantiate a new cloud account creator + + :param str label: The label of the cloud account + :rtype: :class:`dataikuapi.fm.cloudaccounts.FMAWSCloudAccountCreator` + """ + return FMAWSCloudAccountCreator(self, label) + def new_virtual_network_creator(self, label): """ Instantiate a new virtual network creator @@ -371,6 +468,16 @@ def new_virtual_network_creator(self, label): """ return FMAWSVirtualNetworkCreator(self, label) + def new_load_balancer_creator(self, name, virtual_network_id): + """ + Instantiate a new load balancer creator + + :param str name: The name of the load balancer + :param str virtual_network_id: The id of the virtual network + :rtype: :class:`dataikuapi.fm.loadbalancers.FMAWSLoadBalancerCreator` + """ + return FMAWSLoadBalancerCreator(self, name, virtual_network_id) + def new_instance_template_creator(self, label): """ Instantiate a new instance template creator @@ -419,6 +526,15 @@ def __init__( host, api_key_id, api_key_secret, tenant_id, extra_headers, insecure_tls ) + def new_cloud_account_creator(self, label): + """ + Instantiate a new cloud account creator + + :param str label: The label of the cloud account + :rtype: :class:`dataikuapi.fm.cloudaccounts.FMAzureCloudAccountCreator` + """ + return FMAzureCloudAccountCreator(self, label) + def new_virtual_network_creator(self, label): """ Instantiate a new virtual network creator @@ -428,6 +544,16 @@ def new_virtual_network_creator(self, label): """ return FMAzureVirtualNetworkCreator(self, label) + def new_load_balancer_creator(self, name, virtual_network_id): + """ + Instantiate a new Load balancer creator + + :param str name: The name of the load balancer + :param str virtual_network_id: The id of the virtual network + :rtype: :class:`dataikuapi.fm.loadbalancers.FMAzureLoadBalancerCreator` + """ + return FMAzureLoadBalancerCreator(self, name, virtual_network_id) + def new_instance_template_creator(self, label): """ Instantiate a new instance template creator @@ -475,6 +601,15 @@ def __init__( host, api_key_id, api_key_secret, tenant_id, extra_headers, insecure_tls ) + def new_cloud_account_creator(self, label): + """ + Instantiate a new cloud account creator + + :param str label: The label of the cloud account + :rtype: :class:`dataikuapi.fm.cloudaccounts.FMGCPCloudAccountCreator` + """ + return FMGCPCloudAccountCreator(self, label) + def new_virtual_network_creator(self, label): """ Instantiate a new virtual network creator From bdcb87fe49273190f04f472c61751fab8ddb95dd Mon Sep 17 00:00:00 2001 From: Kevin REMY Date: Mon, 9 Dec 2024 09:34:29 +0100 Subject: [PATCH 2/3] Adapt packaging --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 7ac92538..c3b8ca21 100644 --- a/setup.py +++ b/setup.py @@ -12,7 +12,7 @@ version=VERSION, license="Apache Software License", packages=["dataikuapi", "dataikuapi.dss", "dataikuapi.apinode_admin", "dataikuapi.fm", "dataikuapi.iam", - "dataikuapi.govern", "dataikuapi.dss_plugin_mlflow", "dataikuapi.dss.langchain", "dataikuapi.dss.tools"], + "dataikuapi.govern", "dataikuapi.dss_plugin_mlflow", "dataikuapi.dss.langchain", "dataikuapi.dss.tools", "dataikuapi.dss.llm_tracing"], description="Python API client for Dataiku APIs", long_description=long_description, author="Dataiku", From e1d4bd22ce38170bef486058ee69325182525c77 Mon Sep 17 00:00:00 2001 From: Kevin REMY Date: Fri, 20 Dec 2024 10:19:07 +0100 Subject: [PATCH 3/3] Release metadata --- HISTORY.txt | 5 +++++ setup.py | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/HISTORY.txt b/HISTORY.txt index aa6f4f50..34dab416 100644 --- a/HISTORY.txt +++ b/HISTORY.txt @@ -2,6 +2,11 @@ Changelog ========== +13.3.1 (2024-12-20) +--------------------- + +* Initial release for DSS 13.3.1 + 13.2.4 (2024-12-03) --------------------- diff --git a/setup.py b/setup.py index c3b8ca21..9caa8809 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ from setuptools import setup -VERSION = "13.2.4" +VERSION = "13.3.1" long_description = (open('README').read() + '\n\n' + open('HISTORY.txt').read())